#!/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 <div style="display: flex; gap: 10px;"> <a href="https://huggingface.co/docs/transformers/en/model_doc/vitpose"> <img src="https://img.shields.io/badge/Huggingface-FFD21E?style=flat&logo=Huggingface&logoColor=black" alt="Huggingface"> </a> <a href="https://arxiv.org/abs/2204.12484"> <img src="https://img.shields.io/badge/Arvix-B31B1B?style=flat&logo=arXiv&logoColor=white" alt="Paper"> </a> <a href="https://github.com/ViTAE-Transformer/ViTPose"> <img src="https://img.shields.io/badge/Github-100000?style=flat&logo=github&logoColor=white" alt="Github"> </a> </div> ViTPose is a state-of-the-art human pose estimation model based on Vision Transformers (ViT). It employs a standard, non-hierarchical ViT backbone and a simple decoder head to predict keypoint heatmaps from images. Despite its simplicity, ViTPose achieves top results on the MS COCO Keypoint Detection benchmark. ViTPose++ further improves performance with a mixture-of-experts (MoE) module and extensive pre-training. The model is scalable, flexible, and demonstrates strong transferability across pose estimation tasks. **Key features:** - PyTorch implementation - Scalable model size (100M to 1B parameters) - Flexible training and inference - State-of-the-art accuracy on challenging benchmarks """ COLORS = [ "#A351FB", "#FF4040", "#FFA1A0", "#FF7633", "#FFB633", "#D1D435", "#4CFB12", "#94CF1A", "#40DE8A", "#1B9640", "#00D6C1", "#2E9CAA", "#00C4FF", "#364797", "#6675FF", "#0019EF", "#863AFF", ] COLORS = [sv.Color.from_hex(color_hex=c) for c in COLORS] MAX_NUM_FRAMES = 300 keypoint_score = 0.3 enable_labels_annotator = True enable_vertices_annotator = True 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 detect_pose_image( image: PIL.Image.Image, threshold: float = 0.3, enable_labels_annotator: bool = True, enable_vertices_annotator: bool = True, ) -> tuple[PIL.Image.Image, list[dict]]: """Detects persons and estimates their poses in a single image. Args: image (PIL.Image.Image): Input image in which to detect persons and estimate poses. threshold (Float): Confidence threshold for pose keypoints. enable_labels_annotator (bool): Whether to enable annotating labels for pose keypoints. enable_vertices_annotator (bool): Whether to enable annotating vertices for pose keypoints Returns: tuple[PIL.Image.Image, list[dict]]: - Annotated image with bounding boxes and pose keypoints drawn. - List of dictionaries containing human-readable pose estimation results for each detected person. """ 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=threshold ) result = results[0] # take first image results detections = sv.Detections.from_transformers(result) person_detections_xywh = sv.xyxy_to_xywh(detections[detections.class_id == 0].xyxy) inputs = pose_image_processor(image, boxes=[person_detections_xywh], 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_detections_xywh]) image_pose_result = pose_results[0] # results for first image # make results more human-readable human_readable_results = [] person_pose_labels = [] 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()] person_pose_labels.append(keypoint_name) x, y = keypoint data["keypoints"].append({"name": keypoint_name, "x": x.item(), "y": y.item(), "score": score.item()}) human_readable_results.append(data) line_thickness = sv.calculate_optimal_line_thickness(resolution_wh=(image.width, image.height)) text_scale = sv.calculate_optimal_text_scale(resolution_wh=(image.width, image.height)) edge_annotator = sv.EdgeAnnotator(color=sv.Color.WHITE, thickness=line_thickness) vertex_annotator = sv.VertexAnnotator(color=sv.Color.BLUE, radius=3) box_annotator = sv.BoxAnnotator(color=sv.Color.WHITE, color_lookup=sv.ColorLookup.INDEX, thickness=3) vertex_label_annotator = sv.VertexLabelAnnotator( color=COLORS, smart_position=True, border_radius=3, text_thickness=2, text_scale=text_scale ) annotated_frame = box_annotator.annotate(scene=image.copy(), detections=detections) for _, person_pose in enumerate(image_pose_result): person_keypoints = sv.KeyPoints.from_transformers([person_pose]) person_labels = [pose_model.config.id2label[label.item()] for label in person_pose["labels"]] # annotate edges and vertices for this person annotated_frame = edge_annotator.annotate(scene=annotated_frame, key_points=person_keypoints) # annotate labels for this person if enable_labels_annotator: annotated_frame = vertex_label_annotator.annotate( scene=np.array(annotated_frame), key_points=person_keypoints, labels=person_labels ) # annotate vertices for this person if enable_vertices_annotator: annotated_frame = vertex_annotator.annotate(scene=annotated_frame, key_points=person_keypoints) return annotated_frame, human_readable_results # Decorate this function with `@spaces.GPU` to ensure that ZeroGPU is allocated once for the entire video processing. # Although `detect_pose_image` (called per frame) is already decorated, without this decorator, ZeroGPU would be invoked for each frame, # causing significant overhead and slowdowns. By decorating this function, all frames are processed sequentially after a single GPU allocation. @spaces.GPU(duration=90) def detect_pose_video( video_path: str, threshold: float, enable_labels_annotator: bool = True, enable_vertices_annotator: bool = True, progress: gr.Progress = gr.Progress(track_tqdm=True), # noqa: ARG001, B008 ) -> str: """Detects persons and estimates their poses for each frame in a video, saving the annotated video. Args: video_path (str): Path to the input video file. threshold (Float): Confidence threshold for pose keypoints. enable_labels_annotator (bool): Whether to enable annotating labels for pose keypoints. enable_vertices_annotator (bool): Whether to enable annotating vertices for pose keypoints. progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True). Returns: str: Path to the output video file with annotated bounding boxes and pose keypoints. """ 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, _ = detect_pose_image( PIL.Image.fromarray(rgb_frame), threshold=threshold, enable_labels_annotator=enable_labels_annotator, enable_vertices_annotator=enable_vertices_annotator, ) 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) keypoint_score = gr.Slider( minimum=0.0, maximum=1.0, value=0.6, step=0.01, info="Adjust the confidence threshold for keypoint detection.", label="Keypoint Score Threshold", ) enable_labels_annotator = gr.Checkbox(interactive=True, value=True, label="Enable Labels") enable_vertices_annotator = gr.Checkbox(interactive=True, value=True, label="Enable Vertices") 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=[[str(img), 0.5, True, True] for img in sorted(pathlib.Path("images").glob("*.jpg"))], inputs=[input_image, keypoint_score, enable_labels_annotator, enable_vertices_annotator], outputs=[output_image, output_json], fn=detect_pose_image, ) run_button_image.click( fn=detect_pose_image, inputs=[input_image, keypoint_score, enable_labels_annotator, enable_vertices_annotator], 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=[[str(video), 0.5, True, True] for video in sorted(pathlib.Path("videos").glob("*.mp4"))], inputs=[input_video, keypoint_score, enable_labels_annotator, enable_vertices_annotator], outputs=output_video, fn=detect_pose_video, cache_examples=False, ) run_button_video.click( fn=detect_pose_video, inputs=[input_video, keypoint_score, enable_labels_annotator, enable_vertices_annotator], outputs=output_video, ) if __name__ == "__main__": demo.launch(mcp_server=True, ssr_mode=False)