Update app.py

app.py

@@ -3,7 +3,7 @@
 
 import cv2
 from ultralytics import YOLO
-from PIL import Image, ImageDraw, ImageFont
+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')  # اطمینان حاصل کنید که مسیر مدل در محیط Spaces صحیح است
+model = YOLO('/content/best.pt')  # یا '/content/best.pt' بر اساس مدل مورد نظر شما
 
 # تعریف نام کلاس‌ها به انگلیسی و فارسی
 class_names = {
@@ -32,7 +32,7 @@ class_names = {
     14: ('swimming pool', 'استخر شنا')
 }
 
-# رنگ‌ها برای هر کلاس (RGB برای PIL)
+# رنگ‌ها برای هر کلاس (BGR برای OpenCV)
 colors = {
     0: (255, 0, 0),       # قرمز
     1: (0, 255, 0),       # سبز
@@ -44,19 +44,12 @@ colors = {
     7: (255, 165, 0),     # نارنجی
     8: (0, 128, 0),       # سبز تیره
     9: (128, 128, 0),     # زیتونی
-    10: (0, 255, 0),      # سبز روشن برای class_id=10
+    10: (128, 0, 0),      # سرخ کلید
     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:
@@ -64,8 +57,12 @@ def detect_and_draw_image(input_image):
         input_image_np = np.array(input_image)
         print("Image converted to NumPy array.")
 
-        # اجرای مدل روی تصویر با استفاده از آرایه NumPy (RGB)
-        results = model.predict(source=input_image_np, conf=0.3)
+        # تبدیل تصویر به 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)
         print("Model prediction completed.")
 
         # دسترسی به نتایج OBB
@@ -87,37 +84,34 @@ def detect_and_draw_image(input_image):
             return input_image, df
 
         counts = {}
-        draw = ImageDraw.Draw(input_image)
-
         # پردازش نتایج و رسم جعبه‌ها
-        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)
+        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)
-
-            # تبدیل نقاط به لیست برای رسم در 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)))
+            color = colors.get(class_id, (0, 255, 0))
+            cv2.drawContours(image_cv, [box_points], 0, color, 2)
             print(f"Drawn OBB for class_id {class_id} with confidence {confidence}.")
 
             # رسم برچسب
             label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
-            draw.text((int(x_center), int(y_center)), f'{label_en}: {confidence:.2f}', fill=colors.get(class_id, (0, 255, 0)), font=font)
+            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)
+        print("Image converted back to RGB for Gradio.")
+
         # ایجاد DataFrame برای نمایش نتایج
         df = pd.DataFrame({
             'Label (English)': list(counts.keys()),
@@ -126,7 +120,7 @@ def detect_and_draw_image(input_image):
         })
         print("DataFrame created.")
 
-        return input_image, df
+        return output_image, df
 
     except Exception as e:
         print(f"Error in detect_and_draw_image: {e}")
@@ -169,40 +163,28 @@ def detect_and_draw_video(video_path):
                 obb_results = None
 
             if obb_results is not None and len(obb_results.data) > 0:
-                # تبدیل فریم به 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
+                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)
-
-                    # تبدیل نقاط به لیست برای رسم در 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)))
+                    color = colors.get(class_id, (0, 255, 0))
+                    cv2.drawContours(frame, [box_points], 0, color, 2)
                     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', 'ناشناخته'))
-                    draw.text((int(x_center), int(y_center)), f"{label_en}: {confidence:.2f}", fill=colors.get(class_id, (0, 255, 0)), font=font)
+                    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
 
-                # تبدیل فریم به آرایه NumPy برای ذخیره
-                frame = cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR)
-
             frames.append(frame)
             print(f"Frame {frame_count} processed.")