Spaces:
Paused
Paused
File size: 3,221 Bytes
0d66750 b535d9f c9861e0 0d66750 c9861e0 b535d9f 0d66750 fdb72e1 0d66750 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
# YOLO
import cv2
import torch
import numpy as np
from ultralytics import YOLO
from PIL import Image, ImageDraw
model_name = 'yolo11s-pose.pt'
model = YOLO(model_name)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
half = device=='cuda'
connection_color_map = {
(0, 1): 'lightgreen', # Nose to Left Eye
(0, 2): 'lightgreen', # Nose to Right Eye
(1, 2): 'lightgreen', # Left Eye to Right Eye
(1, 3): 'magenta', # Left Eye to Left Ear
(2, 4): 'magenta', # Right Eye to Right Ear
(3, 5): 'lightblue', # Left Ear to Left Shoulder
(4, 6): 'lightblue', # Right Ear to Right Shoulder
(5, 6): 'purple', # Left Shoulder to Right Shoulder
(5, 7): 'blue', # Left Shoulder to Left Elbow
(7, 9): 'blue', # Left Elbow to Left Wrist
(6, 8): 'red', # Right Shoulder to Right Elbow
(8, 10): 'red', # Right Elbow to Right Wrist
(5, 11): 'purple', # Left Shoulder to Left Hip
(6, 12): 'purple', # Right Shoulder to Right Hip
(11, 12): 'purple', # Left Hip to Right Hip
(11, 13): 'orange', # Left Hip to Left Knee
(13, 15): 'orange', # Left Knee to Left Ankle
(12, 14): 'orange', # Right Hip to Right Knee
(14, 16): 'orange', # Right Knee to Right Ankle
}
def detect_pose(src_img):
results = model.predict(
source=src_img,
# conf=conf_threshold,
# iou=iou_threshold,
max_det=1,
show_labels=True,
show_conf=True,
imgsz=32*5,
classes=[0],
# line_width=3,
half=half,
device=device,
)
return results
def draw_pose(skeleton_image:Image.Image, keypoints, threshold=0.7, line_width = 3,circle_radius = 6):
coords = keypoints.cpu().numpy().xy[0]
scores = keypoints.data.squeeze()[:,2].cpu().numpy()
# Create a blank image with the same size as the input image
# skeleton_image = Image.new("RGB", im.size, "black") # White background
draw = ImageDraw.Draw(skeleton_image)
# Draw lines connecting the keypoints
for connection, color in connection_color_map.items():
try:
start_point = coords[connection[0]]
end_point = coords[connection[1]]
if scores[connection[0]] > threshold and scores[connection[1]] > threshold:
draw.line([tuple(start_point), tuple(end_point)], fill=color, width=line_width)
draw.ellipse(
(start_point[0] - circle_radius, start_point[1] - circle_radius,
start_point[0] + circle_radius, start_point[1] + circle_radius),
fill=color
)
draw.ellipse(
(end_point[0] - circle_radius, end_point[1] - circle_radius,
end_point[0] + circle_radius, end_point[1] + circle_radius),
fill=color
)
except IndexError:
pass # Handle cases where keypoints might be missing
# Save the skeleton image
return skeleton_image
|