Update app.py

app.py

@@ -11,7 +11,8 @@ import tempfile
 import os
 
 # بارگذاری مدل آموزش‌دیده شما
-model = YOLO('/content/yolo11n-obb.pt')  # اطمینان حاصل کنید که مسیر مدل صحیح است
+# اطمینان حاصل کنید که فقط یک بار مدل را بارگذاری می‌کنید و مسیر صحیح است
+model = YOLO('/content/yolo11n-obb.pt')  # یا '/content/best.pt' بر اساس مدل مورد نظر شما
 
 # تعریف نام کلاس‌ها به انگلیسی و فارسی
 class_names = {
@@ -53,15 +54,72 @@ colors = {
 
 # تابع برای تشخیص اشیاء در تصاویر
 def detect_and_draw_image(input_image):
-    # تبدیل تصویر PIL به آرایه NumPy
-    input_image_np = np.array(input_image)
+    try:
+        # تبدیل تصویر PIL به آرایه NumPy
+        input_image_np = np.array(input_image)
+
+        # اجرای مدل روی تصویر
+        results = model.predict(source=input_image_np, conf=0.3)
+
+        # دسترسی به نتایج OBB
+        obb_results = results[0].boxes  # یا results[0].obb بر اساس نسخه مدل شما
+
+        # بررسی وجود جعبه‌های شناسایی شده
+        if obb_results is None or len(obb_results) == 0:
+            print("هیچ شیء شناسایی نشده است.")
+            df = pd.DataFrame({
+                'Label (English)': [],
+                'Label (Persian)': [],
+                'Object Count': []
+            })
+            return input_image, df
+
+        # بارگذاری تصویر اصلی به صورت OpenCV برای رسم جعبه‌ها
+        image_cv = cv2.cvtColor(input_image_np, cv2.COLOR_RGB2BGR)
+
+        counts = {}
+        for box in obb_results:
+            # دریافت مختصات جعبه
+            x_center, y_center, width, height, rotation = box.xywhn.tolist()[0]
+            class_id = int(box.cls.tolist()[0])
+            confidence = float(box.conf.tolist()[0])
+
+            # تبدیل مختصات از نرمال شده به پیکسل
+            img_height, img_width, _ = image_cv.shape
+            x_center *= img_width
+            y_center *= img_height
+            width *= img_width
+            height *= img_height
+
+            # دریافت برچسب‌ها
+            label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
+            counts[label_en] = counts.get(label_en, 0) + 1
+
+            # رسم جعبه چرخان با استفاده از OpenCV
+            rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
+            box_points = cv2.boxPoints(rect)
+            box_points = np.int0(box_points)
+            color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
+            cv2.drawContours(image_cv, [box_points], 0, color, 2)
+            cv2.putText(image_cv, f'{label_en}: {confidence:.2f}', 
+                        (int(x_center), int(y_center)), 
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2, cv2.LINE_AA)
+
+        # تبدیل تصویر به RGB برای Gradio
+        image_rgb = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
+        output_image = Image.fromarray(image_rgb)
+
+        # ایجاد DataFrame برای نمایش نتایج
+        df = pd.DataFrame({
+            'Label (English)': list(counts.keys()),
+            'Label (Persian)': [class_names.get(k, ('unknown', 'ناشناخته'))[1] for k in counts.keys()],
+            'Object Count': list(counts.values())
+        })
 
-    # اجرای مدل روی تصویر
-    results = model(input_image_np)
+        return output_image, df
 
-    # بررسی کردن اینکه آیا اشیاء شناسایی شده اند یا خیر
-    if not results or not results[0].boxes:
-        # اگر هیچ شیء شناسایی نشده باشد
+    except Exception as e:
+        print(f"Error in detect_and_draw_image: {e}")
         df = pd.DataFrame({
             'Label (English)': [],
             'Label (Persian)': [],
@@ -69,102 +127,93 @@ def detect_and_draw_image(input_image):
         })
         return input_image, df
 
-    detections = results[0].boxes  # دسترسی به نتایج در YOLOv8
-
-    # تبدیل تصویر به RGB برای رسم
-    image_draw = input_image.convert("RGB")
-    draw = ImageDraw.Draw(image_draw)
-    counts = {}
-    for box in detections:
-        # دریافت مختصات جعبه
-        xmin, ymin, xmax, ymax = box.xyxy.tolist()
-        conf = box.conf.tolist()[0]
-        class_id = int(box.cls.tolist()[0])
-
-        # دریافت برچسب‌ها
-        label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
-        counts[label_en] = counts.get(label_en, 0) + 1
-
-        # رسم مستطیل
-        draw.rectangle([(xmin, ymin), (xmax, ymax)], outline=colors.get(class_id, (255,0,0)), width=2)
-        # رسم برچسب
-        draw.text((xmin, ymin), f"{label_en}: {conf:.2f}", fill="white")
-
-    # ایجاد DataFrame
-    df = pd.DataFrame({
-        'Label (English)': list(counts.keys()),
-        'Label (Persian)': [class_names.get(k, ('unknown', 'ناشناخته'))[1] for k in counts.keys()],
-        'Object Count': list(counts.values())
-    })
-
-    return image_draw, df
-
 # تابع برای تشخیص اشیاء در ویدئوها
 def detect_and_draw_video(video_path):
-    cap = cv2.VideoCapture(video_path)
-    frames = []
-    overall_counts = {}
-    seen_objects = []  # لیست برای دنبال کردن اشیاء شناسایی شده
-
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        # تغییر اندازه فریم
-        frame = cv2.resize(frame, (640, 480))
-        # تبدیل رنگ از BGR به RGB
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-
-        # اجرای مدل روی فریم
-        results = model(frame_rgb)
-
-        if results and results[0].boxes:
-            detections = results[0].boxes
-            for box in detections:
-                # دریافت مختصات جعبه
-                xmin, ymin, xmax, ymax = box.xyxy.tolist()
-                conf = box.conf.tolist()[0]
-                class_id = int(box.cls.tolist()[0])
-                label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
-                current_object = (label_en, int(xmin), int(ymin), int(xmax), int(ymax))
-
-                # بررسی وجود شیء در لیست seen_objects
-                if not any(existing[0] == label_en and
-                           (existing[1] < current_object[3] and existing[3] > current_object[1] and
-                            existing[2] < current_object[4] and existing[4] > current_object[2]) for existing in seen_objects):
-                    seen_objects.append(current_object)
-                    overall_counts[label_en] = overall_counts.get(label_en, 0) + 1
-
-                # رسم مستطیل
-                color = colors.get(class_id, (255, 0, 0))
-                cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, 2)
-                # رسم برچسب
-                cv2.putText(frame, f"{label_en}: {conf:.2f}", (int(xmin), int(ymin) - 10),
-                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
-
-        frames.append(frame)
-
-    cap.release()
-
-    # ذخیره ویدئو پردازش‌شده در یک فایل موقت
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmpfile:
-        output_path = tmpfile.name
-
-    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 20.0, (640, 480))
-
-    for frame in frames:
-        out.write(frame)
-    out.release()
-
-    # ایجاد DataFrame برای ذخیره نتایج
-    df = pd.DataFrame({
-        'Label (English)': list(overall_counts.keys()),
-        'Label (Persian)': [class_names.get(k, ('unknown', 'ناشناخته'))[1] for k in overall_counts.keys()],
-        'Object Count': list(overall_counts.values())
-    })
-
-    return output_path, df
+    try:
+        cap = cv2.VideoCapture(video_path)
+        frames = []
+        overall_counts = {}
+        seen_objects = []  # لیست برای دنبال کردن اشیاء شناسایی شده
+
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+
+            # تغییر اندازه فریم
+            frame = cv2.resize(frame, (640, 480))
+            # اجرای مدل روی فریم
+            results = model.predict(source=frame, conf=0.3)
+            obb_results = results[0].boxes  # یا results[0].obb بر اساس نسخه مدل شما
+
+            if obb_results is not None and len(obb_results) > 0:
+                for box in obb_results:
+                    # دریافت مختصات جعبه
+                    x_center, y_center, width, height, rotation = box.xywhn.tolist()[0]
+                    class_id = int(box.cls.tolist()[0])
+                    confidence = float(box.conf.tolist()[0])
+
+                    # تبدیل مختصات از نرمال شده به پیکسل
+                    img_height, img_width, _ = frame.shape
+                    x_center *= img_width
+                    y_center *= img_height
+                    width *= img_width
+                    height *= img_height
+
+                    # دریافت برچسب‌ها
+                    label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
+                    current_object = (label_en, int(x_center - width / 2), int(y_center - height / 2), 
+                                      int(x_center + width / 2), int(y_center + height / 2))
+
+                    # بررسی وجود شیء در لیست seen_objects
+                    if not any(existing[0] == label_en and
+                               (existing[1] < current_object[3] and existing[3] > current_object[1] and
+                                existing[2] < current_object[4] and existing[4] > current_object[2]) 
+                               for existing in seen_objects):
+                        seen_objects.append(current_object)
+                        overall_counts[label_en] = overall_counts.get(label_en, 0) + 1
+
+                    # رسم جعبه چرخان با استفاده از OpenCV
+                    rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
+                    box_points = cv2.boxPoints(rect)
+                    box_points = np.int0(box_points)
+                    color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
+                    cv2.drawContours(frame, [box_points], 0, color, 2)
+                    cv2.putText(frame, f"{label_en}: {confidence:.2f}", 
+                                (int(x_center), int(y_center)), 
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2, cv2.LINE_AA)
+
+            frames.append(frame)
+
+        cap.release()
+
+        # ذخیره ویدئو پردازش‌شده در یک فایل موقت
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmpfile:
+            output_path = tmpfile.name
+
+        out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 20.0, (640, 480))
+
+        for frame in frames:
+            out.write(frame)
+        out.release()
+
+        # ایجاد DataFrame برای ذخیره نتایج
+        df = pd.DataFrame({
+            'Label (English)': list(overall_counts.keys()),
+            'Label (Persian)': [class_names.get(k, ('unknown', 'ناشناخته'))[1] for k in overall_counts.keys()],
+            'Object Count': list(overall_counts.values())
+        })
+
+        return output_path, df
+
+    except Exception as e:
+        print(f"Error in detect_and_draw_video: {e}")
+        # در صورت بروز خطا، بازگرداندن ویدئوی اصلی بدون تغییر و یک DataFrame خالی
+        return video_path, pd.DataFrame({
+            'Label (English)': [],
+            'Label (Persian)': [],
+            'Object Count': []
+        })
 
 # رابط کاربری تصویر
 image_interface = gr.Interface(