Update app.py

app.py

@@ -11,7 +11,6 @@ import tempfile
 import os
 
 # بارگذاری مدل آموزش‌دیده شما
-# اطمینان حاصل کنید که مسیر مدل صحیح است و فقط یک بار مدل را بارگذاری می‌کنید
 model = YOLO('weights/best.pt')  # یا '/content/best.pt' بر اساس مدل مورد نظر شما
 
 # تعریف نام کلاس‌ها به انگلیسی و فارسی
@@ -51,23 +50,24 @@ colors = {
     13: (75, 0, 130),     # ایندیگو
     14: (199, 21, 133)    # رز متوسط
 }
-
 # تابع برای تشخیص اشیاء در تصاویر
 def detect_and_draw_image(input_image):
     try:
-        # ذخیره تصویر ورودی به یک فایل موقت
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmpfile:
-            input_image.save(tmpfile.name)
-            temp_image_path = tmpfile.name
-        print("Image saved to temporary file:", temp_image_path)
-
-        # اجرای مدل روی تصویر با استفاده از مسیر فایل
-        results = model.predict(source=temp_image_path, conf=0.3)
+        # تبدیل تصویر PIL به آرایه NumPy (RGB)
+        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)
         print("Model prediction completed.")
 
         # دسترسی به نتایج OBB
         if hasattr(results[0], 'obb') and results[0].obb is not None:
-            obb_results = results[0].obb  # استفاده از obb برای مدل OBB
+            obb_results = results[0].obb
             print("Accessed obb_results.")
         else:
             print("No 'obb' attribute found in results[0].")
@@ -81,15 +81,10 @@ def detect_and_draw_image(input_image):
                 'Label (Persian)': [],
                 'Object Count': []
             })
-            # حذف فایل موقت
-            os.remove(temp_image_path)
             return input_image, df
 
-        # بارگذاری تصویر اصلی به صورت OpenCV برای رسم جعبه‌ها (BGR)
-        image_cv = cv2.imread(temp_image_path)
-        print("Image loaded with OpenCV.")
-
         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)
@@ -99,14 +94,14 @@ def detect_and_draw_image(input_image):
             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))  # استفاده از رنگ مشخص برای هر کلاس
+            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', 'ناشناخته'))
-            cv2.putText(image_cv, f'{label_en}: {confidence:.2f}', 
-                        (int(x_center), int(y_center)), 
+            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)
 
             # شمارش اشیاء
@@ -125,9 +120,6 @@ def detect_and_draw_image(input_image):
         })
         print("DataFrame created.")
 
-        # حذف فایل موقت
-        os.remove(temp_image_path)
-
         return output_image, df
 
     except Exception as e:
@@ -145,7 +137,6 @@ def detect_and_draw_video(video_path):
         cap = cv2.VideoCapture(video_path)
         frames = []
         overall_counts = {}
-        seen_objects = []  # لیست برای دنبال کردن اشیاء شناسایی شده
         frame_count = 0
 
         while cap.isOpened():
@@ -159,16 +150,13 @@ 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_rgb, conf=0.3)
+            results = model.predict(source=frame, conf=0.3)
             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  # استفاده از obb برای مدل OBB
+                obb_results = results[0].obb
                 print("Accessed obb_results for frame.")
             else:
                 print("No 'obb' attribute found in results[0] for frame.")
@@ -181,17 +169,17 @@ def detect_and_draw_video(video_path):
                     confidence = float(conf)
 
                     # رسم جعبه چرخان با استفاده از OpenCV
-                    rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
+                    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))  # استفاده از رنگ مشخص برای هر کلاس
+                    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', 'ناشناخته'))
-                    cv2.putText(frame, f"{label_en}: {confidence:.2f}", 
-                                (int(x_center), int(y_center)), 
+                    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)
 
                     # شمارش اشیاء
@@ -244,7 +232,7 @@ image_interface = gr.Interface(
     title="تشخیص اشیاء در تصاویر هوایی",
     description="یک تصویر هوایی بارگذاری کنید تا اشیاء شناسایی شده و تعداد آن‌ها را ببینید.",
     examples=[
-        'Examples/images/areial_car.jpg', 
+        'Examples/images/areial_car.jpg',
         'Examples/images/arieal_car_1.jpg',
         'Examples/images/t.jpg'
     ]
@@ -258,7 +246,7 @@ video_interface = gr.Interface(
     title="تشخیص اشیاء در ویدئوها",
     description="یک ویدئو بارگذاری کنید تا اشیاء شناسایی شده و تعداد آن‌ها را ببینید.",
     examples=[
-        'Examples/video/city.mp4', 
+        'Examples/video/city.mp4',
         'Examples/video/airplane.mp4'
     ]
 )