Update app.py

app.py

@@ -1,6 +1,6 @@
 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,40 +32,121 @@ def detect_and_draw_image(input_image):
     # تبدیل تصویر PIL به آرایه NumPy
     input_image_np = np.array(input_image)
 
-    # اجرای مدل روی تصویر با تنظیمات اطمینان و iou پایین‌تر
-    results = model.predict(source=input_image_np, conf=0.1, iou=0.4)
-    
-    # بررسی وجود اشیاء شناسایی شده
-    if results and results[0].boxes is not None:
-        # دریافت تصویر با نتایج رسم‌شده
-        annotated_image = results[0].plot()
-        
-        # تبدیل تصویر به فرمت PIL برای نمایش در Gradio
-        annotated_image_pil = Image.fromarray(annotated_image)
-
-        # شمارش اشیاء شناسایی‌شده
-        counts = {}
-        for box in results[0].boxes:
-            class_id = int(box.cls[0])
-            label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
-            counts[label_en] = counts.get(label_en, 0) + 1
-
-        # ساخت DataFrame با نام‌ها و تعداد اشیاء
+    # اجرای مدل روی تصویر با سطح اطمینان پایین‌تر برای اطمینان از شناسایی بیشتر اشیاء
+    results = model(input_image_np, conf=0.25)
+
+    # بررسی تعداد نتایج
+    if not results:
+        print("مدل هیچ نتیجه‌ای برنگرداند.")
         df = pd.DataFrame({
-            'Label (English)': list(counts.keys()),
-            'Label (Persian)': [class_names.get(k, ('unknown', 'ناشناخته'))[1] for k in counts.keys()],
-            'Object Count': list(counts.values())
+            'Label (English)': [],
+            'Label (Persian)': [],
+            'Object Count': []
         })
-    else:
-        # اگر هیچ شیء شناسایی نشده باشد، بازگرداندن تصویر اصلی و DataFrame خالی
-        annotated_image_pil = input_image
+        return input_image, df
+
+    # بررسی وجود جعبه‌های شناسایی شده
+    detections = results[0].boxes
+    if detections is None or len(detections) == 0:
+        print("هیچ شیء شناسایی نشده است.")
         df = pd.DataFrame({
             'Label (English)': [],
             'Label (Persian)': [],
             'Object Count': []
         })
+        return input_image, df
+
+    print(f"تعداد اشیاء شناسایی‌شده: {len(detections)}")
+
+    # تبدیل تصویر به PIL برای رسم
+    draw = ImageDraw.Draw(input_image)
+    # بارگذاری یک فونت برای نوشتن متن
+    try:
+        font = ImageFont.truetype("arial.ttf", size=15)
+    except IOError:
+        font = ImageFont.load_default()
+
+    counts = {}
+    for box in detections:
+        # دسترسی به مختصات جعبه و اطمینان
+        xmin, ymin, xmax, ymax = box.xyxy[0].tolist()
+        conf = box.conf[0].item()
+        class_id = int(box.cls[0].item())
+
+        # دریافت برچسب‌های انگلیسی و فارسی
+        label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
+        counts[label_en] = counts.get(label_en, 0) + 1
+
+        # رسم مستطیل
+        draw.rectangle([(xmin, ymin), (xmax, ymax)], outline="red", width=2)
+        # نوشتن برچسب و اطمینان
+        draw.text((xmin, ymin - 10), f"{label_en}: {conf:.2f}", fill="red", font=font)
+
+    # ایجاد DataFrame برای نمایش نتایج
+    df = pd.DataFrame({
+        'Label (English)': list(counts.keys()),
+        'Label (Persian)': [class_names.get(k, ('unknown', 'ناشناخته'))[1] for k in counts.keys()],
+        'Object Count': list(counts.values())
+    })
+
+    return input_image, df
+
+# تابع برای تشخیص اشیاء در ویدئوها
+def detect_and_draw_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    overall_counts = {}
+    seen_objects = []  # لیست برای دنبال کردن اشیاء شناسایی شده
 
-    return annotated_image_pil, df
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        frame = cv2.resize(frame, (640, 480))
+        results = model(frame, conf=0.25)
+
+        detections = results[0].boxes
+
+        if detections is not None and len(detections) > 0:
+            for box in detections:
+                xmin, ymin, xmax, ymax = box.xyxy[0].tolist()
+                conf = box.conf[0].item()
+                class_id = int(box.cls[0].item())
+                label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
+                current_object = (label_en, int(xmin), int(ymin), int(xmax), int(ymax))
+
+                # بررسی وجود شیء در لیست seen_objects
+                if not any(existing[0] == label_en and
+                           (existing[1] < xmax and existing[3] > xmin and
+                            existing[2] < ymax and existing[4] > ymin) for existing in seen_objects):
+                    seen_objects.append(current_object)
+                    overall_counts[label_en] = overall_counts.get(label_en, 0) + 1
+
+                    # رسم مستطیل و نام شیء بر روی فریم
+                    cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (255, 0, 0), 2)
+                    cv2.putText(frame, f"{label_en}: {conf:.2f}", (int(xmin), int(ymin) - 10),
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
+
+        frames.append(frame)
+
+    cap.release()
+
+    output_path = 'output.mp4'
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 20.0, (640, 480))
+
+    for frame in frames:
+        out.write(frame)
+    out.release()
+
+    # ایجاد DataFrame برای ذخیره نتایج
+    df = pd.DataFrame({
+        'Label (English)': list(overall_counts.keys()),
+        'Label (Persian)': [class_names.get(k, ('unknown', 'ناشناخته'))[1] for k in overall_counts.keys()],
+        'Object Count': list(overall_counts.values())
+    })
+
+    return output_path, df
 
 # رابط کاربری تصویر
 image_interface = gr.Interface(
@@ -77,6 +158,16 @@ image_interface = gr.Interface(
     examples=['Examples/images/areial_car.jpg', 'Examples/images/arieal_car_1.jpg','Examples/images/t.jpg']
 )
 
-# اجرای برنامه
-app = gr.Interface(fn=detect_and_draw_image, inputs=gr.Image(type="pil"), outputs=[gr.Image(type="pil"), gr.Dataframe()])
+# رابط کاربری ویدئو
+video_interface = gr.Interface(
+    fn=detect_and_draw_video,
+    inputs=gr.Video(label="بارگذاری ویدئو"),
+    outputs=[gr.Video(label="ویدئوی پردازش شده"), gr.Dataframe(label="تعداد اشیاء")],
+    title="تشخیص اشیاء در ویدئوها",
+    description="یک ویدئو بارگذاری کنید تا اشیاء شناسایی شده و تعداد آن‌ها را ببینید.",
+    examples=['Examples/video/city.mp4', 'Examples/video/airplane.mp4']
+)
+
+# اجرای برنامه با استفاده از رابط کاربری تب‌دار
+app = gr.TabbedInterface([image_interface, video_interface], ["تشخیص تصویر", "تشخیص ویدئو"])
 app.launch(debug=True, share=True)