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

# Load model
model = YOLO("yolov8n.pt")

def detect_faults(video_file):
    # Save uploaded video to a temp file path
    temp_input_path = video_file

    # Create a temp directory to store the output video
    temp_dir = tempfile.mkdtemp()
    output_path = os.path.join(temp_dir, "output.mp4")

    # Open input video
    cap = cv2.VideoCapture(temp_input_path)
    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps    = cap.get(cv2.CAP_PROP_FPS)

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out    = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            break

        results = model.predict(source=frame, conf=0.25, verbose=False)
        annotated = results[0].plot()
        out.write(annotated)

    cap.release()
    out.release()

    return output_path

# Gradio UI
interface = gr.Interface(
    fn=detect_faults,
    inputs=gr.Video(label="Upload Thermal Video (MP4)"),
    outputs=gr.Video(label="Detected Output"),
    title="Thermal Fault Detection",
    description="Upload a thermal drone video to detect overheating or burning areas using YOLOv8."
)

if __name__ == "__main__":
    interface.launch()