app.py

import os
os.environ["YOLO_CONFIG_DIR"] = "/tmp"
os.environ["MPLCONFIGDIR"] = "/tmp"

import spaces
import base64
import io
import math
import re
import cv2
import numpy as np
import gradio as gr
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
from gtts import gTTS
from pydantic import BaseModel
from ultralytics import YOLO
from PIL import Image

try:
    from google import genai
except ImportError:
    genai = None

fastapi_app = FastAPI(title="ErgoAI - YOLOv8 System")

fastapi_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], 
    allow_methods=["*"],
    allow_headers=["*"],
)

desk_model = None
pose_model = None

def get_desk_model():
    global desk_model
    if desk_model is None:
        desk_model = YOLO('yolov8n.pt')
    return desk_model

def get_pose_model():
    global pose_model
    if pose_model is None:
        pose_model = YOLO('yolov8n-pose.pt')
    return pose_model

@spaces.GPU
def infer_desk(img_out):
    model = get_desk_model() 
    return model(img_out, conf=0.20, imgsz=640, verbose=False)

@spaces.GPU
def infer_pose(img):
    model = get_pose_model() 
    return model(img, imgsz=320, verbose=False)

GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "").strip()
gemini_client = None
if GEMINI_API_KEY and genai:
    try:
        gemini_client = genai.Client(api_key=GEMINI_API_KEY)
    except Exception:
        pass

AUDIO_CACHE = {}
MAX_CACHE_SIZE = 100
SCREEN_REAL_WIDTH_CM = 35.0

class ImageData(BaseModel):
    image_base64: str
    is_calibration: bool
    baseline_eye_dist: float = 0.0
    baseline_shoulder_y: float = 0.0

def calculate_distance(p1, p2):
    return math.sqrt((p2[0] - p1[0])**2 + (p2[1] - p1[1])**2)

def cv2_to_base64(image):
    _, buffer = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, 70])
    return f"data:image/jpeg;base64,{base64.b64encode(buffer).decode('utf-8')}"

def clean_text_for_tts(text):
    cleaned = re.sub(r'[\#\*\_\`\>\-\[\]\(\)]', ' ', text)
    return re.sub(r'\s+', ' ', cleaned).strip()

def get_voice_base64(text: str, lang='en') -> str:
    if not text: 
        return ""
    if len(AUDIO_CACHE) > MAX_CACHE_SIZE:
        AUDIO_CACHE.clear()
    cache_key = f"{lang}_{text}"
    if cache_key not in AUDIO_CACHE:
        try:
            tts = gTTS(text=text, lang=lang, slow=False)
            fp = io.BytesIO()
            tts.write_to_fp(fp)
            fp.seek(0)
            AUDIO_CACHE[cache_key] = base64.b64encode(fp.read()).decode('utf-8')
        except Exception:
            return ""
    return AUDIO_CACHE[cache_key]

