from ultralytics import YOLO
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download
from tqdm.auto import tqdm
import os
import tempfile
import cv2  # OpenCV for video processing

model_path = "best_int8_openvino_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


def process_image(pilimg):
    source = pilimg
    result = detection_model.predict(source, conf=0.5)
    img_bgr = result[0].plot()
    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
    return out_pilimg

def process_video(video):
    video_reader = cv2.VideoCapture(video)
    nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
    fps = video_reader.get(cv2.CAP_PROP_FPS)

    temp_dir = tempfile.mkdtemp()
    output_path = os.path.join(temp_dir, "annotated_video.mp4")
    video_writer = cv2.VideoWriter(output_path,
                            cv2.VideoWriter_fourcc(*'mp4v'),
                            fps,
                            (frame_w, frame_h))

    # Loop through the video frames
    for i in tqdm(range(nb_frames)):
        # Read a frame from the video
        success, frame = video_reader.read()
        
        if success:
            # Run YOLO inference on the frame on GPU Device 0
            results = detection_model.predict(frame, conf=0.5)

            # Visualize the results on the frame
            annotated_frame = results[0].plot()

            # Write the annotated frame
            video_writer.write(annotated_frame)

    video_reader.release()
    video_writer.release()
    cv2.destroyAllWindows()
    cv2.waitKey(1)
    return output_path

REPO_ID = "leryongfeng/iti107_assignment"
detection_model = load_model(REPO_ID)

# Define a message to display at the top of the app
message = "Sword and Bow detection model"


# Create the interface for image upload
image_interface = gr.Interface(fn=process_image, 
                               inputs=gr.Image(type="pil"),
                               outputs=gr.Image(type="pil"))

# Create the interface for video upload
video_interface = gr.Interface(fn=process_video, 
                               inputs=gr.Video(label="Upload a Video"), 
                               outputs="video")


# Use gr.Blocks to arrange components and launch the app
with gr.Blocks() as app:
    gr.HTML(message)  # Add the message at the top
    gr.TabbedInterface([image_interface, video_interface], 
                          tab_names=["Image Upload", "Video Upload"])

# Launch the interface
app.launch()