Spaces:
Paused
Paused
| # 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 | |