from PIL import Image
import gradio as gr
import cv2
from ultralytics import ASSETS, YOLO
import tempfile
import numpy as np
import time

def load_model(model_name):
    """Loads the specified YOLO model for either segmentation or detection."""
    if model_name == "yolov9c-seg":
        model_path = "yolov9c-seg.pt"
    elif model_name == "yolov9e-seg":
        model_path = "yolov9e-seg.pt"
    elif model_name == "yolov9c":
        model_path = "yolov9c.pt"
    elif model_name == "yolov9e":
        model_path = "yolov9e.pt"
    else:
        raise ValueError(f"Invalid model name: {model_name}")
    
    return YOLO(model_path)

def predict_image(img, conf_threshold, iou_threshold, task="detection", model_name=None):
    """Predicts and plots results in an image using YOLO model with adjustable confidence and IOU thresholds."""
    if task == "segmentation":
        if not model_name:
            model_name = "yolov9c-seg"
        elif model_name not in ["yolov9c-seg", "yolov9e-seg"]:
            raise ValueError(f"Invalid model name for segmentation: {model_name}")
    elif task == "detection":
        if not model_name:
            model_name = "yolov9c"
        elif model_name not in ["yolov9c", "yolov9e"]:
            raise ValueError(f"Invalid model name for detection: {model_name}")
    else:
        raise ValueError(f"Invalid task: {task}. Choose either 'segmentation' or 'detection'.")

    model = load_model(model_name)
    results = model.predict(
        source=img,
        conf=conf_threshold,
        iou=iou_threshold,
        show_labels=True,
        show_conf=True,
        imgsz=640,
    )

    for r in results:
        im_array = r.plot()
        im = Image.fromarray(im_array[..., ::-1])
    
    return im

def predict_image_with_task(img, conf_threshold, iou_threshold, task, model_name):

    return predict_image(img, conf_threshold, iou_threshold, task, model_name)

image_iface = gr.Interface(
    fn=predict_image_with_task,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
        gr.Dropdown(choices=["detection", "segmentation"], value="detection", label="Task"),
        gr.Dropdown(choices=["yolov9c", "yolov9e", "yolov9c-seg", "yolov9e-seg"], value="yolov9c", label="Model"),
    ],
    outputs=gr.Image(type="pil", label="Result"),
    title="X509",
    description="Upload images for inference. Choose task and corresponding model.",
    examples=[
        ["cars.jpg", 0.25, 0.45, "detection", "yolov9c"],
    ],
)

def predict_video(video_path, conf_threshold, iou_threshold, task="detection", model_name=None):
    """Predicts and processes video frames using YOLO model with adjustable confidence and IOU thresholds."""
    if task == "segmentation":
        if not model_name:
            model_name = "yolov9c-seg"
        elif model_name not in ["yolov9c-seg", "yolov9e-seg"]:
            raise ValueError(f"Invalid model name for segmentation: {model_name}")
    elif task == "detection":
        if not model_name:
            model_name = "yolov9c"
        elif model_name not in ["yolov9c", "yolov9e"]:
            raise ValueError(f"Invalid model name for detection: {model_name}")
    else:
        raise ValueError(f"Invalid task: {task}. Choose either 'segmentation' or 'detection'.")

    model = load_model(model_name)
    cap = cv2.VideoCapture(video_path)

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    temp_video_path = tempfile.mktemp(suffix=".mp4")
    out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height))

    frame_count = 0
    start_time = time.time()

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame_count += 1

        elapsed_time = time.time() - start_time
        current_fps = frame_count / elapsed_time

        pil_img = Image.fromarray(frame[..., ::-1])
        results = model.predict(
            source=pil_img,
            conf=conf_threshold,
            iou=iou_threshold,
            show_labels=True,
            show_conf=True,
            imgsz=640,
        )

        for r in results:
            im_array = r.plot()
            processed_frame = Image.fromarray(im_array[..., ::-1])
            frame = cv2.cvtColor(np.array(processed_frame), cv2.COLOR_RGB2BGR)

            cv2.putText(frame, f"FPS: {current_fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            out.write(frame)
    
    cap.release()
    out.release()

    return temp_video_path

def predict_video_with_task(video_path, conf_threshold, iou_threshold, task, model_name):

    return predict_video(video_path, conf_threshold, iou_threshold, task, model_name)

video_iface = gr.Interface(
    fn=predict_video_with_task,
    inputs=[
        gr.Video(label="Upload Video", interactive=True),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
        gr.Dropdown(choices=["detection", "segmentation"], value="detection", label="Task"),
        gr.Dropdown(choices=["yolov9c", "yolov9e", "yolov9c-seg", "yolov9e-seg"], value="yolov9c", label="Model"),
    ],
    outputs=gr.File(label="Result"),
    title="X509",
    description="Upload video for inference. Choose task and corresponding model.",
    examples=[
        ["VID_20240517112011.mp4", 0.25, 0.45, "detection", "yolov9c"],
    ]
)

production = gr.TabbedInterface([image_iface, video_iface], ["Image Inference", "Video Inference"])

if __name__ == '__main__':
    production.launch()