Update app.py

app.py

@@ -66,11 +66,11 @@ def detect_and_draw_image(input_image):
         print("Model prediction completed.")
 
         # دسترسی به نتایج OBB
-        if hasattr(results[0], 'boxes') and results[0].boxes is not None:
-            obb_results = results[0].boxes  # استفاده از boxes که شامل OBB است
+        if hasattr(results[0], 'obb') and results[0].obb is not None:
+            obb_results = results[0].obb  # استفاده از obb برای مدل OBB
             print("Accessed obb_results.")
         else:
-            print("No 'boxes' attribute found in results[0].")
+            print("No 'obb' attribute found in results[0].")
             obb_results = None
 
         # بررسی وجود جعبه‌های شناسایی شده
@@ -90,14 +90,10 @@ def detect_and_draw_image(input_image):
         print("Image loaded with OpenCV.")
 
         counts = {}
-        for box in obb_results:
-            # دسترسی به داده‌های جعبه
-            if hasattr(box, 'xywhn'):
-                x_center, y_center, width, height, rotation = box.xywhn.tolist()
-            else:
-                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
+        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)  # تبدیل رادیان به درجه
@@ -171,22 +167,18 @@ def detect_and_draw_video(video_path):
             print(f"Model prediction completed for frame {frame_count}.")
 
             # دسترسی به نتایج OBB
-            if hasattr(results[0], 'boxes') and results[0].boxes is not None:
-                obb_results = results[0].boxes  # استفاده از boxes که شامل OBB است
+            if hasattr(results[0], 'obb') and results[0].obb is not None:
+                obb_results = results[0].obb  # استفاده از obb برای مدل OBB
                 print("Accessed obb_results for frame.")
             else:
-                print("No 'boxes' attribute found in results[0] for frame.")
+                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 box in obb_results:
-                    # دسترسی به داده‌های جعبه
-                    if hasattr(box, 'xywhn'):
-                        x_center, y_center, width, height, rotation = box.xywhn.tolist()
-                    else:
-                        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
+                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)  # تبدیل رادیان به درجه