import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import tempfile
from PIL import Image
from collections import Counter

# تحميل نموذج YOLOv8
model = YOLO('yolov8n.pt')

def process_image(image):
    """ معالجة الصور للكشف عن الكائنات """
    results = model(image)
    boxes = results[0].boxes
    detection_info = [f"{results[0].names[int(box.cls[0])]}: {float(box.conf[0]):.2%}" for box in boxes]
    return Image.fromarray(results[0].plot()), "\n".join(detection_info)

def process_video(video_path):
    """ معالجة الفيديو للكشف عن الكائنات """
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    output_path = temp_file.name

    cap = cv2.VideoCapture(video_path)
    width, height = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps, total_frames = int(cap.get(cv2.CAP_PROP_FPS)), int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    detection_summary = []
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        results = model(frame)
        if int(cap.get(cv2.CAP_PROP_POS_FRAMES)) % int(fps) == 0:
            detection_summary.extend([results[0].names[int(box.cls[0])] for box in results[0].boxes])

        out.write(results[0].plot())

    cap.release()
    out.release()

    summary = "\n".join([f"{obj}: {count} occurrences" for obj, count in Counter(detection_summary).most_common()]) if detection_summary else "No objects detected"
    return output_path, summary

def detect_objects(media):
    """ دالة موحدة للتعامل مع الصور والفيديو """
    if media is None:
        return None, None, "Please upload an image or video.", gr.update(visible=True), gr.update(visible=False)

    try:
        if isinstance(media, str) and media.lower().endswith(('.mp4', '.avi', '.mov')):
            output_video, detection_summary = process_video(media)
            return None, output_video, detection_summary, "✅ Video processing complete!", gr.update(visible=False), gr.update(visible=True)
        else:
            image = cv2.imread(media) if isinstance(media, str) else media
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            processed_image, detection_info = process_image(image)
            return processed_image, None, detection_info, "✅ Image processing complete!", gr.update(visible=True), gr.update(visible=False)
    except Exception as e:
        return None, None, f"❌ Error: {str(e)}", gr.update(visible=False), gr.update(visible=False)

# تصميم الواجهة باستخدام Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🔍 Object Detection")

    input_media = gr.File(label="Upload Image/Video (jpg, png, mp4, avi)", file_types=["image", "video"])
    status_text = gr.Textbox(label="Status", value="Waiting for upload...", interactive=False)
    detection_info = gr.Textbox(label="Detection Results", interactive=False)

    with gr.Row():
        with gr.Column(visible=False) as image_column:
            output_image = gr.Image(label="Detected Objects")
        with gr.Column(visible=False) as video_column:
            output_video = gr.Video(label="Processed Video")

    input_media.upload(detect_objects, inputs=[input_media], outputs=[output_image, output_video, detection_info, status_text, image_column, video_column])

# تشغيل التطبيق
if __name__ == "__main__":
    demo.launch(share=True)