app.py

# !pip install ultralytics
# !pip install gradio

import cv2
from ultralytics import YOLO
from PIL import Image, ImageDraw
import gradio as gr
import pandas as pd
import numpy as np
import tempfile
import os

# بارگذاری مدل آموزش‌دیده شما
# اطمینان حاصل کنید که فقط یک بار مدل را بارگذاری می‌کنید و مسیر صحیح است
model = YOLO('/content/yolo11n-obb.pt')  # یا '/content/best.pt' بر اساس مدل مورد نظر شما

# تعریف نام کلاس‌ها به انگلیسی و فارسی
class_names = {
    0: ('plane', 'هواپیما'),
    1: ('ship', 'کشتی'),
    2: ('storage tank', 'مخزن ذخیره'),
    3: ('baseball diamond', 'زمین بیسبال'),
    4: ('tennis court', 'زمین تنیس'),
    5: ('basketball court', 'زمین بسکتبال'),
    6: ('ground track field', 'زمین دو و میدانی'),
    7: ('harbor', 'بندرگاه'),
    8: ('bridge', 'پل'),
    9: ('large vehicle', 'خودرو بزرگ'),
    10: ('small vehicle', 'خودرو کوچک'),
    11: ('helicopter', 'هلیکوپتر'),
    12: ('roundabout', 'میدان'),
    13: ('soccer ball field', 'زمین فوتبال'),
    14: ('swimming pool', 'استخر شنا')
}

# رنگ‌ها برای هر کلاس
colors = {
    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)    # رز متوسط
}

# تابع برای تشخیص اشیاء در تصاویر
def detect_and_draw_image(input_image):
    try:
        # تبدیل تصویر PIL به آرایه NumPy
        input_image_np = np.array(input_image)

        # اجرای مدل روی تصویر
        results = model.predict(source=input_image_np, conf=0.3)

        # دسترسی به نتایج OBB
        obb_results = results[0].boxes  # یا results[0].obb بر اساس نسخه مدل شما

        # بررسی وجود جعبه‌های شناسایی شده
        if obb_results is None or len(obb_results) == 0:
            print("هیچ شیء شناسایی نشده است.")
            df = pd.DataFrame({
                'Label (English)': [],
                'Label (Persian)': [],
                'Object Count': []
            })
            return input_image, df

        # بارگذاری تصویر اصلی به صورت OpenCV برای رسم جعبه‌ها
        image_cv = cv2.cvtColor(input_image_np, cv2.COLOR_RGB2BGR)

        counts = {}
        for box in obb_results:
            # دریافت مختصات جعبه
            x_center, y_center, width, height, rotation = box.xywhn.tolist()[0]
            class_id = int(box.cls.tolist()[0])
            confidence = float(box.conf.tolist()[0])

            # تبدیل مختصات از نرمال شده به پیکسل
            img_height, img_width, _ = image_cv.shape
            x_center *= img_width
            y_center *= img_height
            width *= img_width
            height *= img_height

            # دریافت برچسب‌ها
            label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
            counts[label_en] = counts.get(label_en, 0) + 1

            # رسم جعبه چرخان با استفاده از OpenCV
            rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
            box_points = cv2.boxPoints(rect)
            box_points = np.int0(box_points)
            color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
            cv2.drawContours(image_cv, [box_points], 0, color, 2)
            cv2.putText(image_cv, f'{label_en}: {confidence:.2f}', 
                        (int(x_center), int(y_center)), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2, cv2.LINE_AA)

        # تبدیل تصویر به RGB برای Gradio
        image_rgb = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
        output_image = Image.fromarray(image_rgb)

        # ایجاد 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

    except Exception as e:
        print(f"Error in detect_and_draw_image: {e}")
        df = pd.DataFrame({
            'Label (English)': [],
            'Label (Persian)': [],
            'Object Count': []
        })
        return input_image, df

# تابع برای تشخیص اشیاء در ویدئوها
def detect_and_draw_video(video_path):
    try:
        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)
            obb_results = results[0].boxes  # یا results[0].obb بر اساس نسخه مدل شما

            if obb_results is not None and len(obb_results) > 0:
                for box in obb_results:
                    # دریافت مختصات جعبه
                    x_center, y_center, width, height, rotation = box.xywhn.tolist()[0]
                    class_id = int(box.cls.tolist()[0])
                    confidence = float(box.conf.tolist()[0])

                    # تبدیل مختصات از نرمال شده به پیکسل
                    img_height, img_width, _ = frame.shape
                    x_center *= img_width
                    y_center *= img_height
                    width *= img_width
                    height *= img_height

                    # دریافت برچسب‌ها
                    label_en, label_fa = class_names.get(class_id, ('unknown', 'ناشناخته'))
                    current_object = (label_en, int(x_center - width / 2), int(y_center - height / 2), 
                                      int(x_center + width / 2), int(y_center + height / 2))

                    # بررسی وجود شیء در لیست 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

                    # رسم جعبه چرخان با استفاده از OpenCV
                    rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
                    box_points = cv2.boxPoints(rect)
                    box_points = np.int0(box_points)
                    color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
                    cv2.drawContours(frame, [box_points], 0, color, 2)
                    cv2.putText(frame, f"{label_en}: {confidence:.2f}", 
                                (int(x_center), int(y_center)), 
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2, cv2.LINE_AA)

            frames.append(frame)

        cap.release()

        # ذخیره ویدئو پردازش‌شده در یک فایل موقت
        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:
            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

    except Exception as e:
        print(f"Error in detect_and_draw_video: {e}")
        # در صورت بروز خطا، بازگرداندن ویدئوی اصلی بدون تغییر و یک DataFrame خالی
        return video_path, pd.DataFrame({
            'Label (English)': [],
            'Label (Persian)': [],
            'Object Count': []
        })

# رابط کاربری تصویر
image_interface = gr.Interface(
    fn=detect_and_draw_image,
    inputs=gr.Image(type="pil", label="بارگذاری تصویر"),
    outputs=[gr.Image(type="pil"), gr.Dataframe(label="تعداد اشیاء")],
    title="تشخیص اشیاء در تصاویر هوایی",
    description="یک تصویر هوایی بارگذاری کنید تا اشیاء شناسایی شده و تعداد آن‌ها را ببینید.",
    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)