Update app.py

app.py

@@ -1,13 +1,17 @@
+# !pip install ultralytics
+# !pip install gradio
+
 import cv2
-import matplotlib.pyplot as plt
-import numpy as np
 from ultralytics import YOLO
-from PIL import Image, ImageDraw, ImageFont
-import pandas as pd
+from PIL import Image, ImageDraw
 import gradio as gr
+import pandas as pd
+import numpy as np
+import tempfile
+import os
 
-# بارگذاری مدل
-model = YOLO('yolo11n-obb.pt')  # مدل از پیش آموزش داده شده OBB را بارگذاری کنید
+# بارگذاری مدل آموزش‌دیده شما
+model = YOLO('/content/yolo11n-obb.pt')  # اطمینان حاصل کنید که مسیر مدل صحیح است
 
 # تعریف نام کلاس‌ها به انگلیسی و فارسی
 class_names = {
@@ -30,21 +34,21 @@ class_names = {
 
 # رنگ‌ها برای هر کلاس
 colors = {
-    0: (255, 0, 0),  # Red
-    1: (0, 255, 0),  # Green
-    2: (0, 0, 255),  # Blue
-    3: (255, 255, 0),  # Yellow
-    4: (255, 0, 255),  # Magenta
-    5: (0, 255, 255),  # Cyan
-    6: (128, 0, 128),  # Purple
-    7: (255, 165, 0),  # Orange
-    8: (0, 128, 0),  # Dark Green
-    9: (128, 128, 0),  # Olive
-    10: (128, 0, 0),  # Maroon
-    11: (0, 128, 128),  # Teal
-    12: (0, 0, 128),  # Navy
-    13: (75, 0, 130),  # Indigo
-    14: (199, 21, 133)  # MediumVioletRed
+    0: (255, 0, 0),       # قرمز
+    1: (0, 255, 0),       # سبز
+    2: (0, 0, 255),       # آبی
+    3: (255, 255, 0),     # زرد
+    4: (255, 0, 255),     # مجنتا
+    5: (0, 255, 255),     # فیروزه‌ای
+    6: (128, 0, 128),     # بنفش
+    7: (255, 165, 0),     # نارنجی
+    8: (0, 128, 0),       # سبز تیره
+    9: (128, 128, 0),     # زیتونی
+    10: (128, 0, 0),      # سرخ کلید
+    11: (0, 128, 128),    # سبز نفتی
+    12: (0, 0, 128),      # نیوی
+    13: (75, 0, 130),     # ایندیگو
+    14: (199, 21, 133)    # رز متوسط
 }
 
 # تابع برای تشخیص اشیاء در تصاویر
@@ -52,10 +56,12 @@ def detect_and_draw_image(input_image):
     # تبدیل تصویر PIL به آرایه NumPy
     input_image_np = np.array(input_image)
 
-    # اجرای مدل روی تصویر با سطح اطمینان پایین‌تر برای اطمینان از شناسایی بیشتر اشیاء
-    results = model.predict(source=input_image_np, conf=0.3)
-    if not results or not hasattr(results[0], 'boxes') or results[0].boxes is None:
-        print("هیچ شیء شناسایی نشده است.")
+    # اجرای مدل روی تصویر
+    results = model(input_image_np)
+
+    # بررسی کردن اینکه آیا اشیاء شناسایی شده اند یا خیر
+    if not results or not results[0].boxes:
+        # اگر هیچ شیء شناسایی نشده باشد
         df = pd.DataFrame({
             'Label (English)': [],
             'Label (Persian)': [],
@@ -63,73 +69,88 @@ def detect_and_draw_image(input_image):
         })
         return input_image, df
 
-    # بارگذاری تصویر اصلی به صورت OpenCV برای رسم جعبه‌ها
-    image_np = np.array(input_image.convert('RGB'))[:, :, ::-1]  # تبدیل PIL به OpenCV
+    detections = results[0].boxes  # دسترسی به نتایج در YOLOv8
 
+    # تبدیل تصویر به RGB برای رسم
+    image_draw = input_image.convert("RGB")
+    draw = ImageDraw.Draw(image_draw)
     counts = {}
-    for box in results[0].boxes:
-        # دسترسی به مختصات جعبه و اطمینان
-        xmin, ymin, xmax, ymax, conf, class_id = box.cpu().numpy()
-        class_id = int(class_id)
-        confidence = float(conf)
+    for box in detections:
+        # دریافت مختصات جعبه
+        xmin, ymin, xmax, ymax = box.xyxy.tolist()
+        conf = box.conf.tolist()[0]
+        class_id = int(box.cls.tolist()[0])
 
-        # دریافت برچسب‌های انگلیسی و فارسی
+        # دریافت برچسب‌ها
         label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
         counts[label_en] = counts.get(label_en, 0) + 1
 
-        # رسم مستطیل با استفاده از OpenCV
-        color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
-        cv2.rectangle(image_np, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, 2)
-        cv2.putText(image_np, f'{label_en}: {confidence:.2f}', (int(xmin), int(ymin) - 10), 
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)
+        # رسم مستطیل
+        draw.rectangle([(xmin, ymin), (xmax, ymax)], outline=colors.get(class_id, (255,0,0)), width=2)
+        # رسم برچسب
+        draw.text((xmin, ymin), f"{label_en}: {conf:.2f}", fill="white")
 
-    # تبدیل تصویر به RGB برای Gradio
-    image_rgb = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
-    output_image = Image.fromarray(image_rgb)
-
-    # ایجاد DataFrame برای نمایش نتایج
+    # ایجاد 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 output_image, df
+    return image_draw, df
 
 # تابع برای تشخیص اشیاء در ویدئوها
 def detect_and_draw_video(video_path):
     cap = cv2.VideoCapture(video_path)
     frames = []
     overall_counts = {}
+    seen_objects = []  # لیست برای دنبال کردن اشیاء شناسایی شده
 
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
 
+        # تغییر اندازه فریم
         frame = cv2.resize(frame, (640, 480))
-        results = model.predict(source=frame, conf=0.3)
-        if not results or not hasattr(results[0], 'boxes') or results[0].boxes is None:
-            continue
-
-        for box in results[0].boxes:
-            xmin, ymin, xmax, ymax, conf, class_id = box.cpu().numpy()
-            class_id = int(class_id)
-            confidence = float(conf)
-            label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
-            overall_counts[label_en] = overall_counts.get(label_en, 0) + 1
-
-            # رسم مستطیل و نام شیء بر روی فریم
-            color = colors.get(class_id, (0, 255, 0))
-            cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, 2)
-            cv2.putText(frame, f'{label_en}: {confidence:.2f}', (int(xmin), int(ymin) - 10),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
+        # تبدیل رنگ از BGR به RGB
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+        # اجرای مدل روی فریم
+        results = model(frame_rgb)
+
+        if results and results[0].boxes:
+            detections = results[0].boxes
+            for box in detections:
+                # دریافت مختصات جعبه
+                xmin, ymin, xmax, ymax = box.xyxy.tolist()
+                conf = box.conf.tolist()[0]
+                class_id = int(box.cls.tolist()[0])
+                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] < current_object[3] and existing[3] > current_object[1] and
+                            existing[2] < current_object[4] and existing[4] > current_object[2]) for existing in seen_objects):
+                    seen_objects.append(current_object)
+                    overall_counts[label_en] = overall_counts.get(label_en, 0) + 1
+
+                # رسم مستطیل
+                color = colors.get(class_id, (255, 0, 0))
+                cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, 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'
+    # ذخیره ویدئو پردازش‌شده در یک فایل موقت
+    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))
 
     for frame in frames: