Update app.py

app.py

@@ -3,7 +3,7 @@
 
 import cv2
 from ultralytics import YOLO
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 import gradio as gr
 import pandas as pd
 import numpy as np
@@ -32,7 +32,7 @@ class_names = {
     14: ('swimming pool', 'استخر شنا')
 }
 
-# رنگ‌ها برای هر کلاس (BGR برای OpenCV)
+# رنگ‌ها برای هر کلاس (RGB برای PIL)
 colors = {
     0: (255, 0, 0),       # قرمز
     1: (0, 255, 0),       # سبز
@@ -44,13 +44,19 @@ colors = {
     7: (255, 165, 0),     # نارنجی
     8: (0, 128, 0),       # سبز تیره
     9: (128, 128, 0),     # زیتونی
-    10: (0, 255, 0),      # تغییر به سبز روشن برای class_id=10
+    10: (0, 255, 0),      # سبز روشن برای class_id=10
     11: (0, 128, 128),    # سبز نفتی
     12: (0, 0, 128),      # نیوی
     13: (75, 0, 130),     # ایندیگو
     14: (199, 21, 133)    # رز متوسط
 }
 
+# فونت برای نوشتن برچسب‌ها
+try:
+    font = ImageFont.truetype("arial.ttf", 15)
+except IOError:
+    font = ImageFont.load_default()
+
 # تابع برای تشخیص اشیاء در تصاویر
 def detect_and_draw_image(input_image):
     try:
@@ -58,20 +64,16 @@ def detect_and_draw_image(input_image):
         input_image_np = np.array(input_image)
         print("Image converted to NumPy array.")
 
-        # تبدیل تصویر به BGR برای OpenCV
-        image_cv = cv2.cvtColor(input_image_np, cv2.COLOR_RGB2BGR)
-        print("Image converted to OpenCV format.")
-
         # اجرای مدل روی تصویر با استفاده از آرایه NumPy
-        results = model.predict(source=image_cv, conf=0.3)
+        results = model.predict(source=input_image_np, conf=0.3)
         print("Model prediction completed.")
 
         # دسترسی به نتایج OBB
-        if hasattr(results[0], 'obb') and results[0].obb is not None:
-            obb_results = results[0].obb
+        if hasattr(results[0], 'boxes') and results[0].boxes is not None:
+            obb_results = results[0].boxes
             print("Accessed obb_results.")
         else:
-            print("No 'obb' attribute found in results[0].")
+            print("No 'boxes' attribute found in results[0].")
             obb_results = None
 
         # بررسی وجود جعبه‌های شناسایی شده
@@ -85,34 +87,37 @@ def detect_and_draw_image(input_image):
             return input_image, df
 
         counts = {}
-        # پردازش نتایج و رسم جعبه‌ها
-        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)
+        draw = ImageDraw.Draw(input_image)
 
-            # رسم جعبه چرخان با استفاده از OpenCV
-            rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
+        # پردازش نتایج و رسم جعبه‌ها
+        for box in obb_results:
+            # دسترسی به داده‌های جعبه
+            try:
+                x_center, y_center, width, height, rotation = box.xywh.tolist()
+            except:
+                x_center, y_center, width, height, rotation = box.tolist()[:5]
+            class_id = int(box.cls.tolist()[0]) if hasattr(box, 'cls') else 0
+            confidence = float(box.conf.tolist()[0]) if hasattr(box, 'conf') else 0.0
+
+            # محاسبه مختصات گوشه‌ها با استفاده از 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, 3)  # افزایش ضخامت جعبه
+
+            # تبدیل نقاط به لیست برای رسم در PIL
+            box_points = [(int(point[0]), int(point[1])) for point in box_points]
+
+            # رسم جعبه
+            draw.polygon(box_points, outline=colors.get(class_id, (0, 255, 0)))
             print(f"Drawn OBB for class_id {class_id} with confidence {confidence}.")
 
             # رسم برچسب
             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)
+            draw.text((int(x_center), int(y_center)), f'{label_en}: {confidence:.2f}', fill=colors.get(class_id, (0, 255, 0)), font=font)
 
             # شمارش اشیاء
             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)
-        print("Image converted back to RGB for Gradio.")
-
         # ایجاد DataFrame برای نمایش نتایج
         df = pd.DataFrame({
             'Label (English)': list(counts.keys()),
@@ -121,7 +126,7 @@ def detect_and_draw_image(input_image):
         })
         print("DataFrame created.")
 
-        return output_image, df
+        return input_image, df
 
     except Exception as e:
         print(f"Error in detect_and_draw_image: {e}")
@@ -156,36 +161,48 @@ def detect_and_draw_video(video_path):
             print(f"Model prediction completed for frame {frame_count}.")
 
             # دسترسی به نتایج OBB
-            if hasattr(results[0], 'obb') and results[0].obb is not None:
-                obb_results = results[0].obb
+            if hasattr(results[0], 'boxes') and results[0].boxes is not None:
+                obb_results = results[0].boxes
                 print("Accessed obb_results for frame.")
             else:
                 print("No 'obb' attribute found in results[0] for frame.")
                 obb_results = None
 
             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
+                # تبدیل فریم به PIL Image برای رسم
+                frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+                draw = ImageDraw.Draw(frame_pil)
+
+                for box in obb_results:
+                    try:
+                        x_center, y_center, width, height, rotation = box.xywh.tolist()
+                    except:
+                        x_center, y_center, width, height, rotation = box.tolist()[:5]
+                    class_id = int(box.cls.tolist()[0]) if hasattr(box, 'cls') else 0
+                    confidence = float(box.conf.tolist()[0]) if hasattr(box, 'conf') else 0.0
+
+                    # محاسبه مختصات گوشه‌ها با استفاده از 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, 3)  # افزایش ضخامت جعبه
+
+                    # تبدیل نقاط به لیست برای رسم در PIL
+                    box_points = [(int(point[0]), int(point[1])) for point in box_points]
+
+                    # رسم جعبه
+                    draw.polygon(box_points, outline=colors.get(class_id, (0, 255, 0)))
                     print(f"Drawn OBB for class_id {class_id} with confidence {confidence} in frame {frame_count}.")
 
                     # رسم برچسب
                     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)
+                    draw.text((int(x_center), int(y_center)), f"{label_en}: {confidence:.2f}", fill=colors.get(class_id, (0, 255, 0)), font=font)
 
                     # شمارش اشیاء
                     overall_counts[label_en] = overall_counts.get(label_en, 0) + 1
 
+                # تبدیل فریم به آرایه NumPy برای ذخیره
+                frame = cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR)
+
             frames.append(frame)
             print(f"Frame {frame_count} processed.")