Update app.py

app.py

@@ -3,7 +3,7 @@
 
 import cv2
 from ultralytics import YOLO
-from PIL import Image, ImageDraw
+from PIL import Image
 import gradio as gr
 import pandas as pd
 import numpy as np
@@ -11,7 +11,7 @@ import tempfile
 import os
 
 # بارگذاری مدل آموزش‌دیده شما
-# اطمینان حاصل کنید که فقط یک بار مدل را بارگذاری می‌کنید و مسیر صحیح است
+# اطمینان حاصل کنید که مسیر مدل صحیح است
 model = YOLO('weights/best.pt')  # یا '/content/best.pt' بر اساس مدل مورد نظر شما
 
 # تعریف نام کلاس‌ها به انگلیسی و فارسی
@@ -33,7 +33,7 @@ class_names = {
     14: ('swimming pool', 'استخر شنا')
 }
 
-# رنگ‌ها برای هر کلاس
+# رنگ‌ها برای هر کلاس (BGR برای OpenCV)
 colors = {
     0: (255, 0, 0),       # قرمز
     1: (0, 255, 0),       # سبز
@@ -55,17 +55,17 @@ colors = {
 # تابع برای تشخیص اشیاء در تصاویر
 def detect_and_draw_image(input_image):
     try:
-        # تبدیل تصویر PIL به آرایه NumPy
+        # تبدیل تصویر PIL به آرایه NumPy (RGB)
         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 بر اساس نسخه مدل شما
+        obb_results = results[0].obb
 
         # بررسی وجود جعبه‌های شناسایی شده
-        if obb_results is None or len(obb_results) == 0:
+        if obb_results is None or len(obb_results.data) == 0:
             print("هیچ شیء شناسایی نشده است.")
             df = pd.DataFrame({
                 'Label (English)': [],
@@ -74,37 +74,31 @@ def detect_and_draw_image(input_image):
             })
             return input_image, df
 
-        # بارگذاری تصویر اصلی به صورت OpenCV برای رسم جعبه‌ها
+        # بارگذاری تصویر اصلی به صورت OpenCV برای رسم جعبه‌ها (BGR)
         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
+        for obb, conf, cls in zip(obb_results.data.cpu().numpy(), obb_results.conf.cpu().numpy(), obb_results.cls.cpu().numpy()):
+            x_center, y_center, width, height, rotation = obb[:5]
+            class_id = int(cls)
+            confidence = float(conf)
 
             # رسم جعبه چرخان با استفاده از OpenCV
             rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
-            box_points = cv2.boxPoints(rect)
-            box_points = np.int0(box_points)
+            box = cv2.boxPoints(rect)
+            box = np.int0(box)
             color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
-            cv2.drawContours(image_cv, [box_points], 0, color, 2)
+            cv2.drawContours(image_cv, [box], 0, color, 2)
+
+            # رسم برچسب
+            label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
             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)
 
+            # شمارش اشیاء
+            counts[label_en] = counts.get(label_en, 0) + 1
+
         # تبدیل تصویر به RGB برای Gradio
         image_rgb = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
         output_image = Image.fromarray(image_rgb)
@@ -142,47 +136,38 @@ def detect_and_draw_video(video_path):
 
             # تغییر اندازه فریم
             frame = cv2.resize(frame, (640, 480))
+
+            # تبدیل BGR به RGB برای مدل
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
             # اجرای مدل روی فریم
-            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))
+            results = model.predict(source=frame_rgb, conf=0.3)
+
+            # دسترسی به نتایج OBB
+            obb_results = results[0].obb
 
-                    # بررسی وجود شیء در لیست 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
+            if obb_results is not None and len(obb_results.data) > 0:
+                for obb, conf, cls in zip(obb_results.data.cpu().numpy(), obb_results.conf.cpu().numpy(), obb_results.cls.cpu().numpy()):
+                    x_center, y_center, width, height, rotation = obb[:5]
+                    class_id = int(cls)
+                    confidence = float(conf)
 
                     # رسم جعبه چرخان با استفاده از OpenCV
                     rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
-                    box_points = cv2.boxPoints(rect)
-                    box_points = np.int0(box_points)
+                    box = cv2.boxPoints(rect)
+                    box = np.int0(box)
                     color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
-                    cv2.drawContours(frame, [box_points], 0, color, 2)
+                    cv2.drawContours(frame, [box], 0, color, 2)
+
+                    # رسم برچسب
+                    label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
                     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)
 
+                    # شمارش اشیاء
+                    overall_counts[label_en] = overall_counts.get(label_en, 0) + 1
+
             frames.append(frame)
 
         cap.release()
@@ -215,24 +200,31 @@ def detect_and_draw_video(video_path):
             'Object Count': []
         })
 
-# رابط کاربری تصویر
+# رابط کاربری Gradio برای تصاویر
 image_interface = gr.Interface(
     fn=detect_and_draw_image,
     inputs=gr.Image(type="pil", label="بارگذاری تصویر"),
-    outputs=[gr.Image(type="pil"), gr.Dataframe(label="تعداد اشیاء")],
+    outputs=[gr.Image(type="pil", label="تصویر پردازش شده"), gr.Dataframe(label="تعداد اشیاء")],
     title="تشخیص اشیاء در تصاویر هوایی",
     description="یک تصویر هوایی بارگذاری کنید تا اشیاء شناسایی شده و تعداد آن‌ها را ببینید.",
-    examples=['Examples/images/areial_car.jpg', 'Examples/images/arieal_car_1.jpg','Examples/images/t.jpg']
+    examples=[
+        'Examples/images/areial_car.jpg', 
+        'Examples/images/arieal_car_1.jpg',
+        'Examples/images/t.jpg'
+    ]
 )
 
-# رابط کاربری ویدئو
+# رابط کاربری Gradio برای ویدئوها
 video_interface = gr.Interface(
     fn=detect_and_draw_video,
     inputs=gr.Video(label="بارگذاری ویدئو"),
     outputs=[gr.Video(label="ویدئوی پردازش شده"), gr.Dataframe(label="تعداد اشیاء")],
     title="تشخیص اشیاء در ویدئوها",
     description="یک ویدئو بارگذاری کنید تا اشیاء شناسایی شده و تعداد آن‌ها را ببینید.",
-    examples=['Examples/video/city.mp4', 'Examples/video/airplane.mp4']
+    examples=[
+        'Examples/video/city.mp4', 
+        'Examples/video/airplane.mp4'
+    ]
 )
 
 # اجرای برنامه با استفاده از رابط کاربری تب‌دار