Create try_1.py

try_1.py

@@ -0,0 +1,86 @@
+import cv2
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig
+
+# Load the ronka/postureDetection-Mediapipe model
+import mediapipe as mp
+mp_pose = mp.solutions.pose
+pose = mp_pose.Pose()
+lmPose = mp_pose.PoseLandmark
+
+class PostureDetectionConfig(PretrainedConfig):
+    model_type = "posture_detection"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+class PostureDetectionModel(PreTrainedModel):
+    config_class = PostureDetectionConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.pose = pose
+
+    def forward(self, pixel_values):
+        images = [cv2.cvtColor(pixel_value.numpy(), cv2.COLOR_RGB2BGR) for pixel_value in pixel_values]
+        keypoints = [self.pose.process(image) for image in images]
+
+        results = []
+        for image, keypoint in zip(images, keypoints):
+            lm = keypoint.pose_landmarks
+            if lm is None:
+                results.append(None)
+                continue
+
+            # Acquire the landmark coordinates.
+            # Left shoulder.
+            l_shldr_x = int(lm.landmark[lmPose.LEFT_SHOULDER].x * image.shape[1])
+            l_shldr_y = int(lm.landmark[lmPose.LEFT_SHOULDER].y * image.shape[0])
+
+            # Right shoulder
+            r_shldr_x = int(lm.landmark[lmPose.RIGHT_SHOULDER].x * image.shape[1])
+            r_shldr_y = int(lm.landmark[lmPose.RIGHT_SHOULDER].y * image.shape[0])
+
+            # Left ear.
+            l_ear_x = int(lm.landmark[lmPose.LEFT_EAR].x * image.shape[1])
+            l_ear_y = int(lm.landmark[lmPose.LEFT_EAR].y * image.shape[0])
+
+            # Left hip.
+            l_hip_x = int(lm.landmark[lmPose.LEFT_HIP].x * image.shape[1])
+            l_hip_y = int(lm.landmark[lmPose.LEFT_HIP].y * image.shape[0])
+
+            # Calculate distance between left shoulder and right shoulder points.
+            offset = self.findDistance(l_shldr_x, l_shldr_y, r_shldr_x, r_shldr_y)
+
+            # Calculate angles.
+            neck_inclination = self.findAngle(l_shldr_x, l_shldr_y, l_ear_x, l_ear_y)
+            torso_inclination = self.findAngle(l_hip_x, l_hip_y, l_shldr_x, l_shldr_y)
+
+            output = {
+                "offset": offset,
+                "neck_inclination": neck_inclination,
+                "torso_inclination": torso_inclination
+            }
+
+            results.append(output)
+
+        return results
+
+    def findDistance(self, x1, y1, x2, y2):
+        dist = torch.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+        return dist
+
+    def findAngle(self, x1, y1, x2, y2):
+        theta = torch.acos((y2 - y1) * (-y1) / (torch.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) * y1))
+        degree = (180 / torch.pi) * theta
+        return degree
+
+# Register the model and configuration
+MODEL_MAPPING = {
+    PostureDetectionConfig.model_type: PostureDetectionModel
+}
+
+CONFIG_MAPPING = {
+    PostureDetectionConfig.model_type: PostureDetectionConfig
+}