yashrajsinha/postSURE
0
1import torch2import torch.nn as nn3import numpy as np4 5 6class PostureClassifier(nn.Module):7 """Simple feedforward network for posture classification (legacy)"""8 def __init__(self, input_size=33*2): # 33 landmarks x 2 coordinates (x, y)9 super(PostureClassifier, self).__init__()10 self.network = nn.Sequential(11 nn.Linear(input_size, 128),12 nn.ReLU(),13 nn.Dropout(0.3),14 nn.Linear(128, 64),15 nn.ReLU(),16 nn.Dropout(0.3),17 nn.Linear(64, 32),18 nn.ReLU(),19 nn.Linear(32, 2) # Good vs Bad posture20 )21 22 def forward(self, x):23 return self.network(x)24 25 26class ResidualBlock(nn.Module):27 """Residual block with batch normalization for stable training"""28 def __init__(self, in_features, out_features, dropout=0.3):29 super(ResidualBlock, self).__init__()30 self.block = nn.Sequential(31 nn.Linear(in_features, out_features),32 nn.BatchNorm1d(out_features),33 nn.ReLU(),34 nn.Dropout(dropout),35 nn.Linear(out_features, out_features),36 nn.BatchNorm1d(out_features),37 )38 # Projection shortcut if dimensions don't match39 self.shortcut = nn.Identity() if in_features == out_features else nn.Linear(in_features, out_features)40 self.relu = nn.ReLU()41 42 def forward(self, x):43 residual = self.shortcut(x)44 out = self.block(x)45 out = out + residual # Skip connection46 return self.relu(out)47 48 49class ImprovedPostureClassifier(nn.Module):50 """51 Enhanced posture classifier with:52 - Batch normalization for training stability53 - Residual connections for better gradient flow54 - Configurable architecture55 - Support for normalized landmark features56 """57 def __init__(self, input_size=66, num_classes=2, hidden_dims=[128, 64, 32], dropout=0.3):58 super(ImprovedPostureClassifier, self).__init__()59 60 self.input_size = input_size61 self.num_classes = num_classes62 63 # Input normalization layer64 self.input_norm = nn.BatchNorm1d(input_size)65 66 # Build residual blocks67 layers = []68 in_dim = input_size69 for hidden_dim in hidden_dims:70 layers.append(ResidualBlock(in_dim, hidden_dim, dropout))71 in_dim = hidden_dim72 73 self.feature_extractor = nn.Sequential(*layers)74 75 # Classification head76 self.classifier = nn.Sequential(77 nn.Linear(hidden_dims[-1], 16),78 nn.ReLU(),79 nn.Dropout(dropout / 2),80 nn.Linear(16, num_classes)81 )82 83 def forward(self, x):84 # Normalize input85 x = self.input_norm(x)86 # Extract features through residual blocks87 features = self.feature_extractor(x)88 # Classify89 return self.classifier(features)90 91 def get_confidence(self, x):92 """Get prediction with confidence score"""93 with torch.no_grad():94 logits = self.forward(x)95 probs = torch.softmax(logits, dim=1)96 confidence, prediction = torch.max(probs, dim=1)97 return prediction, confidence, probs98 99 100class LandmarkNormalizer:101 """102 Normalize MediaPipe landmarks to be position and scale invariant.103 104 This makes the model robust to:105 - Different positions in the frame (sitting left vs right)106 - Different distances from camera (close vs far)107 - Different body sizes108 """109 110 # Key landmark indices from MediaPipe Pose111 LEFT_SHOULDER = 11112 RIGHT_SHOULDER = 12113 LEFT_HIP = 23114 RIGHT_HIP = 24115 NOSE = 0116 117 @staticmethod118 def normalize(landmarks: np.ndarray) -> np.ndarray:119 """120 Normalize 66-element landmark array (33 landmarks x 2 coords).121 122 Args:123 landmarks: Raw [x0, y0, x1, y1, ..., x32, y32] from MediaPipe124 125 Returns:126 Normalized landmarks centered at hip midpoint, scaled by shoulder width127 """128 if len(landmarks) != 66:129 raise ValueError(f"Expected 66 landmarks, got {len(landmarks)}")130 131 # Reshape to (33, 2) for easier manipulation132 points = landmarks.reshape(33, 2)133 134 # Calculate hip center (origin point)135 left_hip = points[LandmarkNormalizer.LEFT_HIP]136 right_hip = points[LandmarkNormalizer.RIGHT_HIP]137 hip_center = (left_hip + right_hip) / 2138 139 # Calculate shoulder width for scale normalization140 left_shoulder = points[LandmarkNormalizer.LEFT_SHOULDER]141 right_shoulder = points[LandmarkNormalizer.RIGHT_SHOULDER]142 shoulder_width = np.linalg.norm(left_shoulder - right_shoulder)143 144 # Avoid division by zero145 if shoulder_width < 0.01:146 shoulder_width = 0.01147 148 # Center and scale149 normalized = (points - hip_center) / shoulder_width150 151 return normalized.flatten().astype(np.float32)152 153 @staticmethod154 def compute_posture_angles(landmarks: np.ndarray) -> np.ndarray:155 """156 Compute meaningful angles that indicate posture quality.157 158 Returns angles for:159 - Head tilt (nose relative to shoulder midpoint)160 - Shoulder alignment (left vs right shoulder height)161 - Spine angle (shoulder center to hip center)162 """163 points = landmarks.reshape(33, 2)164 165 # Get key points166 nose = points[LandmarkNormalizer.NOSE]167 left_shoulder = points[LandmarkNormalizer.LEFT_SHOULDER]168 right_shoulder = points[LandmarkNormalizer.RIGHT_SHOULDER]169 left_hip = points[LandmarkNormalizer.LEFT_HIP]170 right_hip = points[LandmarkNormalizer.RIGHT_HIP]171 172 shoulder_center = (left_shoulder + right_shoulder) / 2173 hip_center = (left_hip + right_hip) / 2174 175 # Head forward tilt angle (how far nose is forward from shoulder line)176 head_forward = nose[1] - shoulder_center[1] # y difference177 178 # Shoulder level difference (should be ~0 for good posture)179 shoulder_tilt = left_shoulder[1] - right_shoulder[1]180 181 # Spine angle (vertical alignment from hips to shoulders)182 spine_vector = shoulder_center - hip_center183 spine_angle = np.arctan2(spine_vector[0], spine_vector[1]) # Angle from vertical184 185 # Neck angle (nose to shoulder center vs vertical)186 neck_vector = nose - shoulder_center187 neck_angle = np.arctan2(neck_vector[0], neck_vector[1])188 189 return np.array([head_forward, shoulder_tilt, spine_angle, neck_angle], dtype=np.float32)