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
import gradio as gr

# بارگذاری مدل
model = YOLO('yolo11n-obb.pt')  # مدل از پیش آموزش داده شده OBB را بارگذاری کنید

# تعریف نام کلاس‌ها به انگلیسی و فارسی
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),  # 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
}

# تابع برای تشخیص اشیاء در تصاویر
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("هیچ شیء شناسایی نشده است.")
        df = pd.DataFrame({
            'Label (English)': [],
            'Label (Persian)': [],
            'Object Count': []
        })
        return input_image, df

    # بارگذاری تصویر اصلی به صورت OpenCV برای رسم جعبه‌ها
    image_np = np.array(input_image.convert('RGB'))[:, :, ::-1]  # تبدیل PIL به OpenCV

    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)

        # دریافت برچسب‌های انگلیسی و فارسی
        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)

    # تبدیل تصویر به RGB برای Gradio
    image_rgb = cv2.cvtColor(image_np, 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

# تابع برای تشخیص اشیاء در ویدئوها
def detect_and_draw_video(video_path):
    cap = cv2.VideoCapture(video_path)
    frames = []
    overall_counts = {}

    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)

        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,
    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)