Update app.py

app.py

@@ -11,7 +11,7 @@ import tempfile
 import os
 
 # بارگذاری مدل آموزش‌دیده شما
-# اطمینان حاصل کنید که مسیر مدل صحیح است
+# اطمینان حاصل کنید که مسیر مدل صحیح است و فقط یک بار مدل را بارگذاری می‌کنید
 model = YOLO('weights/best.pt')  # یا '/content/best.pt' بر اساس مدل مورد نظر شما
 
 # تعریف نام کلاس‌ها به انگلیسی و فارسی
@@ -57,12 +57,19 @@ def detect_and_draw_image(input_image):
     try:
         # تبدیل تصویر PIL به آرایه NumPy (RGB)
         input_image_np = np.array(input_image)
+        print("Image converted to NumPy array.")
 
         # اجرای مدل روی تصویر
         results = model.predict(source=input_image_np, conf=0.3)
+        print("Model prediction completed.")
 
         # دسترسی به نتایج OBB
-        obb_results = results[0].obb
+        if hasattr(results[0], 'obb') and results[0].obb is not None:
+            obb_results = results[0].obb
+            print("Accessed obb_results.")
+        else:
+            print("No 'obb' attribute found in results[0].")
+            obb_results = None
 
         # بررسی وجود جعبه‌های شناسایی شده
         if obb_results is None or len(obb_results.data) == 0:
@@ -76,6 +83,7 @@ def detect_and_draw_image(input_image):
 
         # بارگذاری تصویر اصلی به صورت OpenCV برای رسم جعبه‌ها (BGR)
         image_cv = cv2.cvtColor(input_image_np, cv2.COLOR_RGB2BGR)
+        print("Image converted to BGR for OpenCV.")
 
         counts = {}
         for obb, conf, cls in zip(obb_results.data.cpu().numpy(), obb_results.conf.cpu().numpy(), obb_results.cls.cpu().numpy()):
@@ -89,6 +97,7 @@ def detect_and_draw_image(input_image):
             box = np.int0(box)
             color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
             cv2.drawContours(image_cv, [box], 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', 'ناشناخته'))
@@ -102,6 +111,7 @@ def detect_and_draw_image(input_image):
         # تبدیل تصویر به 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({
@@ -109,6 +119,7 @@ def detect_and_draw_image(input_image):
             'Label (Persian)': [class_names.get(k, ('unknown', 'ناشناخته'))[1] for k in counts.keys()],
             'Object Count': list(counts.values())
         })
+        print("DataFrame created.")
 
         return output_image, df
 
@@ -128,12 +139,16 @@ def detect_and_draw_video(video_path):
         frames = []
         overall_counts = {}
         seen_objects = []  # لیست برای دنبال کردن اشیاء شناسایی شده
+        frame_count = 0
 
         while cap.isOpened():
             ret, frame = cap.read()
             if not ret:
                 break
 
+            frame_count +=1
+            print(f"Processing frame {frame_count}")
+
             # تغییر اندازه فریم
             frame = cv2.resize(frame, (640, 480))
 
@@ -142,9 +157,15 @@ def detect_and_draw_video(video_path):
 
             # اجرای مدل روی فریم
             results = model.predict(source=frame_rgb, conf=0.3)
+            print(f"Model prediction completed for frame {frame_count}.")
 
             # دسترسی به نتایج OBB
-            obb_results = results[0].obb
+            if hasattr(results[0], 'obb') and results[0].obb is not None:
+                obb_results = results[0].obb
+                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()):
@@ -158,6 +179,7 @@ def detect_and_draw_video(video_path):
                     box = np.int0(box)
                     color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
                     cv2.drawContours(frame, [box], 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', 'ناشناخته'))
@@ -169,18 +191,25 @@ def detect_and_draw_video(video_path):
                     overall_counts[label_en] = overall_counts.get(label_en, 0) + 1
 
             frames.append(frame)
+            print(f"Frame {frame_count} processed.")
 
         cap.release()
+        print("Video processing completed.")
 
         # ذخیره ویدئو پردازش‌شده در یک فایل موقت
         with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmpfile:
             output_path = tmpfile.name
 
         out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 20.0, (640, 480))
+        print(f"Saving processed video to {output_path}")
 
-        for frame in frames:
+        for idx, frame in enumerate(frames):
             out.write(frame)
+            if idx % 100 == 0:
+                print(f"Writing frame {idx} to video.")
+
         out.release()
+        print("Video saved.")
 
         # ایجاد DataFrame برای ذخیره نتایج
         df = pd.DataFrame({
@@ -188,6 +217,7 @@ def detect_and_draw_video(video_path):
             'Label (Persian)': [class_names.get(k, ('unknown', 'ناشناخته'))[1] for k in overall_counts.keys()],
             'Object Count': list(overall_counts.values())
         })
+        print("DataFrame created.")
 
         return output_path, df