#!/usr/bin/env python

import pathlib
import tempfile

import cv2
import gradio as gr
import numpy as np
import PIL.Image
import spaces
import supervision as sv
import torch
import tqdm
from transformers import AutoProcessor, RTDetrForObjectDetection, VitPoseForPoseEstimation

DESCRIPTION = "# ViTPose"

MAX_NUM_FRAMES = 300

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

person_detector_name = "PekingU/rtdetr_r50vd_coco_o365"
person_image_processor = AutoProcessor.from_pretrained(person_detector_name)
person_model = RTDetrForObjectDetection.from_pretrained(person_detector_name, device_map=device)

pose_model_name = "usyd-community/vitpose-base-simple"
pose_image_processor = AutoProcessor.from_pretrained(pose_model_name)
pose_model = VitPoseForPoseEstimation.from_pretrained(pose_model_name, device_map=device)


@spaces.GPU(duration=5)
@torch.inference_mode()
def process_image(image: PIL.Image.Image) -> tuple[PIL.Image.Image, list[dict]]:
    inputs = person_image_processor(images=image, return_tensors="pt").to(device)
    outputs = person_model(**inputs)
    results = person_image_processor.post_process_object_detection(
        outputs, target_sizes=torch.tensor([(image.height, image.width)]), threshold=0.3
    )
    result = results[0]

    person_boxes_xyxy = result["boxes"][result["labels"] == 0]
    person_boxes_xyxy = person_boxes_xyxy.cpu().numpy()

    person_boxes = person_boxes_xyxy.copy()
    person_boxes[:, 2] = person_boxes[:, 2] - person_boxes[:, 0]
    person_boxes[:, 3] = person_boxes[:, 3] - person_boxes[:, 1]

    inputs = pose_image_processor(image, boxes=[person_boxes], return_tensors="pt").to(device)

    if pose_model.config.backbone_config.num_experts > 1:
        dataset_index = torch.tensor([0] * len(inputs["pixel_values"]))
        dataset_index = dataset_index.to(inputs["pixel_values"].device)
        inputs["dataset_index"] = dataset_index

    outputs = pose_model(**inputs)

    pose_results = pose_image_processor.post_process_pose_estimation(outputs, boxes=[person_boxes])
    image_pose_result = pose_results[0]

    human_readable_results = []
    for i, person_pose in enumerate(image_pose_result):
        data = {
            "person_id": i,
            "bbox": person_pose["bbox"].numpy().tolist(),
            "keypoints": [],
        }
        for keypoint, label, score in zip(
            person_pose["keypoints"], person_pose["labels"], person_pose["scores"], strict=True
        ):
            keypoint_name = pose_model.config.id2label[label.item()]
            x, y = keypoint
            data["keypoints"].append({"name": keypoint_name, "x": x.item(), "y": y.item(), "score": score.item()})
        human_readable_results.append(data)

    xy = [pose_result["keypoints"] for pose_result in image_pose_result]
    xy = torch.stack(xy).cpu().numpy()

    scores = [pose_result["scores"] for pose_result in image_pose_result]
    scores = torch.stack(scores).cpu().numpy()

    keypoints = sv.KeyPoints(xy=xy, confidence=scores)
    detections = sv.Detections(xyxy=person_boxes_xyxy)

    edge_annotator = sv.EdgeAnnotator(color=sv.Color.GREEN, thickness=1)
    vertex_annotator = sv.VertexAnnotator(color=sv.Color.RED, radius=2)
    bounding_box_annotator = sv.BoxAnnotator(color=sv.Color.WHITE, color_lookup=sv.ColorLookup.INDEX, thickness=1)

    annotated_frame = image.copy()

    annotated_frame = bounding_box_annotator.annotate(scene=image.copy(), detections=detections)

    annotated_frame = edge_annotator.annotate(scene=annotated_frame, key_points=keypoints)
    return vertex_annotator.annotate(scene=annotated_frame, key_points=keypoints), human_readable_results


@spaces.GPU(duration=90)
def process_video(
    video_path: str,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> str:
    cap = cv2.VideoCapture(video_path)

    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    fps = cap.get(cv2.CAP_PROP_FPS)
    num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as out_file:
        writer = cv2.VideoWriter(out_file.name, fourcc, fps, (width, height))
        for _ in tqdm.auto.tqdm(range(min(MAX_NUM_FRAMES, num_frames))):
            ok, frame = cap.read()
            if not ok:
                break
            rgb_frame = frame[:, :, ::-1]
            annotated_frame, _ = process_image(PIL.Image.fromarray(rgb_frame))
            writer.write(np.asarray(annotated_frame)[:, :, ::-1])
        writer.release()
    cap.release()
    return out_file.name


with gr.Blocks(css_paths="style.css") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Tabs():
        with gr.Tab("Image"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(label="Input Image", type="pil")
                    run_button_image = gr.Button()
                with gr.Column():
                    output_image = gr.Image(label="Output Image")
                    output_json = gr.JSON(label="Output JSON")
            gr.Examples(
                examples=sorted(pathlib.Path("images").glob("*.jpg")),
                inputs=input_image,
                outputs=[output_image, output_json],
                fn=process_image,
            )

            run_button_image.click(
                fn=process_image,
                inputs=input_image,
                outputs=[output_image, output_json],
            )

        with gr.Tab("Video"):
            gr.Markdown(f"The input video will be truncated to {MAX_NUM_FRAMES} frames.")

            with gr.Row():
                with gr.Column():
                    input_video = gr.Video(label="Input Video")
                    run_button_video = gr.Button()
                with gr.Column():
                    output_video = gr.Video(label="Output Video")

            gr.Examples(
                examples=sorted(pathlib.Path("videos").glob("*.mp4")),
                inputs=input_video,
                outputs=output_video,
                fn=process_video,
                cache_examples=False,
            )
            run_button_video.click(
                fn=process_video,
                inputs=input_video,
                outputs=output_video,
            )


if __name__ == "__main__":
    demo.queue(max_size=20).launch()