HIMANEESH123/posture_correction

1"""2Pose Generator Module3Extracts and processes pose data from images to generate reference poses4"""5 6import os7import json8import logging9import numpy as np10import cv211import mediapipe as mp12from datetime import datetime13from typing import Dict, List, Optional, Tuple14import uuid15 16logger = logging.getLogger(__name__)17 18mp_pose = mp.solutions.pose19mp_drawing = mp.solutions.drawing_utils20 21class PoseGenerator:22    """Generate pose references from uploaded images."""23    24    def __init__(self):25        self.pose = mp_pose.Pose(26            static_image_mode=True,27            model_complexity=0,28            min_detection_confidence=0.5,29            enable_segmentation=False,30        )31        self.joint_pairs = [32            (11, 13, 15), (12, 14, 16), (11, 23, 25), (12, 24, 26),33            (23, 25, 27), (24, 26, 28), (11, 23, 24), (12, 24, 23),34            (11, 0, 12), (13, 11, 23), (14, 12, 24), (15, 13, 11),35            (16, 14, 12), (25, 23, 24), (26, 24, 23),36        ]37        self.default_weights = [1.0, 1.0, 1.2, 1.2, 1.0, 1.0, 0.8, 0.8, 38                                0.8, 1.0, 1.0, 0.7, 0.7, 0.8, 0.8]39    40    def extract_angles(self, landmarks: np.ndarray) -> Optional[np.ndarray]:41        def calculate_angle(a, b, c):42            a = np.array(a)43            b = np.array(b)44            c = np.array(c)45            radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])46            angle = np.abs(radians * 180.0 / np.pi)47            if angle > 180:48                angle = 360 - angle49            return angle50        51        angles = []52        try:53            for a, b, c in self.joint_pairs:54                if a >= len(landmarks) or b >= len(landmarks) or c >= len(landmarks):55                    return None56                angles.append(calculate_angle(landmarks[a], landmarks[b], landmarks[c]))57        except (IndexError, ValueError):58            return None59        return np.array(angles)60    61    def extract_landmarks(self, image_path: str) -> Optional[Dict]:62        logger.info(f"🔍 Extracting landmarks from {image_path}")63        try:64            img = cv2.imread(image_path)65            if img is None:66                logger.error(f"❌ Could not read image: {image_path}")67                return None68            69            rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)70            results = self.pose.process(rgb)71            72            if not results.pose_landmarks:73                logger.warning(f"⚠️ No pose detected in {image_path}")74                return None75            76            landmarks = []77            for lm in results.pose_landmarks.landmark:78                landmarks.append([lm.x, lm.y])79            80            landmarks_np = np.array(landmarks)81            logger.info(f"✅ Detected {len(landmarks_np)} landmarks")82            83            angles = self.extract_angles(landmarks_np)84            if angles is None:85                logger.warning(f"⚠️ Failed to extract angles from {image_path}")86                return None87            88            logger.info(f"✅ Extracted {len(angles)} angles")89            return {90                'landmarks': landmarks_np.tolist(),91                'angles': angles.tolist(),92                'detection_confidence': results.pose_landmarks.landmark[0].visibility93            }94            95        except Exception as e:96            logger.error(f"❌ Error processing image {image_path}: {e}")97            return None98    99    def calculate_tolerances(self, angle_lists: List[List[float]]) -> List[float]:100        angle_array = np.array(angle_lists)101        std_devs = np.std(angle_array, axis=0)102        tolerances = []103        for std in std_devs:104            tol = max(5, min(25, std * 2))105            tolerances.append(float(tol))106        return tolerances107    108    def compare_landmark_similarity(self, landmarks_list: List[List[List[float]]]) -> float:109        if len(landmarks_list) < 2:110            return 1.0111        112        arrays = [np.array(lm) for lm in landmarks_list]113        similarities = []114        for i in range(len(arrays)):115            for j in range(i + 1, len(arrays)):116                if arrays[i].shape != arrays[j].shape:117                    continue118                diff = np.mean((arrays[i] - arrays[j]) ** 2)119                sim = max(0, 1 - (diff * 5))120                similarities.append(sim)121        122        return float(np.mean(similarities)) if similarities else 0.0123    124    def generate_pose_reference(125        self,126        image_paths: List[str],127        pose_name: str,128        difficulty: str = "Beginner",129        category: str = "General"130    ) -> Optional[Dict]:131        logger.info(f"🔄 Generating reference for '{pose_name}' with {len(image_paths)} images")132        if len(image_paths) != 5:133            logger.error(f"Expected 5 images, got {len(image_paths)}")134            return None135        136        extracted_data = []137        landmarks_list = []138        angle_lists = []139        140        for i, path in enumerate(image_paths):141            data = self.extract_landmarks(path)142            if data is None:143                logger.error(f"Failed to extract landmarks from {path}")144                return {'error': f'Failed to extract pose from image {i+1}'}145            146            extracted_data.append(data)147            landmarks_list.append(data['landmarks'])148            angle_lists.append(data['angles'])149        150        similarity = self.compare_landmark_similarity(landmarks_list)151        if similarity < 0.5:152            logger.warning(f"Pose similarity too low: {similarity:.2f}")153            return {154                'error': f'The uploaded images show different poses (similarity: {similarity:.2f}). Please upload 5 similar images of the same pose.',155                'similarity': similarity156            }157        158        angle_array = np.array(angle_lists)159        avg_angles = np.mean(angle_array, axis=0).tolist()160        tolerances = self.calculate_tolerances(angle_lists)161        weights = self.default_weights.copy()162        163        if len(avg_angles) != 15:164            logger.error(f"Expected 15 angles, got {len(avg_angles)}")165            return None166        167        pose_id = pose_name.lower().replace(' ', '-').replace("'", "").replace('(', '').replace(')', '')168        logger.info(f"📝 Average angles: {len(avg_angles)} angles, tolerances: {len(tolerances)}")169        170        reference = {171            'name': pose_name,172            'difficulty': difficulty,173            'category': category,174            'angles': avg_angles,175            'tolerances': tolerances,176            'weights': weights,177            'sample_count': len(image_paths),178            'generated_at': datetime.now().isoformat(),179            'similarity_score': similarity,180            'landmark_samples': landmarks_list,181        }182        183        return {184            'pose_id': pose_id,185            'reference': reference,186            'similarity': similarity,187            'avg_angles': avg_angles,188            'tolerances': tolerances,189            'sample_count': len(image_paths)190        }191 192    def cleanup(self):193        if hasattr(self, 'pose'):194            self.pose.close()195