Prajith04/ergonomics

1import cv22import math3import base644import numpy as np5import mediapipe as mp6from io import BytesIO7from fastapi import FastAPI, File, UploadFile8from fastapi.responses import Response9from fastapi.middleware.cors import CORSMiddleware  # Add CORS support10from PIL import Image11 12# Initialize FastAPI app13app = FastAPI()14 15# Add CORS middleware16app.add_middleware(17    CORSMiddleware,18    allow_origins=["*"],19    allow_credentials=True,20    allow_methods=["*"],21    allow_headers=["*"],22)23 24# Initialize Mediapipe Pose model25mp_pose = mp.solutions.pose26pose = mp_pose.Pose(27    static_image_mode=False,28    min_detection_confidence=0.5,29    min_tracking_confidence=0.530)31 32# Function to calculate angles between three points33def calculate_angle(a, b, c):34    ab = (b[0] - a[0], b[1] - a[1])35    bc = (c[0] - b[0], c[1] - b[1])36    37    dot_product = ab[0] * bc[0] + ab[1] * bc[1]38    magnitude_ab = math.sqrt(ab[0]**2 + ab[1]**2)39    magnitude_bc = math.sqrt(bc[0]**2 + bc[1]**2)40    41    # To avoid division by zero42    if magnitude_ab * magnitude_bc == 0:43        return 0.044 45    # Clamp the cosine value to the [-1, 1] range to avoid numerical errors46    cosine_angle = max(min(dot_product / (magnitude_ab * magnitude_bc), 1), -1)47    angle_radians = math.acos(cosine_angle)48    angle_degrees = math.degrees(angle_radians)49    50    return angle_degrees51 52# Function to calculate a simplified REBA score based on trunk (hip) and neck angles.53def calculate_reba(trunk_angle, neck_angle):54    """55    This is a simplified approach:56      - For the trunk (approximated by the hip angle), a nearly upright posture (angle >= 160°) is scored as 1,57        a moderately bent posture (angle between 140° and 160°) is scored as 2, and a severely bent posture (<140°) is scored as 3.58      - Similarly for the neck angle.59      - The REBA score is the sum of these scores.60      - Finally, we define a risk level based on the total score.61    """62    # Determine trunk score (using the hip angle)63    if trunk_angle >= 160:64        trunk_score = 165    elif trunk_angle >= 140:66        trunk_score = 267    else:68        trunk_score = 369 70    # Determine neck score71    if neck_angle >= 150:72        neck_score = 173    elif neck_angle >= 130:74        neck_score = 275    else:76        neck_score = 377 78    # Simplified REBA group A score (normally REBA also considers legs, arms, load, etc.)79    reba_score = trunk_score + neck_score80 81    # Define risk levels based on the score82    if reba_score <= 2:83        risk = "Negligible"84    elif reba_score <= 4:85        risk = "Low"86    elif reba_score <= 6:87        risk = "Medium"88    else:89        risk = "High"90 91    return reba_score, risk92 93# Process image with Mediapipe Pose Estimation and analyze posture using REBA score94def process_frame(image):95    h, w, _ = image.shape96    97    # Convert to RGB98    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)99    image_rgb.flags.writeable = False100    results = pose.process(image_rgb)101    image_rgb.flags.writeable = True102    103    # Convert back to BGR for display104    image = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)105    106    if results.pose_landmarks:107        # Get key landmarks from the right side108        right_shoulder = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER]109        right_hip = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_HIP]110        right_ear = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_EAR]111        112        # Convert normalized coordinates to pixel coordinates113        cx_rs, cy_rs = int(right_shoulder.x * w), int(right_shoulder.y * h)114        cx_rh, cy_rh = int(right_hip.x * w), int(right_hip.y * h)115        cx_re, cy_re = int(right_ear.x * w), int(right_ear.y * h)116        117        # Create reference points by applying an offset (helps approximate vertical)118        offset = 60119        upper_shoulder = (cx_rs, max(0, cy_rs - offset))120        upper_hip = (cx_rh, max(0, cy_rh - offset))121        122        # Draw reference landmarks on the image123        cv2.circle(image, upper_shoulder, 5, (0, 255, 0), -1)124        cv2.circle(image, upper_hip, 5, (0, 255, 0), -1)125        126        # Draw lines connecting key points127        cv2.line(image, (cx_rh, cy_rh), (cx_rs, cy_rs), (255, 0, 255), 2)  # Hip to shoulder128        cv2.line(image, (cx_rs, cy_rs), (cx_re, cy_re), (255, 255, 0), 2)  # Shoulder to ear129        cv2.line(image, (cx_rh, cy_rh), upper_hip, (0, 165, 255), 2)       # Hip to upper hip130        cv2.line(image, (cx_rs, cy_rs), upper_shoulder, (0, 255, 255), 2)  # Shoulder to upper shoulder131        132        # Calculate angles using the defined reference points133        angle_hip = calculate_angle(upper_hip, (cx_rh, cy_rh), (cx_rs, cy_rs))134        angle_neck = calculate_angle((cx_rs, cy_rs), (cx_re, cy_re), upper_shoulder)135        136        # Compute the simplified REBA score and corresponding risk level137        reba_score, risk = calculate_reba(angle_hip, angle_neck)138        139        # Display the calculated angles on the image140        cv2.putText(image, f"Hip Angle: {angle_hip:.1f}", (10, 60), 141                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 255), 2)142        cv2.putText(image, f"Neck Angle: {angle_neck:.1f}", (10, 90), 143                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)144        145        # Display the simplified REBA score and risk level on the image146        cv2.putText(image, f"REBA Score: {reba_score} ({risk})", (10, 120), 147                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)148 149    return image150 151# API Route to receive an image and return the processed image with REBA analysis152@app.post("/upload")153async def upload_image(file: UploadFile = File(...)):154    contents = await file.read()155    image = Image.open(BytesIO(contents))156    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)157    158    # Process the image (the processing function now includes REBA score analysis)159    processed_image = process_frame(image)160    161    # Encode the processed image to return it as JPEG162    _, buffer = cv2.imencode(".jpg", processed_image)163    return Response(content=buffer.tobytes(), media_type="image/jpeg")164