Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| from .config import WIDTH, HEIGHT, FPS, POSE_CONNECTIONS | |
| # ========================= | |
| # COLORS | |
| # ========================= | |
| WHITE = (255, 255, 255) | |
| # Background (#190124 in HEX) | |
| BACKGROUND = (36, 1, 25) | |
| # Finger colors (BGR) | |
| FINGER_COLORS = { | |
| "thumb": (0, 0, 255), # Red | |
| "index": (0, 255, 0), # Green | |
| "middle": (255, 0, 0), # Blue | |
| "ring": (0, 255, 255), # Yellow | |
| "pinky": (255, 0, 255) # Magenta | |
| } | |
| # ========================= | |
| # HAND | |
| # ========================= | |
| def draw_hand_stylish(img, hand_points, thickness=5): | |
| h, w = img.shape[:2] | |
| pts = [(int(x * w), int(y * h)) for x, y, _ in hand_points] | |
| fingers = { | |
| "thumb": [0, 1, 2, 3, 4], | |
| "index": [0, 5, 6, 7, 8], | |
| "middle": [0, 9, 10, 11, 12], | |
| "ring": [0, 13, 14, 15, 16], | |
| "pinky": [0, 17, 18, 19, 20] | |
| } | |
| for finger_name, idxs in fingers.items(): | |
| color = FINGER_COLORS[finger_name] | |
| for i in range(len(idxs) - 1): | |
| x1, y1 = pts[idxs[i]] | |
| x2, y2 = pts[idxs[i + 1]] | |
| cv2.line( | |
| img, | |
| (x1, y1), | |
| (x2, y2), | |
| color, | |
| thickness, | |
| cv2.LINE_AA | |
| ) | |
| # ========================= | |
| # POSE | |
| # ========================= | |
| def draw_pose_lines(img, points, connections, thickness=4): | |
| h, w = img.shape[:2] | |
| pts = [(int(x * w), int(y * h)) for x, y, _ in points] | |
| for i, j in connections: | |
| if i < len(pts) and j < len(pts): | |
| cv2.line( | |
| img, | |
| pts[i], | |
| pts[j], | |
| WHITE, | |
| thickness, | |
| cv2.LINE_AA | |
| ) | |
| # ========================= | |
| # FACE LANDMARKS | |
| # ========================= | |
| def draw_points_and_connections(img, points): | |
| h, w = img.shape[:2] | |
| pts = [(int(x * w), int(y * h)) for x, y, _ in points] | |
| for (x, y) in pts: | |
| if 0 <= x < w and 0 <= y < h: | |
| cv2.circle( | |
| img, | |
| (x, y), | |
| 2, | |
| WHITE, | |
| -1 | |
| ) | |
| # ========================= | |
| # MAIN DRAW | |
| # ========================= | |
| def draw_skeleton(frame, img): | |
| # Pose | |
| draw_pose_lines( | |
| img, | |
| frame[42:75], | |
| POSE_CONNECTIONS, | |
| thickness=4 | |
| ) | |
| # Left hand | |
| draw_hand_stylish( | |
| img, | |
| frame[0:21], | |
| thickness=4 | |
| ) | |
| # Right hand | |
| draw_hand_stylish( | |
| img, | |
| frame[21:42], | |
| thickness=4 | |
| ) | |
| # Face | |
| draw_points_and_connections( | |
| img, | |
| frame[75:553] | |
| ) | |
| # ========================= | |
| # RENDER | |
| # ========================= | |
| def render(frames, output="output.mp4"): | |
| writer = cv2.VideoWriter( | |
| output, | |
| cv2.VideoWriter_fourcc(*"mp4v"), | |
| FPS, | |
| (WIDTH, HEIGHT) | |
| ) | |
| for frame in frames: | |
| img = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8) | |
| img[:] = BACKGROUND | |
| draw_skeleton(frame, img) | |
| writer.write(img) | |
| writer.release() | |
| return output |