import gradio as gr
from ultralytics import YOLO
import cv2
import tempfile

# Load your trained OpenVINO model
model = YOLO("best_int8_openvino_model/")

def predict_image(image):
    """Run detection on an uploaded image."""
    results = model(image)
    annotated = results[0].plot()
    return annotated

def predict_video(video):
    """Process an uploaded video frame by frame."""
    temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    
    cap = cv2.VideoCapture(video)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(temp_output.name, fourcc, fps, (width, height))
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = model(frame)
        annotated_frame = results[0].plot()
        out.write(annotated_frame)
    
    cap.release()
    out.release()
    return temp_output.name

# Create Gradio interface
with gr.Blocks(title="Construction Safety Detector") as demo:
    gr.Markdown("""
    # 🚧 Construction Safety Detector
    Detects **Reflective Jackets** and **Safety Helmets** in construction site images and videos.
    """)
    
    with gr.Tab("📷 Image Detection"):
        img_input = gr.Image(type="numpy", label="Upload Image")
        img_output = gr.Image(label="Detection Results")
        img_button = gr.Button("Detect Objects")
        img_button.click(fn=predict_image, inputs=img_input, outputs=img_output)
    
    with gr.Tab("🎥 Video Detection"):
        vid_input = gr.Video(label="Upload Video")
        vid_output = gr.Video(label="Processed Video")
        vid_button = gr.Button("Process Video")
        vid_button.click(fn=predict_video, inputs=vid_input, outputs=vid_output)

demo.launch()