|
|
|
|
| from PIL import Image
|
| import mediapipe as mp
|
| import numpy as np
|
|
|
| mp_pose = mp.solutions.pose
|
| mp_drawing = mp.solutions.drawing_utils
|
|
|
| LANDMARK = mp_pose.PoseLandmark
|
| TORSO_INDICES = [
|
| LANDMARK.LEFT_SHOULDER.value,
|
| LANDMARK.RIGHT_SHOULDER.value,
|
| LANDMARK.LEFT_HIP.value,
|
| LANDMARK.RIGHT_HIP.value,
|
| ]
|
|
|
|
|
| def extract_pose_mediapipe(image_pil: Image.Image):
|
| img_rgb = np.array(image_pil.convert("RGB"))
|
| height, width = img_rgb.shape[:2]
|
|
|
| with mp_pose.Pose(
|
| static_image_mode=True,
|
| model_complexity=1,
|
| enable_segmentation=False,
|
| min_detection_confidence=0.5,
|
| min_tracking_confidence=0.5,
|
| ) as pose:
|
| results = pose.process(img_rgb)
|
|
|
| if not results.pose_landmarks:
|
| return None
|
|
|
| landmarks_raw = results.pose_landmarks.landmark
|
|
|
| landmarks = {}
|
| for idx, name in enumerate(LANDMARK):
|
| px = int(landmarks_raw[idx].x * width)
|
| py = int(landmarks_raw[idx].y * height)
|
| landmarks[name.name] = (px, py)
|
|
|
| skeleton_canvas = np.zeros_like(img_rgb)
|
| mp_drawing.draw_landmarks(
|
| skeleton_canvas,
|
| results.pose_landmarks,
|
| mp_pose.POSE_CONNECTIONS,
|
| landmark_drawing_spec=mp_drawing.DrawingSpec(
|
| color=(0, 255, 200), thickness=2, circle_radius=4
|
| ),
|
| connection_drawing_spec=mp_drawing.DrawingSpec(color=(100, 180, 255), thickness=2),
|
| )
|
|
|
| torso_x = [landmarks_raw[idx].x * width for idx in TORSO_INDICES]
|
| torso_y = [landmarks_raw[idx].y * height for idx in TORSO_INDICES]
|
|
|
| pad_x = int((max(torso_x) - min(torso_x)) * 0.12)
|
| pad_y_top = int((max(torso_y) - min(torso_y)) * 0.08)
|
| pad_y_bottom = int((max(torso_y) - min(torso_y)) * 0.05)
|
|
|
| x1 = max(0, int(min(torso_x)) - pad_x)
|
| y1 = max(0, int(min(torso_y)) - pad_y_top)
|
| x2 = min(width, int(max(torso_x)) + pad_x)
|
| y2 = min(height, int(max(torso_y)) + pad_y_bottom)
|
|
|
| return {
|
| "landmarks": landmarks,
|
| "skeleton_pil": Image.fromarray(skeleton_canvas),
|
| "torso_box": (x1, y1, x2, y2),
|
| "image_size": (width, height),
|
| }
|
|
|
|
|
| def draw_skeleton_overlay(image_pil: Image.Image, pose_data: dict, alpha: float = 0.6):
|
| base = np.array(image_pil.convert("RGB")).astype(float)
|
| skeleton = np.array(pose_data["skeleton_pil"].convert("RGB")).astype(float)
|
|
|
| skeleton_mask = (skeleton.sum(axis=2) > 0).astype(float)[:, :, np.newaxis]
|
| blended = base * (1 - skeleton_mask * alpha) + skeleton * skeleton_mask * alpha
|
|
|
| return Image.fromarray(blended.astype(np.uint8))
|
|
|