def process_single_image(img, user_height):
    h_img, w_img, _ = img.shape
    img_out = img.copy()

    results = infer_desk(img_out)
    
    object_coords = {}
    for r in results:
        for box in r.boxes:
            label = get_desk_model().names[int(box.cls[0])].lower()
            x_center, y_center, width, height = box.xywh[0].tolist()
            if label not in object_coords:
                object_coords[label] = []
            object_coords[label].append({'center': (x_center, y_center), 'width_px': width})

    x_origin, y_origin, w_origin_px = 0, 0, 0
    if 'laptop' in object_coords:
        x_origin, y_origin = object_coords['laptop'][0]['center']
        w_origin_px = object_coords['laptop'][0]['width_px']
    elif 'tvmonitor' in object_coords:
        x_origin, y_origin = object_coords['tvmonitor'][0]['center']
        w_origin_px = object_coords['tvmonitor'][0]['width_px']
    else:
        x_origin, y_origin = w_img / 2, h_img * 0.4
        w_origin_px = w_img * 0.35

    raw_r_normal = w_origin_px * ((user_height * 0.24) / SCREEN_REAL_WIDTH_CM)
    raw_r_max = w_origin_px * ((user_height * 0.35) / SCREEN_REAL_WIDTH_CM)
    max_limit = min(w_img, h_img) * 0.35
    R_NORMAL_PX = int(min(raw_r_normal, max_limit))
    R_MAX_PX = int(min(raw_r_max, max_limit * 1.3))

    overlay = img_out.copy()
    cv2.circle(overlay, (int(x_origin), int(y_origin)), R_NORMAL_PX, (0, 255, 0), -1)
    cv2.circle(img_out, (int(x_origin), int(y_origin)), R_MAX_PX, (0, 0, 255), 3)
    cv2.addWeighted(overlay, 0.20, img_out, 0.80, 0, img_out)

    violations = []
    deductions = 0
    for label_name, object_list in object_coords.items():
        for obj in object_list:
            x_obj, y_obj = obj['center']
            if abs(x_obj - x_origin) < 20 and abs(y_obj - y_origin) < 20: 
                continue
            
            distance_px = math.sqrt((x_obj - x_origin)**2 + (y_obj - y_origin)**2)
            if distance_px == 0: 
                continue
            distance_cm = (distance_px / w_origin_px) * SCREEN_REAL_WIDTH_CM

            if label_name in ['cup', 'bottle'] and distance_px < R_MAX_PX:
                violations.append(f"**{label_name.title()}**: In danger zone ({distance_cm:.1f}cm)")
                deductions += 15
            elif label_name in ['mouse', 'keyboard'] and distance_px > R_NORMAL_PX:
                violations.append(f"**{label_name.title()}**: Too far ({distance_cm:.1f}cm)")
                deductions += 10

    score = max(0, 100 - deductions)
    return img_out, list(dict.fromkeys(violations)), score

@fastapi_app.get("/")
def read_root():
    return {"status": "online", "message": "ErgoAI Production Backend is running!"}

@fastapi_app.post("/api/assess_desk")
def assess_desk(
    file: UploadFile = File(...),
    user_height: float = Form(170.0)
):
    contents = file.file.read()
    nparr = np.frombuffer(contents, np.uint8)
    img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    annotated_bgr, violations, score = process_single_image(img_bgr, user_height)
    
    violations_text = "\n".join([f"* {v}" for v in violations]) if violations else "* Setup is optimal."
    spatial_section = f"### Spatial Alerts\n{violations_text}\n"

    gemini_insights = ""
    if gemini_client:
        try:
            img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
            raw_img = Image.fromarray(img_rgb)
            raw_img.thumbnail((1024, 1024))
            
            prompt = f"User Height: {user_height}cm. Score: {score}. Issues: {violations_text}. Be concise in English."
            response = gemini_client.models.generate_content(model='gemini-2.5-flash', contents=[prompt, raw_img])
            gemini_insights = response.text if response.text else ""
        except Exception:
            pass

    final_report = f"### Ergonomic Score: **{score}/100**\n\n{spatial_section}\n---\n{gemini_insights}"
    tts_text = clean_text_for_tts(gemini_insights if gemini_insights else "Analysis complete.")
    audio_base64_str = get_voice_base64(tts_text, lang='en')

    return {
        "feedback": final_report,
        "processed_image": cv2_to_base64(annotated_bgr),
        "audio_base64": f"data:audio/mp3;base64,{audio_base64_str}" if audio_base64_str else ""
    }

