Update app.py

app.py

@@ -1,12 +1,13 @@
 import cv2
+import matplotlib.pyplot as plt
+import numpy as np
 from ultralytics import YOLO
 from PIL import Image, ImageDraw, ImageFont
-import gradio as gr
 import pandas as pd
-import numpy as np
+import gradio as gr
 
-# بارگذاری مدل آموزش‌دیده شما
-model = YOLO('weights/best.pt')
+# بارگذاری مدل
+model = YOLO('yolo11n-obb.pt')  # مدل از پیش آموزش داده شده OBB را بارگذاری کنید
 
 # تعریف نام کلاس‌ها به انگلیسی و فارسی
 class_names = {
@@ -33,30 +34,8 @@ def detect_and_draw_image(input_image):
     input_image_np = np.array(input_image)
 
     # اجرای مدل روی تصویر با سطح اطمینان پایین‌تر برای اطمینان از شناسایی بیشتر اشیاء
-    results = model(input_image_np, conf=0.25)
-
-    # بررسی تعداد نتایج
-    if not results:
-        print("مدل هیچ نتیجه‌ای برنگرداند.")
-        df = pd.DataFrame({
-            'Label (English)': [],
-            'Label (Persian)': [],
-            'Object Count': []
-        })
-        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)}")
+    results = model.predict(source=input_image_np, conf=0.3)
+    obb_results = results[0].obb
 
     # تبدیل تصویر به PIL برای رسم
     draw = ImageDraw.Draw(input_image)
@@ -67,20 +46,23 @@ def detect_and_draw_image(input_image):
         font = ImageFont.load_default()
 
     counts = {}
-    for box in detections:
+    for obb, conf, cls in zip(obb_results.data.cpu().numpy(), obb_results.conf.cpu().numpy(), obb_results.cls.cpu().numpy()):
         # دسترسی به مختصات جعبه و اطمینان
-        xmin, ymin, xmax, ymax = box.xyxy[0].tolist()
-        conf = box.conf[0].item()
-        class_id = int(box.cls[0].item())
+        x_center, y_center, width, height, rotation = obb[:5]
+        class_id = int(cls)
+        confidence = float(conf)
 
         # دریافت برچسب‌های انگلیسی و فارسی
         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)
+        # رسم جعبه چرخان با استفاده از OpenCV
+        rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
+        box = cv2.boxPoints(rect)
+        box = np.int0(box)
+        draw.polygon([tuple(point) for point in box], outline="red")
         # نوشتن برچسب و اطمینان
-        draw.text((xmin, ymin - 10), f"{label_en}: {conf:.2f}", fill="red", font=font)
+        draw.text((x_center, y_center - 10), f"{label_en}: {confidence:.2f}", fill="red", font=font)
 
     # ایجاد DataFrame برای نمایش نتایج
     df = pd.DataFrame({
@@ -91,63 +73,6 @@ def detect_and_draw_image(input_image):
 
     return input_image, 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(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(
     fn=detect_and_draw_image,
@@ -158,16 +83,5 @@ image_interface = gr.Interface(
     examples=['Examples/images/areial_car.jpg', 'Examples/images/arieal_car_1.jpg','Examples/images/t.jpg']
 )
 
-# رابط کاربری ویدئو
-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)
+# اجرای برنامه با استفاده از رابط کاربری
+image_interface.launch(debug=True, share=True)