import torch
import requests
import numpy as np

from PIL import Image

from transformers import (
    AutoProcessor,
    RTDetrForObjectDetection,
    VitPoseForPoseEstimation,
)

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

url = "http://images.cocodataset.org/val2017/000000000139.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# ------------------------------------------------------------------------
# Stage 1. Detect humans on the image
# ------------------------------------------------------------------------

# You can choose detector by your choice
person_image_processor = AutoProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
person_model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365", device_map=device)

inputs = person_image_processor(images=image, return_tensors="pt").to(device)

with torch.no_grad():
    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 = result["boxes"][result["labels"] == 0]
person_boxes = person_boxes.cpu().numpy()

# Convert boxes from VOC (x1, y1, x2, y2) to COCO (x1, y1, w, h) format
person_boxes[:, 2] = person_boxes[:, 2] - person_boxes[:, 0]
person_boxes[:, 3] = person_boxes[:, 3] - person_boxes[:, 1]

# ------------------------------------------------------------------------
# Stage 2. Detect keypoints for each person found
# ------------------------------------------------------------------------

image_processor = AutoProcessor.from_pretrained("usyd-community/vitpose-plus-base")
model = VitPoseForPoseEstimation.from_pretrained("usyd-community/vitpose-plus-base", device_map=device)

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

# This is MOE architecture, we should specify dataset indexes for each image in range 0..5
inputs["dataset_index"] = torch.tensor([0], device=device)

with torch.no_grad():
    outputs = model(**inputs)

pose_results = image_processor.post_process_pose_estimation(outputs, boxes=[person_boxes], threshold=0.3)
image_pose_result = pose_results[0]  # results for first image

for i, person_pose in enumerate(image_pose_result):
    print(f"Person #{i}")
    for keypoint, label, score in zip(
        person_pose["keypoints"], person_pose["labels"], person_pose["scores"]
    ):
        keypoint_name = model.config.id2label[label.item()]
        x, y = keypoint
        print(f" - {keypoint_name}: x={x.item():.2f}, y={y.item():.2f}, score={score.item():.2f}")

Person #0
 - Nose: x=428.81, y=171.53, score=0.92
 - L_Eye: x=429.32, y=168.30, score=0.92
 - R_Eye: x=428.84, y=168.47, score=0.82
 - L_Ear: x=434.60, y=166.54, score=0.90
 - R_Ear: x=440.14, y=165.80, score=0.80
 - L_Shoulder: x=440.74, y=176.95, score=0.96
 - R_Shoulder: x=444.06, y=177.52, score=0.68
 - L_Elbow: x=436.30, y=197.08, score=0.91
 - R_Elbow: x=432.29, y=201.22, score=0.79
 - L_Wrist: x=429.91, y=217.90, score=0.84
 - R_Wrist: x=421.08, y=212.72, score=0.90
 - L_Hip: x=446.15, y=223.88, score=0.74
 - R_Hip: x=449.32, y=223.45, score=0.65
 - L_Knee: x=443.73, y=255.72, score=0.76
 - R_Knee: x=450.72, y=255.21, score=0.73
 - L_Ankle: x=452.14, y=287.30, score=0.66
 - R_Ankle: x=456.02, y=285.99, score=0.72
Person #1
 - Nose: x=398.22, y=181.60, score=0.88
 - L_Eye: x=398.67, y=179.84, score=0.87
 - R_Eye: x=396.07, y=179.44, score=0.87
 - R_Ear: x=388.94, y=180.38, score=0.87
 - L_Shoulder: x=397.11, y=194.19, score=0.71
 - R_Shoulder: x=384.75, y=190.74, score=0.55