yashrajsinha/postSURE
0
1# 1. Captures video from webcam │2# │ 2. Sends each frame to MediaPipe │3# │ 3. Extracts the 33 landmark coordinates │4# │ 4. Feeds landmarks to PyTorch model │5# │ 5. Gets prediction: Good (1) or Bad (0) │6# │ 6. Smooths predictions to avoid jitter │7# │ 7. Displays result (green/red border)8 9 10import cv211import numpy as np12import torch13import torch.nn as nn14import mediapipe as mp15from collections import deque16import json17import os18import random19from classifier import PostureClassifier20 21class PostureMonitor:22 def __init__(self, model_path=None, test_mode=False):23 # Initialize MediaPipe Pose24 self.mp_pose = mp.solutions.pose25 self.pose = self.mp_pose.Pose(26 min_detection_confidence=0.5,27 min_tracking_confidence=0.528 )29 self.mp_drawing = mp.solutions.drawing_utils30 31 # TEST MODE: Skip model loading entirely32 self.test_mode = test_mode33 34 if not test_mode:35 # Initialize PyTorch model36 self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')37 self.model = PostureClassifier().to(self.device)38 39 if model_path and os.path.exists(model_path):40 self.model.load_state_dict(torch.load(model_path, map_location=self.device))41 self.model.eval()42 self.is_calibrated = True43 else:44 self.is_calibrated = False45 else:46 self.is_calibrated = True # Pretend we're calibrated in test mode47 print("🧪 TEST MODE: Using random predictions (no model needed)")48 49 # Smoothing buffer for predictions50 self.prediction_buffer = deque(maxlen=10)51 52 def extract_landmarks(self, frame):53 """Extract pose landmarks from frame"""54 # Convert BGR to RGB55 image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)56 results = self.pose.process(image_rgb)57 58 if results.pose_landmarks:59 # Extract x, y coordinates (ignore z for now)60 landmarks = []61 for landmark in results.pose_landmarks.landmark:62 landmarks.extend([landmark.x, landmark.y])63 return np.array(landmarks, dtype=np.float32), results.pose_landmarks64 65 return None, None66 67 def predict_posture(self, landmarks):68 """Predict posture from landmarks"""69 if landmarks is None:70 return None71 72 # TEST MODE: Generate random predictions73 if self.test_mode:74 print(f"📊 Landmarks shape: {landmarks.shape} (66 values)")75 76 # Generate random prediction (0 or 1)77 prediction = random.choice([0, 1])78 confidence = random.uniform(0.7, 0.95) # Random confidence between 70-95%79 80 print(f"🎲 Random prediction: {prediction} ({'Good' if prediction == 1 else 'Bad'}) with {confidence:.1%} confidence")81 82 # Still smooth the predictions83 self.prediction_buffer.append(prediction)84 smoothed_prediction = round(np.mean(self.prediction_buffer))85 86 return smoothed_prediction, confidence87 88 # NORMAL MODE: Use the actual model89 if not self.is_calibrated:90 return None91 92 # Convert to tensor and predict93 with torch.no_grad():94 landmark_tensor = torch.FloatTensor(landmarks).unsqueeze(0).to(self.device)95 print(f"📊 Tensor shape going into model: {landmark_tensor.shape}")96 97 output = self.model(landmark_tensor)98 print(f"🧠 Model raw output: {output}")99 100 probabilities = torch.softmax(output, dim=1)101 print(f"📈 Probabilities: {probabilities} (Bad={probabilities[0][0]:.1%}, Good={probabilities[0][1]:.1%})")102 103 prediction = torch.argmax(probabilities, dim=1).item()104 confidence = probabilities[0][prediction].item()105 106 print(f"✅ Final prediction: {prediction} ({'Good' if prediction == 1 else 'Bad'}) with {confidence:.1%} confidence")107 108 # Smooth predictions109 self.prediction_buffer.append(prediction)110 smoothed_prediction = round(np.mean(self.prediction_buffer))111 112 return smoothed_prediction, confidence113 114 def run(self):115 """Main monitoring loop"""116 cap = cv2.VideoCapture(0)117 118 if not cap.isOpened():119 print("Error: Could not open camera")120 return121 122 print("\n" + "="*60)123 print("Posture Monitor Started!")124 print("="*60)125 126 if self.test_mode:127 print("🧪 TEST MODE ACTIVE")128 print(" - MediaPipe will detect your pose")129 print(" - Predictions are RANDOM (for demo)")130 print(" - Watch the console to see the data flow!")131 elif not self.is_calibrated:132 print("⚠️ Model not calibrated. Please run calibration first.")133 else:134 print("✓ Model loaded and ready")135 136 print("\nPress 'q' to quit")137 print("="*60 + "\n")138 139 frame_count = 0140 141 while cap.isOpened():142 ret, frame = cap.read()143 if not ret:144 break145 146 frame_count += 1147 148 # Flip frame horizontally for mirror view149 frame = cv2.flip(frame, 1)150 151 # Extract landmarks152 print(f"\n--- Frame {frame_count} ---")153 print("1️⃣ Capturing webcam frame...")154 155 landmarks, pose_landmarks = self.extract_landmarks(frame)156 157 if landmarks is not None:158 print(f"2️⃣ MediaPipe detected pose! Found {len(landmarks)//2} landmarks")159 else:160 print("❌ No pose detected in this frame")161 162 # Default border color163 border_color = (128, 128, 128) # Gray for uncalibrated164 status_text = "Not Calibrated"165 166 if self.is_calibrated and landmarks is not None:167 # Predict posture168 print("3️⃣ Feeding landmarks to model...")169 result = self.predict_posture(landmarks)170 171 if result:172 prediction, confidence = result173 print(f"4️⃣ Smoothed prediction: {prediction} ({'Good' if prediction == 1 else 'Bad'})")174 175 if prediction == 1: # Good posture176 border_color = (0, 255, 0) # Green177 status_text = f"Good Posture ({confidence:.1%})"178 print("5️⃣ Displaying: GREEN BORDER ✓")179 else: # Bad posture180 border_color = (0, 0, 255) # Red181 status_text = f"Bad Posture ({confidence:.1%})"182 print("5️⃣ Displaying: RED BORDER ✗")183 184 # Draw pose landmarks185 if pose_landmarks:186 self.mp_drawing.draw_landmarks(187 frame, 188 pose_landmarks, 189 self.mp_pose.POSE_CONNECTIONS190 )191 192 # Draw border193 border_thickness = 15194 cv2.rectangle(195 frame, 196 (0, 0), 197 (frame.shape[1], frame.shape[0]), 198 border_color, 199 border_thickness200 )201 202 # Draw status text203 cv2.putText(204 frame, 205 status_text, 206 (20, 50), 207 cv2.FONT_HERSHEY_SIMPLEX, 208 1.2, 209 border_color, 210 2211 )212 213 # Add test mode indicator214 if self.test_mode:215 cv2.putText(216 frame,217 "TEST MODE - Random Predictions",218 (20, frame.shape[0] - 20),219 cv2.FONT_HERSHEY_SIMPLEX,220 0.7,221 (255, 255, 0),222 2223 )224 225 cv2.imshow('Posture Monitor', frame)226 227 if cv2.waitKey(1) & 0xFF == ord('q'):228 break229 230 cap.release()231 cv2.destroyAllWindows()232 233 234 235if __name__ == "__main__":236 import sys237 238 # Check if user wants test mode239 if len(sys.argv) > 1 and sys.argv[1] == '--test':240 print("\n🧪 Starting in TEST MODE")241 print("This will show you how the system works WITHOUT training a model\n")242 monitor = PostureMonitor(test_mode=True)243 else:244 print("\n💡 TIP: Run with --test flag to test without a trained model:")245 print(" python monitor.py --test\n")246 monitor = PostureMonitor(model_path='models/posture_model.pth')247 248 monitor.run()