Create try_1.py

9777f9d verified over 2 years ago

2.95 kB

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
}

	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
	}