Spaces:
Sleeping
Sleeping
File size: 2,271 Bytes
f20fe1f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import cv2
import argparse
from movenet_analyzer import MoveNetAnalyzer
def main():
parser = argparse.ArgumentParser(description='MoveNet Pose Analysis Demo')
parser.add_argument('--video', type=str, help='Path to video file (optional)')
parser.add_argument('--model', type=str, default='lightning', choices=['lightning', 'thunder'],
help='MoveNet model variant (lightning or thunder)')
args = parser.parse_args()
# Initialize the MoveNet analyzer
analyzer = MoveNetAnalyzer(model_name=args.model)
# Initialize video capture
if args.video:
cap = cv2.VideoCapture(args.video)
else:
cap = cv2.VideoCapture(0) # Use webcam if no video file provided
if not cap.isOpened():
print("Error: Could not open video source")
return
while True:
ret, frame = cap.read()
if not ret:
break
# Process frame
frame_with_pose, analysis = analyzer.process_frame(frame)
# Display analysis results
if 'error' not in analysis:
# Display pose type
cv2.putText(frame_with_pose, f"Pose: {analysis['pose_type']}",
(10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
# Display angles
y_offset = 60
for joint, angle in analysis['angles'].items():
cv2.putText(frame_with_pose, f"{joint}: {angle:.1f}°",
(10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
y_offset += 30
# Display corrections
for correction in analysis['corrections']:
cv2.putText(frame_with_pose, correction,
(10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
y_offset += 30
else:
cv2.putText(frame_with_pose, analysis['error'],
(10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
# Display the frame
cv2.imshow('MoveNet Pose Analysis', frame_with_pose)
# Break the loop if 'q' is pressed
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
if __name__ == '__main__':
main() |