@fastapi_app.post("/api/analyze_frame")
def analyze_frame(data: ImageData):
    try:
        encoded_data = data.image_base64.split(',')[1] if ',' in data.image_base64 else data.image_base64
        nparr = np.frombuffer(base64.b64decode(encoded_data), np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError("Invalid Image Data")
    except Exception:
        return {"status": "Invalid Image", "is_bad_posture": False, "eye_dist": 0, "shoulder_y": 0, "is_success": False}

    results = infer_pose(img)
    
    status = "No person detected"
    current_eye_dist = 0.0
    current_shoulder_y = 0.0
    color = (0, 0, 255)
    is_success = False
    is_bad_posture = False

    if len(results) > 0 and results[0].keypoints is not None and len(results[0].keypoints.xy) > 0:
        kpts = results[0].keypoints.xy[0].cpu().numpy()
        confs = results[0].keypoints.conf[0].cpu().numpy() if results[0].keypoints.conf is not None else np.ones(len(kpts))
        
        CONF_THRESHOLD = 0.5

        if len(kpts) >= 7:
            left_eye, right_eye = kpts[1], kpts[2]
            left_shoulder, right_shoulder = kpts[5], kpts[6]

            if (confs[1] > CONF_THRESHOLD and confs[2] > CONF_THRESHOLD and
                confs[5] > CONF_THRESHOLD and confs[6] > CONF_THRESHOLD and
                np.any(left_eye) and np.any(right_eye) and 
                np.any(left_shoulder) and np.any(right_shoulder)):

                current_eye_dist = calculate_distance(left_eye, right_eye)
                current_shoulder_y = float((left_shoulder[1] + right_shoulder[1]) / 2)

                shoulder_tilt = abs(left_shoulder[1] - right_shoulder[1])
                eye_tilt = abs(left_eye[1] - right_eye[1])

                if data.is_calibration:
                    if current_eye_dist > 60:
                        status = "Error: Too CLOSE! Move back 50-70cm and recalibrate."
                        color = (0, 0, 255)
                    elif current_eye_dist < 30:
                        status = "Error: Too FAR! Move closer 50-70cm and recalibrate."
                        color = (0, 0, 255)
                    elif shoulder_tilt > 12:
                        status = "Error: SHOULDERS TILTED! Please sit straight."
                        color = (0, 0, 255)
                    else:
                        status = "Success: Standard posture saved!"
                        color = (0, 255, 0)
                        is_success = True
                else:
                    if data.baseline_eye_dist > 0 and current_eye_dist > (data.baseline_eye_dist * 1.25):
                        status = "Warning: Too close to screen!"
                        is_bad_posture = True
                        color = (0, 0, 255)
                    elif data.baseline_shoulder_y > 0 and current_shoulder_y > (data.baseline_shoulder_y + 15):
                        status = "Warning: Bad posture (Slouching)!"
                        is_bad_posture = True
                        color = (0, 0, 255)
                    elif shoulder_tilt > 18 or eye_tilt > 15:
                        status = "Warning: Bad posture (Leaning)!"
                        is_bad_posture = True
                        color = (0, 0, 255)
                    else:
                        status = "Good posture"
                        is_bad_posture = False
                        color = (0, 255, 0)

                cv2.putText(img, status, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
                cv2.circle(img, (int(left_eye[0]), int(left_eye[1])), 4, (255, 0, 0), -1)
                cv2.circle(img, (int(right_eye[0]), int(right_eye[1])), 4, (255, 0, 0), -1)
                cv2.line(img, (int(left_shoulder[0]), int(left_shoulder[1])), 
                         (int(right_shoulder[0]), int(right_shoulder[1])), (0, 255, 0), 2)

    _, buffer = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 70])
    processed_base64 = base64.b64encode(buffer).decode('utf-8')

    return {
        "image": f"data:image/jpeg;base64,{processed_base64}",
        "status": status,
        "is_bad_posture": is_bad_posture,
        "eye_dist": current_eye_dist,
        "shoulder_y": current_shoulder_y,
        "is_success": is_success
    }

with gr.Blocks() as dummy_ui:
    gr.Markdown("## ErgoAI Backend is Running")
    gr.Markdown("API and ZeroGPU are active.")
    dummy_img = gr.Image(visible=False)
    dummy_btn = gr.Button("Init", visible=False)
    dummy_btn.click(fn=infer_desk, inputs=dummy_img, outputs=None)
    dummy_btn.click(fn=infer_pose, inputs=dummy_img, outputs=None)

app = gr.mount_gradio_app(fastapi_app, dummy_ui, path="/gradio")