from PIL import Image, ImageDraw, ImageFont
import cv2
import numpy as np
from transformers import AutoTokenizer, PaliGemmaForConditionalGeneration, PaliGemmaProcessor
import torch
import spaces
import gradio as gr
import jax
import jax.numpy as jnp

# Load PaliGemma
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_id = "google/paligemma-3b-mix-224"
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(device)
processor = PaliGemmaProcessor.from_pretrained(model_id)

# Function to draw bounding boxes (your original code)
def draw_bounding_box(draw, coordinates, label, width, height):
    y1, x1, y2, x2 = coordinates
    y1, x1, y2, x2 = map(round, (y1*height, x1*width, y2*height, x2*width))

    text_width, text_height = draw.textsize(label)
    draw.rectangle([(x1, y1 - text_height - 2), (x1 + text_width + 4, y1)], fill="red")

    # Draw label text
    draw.text((x1 + 2, y1 - text_height - 2), label, fill="white")

    # Draw bounding box
    draw.rectangle([(x1, y1), (x2, y2)], outline="red", width=2)

@spaces.GPU
def process_video(video_path, input_text):
    cap = cv2.VideoCapture(video_path)
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter('output_paligemma_keras.avi', fourcc, 20.0, (int(cap.get(3)), int(cap.get(4))))

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

        # Convert the frame to a PIL Image
        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

        # Send text prompt and image as input.
        inputs = processor(text=input_text, images=img,
                            padding="longest", do_convert_rgb=True, return_tensors="pt").to("cuda")
        inputs = inputs.to(dtype=model.dtype)

        # Get output.
        with torch.no_grad():
            output = model.generate(**inputs, max_length=496)

        paligemma_response = processor.decode(output[0], skip_special_tokens=True)[len(input_text):].lstrip("\n")
        # print(paligemma_response)  # For debugging

        detections = paligemma_response.split(" ; ")

        # Parse the output bounding box coordinates
        parsed_coordinates = []
        labels = []

        for item in detections:
            # Remove '<loc>' tags and split the string
            # print(item)
            detection = item.replace("<loc", "").split()

            if len(detection) >= 2:
              coordinates_str = detection[0]
              label = detection[1]
              labels.append(label)
            else:
              # No label detected, skip the iteration.
              continue

            # Split the coordinates string by '>' to get individual coordinates
            coordinates = coordinates_str.split(">")
            coordinates = coordinates[:4]  # Slicing to ensure only 4 values

            if coordinates[-1] == '':
                coordinates = coordinates[:-1]
                # print(coordinates)

            coordinates = [int(coord)/1024 for coord in coordinates]
            # location_values = [int(loc) for loc in re.findall(r'\d{4}', coordinates)]
            # y1, x1, y2, x2 = [value / 1024 for value in location_values]
            parsed_coordinates.append(coordinates)

        width = img.size[0]
        height = img.size[1]

        # Draw bounding boxes on the frame using PIL
        draw = ImageDraw.Draw(img)
        for coordinates, label in zip(parsed_coordinates, labels):
            draw_bounding_box(draw, coordinates, label, width=width, height=height)

        # Convert the PIL Image back to OpenCV format
        frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

        # Write the frame to the output video
        out.write(frame)

    cap.release()
    out.release()

    return "output_paligemma_keras.avi"

with gr.Blocks() as demo:
    gr.Markdown("## Zero-shot Object Tracking with PaliGemma")
    gr.Markdown("This is a demo for zero-shot object tracking using [PaliGemma](https://huggingface.co/google/paligemma-3b-mix-448) vision language model by Google.")
    gr.Markdown("Simply upload a video and enter the candidate labels, or try the example below. Text input should be ; separated. 👇")
    with gr.Tab(label="Video"):
        with gr.Row():
            input_video = gr.Video(label='Input Video')
            output_video = gr.Video(label='Output Video')
        with gr.Row():
            candidate_labels = gr.Textbox(
                label='Labels',
                placeholder='Labels separated by a comma',
            )
            submit = gr.Button()
        gr.Examples(
            fn=process_video,
            examples=[["./input.mp4", "detect person"]],
            inputs=[
                input_video,
                candidate_labels,
                
            ],
            outputs=output_video
        )

    submit.click(fn=process_video,
        inputs=[input_video, candidate_labels],
        outputs=output_video
        )

demo.launch(debug=False, show_error=True)