Spaces:
Sleeping
Sleeping
| """ | |
| Example usage of MediaPipePoseEstimator. | |
| This script demonstrates how to use the MediaPipePoseEstimator to detect | |
| pose landmarks in an image and visualize the results. | |
| """ | |
| import cv2 | |
| import numpy as np | |
| from mediapipe_pose_estimator import MediaPipePoseEstimator | |
| def main(): | |
| # Initialize the estimator with the model file | |
| # Note: You need to have the pose_landmarker_lite.task file in the current directory | |
| try: | |
| estimator = MediaPipePoseEstimator(model_asset_path='pose_landmarker_lite.task') | |
| except Exception as e: | |
| print(f"Error initializing estimator: {e}") | |
| print("Make sure you have downloaded the pose_landmarker_lite.task model file") | |
| print("Run: python download_model.py") | |
| return | |
| # Create a sample image or load an existing one | |
| # For this example, we'll create a blank image | |
| sample_image = np.zeros((480, 640, 3), dtype=np.uint8) | |
| # Add a simple stick figure for testing (optional) | |
| # Head | |
| cv2.circle(sample_image, (320, 100), 30, (255, 255, 255), -1) | |
| # Body | |
| cv2.line(sample_image, (320, 130), (320, 250), (255, 255, 255), 2) | |
| # Arms | |
| cv2.line(sample_image, (270, 180), (370, 180), (255, 255, 255), 2) | |
| # Legs | |
| cv2.line(sample_image, (320, 250), (280, 320), (255, 255, 255), 2) | |
| cv2.line(sample_image, (320, 250), (360, 320), (255, 255, 255), 2) | |
| print("Detecting pose in sample image...") | |
| # Detect pose | |
| result = estimator.detect_pose(sample_image) | |
| print(f"Inference time: {result['inference_time_ms']:.2f} ms") | |
| print(f"Number of keypoints: {len(result['keypoints'])}") | |
| # Print detected keypoints with confidence > 0.1 | |
| print("\nDetected keypoints (confidence > 0.1):") | |
| for name, kp in result['keypoints'].items(): | |
| if kp['confidence'] > 0.1: | |
| print(f" {name}: ({kp['x']:.3f}, {kp['y']:.3f}) conf={kp['confidence']:.3f}") | |
| # Draw keypoints on the image | |
| annotated_image = estimator.draw_keypoints(sample_image, result) | |
| # Display the result | |
| cv2.imshow('Original Image', sample_image) | |
| cv2.imshow('Annotated Image', annotated_image) | |
| print("\nPress any key to close windows...") | |
| cv2.waitKey(0) | |
| cv2.destroyAllWindows() | |
| # Save the annotated image | |
| cv2.imwrite('annotated_sample.jpg', annotated_image) | |
| print("Annotated image saved as 'annotated_sample.jpg'") | |
| if __name__ == '__main__': | |
| main() | |