| import cv2 |
| import mediapipe as mp |
| import csv |
| import os |
|
|
| mp_drawing = mp.solutions.drawing_utils |
| mp_pose = mp.solutions.pose |
|
|
| cap = cv2.VideoCapture(0) |
| count = 0 |
|
|
| with mp_pose.Pose( |
| min_detection_confidence=0.5, |
| min_tracking_confidence=0.5) as pose: |
| fields = ['frame', 'nose_x', 'nose_y', 'left_shoulder_x', 'left_shoulder_y', 'right_shoulder_x', 'right_shoulder_y'] |
| filename = 'pose_data.csv' |
| with open(filename, 'w') as csvfile: |
| csvwriter = csv.writer(csvfile) |
| csvwriter.writerow(fields) |
| while cap.isOpened(): |
| ret, frame = cap.read() |
| if not ret: |
| break |
| image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| image.flags.writeable = False |
| results = pose.process(image) |
| image.flags.writeable = True |
| image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) |
|
|
| if results.pose_landmarks is not None: |
| mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS) |
|
|
| cv2.imshow('Pose Estimation', image) |
| for result.pose.landmarks in fields: |
| print(results.pose.landmarks.landmark.x,results.pose.landmarks.landmark.y ) |
|
|
| if (count % 10 == 0): |
| num_coords = len(results.pose_landmarks.landmark) |
| print(results.pose_landmarks.landmark) |
| row = [count] + [ |
| results.pose_landmarks.landmark[i].x if results.pose_landmarks.landmark[i] is not None else float('nan') |
| for i in range(num_coords)] + \ |
| [results.pose_landmarks.landmark[i].y if results.pose_landmarks.landmark[i] is not None else float( |
| 'nan') for i in range(num_coords)] |
| with open(filename, 'a') as csvfile: |
| csvwriter = csv.writer(csvfile) |
| csvwriter.writerow(row) |
| count += 1 |
|
|
| if cv2.waitKey(1) & 0xFF == ord('q'): |
| break |
|
|
| cap.release() |
| cv2.destroyAllWindows() |
|
|
