import gradio as gr
import cv2
from PIL import Image, ImageDraw, ImageFont
import torch
from transformers import Owlv2Processor, Owlv2ForObjectDetection
import numpy as np
import os

# Check if CUDA is available, otherwise use CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'

processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16").to(device)

def detect_objects_in_frame(image, target):
    draw = ImageDraw.Draw(image)
    texts = [[target]]
    inputs = processor(text=texts, images=image, return_tensors="pt", padding=True).to(device)
    outputs = model(**inputs)

    target_sizes = torch.Tensor([image.size[::-1]])
    results = processor.post_process_object_detection(outputs=outputs, threshold=0.1, target_sizes=target_sizes)

    color_map = {target: "red"}

    try:
        font = ImageFont.truetype("arial.ttf", 15)
    except IOError:
        font = ImageFont.load_default()

    i = 0
    text = texts[i]
    boxes, scores, labels = results[i]["boxes"], results[i]["scores"], results[i]["labels"]

    for box, score, label in zip(boxes, scores, labels):
        if score.item() >= 0.25:
            box = [round(i, 2) for i in box.tolist()]
            object_label = text[label]
            confidence = round(score.item(), 3)
            annotation = f"{object_label}: {confidence}"

            draw.rectangle(box, outline=color_map.get(object_label, "red"), width=2)
            text_position = (box[0], box[1] - 10)
            draw.text(text_position, annotation, fill="white", font=font)

    return image

def process_video(video_path, target, progress=gr.Progress()):
    if video_path is None:
        return None, "Error: No video uploaded"

    if not os.path.exists(video_path):
        return None, f"Error: Video file not found at {video_path}"

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None, f"Error: Unable to open video file at {video_path}"

    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    original_fps = int(cap.get(cv2.CAP_PROP_FPS))
    original_duration = frame_count / original_fps
    output_fps = 5

    output_path = "output_video.mp4"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, output_fps, (int(cap.get(3)), int(cap.get(4))))

    batch_size = 64
    frames = []

    for frame in progress.tqdm(range(frame_count)):
        ret, img = cap.read()
        if not ret:
            break

        if frame % (original_fps // output_fps) != 0:
            continue

        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        frames.append(pil_img)

        if len(frames) == batch_size or frame == frame_count - 1:
            annotated_frames = [detect_objects_in_frame(frame, target) for frame in frames]
            for annotated_img in annotated_frames:
                annotated_frame = cv2.cvtColor(np.array(annotated_img), cv2.COLOR_RGB2BGR)
                out.write(annotated_frame)
            frames = []

    cap.release()
    out.release()

    return output_path, None

def load_sample_frame(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    ret, frame = cap.read()
    cap.release()
    if not ret:
        return None
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    return frame_rgb

def gradio_app():
    with gr.Blocks() as app:
        gr.Markdown("# Video Object Detection with Owlv2")

        video_input = gr.Video(label="Upload Video")
        target_input = gr.Textbox(label="Target Object")
        output_video = gr.Video(label="Output Video")
        error_output = gr.Textbox(label="Error Messages", visible=False)
        sample_video_frame = gr.Image(value=load_sample_frame("Drone Video of African Wildlife Wild Botswan.mp4"), label="Sample Video Frame")
        use_sample_button = gr.Button("Use Sample Video")
        progress_bar = gr.Progress()

        video_path = gr.State(None)
        def process_and_update(video, target):
            output_video_path, error = process_video(video, target, progress_bar)
            if error:
                error_output.visible = True
            else:
                error_output.visible = False
            return output_video_path, error

        video_input.upload(process_and_update, 
                           inputs=[video_input, target_input], 
                           outputs=[output_video, error_output])

        def use_sample_video():
            sample_video_path = "Drone Video of African Wildlife Wild Botswan.mp4"
            return process_and_update(sample_video_path, "animal")

        use_sample_button.click(use_sample_video, 
                                inputs=None, 
                                outputs=[output_video, error_output])

    return app

if __name__ == "__main__":
    app = gradio_app()
    app.launch(share=True)