app.py

import os
from typing import Optional, Dict, Any, List

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from PIL import Image
from io import BytesIO

from transformers import pipeline

# =========================
# إعداد أسماء الموديلات (تقدر تغيّرها من Variables في السبيس لو حاب)
# =========================

DETECTOR_MODEL = os.getenv("DETECTOR_MODEL", "Tinny-Robot/acne")
SEVERITY_MODEL = os.getenv("SEVERITY_MODEL", "imfarzanansari/skintelligent-acne")
CONDITION_MODEL = os.getenv("CONDITION_MODEL", "Tanishq77/skin-condition-classifier")

# =========================
# تحميل الموديلات محلياً (مرّة وحدة عند تشغيل السبيس)
# =========================

print(f"Loading detector model: {DETECTOR_MODEL}")
detector = pipeline("object-detection", model=DETECTOR_MODEL)

print(f"Loading severity model: {SEVERITY_MODEL}")
severity_model = pipeline("image-classification", model=SEVERITY_MODEL)

print(f"Loading condition model: {CONDITION_MODEL}")
condition_model = pipeline("image-classification", model=CONDITION_MODEL)


# =========================
# Schemas
# =========================

class SeverityOut(BaseModel):
    raw: str
    label_ar: str
    score: float


class ConditionOut(BaseModel):
    label: str
    score: float


class AnalysisResponse(BaseModel):
    num_lesions: int
    severity: SeverityOut
    condition: ConditionOut
    meta: Dict[str, Any]


# =========================
# Utilities
# =========================

def map_severity_label_ar(label: str) -> str:
    m = {
        "clear": "صافية",
        "mild": "خفيفة",
        "moderate": "متوسطة",
        "severe": "شديدة",
    }
    return m.get(label.lower(), label)


def read_image_from_bytes(image_bytes: bytes) -> Image.Image:
    try:
        img = Image.open(BytesIO(image_bytes)).convert("RGB")
        return img
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid image file: {e}")


# =========================
# FastAPI app
# =========================

app = FastAPI(title="Acne Orchestrator (Local Models)", version="0.2.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # عدّلها لو تبغى Origins معيّنة
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/")
async def root():
    return {
        "message": "Acne Orchestrator is running (local models).",
        "models": {
            "detector": DETECTOR_MODEL,
            "severity": SEVERITY_MODEL,
            "condition": CONDITION_MODEL,
        },
    }


@app.post("/analyze", response_model=AnalysisResponse)
async def analyze(
    file: UploadFile = File(...),
    age: Optional[int] = Form(None),
    skin_type: Optional[str] = Form(None),
    notes: Optional[str] = Form(None),
):
    """
    يستقبل صورة + معلومات بسيطة، ويستخدم 3 موديلات محلياً:
    - detector: يحدد عدد البثور (object detection)
    - severity_model: يقدّر شدة حب الشباب
    - condition_model: يحدّد نوع الحالة الجلدية العامة (acne, eczema, rosacea, ...)
    """

    image_bytes = await file.read()
    if not image_bytes:
        raise HTTPException(status_code=400, detail="Empty image file.")

    # قراءة الصورة كـ PIL
    image = read_image_from_bytes(image_bytes)

    # 1) كشف البثور
    det_results: List[Dict[str, Any]] = detector(image)
    # فلتره بسيطة للثقة
    det_results = [d for d in det_results if d.get("score", 0) > 0.3]
    num_lesions = len(det_results)

    # 2) شدة الحالة
    sev_preds: List[Dict[str, Any]] = severity_model(image)
    if not sev_preds:
        raise HTTPException(status_code=500, detail="Empty response from severity model.")
    sev_top = max(sev_preds, key=lambda x: x.get("score", 0))
    severity_raw = str(sev_top.get("label", "unknown"))
    severity_score = float(sev_top.get("score", 0.0))
    severity_ar = map_severity_label_ar(severity_raw)

    severity_obj = SeverityOut(
        raw=severity_raw,
        label_ar=severity_ar,
        score=severity_score,
    )

    # 3) نوع الحالة الجلدية
    cond_preds: List[Dict[str, Any]] = condition_model(image)
    if not cond_preds:
        raise HTTPException(status_code=500, detail="Empty response from condition model.")
    cond_top = max(cond_preds, key=lambda x: x.get("score", 0))
    condition_label = str(cond_top.get("label", "unknown"))
    condition_score = float(cond_top.get("score", 0.0))

    condition_obj = ConditionOut(
        label=condition_label,
        score=condition_score,
    )

    meta = {
        "age": age,
        "skin_type": skin_type,
        "notes": notes,
        "filename": file.filename,
        "content_type": file.content_type,
    }

    return AnalysisResponse(
        num_lesions=num_lesions,
        severity=severity_obj,
        condition=condition_obj,
        meta=meta,
    )