app.py

import cv2
import numpy as np
import mediapipe as mp
import tempfile
import os
import json
import re
from typing import List, Tuple, Dict
import math
from pathlib import Path
import warnings
from fastapi import FastAPI, UploadFile, File
from pydantic import BaseModel
import time
import google.generativeai as genai

# Suppress MediaPipe GPU warnings
warnings.filterwarnings("ignore", category=UserWarning)
os.environ['MEDIAPIPE_DISABLE_GPU'] = '1'
os.environ['MESA_GL_VERSION_OVERRIDE'] = '3.3'
os.environ['LIBGL_ALWAYS_SOFTWARE'] = '1'

# Configure the Gemini API
genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))

# Initialize the FastAPI app
app = FastAPI(
    title="Squat Form Analyzer API",
    description="API for analyzing squat form from video files."
)


class SquatFormAnalyzer:
    def __init__(self):
        self.mp_pose = mp.solutions.pose
        self.pose = self.mp_pose.Pose(
            static_image_mode=False,
            model_complexity=1,
            enable_segmentation=False,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5
        )
        self.mp_drawing = mp.solutions.drawing_utils

    def calculate_angle(self, a, b, c):
        """Calculate angle between three points"""
        a = np.array(a)
        b = np.array(b)
        c = np.array(c)
        radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
        angle = np.abs(radians * 180.0 / np.pi)
        if angle > 180.0:
            angle = 360 - angle
        return angle

    def analyze_knee_alignment(self, landmarks, frame_width, frame_height):
        """Check for knees caving inward and return angles"""
        try:
            left_hip = [landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value].x * frame_width,
                        landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value].y * frame_height]
            right_hip = [landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value].x * frame_width,
                         landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value].y * frame_height]
            left_knee = [landmarks[self.mp_pose.PoseLandmark.LEFT_KNEE.value].x * frame_width,
                         landmarks[self.mp_pose.PoseLandmark.LEFT_KNEE.value].y * frame_height]
            right_knee = [landmarks[self.mp_pose.PoseLandmark.RIGHT_KNEE.value].x * frame_width,
                          landmarks[self.mp_pose.PoseLandmark.RIGHT_KNEE.value].y * frame_height]
            left_ankle = [landmarks[self.mp_pose.PoseLandmark.LEFT_ANKLE.value].x * frame_width,
                          landmarks[self.mp_pose.PoseLandmark.LEFT_ANKLE.value].y * frame_height]
            right_ankle = [landmarks[self.mp_pose.PoseLandmark.RIGHT_ANKLE.value].x * frame_width,
                           landmarks[self.mp_pose.PoseLandmark.RIGHT_ANKLE.value].y * frame_height]
            left_knee_angle = self.calculate_angle(left_hip, left_knee, left_ankle)
            right_knee_angle = self.calculate_angle(right_hip, right_knee, right_ankle)
            hip_width = abs(left_hip[0] - right_hip[0])
            knee_width = abs(left_knee[0] - right_knee[0])
            ankle_width = abs(left_ankle[0] - right_ankle[0])
            knee_ratio = knee_width / hip_width if hip_width > 0 else 1
            ankle_ratio = knee_width / ankle_width if ankle_width > 0 else 1
            knee_valgus = (knee_ratio < 0.75 or ankle_ratio < 0.9 or
                           left_knee_angle < 160 or right_knee_angle < 160)
            return knee_valgus, left_knee_angle, right_knee_angle
        except (AttributeError, IndexError, ZeroDivisionError):
            return False, None, None

    def analyze_forward_lean(self, landmarks):
        """Check for excessive forward lean and return angle"""
        try:
            left_shoulder = [landmarks[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value].x,
                             landmarks[self.mp_pose.PoseLandmark.LEFT_SHOULDER.value].y]
            right_shoulder = [landmarks[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value].x,
                              landmarks[self.mp_pose.PoseLandmark.RIGHT_SHOULDER.value].y]
            left_hip = [landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value].x,
                        landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value].y]
            right_hip = [landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value].x,
                         landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value].y]
            avg_shoulder = [(left_shoulder[0] + right_shoulder[0]) / 2,
                            (left_shoulder[1] + right_shoulder[1]) / 2]
            avg_hip = [(left_hip[0] + right_hip[0]) / 2,
                       (left_hip[1] + right_hip[1]) / 2]
            horizontal_distance = abs(avg_shoulder[0] - avg_hip[0])
            vertical_distance = abs(avg_shoulder[1] - avg_hip[1])
            if vertical_distance > 0:
                lean_angle = math.degrees(math.atan(horizontal_distance / vertical_distance))
                return lean_angle > 15, lean_angle
            return False, None
        except (AttributeError, IndexError, ZeroDivisionError):
            return False, None

    def analyze_depth(self, landmarks):
        """Check if squat depth is adequate"""
        try:
            left_hip = [landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value].x,
                        landmarks[self.mp_pose.PoseLandmark.LEFT_HIP.value].y]
            right_hip = [landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value].x,
                         landmarks[self.mp_pose.PoseLandmark.RIGHT_HIP.value].y]
            left_knee = [landmarks[self.mp_pose.PoseLandmark.LEFT_KNEE.value].x,
                         landmarks[self.mp_pose.PoseLandmark.LEFT_KNEE.value].y]
            right_knee = [landmarks[self.mp_pose.PoseLandmark.RIGHT_KNEE.value].x,
                          landmarks[self.mp_pose.PoseLandmark.RIGHT_KNEE.value].y]
            avg_hip = [(left_hip[0] + right_hip[0]) / 2, (left_hip[1] + right_hip[1]) / 2]
            avg_knee = [(left_knee[0] + right_knee[0]) / 2, (left_knee[1] + right_knee[1]) / 2]
            return avg_hip[1] > avg_knee[1]
        except (AttributeError, IndexError):
            return False

    def process_video_metrics(self, video_path):
        """Process video and analyze squat form, returning metrics"""
        print("Starting video processing...")
        cap = cv2.VideoCapture(video_path)
        frame_count = 0
        analysis_results = {
            'knee_inward': [],
            'forward_lean': [],
            'depth_issues': [],
            'torso_angles': [],
            'left_knee_angles': [],
            'right_knee_angles': [],
            'total_frames': 0
        }
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            frame_count += 1
            if frame_count % 30 == 0:
                print(f"Processing frame {frame_count}...")
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = self.pose.process(rgb_frame)
            if results.pose_landmarks:
                knee_error, left_knee_angle, right_knee_angle = self.analyze_knee_alignment(
                    results.pose_landmarks.landmark, frame.shape[1], frame.shape[0])
                lean_error, lean_angle = self.analyze_forward_lean(results.pose_landmarks.landmark)
                depth_ok = self.analyze_depth(results.pose_landmarks.landmark)
                analysis_results['knee_inward'].append(knee_error)
                analysis_results['forward_lean'].append(lean_error)
                analysis_results['depth_issues'].append(not depth_ok)
                if lean_angle is not None:
                    analysis_results['torso_angles'].append(lean_angle)
                if left_knee_angle is not None:
                    analysis_results['left_knee_angles'].append(left_knee_angle)
                if right_knee_angle is not None:
                    analysis_results['right_knee_angles'].append(right_knee_angle)
            analysis_results['total_frames'] = frame_count
        cap.release()
        print("Video processing finished.")
        return analysis_results

    def generate_summary(self, metrics):
        """Generate analysis summary based on metrics"""
        summary_lines = []
        summary_lines.append(f"## Squat Form Analysis Results\n")
        summary_lines.append(f"**Total Frames Analyzed:** {metrics['total_frames_analyzed']}")
        depth_percentage = metrics['squat_depth_issues_percentage']
        summary_lines.append(f"\n### Squat Depth: {depth_percentage:.1f}% shallow")
        if depth_percentage > 80:
            summary_lines.append("❌ Significant issue: The hips did not reach below the knees in most of the attempts. Focus on achieving full depth.")
        elif depth_percentage > 40:
            summary_lines.append("⚠️ Moderate issue: The squat depth is inconsistent. Try to go deeper on each repetition.")
        else:
            summary_lines.append("✅ Good depth: You are consistently hitting adequate squat depth.")
        lean_percentage = metrics['forward_lean_percentage']
        summary_lines.append(f"\n### Torso Angle (Forward Lean): {lean_percentage:.1f}% with excessive lean")
        if lean_percentage > 30:
            summary_lines.append("❌ Significant issue: Excessive forward lean detected. This can put strain on your lower back. Focus on keeping your chest up and torso upright.")
        elif lean_percentage > 10:
            summary_lines.append("⚠️ Minor issue: Some forward lean was detected. Pay attention to keeping your chest up and core engaged.")
        else:
            summary_lines.append("✅ Good posture: You maintained a stable and upright torso throughout the squat.")
        knee_valgus_percentage = metrics['knee_valgus_percentage']
        summary_lines.append(f"\n### Knee Valgus (Knees Caving Inward): {knee_valgus_percentage:.1f}% of frames with knee valgus")
        if knee_valgus_percentage > 50:
            summary_lines.append("❌ Signific    ant issue: Your knees are collapsing inward. This is a common but dangerous issue. Focus on pushing your knees out in line with your toes.")
        elif knee_valgus_percentage > 10:
            summary_lines.append("⚠️ Minor issue: Your knees are tracking inward slightly. Actively work on external hip rotation and glute strength.")
        else:
            summary_lines.append("✅ Good knee alignment: Your knees are tracking well and are in line with your feet.")
        return "\n".join(summary_lines)

    def get_gemini_analysis(self, metrics_text):
        """Generate a final verdict, short summary, and detailed description using the Gemini API."""
        try:
            model = genai.GenerativeModel('gemini-1.5-pro-latest')
            prompt = (
                "You are an expert physical therapist and certified strength coach. You are analyzing an overhead squat from a front view. Your analysis should adhere to the professional standards of the National Academy of Sports Medicine (NASM) and the Functional Movement Screen (FMS). A 'FAIL' verdict should only be given for significant, repeatable compensations that pose a high risk of injury. A perfect squat is not required for a 'PASS'.\n\n"
                "Below are the raw metrics and a summary generated from a video analysis. Don't use this for determining pass or fail. Just show in output "
                "Provide a final, human-readable verdict (e.g., 'PASS' or 'FAIL'). "
                "Then, provide a 2-3 line short summary explaining why. "
                "Finally, provide a detailed, actionable description of the squat form. "
                "Explain what the user is doing well and what needs significant work, with clear recommendations. "
                "Return the response as a single JSON object with the keys 'verdict', 'short_summary', and 'detailed_analysis'. "
                "Ensure the entire response is a valid JSON object. "
                f"Raw Metrics: {metrics_text}\n\n"
            )
            response = model.generate_content(prompt)
            return response.text
        except Exception as e:
            print(f"Error calling Gemini API: {e}")
            return json.dumps({
                "verdict": "ERROR",
                "short_summary": "Could not generate AI analysis due to an error.",
                "detailed_analysis": ""
            })


