import gradio as gr
from huggingface_hub import snapshot_download
from ultralytics import YOLO
import os
from PIL import Image
import cv2
import numpy as np
import tempfile

#public model path location
#MODEL_REPO_ID = "mintheinwin/3907578Y"

#Organizations model path location
MODEL_REPO_ID = "ITI107-2024S2/3907578Y"

# Load model
def load_model(repo_id):
    download_dir = snapshot_download(repo_id)
    path = os.path.join(download_dir, "best_int8_openvino_model")
    detection_model = YOLO(path, task="detect")
    return detection_model


detection_model = load_model(MODEL_REPO_ID)

#Student ID
student_info = "Student Id: 3907578Y, Name: Min Thein Win"

#Prediction for images
def predict_image(pil_img):
    result = detection_model.predict(pil_img, conf=0.5, iou=0.5)
    img_bgr = result[0].plot()  # Annotated image
    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert to RGB PIL image
    return out_pilimg

#Prediction for videos
def predict_video(video):
    cap = cv2.VideoCapture(video)
    frames = []
    temp_dir = tempfile.mkdtemp()

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

        # Detection
        result = detection_model.predict(frame, conf=0.5, iou=0.5)
        annotated_frame = result[0].plot()
        frames.append(annotated_frame)

    cap.release()

    # Save annotated video
    height, width, _ = frames[0].shape
    output_path = os.path.join(temp_dir, "annotated_video.mp4")
    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), 20, (width, height))

    for frame in frames:
        out.write(frame)

    out.release()
    return output_path

# Unified prediction function
def unified_predict(file):
    if isinstance(file, Image.Image):
        # If the input is a PIL Image, treat it as an image
        return predict_image(file)
    elif isinstance(file, str) and file.endswith(('.mp4', '.avi', '.mov')):
        # If the input is a video file path, treat it as a video
        return predict_video(file)
    else:
        raise ValueError("Unsupported file type. Please upload an image or a video.")

# UI Interface
with gr.Blocks() as interface:
    gr.Markdown("# Wild Animal Detection (Tiger/Lion)")
    gr.Markdown(student_info)

    # Unified Section
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Upload an Image or Video:")
            input_file = gr.File(label="Input File")

        with gr.Column():
            gr.Markdown("### Output Results:")
            output_display = gr.Output(label="Output")
            
    clear_btn= gr.Button("CLEAR")
    submit_btn = gr.Button("SUBMIT")

    def process_file(file):
        if file.name.endswith((".jpg", ".jpeg", ".png")):
            pil_image = Image.open(file.name)
            return predict_image(pil_image)
        elif file.name.endswith((".mp4", ".avi", ".mov")):
            return predict_video(file.name)
        else:
            return "Unsupported file type. Please upload an image or a video."

    def clear_all():
        return None, ""
        
    submit_btn.click(fn=process_file, inputs=input_file, outputs=output_display)
    clear_btn.click(fn=clear_all, inputs=None, outputs=[input_file, output_display])

# Launch app
interface.launch(share=True)