import gradio as gr
from ultralytics import YOLO
from PIL import Image
import tempfile
import cv2
import os

# ---------------------------
# Load models EARLY
# ---------------------------
yolo_model = YOLO("yolov8n-pose.pt")

# ---------------------------
# Prediction for IMAGE
# ---------------------------
def predict_image(image):
    results = yolo_model.predict(source=image, show=False, conf=0.6)
    results_img = results[0].plot()
    return Image.fromarray(results_img)

# ---------------------------
# Prediction for VIDEO
# ---------------------------
def predict_video(video):
    temp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    cap = cv2.VideoCapture(video)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width, height = int(cap.get(3)), int(cap.get(4))
    out = cv2.VideoWriter(temp_out.name, fourcc, fps, (width, height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = yolo_model.predict(source=frame, conf=0.6, verbose=False)
        annotated_frame = results[0].plot()
        out.write(annotated_frame)

    cap.release()
    out.release()
    return temp_out.name

# ---------------------------
# Gradio UI
# ---------------------------
with gr.Blocks(css="""
    body {background: linear-gradient(135deg, #1f1c2c, #928DAB);}
    .gradio-container {font-family: 'Segoe UI', sans-serif;}
    h1 {text-align: center; color: white; padding: 20px; font-size: 2.5em;}
    .tabs {margin-top: 20px;}
    .footer {text-align:center; color:#eee; font-size:14px; margin-top:25px;}
    .gr-button {border-radius:12px; font-weight:bold; padding:10px 18px;}
""") as demo:

    gr.HTML("<h1>🚨 Suspicious Activity Detection</h1>")

    with gr.Tab("📷 Image Detection"):
        with gr.Row():
            with gr.Column(scale=1):
                img_input = gr.Image(type="pil", label="Upload Image")
                img_btn = gr.Button("🔍 Detect Suspicious Activity")
            with gr.Column(scale=1):
                img_output = gr.Image(type="pil", label="Detection Result")
        img_btn.click(predict_image, inputs=img_input, outputs=img_output)

    with gr.Tab("🎥 Video Detection"):
        with gr.Row():
            with gr.Column(scale=1):
                vid_input = gr.Video(label="Upload Video")
                vid_btn = gr.Button("🎬 Detect in Video")
            with gr.Column(scale=1):
                vid_output = gr.Video(label="Processed Video")
        vid_btn.click(predict_video, inputs=vid_input, outputs=vid_output)

    gr.HTML("<div class='footer'>Made with ❤️ using YOLO + Gradio</div>")

# ---------------------------
# Launch App
# ---------------------------
demo.launch(share=True)