# Initialize the analyzer
analyzer = SquatFormAnalyzer()


class MetricsResponse(BaseModel):
    total_frames_analyzed: int
    frames_with_pose_detected: int
    squat_depth_issues_percentage: float
    forward_lean_percentage: float
    knee_valgus_percentage: float
    average_torso_lean_angle: float
    average_left_knee_angle: float
    average_right_knee_angle: float


class AnalysisResponse(BaseModel):
    metrics: MetricsResponse
    verdict: str
    short_summary: str
    gemini_verdict: str


@app.post("/analyze-video/", response_model=AnalysisResponse)
async def analyze_squat_video_api(video_file: UploadFile = File(...)):
    """API endpoint to analyze an uploaded squat video and return metrics, analysis, and a Gemini verdict."""
    try:
        print("API call received.")
        start_time = time.time()
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_file:
            contents = await video_file.read()
            tmp_file.write(contents)
            tmp_file_path = tmp_file.name
        print(f"File saved to temporary path: {tmp_file_path}")
        results = analyzer.process_video_metrics(tmp_file_path)
        os.remove(tmp_file_path)
        print("Temporary file cleaned up.")

        total_frames = results['total_frames']
        frames_with_pose = len(results['depth_issues'])

        if frames_with_pose == 0:
            summary = "No human pose could be detected in the video. Please ensure the video has a clear view of the person's body."
            return AnalysisResponse(
                metrics=MetricsResponse(
                    total_frames_analyzed=total_frames,
                    frames_with_pose_detected=0,
                    squat_depth_issues_percentage=0.0,
                    forward_lean_percentage=0.0,
                    knee_valgus_percentage=0.0,
                    average_torso_lean_angle=0.0,
                    average_left_knee_angle=0.0,
                    average_right_knee_angle=0.0
                ),
                verdict="N/A",
                short_summary="No pose detected.",
                gemini_verdict="No analysis could be performed."
            )

        depth_issues_count = sum(results['depth_issues'])
        lean_issues_count = sum(results['forward_lean'])
        knee_valgus_count = sum(results['knee_inward'])

        depth_percentage = (depth_issues_count / frames_with_pose) * 100
        lean_percentage = (lean_issues_count / frames_with_pose) * 100
        knee_valgus_percentage = (knee_valgus_count / frames_with_pose) * 100

        avg_torso_angle = np.mean(results['torso_angles']) if results['torso_angles'] else 0.0
        avg_left_knee_angle = np.mean(results['left_knee_angles']) if results['left_knee_angles'] else 0.0
        avg_right_knee_angle = np.mean(results['right_knee_angles']) if results['right_knee_angles'] else 0.0

        metrics_response = MetricsResponse(
            total_frames_analyzed=total_frames,
            frames_with_pose_detected=frames_with_pose,
            squat_depth_issues_percentage=round(depth_percentage, 2),
            forward_lean_percentage=round(lean_percentage, 2),
            knee_valgus_percentage=round(knee_valgus_percentage, 2),
            average_torso_lean_angle=round(avg_torso_angle, 2),
            average_left_knee_angle=round(avg_left_knee_angle, 2),
            average_right_knee_angle=round(avg_right_knee_angle, 2)
        )

        analysis_summary = analyzer.generate_summary(metrics_response.dict())
        gemini_response_text = analyzer.get_gemini_analysis(
            metrics_text=metrics_response.json()
        )

        # Use a regex to extract the JSON from the Markdown code block
        match = re.search(r'```(?:json)?\n(.*)```', gemini_response_text, re.DOTALL)
        if match:
            json_str = match.group(1)
            try:
                gemini_json = json.loads(json_str)
                verdict = gemini_json.get("verdict", "N/A")
                short_summary = gemini_json.get("short_summary", "Summary not available.")
                gemini_verdict = gemini_json.get("detailed_analysis", "Detailed analysis not available.")
            except json.JSONDecodeError:
                verdict = "ERROR"
                short_summary = "Failed to parse AI response."
                gemini_verdict = gemini_response_text
        else:
            verdict = "ERROR"
            short_summary = "Failed to extract JSON from AI response."
            gemini_verdict = gemini_response_text


        end_time = time.time()
        processing_time = end_time - start_time
        print(f"Video analysis completed in {processing_time:.2f} seconds.")

        return AnalysisResponse(metrics=metrics_response, verdict=verdict, short_summary=short_summary, gemini_verdict=gemini_verdict)

    except Exception as e:
        print(f"An error occurred: {e}")
        return {"error": str(e)}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)