commit

app.py

@@ -28,7 +28,8 @@ pose = mp_pose.Pose(
     min_detection_confidence=0.5,
     min_tracking_confidence=0.5
 )
-# Function to calculate angles between points
+
+# Function to calculate angles between three points
 def calculate_angle(a, b, c):
     ab = (b[0] - a[0], b[1] - a[1])
     bc = (c[0] - b[0], c[1] - b[1])
@@ -37,12 +38,59 @@ def calculate_angle(a, b, c):
     magnitude_ab = math.sqrt(ab[0]**2 + ab[1]**2)
     magnitude_bc = math.sqrt(bc[0]**2 + bc[1]**2)
     
-    angle_radians = math.acos(dot_product / (magnitude_ab * magnitude_bc))
+    # To avoid division by zero
+    if magnitude_ab * magnitude_bc == 0:
+        return 0.0
+
+    # Clamp the cosine value to the [-1, 1] range to avoid numerical errors
+    cosine_angle = max(min(dot_product / (magnitude_ab * magnitude_bc), 1), -1)
+    angle_radians = math.acos(cosine_angle)
     angle_degrees = math.degrees(angle_radians)
     
     return angle_degrees
 
-# Process image with Mediapipe Pose Estimation
+# Function to calculate a simplified REBA score based on trunk (hip) and neck angles.
+def calculate_reba(trunk_angle, neck_angle):
+    """
+    This is a simplified approach:
+      - For the trunk (approximated by the hip angle), a nearly upright posture (angle >= 160°) is scored as 1,
+        a moderately bent posture (angle between 140° and 160°) is scored as 2, and a severely bent posture (<140°) is scored as 3.
+      - Similarly for the neck angle.
+      - The REBA score is the sum of these scores.
+      - Finally, we define a risk level based on the total score.
+    """
+    # Determine trunk score (using the hip angle)
+    if trunk_angle >= 160:
+        trunk_score = 1
+    elif trunk_angle >= 140:
+        trunk_score = 2
+    else:
+        trunk_score = 3
+
+    # Determine neck score
+    if neck_angle >= 150:
+        neck_score = 1
+    elif neck_angle >= 130:
+        neck_score = 2
+    else:
+        neck_score = 3
+
+    # Simplified REBA group A score (normally REBA also considers legs, arms, load, etc.)
+    reba_score = trunk_score + neck_score
+
+    # Define risk levels based on the score
+    if reba_score <= 2:
+        risk = "Negligible"
+    elif reba_score <= 4:
+        risk = "Low"
+    elif reba_score <= 6:
+        risk = "Medium"
+    else:
+        risk = "High"
+
+    return reba_score, risk
+
+# Process image with Mediapipe Pose Estimation and analyze posture using REBA score
 def process_frame(image):
     h, w, _ = image.shape
     
@@ -56,59 +104,60 @@ def process_frame(image):
     image = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
     
     if results.pose_landmarks:
-        # Get landmarks
+        # Get key landmarks from the right side
         right_shoulder = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER]
         right_hip = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_HIP]
         right_ear = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_EAR]
         
-        # Convert to pixel coordinates
+        # Convert normalized coordinates to pixel coordinates
         cx_rs, cy_rs = int(right_shoulder.x * w), int(right_shoulder.y * h)
         cx_rh, cy_rh = int(right_hip.x * w), int(right_hip.y * h)
         cx_re, cy_re = int(right_ear.x * w), int(right_ear.y * h)
         
-        # Create upper reference points
+        # Create reference points by applying an offset (helps approximate vertical)
         offset = 60
         upper_shoulder = (cx_rs, max(0, cy_rs - offset))
         upper_hip = (cx_rh, max(0, cy_rh - offset))
         
-        # Draw landmarks
+        # Draw reference landmarks on the image
         cv2.circle(image, upper_shoulder, 5, (0, 255, 0), -1)
         cv2.circle(image, upper_hip, 5, (0, 255, 0), -1)
         
-        # Draw lines
+        # Draw lines connecting key points
         cv2.line(image, (cx_rh, cy_rh), (cx_rs, cy_rs), (255, 0, 255), 2)  # Hip to shoulder
         cv2.line(image, (cx_rs, cy_rs), (cx_re, cy_re), (255, 255, 0), 2)  # Shoulder to ear
         cv2.line(image, (cx_rh, cy_rh), upper_hip, (0, 165, 255), 2)       # Hip to upper hip
         cv2.line(image, (cx_rs, cy_rs), upper_shoulder, (0, 255, 255), 2)  # Shoulder to upper shoulder
         
-        # Calculate angles
+        # Calculate angles using the defined reference points
         angle_hip = calculate_angle(upper_hip, (cx_rh, cy_rh), (cx_rs, cy_rs))
         angle_neck = calculate_angle((cx_rs, cy_rs), (cx_re, cy_re), upper_shoulder)
         
-        # Determine posture status
-        hip_posture = "Good" if 160 <= angle_hip <= 180 else "Poor"
-        neck_posture = "Good" if 150 <= angle_neck <= 180 else "Poor"
-        hip_color = (0, 255, 0) if hip_posture == "Good" else (0, 0, 255)
-        neck_color = (0, 255, 0) if neck_posture == "Good" else (0, 0, 255)
+        # Compute the simplified REBA score and corresponding risk level
+        reba_score, risk = calculate_reba(angle_hip, angle_neck)
+        
+        # Display the calculated angles on the image
+        cv2.putText(image, f"Hip Angle: {angle_hip:.1f}", (10, 60), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 255), 2)
+        cv2.putText(image, f"Neck Angle: {angle_neck:.1f}", (10, 90), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
         
-        # Display angles
-        cv2.putText(image, f"Hip Angle: {angle_hip:.1f} ({hip_posture})", (10, 60), 
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, hip_color, 2)
-        cv2.putText(image, f"Neck Angle: {angle_neck:.1f} ({neck_posture})", (10, 90), 
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, neck_color, 2)
+        # Display the simplified REBA score and risk level on the image
+        cv2.putText(image, f"REBA Score: {reba_score} ({risk})", (10, 120), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
 
     return image
 
-# API Route to receive an image and return processed image
+# API Route to receive an image and return the processed image with REBA analysis
 @app.post("/upload")
 async def upload_image(file: UploadFile = File(...)):
     contents = await file.read()
     image = Image.open(BytesIO(contents))
     image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
     
-    # Process the image in async context
+    # Process the image (the processing function now includes REBA score analysis)
     processed_image = process_frame(image)
     
-    # Encode processed image to return
+    # Encode the processed image to return it as JPEG
     _, buffer = cv2.imencode(".jpg", processed_image)
-    return Response(content=buffer.tobytes(), media_type="image/jpeg")
+    return Response(content=buffer.tobytes(), media_type="image/jpeg")