| import cv2
|
| import csv
|
| import os
|
| import mediapipe as mp
|
| import json
|
| import pandas as pd
|
| import matplotlib.pyplot as plt
|
| import numpy as np
|
| import math
|
| import logging
|
|
|
| os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'
|
| logging.getLogger('mediapipe').setLevel(logging.ERROR)
|
|
|
| class PoseEstimationModel:
|
| def __init__(self):
|
| self.mp_pose = mp.solutions.pose
|
| self.pose_video = self.mp_pose.Pose(smooth_landmarks=True)
|
|
|
| def detect_pose(self, image, pose, display=True):
|
| output_image = image.copy()
|
| imageRGB = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
| results = pose.process(imageRGB)
|
| height, width, _ = image.shape
|
| landmarks = []
|
| if results.pose_landmarks:
|
| for landmark in results.pose_landmarks.landmark:
|
| landmarks.append((int(landmark.x * width), int(landmark.y * height),
|
| (landmark.z * width)))
|
| if display:
|
| plt.figure(figsize=[22,22])
|
| plt.subplot(121);plt.imshow(image[:,:,::-1]);plt.title("Original Image");plt.axis('off');
|
| plt.subplot(122);plt.imshow(output_image[:,:,::-1]);plt.title("Output Image");plt.axis('off');
|
| plt.show()
|
| else:
|
| return output_image, landmarks
|
|
|
| def calculate_angle(self, landmark1, landmark2, landmark3):
|
| x1, y1, _ = landmark1
|
| x2, y2, _ = landmark2
|
| x3, y3, _ = landmark3
|
| angle = math.degrees(math.atan2(y3 - y2, x3 - x2) - math.atan2(y1 - y2, x1 - x2))
|
| if angle < 0:
|
| angle += 360
|
| if angle > 180:
|
| angle = 360 - angle
|
|
|
| return round(angle, 2)
|
|
|
| def calculate_distance(self, point1, point2):
|
| length = np.linalg.norm(np.array(point1) - np.array(point2))
|
| return round(length, 2)
|
|
|
| def body_angles(self, landmarks):
|
| left_elbow_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_ELBOW.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_WRIST.value])
|
| right_elbow_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_ELBOW.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_WRIST.value])
|
| left_shoulder_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.LEFT_ELBOW.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value])
|
| right_shoulder_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_ELBOW.value])
|
| left_knee_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_KNEE.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_ANKLE.value])
|
| right_knee_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_KNEE.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_ANKLE.value])
|
|
|
| left_ankle_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.LEFT_KNEE.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_ANKLE.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_FOOT_INDEX.value])
|
|
|
| right_ankle_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.RIGHT_KNEE.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_ANKLE.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_FOOT_INDEX.value])
|
|
|
| right_hip_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_KNEE.value])
|
|
|
| left_hip_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_KNEE.value])
|
|
|
| right_hip_to_hip = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.RIGHT_KNEE.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value])
|
|
|
| left_hip_to_hip = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_KNEE.value])
|
|
|
| left_wrist_pinky_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.LEFT_ELBOW.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_WRIST.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_PINKY.value])
|
|
|
| right_wrist_pinky_angle = self.calculate_angle(landmarks[self.mp_pose.PoseLandmark.RIGHT_ELBOW.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_WRIST.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_PINKY.value])
|
|
|
| left_leg_length = self.calculate_distance(
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_ANKLE.value]
|
| )
|
| right_leg_length = self.calculate_distance(
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_ANKLE.value]
|
| )
|
| shoulder_width = self.calculate_distance(
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value]
|
| )
|
| hip_width = self.calculate_distance(
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value],
|
| landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value]
|
| )
|
| torso_height = self.calculate_distance(
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value],
|
| landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value]
|
| )
|
|
|
| return {
|
| "elbow_angles": [left_elbow_angle, right_elbow_angle],
|
| "shoulder_angles": [left_shoulder_angle, right_shoulder_angle],
|
| "knee_angles": [left_knee_angle, right_knee_angle],
|
| "ankle_angles": [left_ankle_angle, right_ankle_angle],
|
| "hip_angles": [left_hip_angle, right_hip_angle],
|
| "wrist_pinky_angles": [left_wrist_pinky_angle, right_wrist_pinky_angle],
|
| "hip_to_hip": [left_hip_to_hip, right_hip_to_hip],
|
| "leg_lengths": [left_leg_length, right_leg_length],
|
| "body_dimensions": {
|
| "shoulder_width": shoulder_width,
|
| "hip_width": hip_width,
|
| "torso_height": torso_height
|
| }
|
| }
|
|
|
| def process_video(self, video_path: str, log_file: str):
|
| provided_video = cv2.VideoCapture(video_path)
|
| if not provided_video.isOpened():
|
| raise Exception("Cannot open video file.")
|
|
|
| if os.path.exists(log_file):
|
| os.remove(log_file)
|
|
|
| while provided_video.isOpened():
|
| ok, frame = provided_video.read()
|
| if not ok:
|
| break
|
| frame, landmarks = self.detect_pose(frame, self.pose_video, display=False)
|
| if landmarks:
|
| body_points = self.body_angles(landmarks)
|
| self.log_landmarks(body_points, log_file)
|
|
|
| provided_video.release()
|
|
|
| def log_landmarks(self, body_points: dict, log_file: str):
|
| log_entry = {
|
| "left_elbow_angles": body_points["elbow_angles"][0],
|
| "right_elbow_angles": body_points["elbow_angles"][1],
|
| "left_shoulder_angles": body_points["shoulder_angles"][0],
|
| "right_shoulder_angles": body_points["shoulder_angles"][1],
|
| "left_knee_angles": body_points["knee_angles"][0],
|
| "right_knee_angles": body_points["knee_angles"][1],
|
| "left_ankle_angles": body_points["ankle_angles"][0],
|
| "right_ankle_angles": body_points["ankle_angles"][1],
|
| "left_hip_angles": body_points["hip_angles"][0],
|
| "right_hip_angles": body_points["hip_angles"][1],
|
| "left_wrist_pinky_angle": body_points["wrist_pinky_angles"][0],
|
| "right_wrist_pinky_angle": body_points["wrist_pinky_angles"][1],
|
| "left_hip_to_hip": body_points["hip_to_hip"][0],
|
| "right_hip_to_hip": body_points["hip_to_hip"][1]
|
| }
|
|
|
| with open(log_file, "a", newline='') as file:
|
| writer = csv.writer(file)
|
| if file.tell() == 0:
|
| writer.writerow(log_entry.keys())
|
| writer.writerow(log_entry.values())
|
|
|
| def process_video_and_scale(self, video_path: str, csv_path: str, json_path: str):
|
| self.process_video(video_path, csv_path)
|
| self.std_scaler(csv_path, json_path)
|
|
|
| return {"csv_path": csv_path, "json_path": json_path}
|
|
|
| def predict(self, instances):
|
| if not instances or "video_path" not in instances[0]:
|
| raise ValueError("Invalid input format. Expected a dictionary with 'video_path', 'csv_path', and 'json_path' keys.")
|
|
|
| video_path = instances[0].get("video_path")
|
| csv_path = instances[0].get("csv_path")
|
| json_path = instances[0].get("json_path")
|
|
|
| print(f"Processing video: {video_path}, saving to {csv_path}, normalizing data in {json_path}")
|
|
|
| result = self.process_video_and_scale(video_path, csv_path, json_path)
|
| return {"predictions": result}
|
|
|
| def save_model_config(self, filename):
|
| config = {"smooth_landmarks": True}
|
| with open(filename, "w") as file:
|
| json.dump(config, file)
|
| print(f"Model configuration saved to {filename}")
|
|
|
| @staticmethod
|
| def load_model_config(filename):
|
| with open(filename, "r") as file:
|
| config = json.load(file)
|
| model = PoseEstimationModel()
|
| print(f"Model configuration loaded from {filename}")
|
| return model
|
|
|
| if __name__ == "__main__":
|
| model = PoseEstimationModel()
|
| model.save_model_config('model_config.json')
|
| loaded_model = PoseEstimationModel.load_model_config('model_config.json')
|
| video_path = 'input_video.mp4'
|
| log_file = 'output_log.csv'
|
| loaded_model.process_video(video_path, log_file)
|
| print("Video processed successfully! Pose data saved in:", log_file) |
