diff --git a/keypoint_demo.py b/keypoint_demo.py new file mode 100644 index 0000000..3a9ea2e --- /dev/null +++ b/keypoint_demo.py @@ -0,0 +1,119 @@ +# keypoint_demo.py +import cv2 +import numpy as np +from mmpose.apis import MMPoseInferencer + +def draw_keypoints(frame, keypoints, keypoint_scores, threshold=0.3): + """Draw COCO 17-keypoint skeleton on frame.""" + + # COCO 17 skeleton connections (joint pairs) + SKELETON = [ + (0, 1), (0, 2), # nose to eyes + (1, 3), (2, 4), # eyes to ears + (5, 6), # shoulders + (5, 7), (7, 9), # left arm + (6, 8), (8, 10), # right arm + (5, 11), (6, 12), # shoulders to hips + (11, 12), # hips + (11, 13), (13, 15), # left leg + (12, 14), (14, 16), # right leg + ] + + KEYPOINT_NAMES = [ + 'nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear', + 'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow', + 'left_wrist', 'right_wrist', 'left_hip', 'right_hip', + 'left_knee', 'right_knee', 'left_ankle', 'right_ankle' + ] + + h, w = frame.shape[:2] + + for person_kps, person_scores in zip(keypoints, keypoint_scores): + # Draw skeleton lines + for (j1, j2) in SKELETON: + if person_scores[j1] > threshold and person_scores[j2] > threshold: + pt1 = (int(person_kps[j1][0]), int(person_kps[j1][1])) + pt2 = (int(person_kps[j2][0]), int(person_kps[j2][1])) + # Clamp to frame bounds + if (0 <= pt1[0] < w and 0 <= pt1[1] < h and + 0 <= pt2[0] < w and 0 <= pt2[1] < h): + cv2.line(frame, pt1, pt2, (0, 255, 0), 2) + + # Draw keypoints + for idx, (kp, score) in enumerate(zip(person_kps, person_scores)): + if score > threshold: + x, y = int(kp[0]), int(kp[1]) + if 0 <= x < w and 0 <= y < h: + # Color by confidence: green=high, yellow=medium + color = (0, 255, 0) if score > 0.7 else (0, 200, 200) + cv2.circle(frame, (x, y), 4, color, -1) + cv2.circle(frame, (x, y), 5, (255, 255, 255), 1) # white border + + return frame + + +def main(): + # Initialize inferencer with RTMPose + # det_model: person detector + # pose2d: pose estimator + inferencer = MMPoseInferencer( + pose2d='human', # uses the built-in alias — downloads everything automatically + device='cuda:0' + ) + + cap = cv2.VideoCapture(0) # 0 = default webcam; use 1, 2, etc. for others + cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) + cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) + + print("Press 'q' to quit, 's' to save a frame") + + while True: + ret, frame = cap.read() + if not ret: + break + + # Run inference + # scope='mmpose' is important to avoid namespace issues + result_generator = inferencer( + frame, + show=False, # we handle display ourselves + return_datasamples=False + ) + results = next(result_generator) + + # Extract keypoints from results + predictions = results.get('predictions', [[]])[0] + + all_keypoints = [] + all_scores = [] + + for pred in predictions: + kps = pred.get('keypoints', []) + scores = pred.get('keypoint_scores', []) + if len(kps) > 0: + all_keypoints.append(np.array(kps)) + all_scores.append(np.array(scores)) + + # Draw on frame + if all_keypoints: + frame = draw_keypoints(frame, all_keypoints, all_scores) + + # Show FPS + cv2.putText(frame, f'Subjects: {len(all_keypoints)}', + (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2) + + cv2.imshow('MMPose RTMPose Keypoints', frame) + + key = cv2.waitKey(1) & 0xFF + if key == ord('q'): + break + elif key == ord('s'): + cv2.imwrite('keypoint_capture.jpg', frame) + print("Frame saved.") + + cap.release() + cv2.destroyAllWindows() + + +if __name__ == '__main__': + main() \ No newline at end of file