#!/usr/bin/env python
"""A demo of the VitPose model.
This code is based on the implementation from the Colab notebook:
https://colab.research.google.com/drive/1e8fcby5rhKZWcr9LSN8mNbQ0TU4Dxxpo
"""
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] # take first image results
# Human label refers 0 index in COCO dataset
person_boxes_xyxy = result["boxes"][result["labels"] == 0]
person_boxes_xyxy = person_boxes_xyxy.cpu().numpy()
# Convert boxes from VOC (x1, y1, x2, y2) to COCO (x1, y1, w, h) format
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)
# for vitpose-plus-base checkpoint we should additionally provide dataset_index
# to specify which MOE experts to use for inference
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] # results for first image
# make results more human-readable
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)
# preprocess to torch tensor of shape (n_objects, n_keypoints, 2)
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()
# annotate bounding boxes
annotated_frame = bounding_box_annotator.annotate(scene=image.copy(), detections=detections)
# annotate edges and vertices
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), # noqa: ARG001, B008
) -> 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()