app.py

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', 'استخر شنا')
}

# تابع برای تشخیص اشیاء در تصاویر
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)
    obb_results = results[0].obb

    # تبدیل تصویر به PIL برای رسم
    draw = ImageDraw.Draw(input_image)
    # بارگذاری یک فونت برای نوشتن متن
    try:
        font = ImageFont.truetype("arial.ttf", size=15)
    except IOError:
        font = ImageFont.load_default()

    counts = {}
    for obb, conf, cls in zip(obb_results.data.cpu().numpy(), obb_results.conf.cpu().numpy(), obb_results.cls.cpu().numpy()):
        # دسترسی به مختصات جعبه و اطمینان
        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

        # رسم جعبه چرخان با استفاده از 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((x_center, y_center - 10), f"{label_en}: {confidence:.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

# رابط کاربری تصویر
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']
)

# اجرای برنامه با استفاده از رابط کاربری
image_interface.launch(debug=True, share=True)