app.py

import os
import urllib.request
from contextlib import asynccontextmanager
from dataclasses import dataclass
from threading import Lock
from typing import Any, List, Optional

import cv2
import mediapipe as mp
import numpy as np
import requests as http_requests
from fastapi import Body, FastAPI, HTTPException, Request
from mediapipe.tasks import python as mp_python
from mediapipe.tasks.python import vision as mp_vision
from PIL import Image
from io import BytesIO
from pydantic import BaseModel

# ---------------------------------------------------------------------------
# Model download
# ---------------------------------------------------------------------------
MODEL_PATH = "face_landmarker_v2.task"
MODEL_URL = (
    "https://storage.googleapis.com/mediapipe-models/"
    "face_landmarker/face_landmarker/float16/latest/face_landmarker.task"
)

if not os.path.exists(MODEL_PATH):
    print(f"Downloading face landmarker model to {MODEL_PATH} …")
    urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
    print("Download complete.")

# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass
class OrientationResult:
    orientation: str
    confidence: float
    face_visibility: float
    nose_depth_signal: float
    shoulder_score: float
    combined_score: float
    details: str


# ---------------------------------------------------------------------------
# Detector
# ---------------------------------------------------------------------------

class FrontBackDetector:
    """
    Front/back classification based on eye visibility.

    Logic:
      - Both eyes visible  → FRONT
      - No eyes visible    → BACK
      - One eye visible    → SIDE / ANGLED
    """

    LEFT_EYE_OUTER = 33
    LEFT_EYE_INNER = 133
    RIGHT_EYE_OUTER = 263
    RIGHT_EYE_INNER = 362
    NOSE_TIP = 1

    def __init__(self):
        options = mp_vision.FaceLandmarkerOptions(
            base_options=mp_python.BaseOptions(model_asset_path=MODEL_PATH),
            running_mode=mp_vision.RunningMode.IMAGE,
            num_faces=1,
            min_face_detection_confidence=0.3,
            min_face_presence_confidence=0.3,
            min_tracking_confidence=0.3,
        )
        self.landmarker = mp_vision.FaceLandmarker.create_from_options(options)

    # ---- helpers --------------------------------------------------------

    def _eye_visibility(self, face_landmarks, img_w, img_h):
        lm = face_landmarks

        le_outer = lm[self.LEFT_EYE_OUTER]
        le_inner = lm[self.LEFT_EYE_INNER]
        re_outer = lm[self.RIGHT_EYE_OUTER]
        re_inner = lm[self.RIGHT_EYE_INNER]

        left_eye_width = abs(le_inner.x - le_outer.x) * img_w
        right_eye_width = abs(re_inner.x - re_outer.x) * img_w
        face_width = abs(le_outer.x - re_outer.x) * img_w

        if face_width < 5:
            return False, False, 0.0, 0.0

        left_ratio = left_eye_width / face_width
        right_ratio = right_eye_width / face_width

        left_score = min(left_ratio / 0.20, 1.0)
        right_score = min(right_ratio / 0.20, 1.0)

        left_in_bounds = 0.02 < le_outer.x < 0.98 and 0.02 < le_inner.x < 0.98
        right_in_bounds = 0.02 < re_outer.x < 0.98 and 0.02 < re_inner.x < 0.98

        left_visible = left_score > 0.35 and left_in_bounds
        right_visible = right_score > 0.35 and right_in_bounds

        return left_visible, right_visible, left_score, right_score

    # ---- main -----------------------------------------------------------

    def detect(self, image_bgr: np.ndarray) -> Optional[OrientationResult]:
        rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        h, w = rgb.shape[:2]
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)

        result = self.landmarker.detect(mp_image)

        if not result.face_landmarks:
            return OrientationResult(
                orientation="BACK",
                confidence=0.95,
                face_visibility=0.0,
                nose_depth_signal=0.0,
                shoulder_score=0.0,
                combined_score=-0.95,
                details="no_face_detected → BACK",
            )

        face_lm = result.face_landmarks[0]
        left_vis, right_vis, left_score, right_score = self._eye_visibility(
            face_lm, w, h
        )
        eyes_visible = int(left_vis) + int(right_vis)

        if eyes_visible == 2:
            confidence = 0.7 + min((left_score + right_score) / 2, 1.0) * 0.3
            orientation = "FRONT"
            combined = confidence
        elif eyes_visible == 1:
            orientation = "SIDE"
            confidence = 0.6
            combined = 0.0
        else:
            orientation = "BACK"
            confidence = 0.75
            combined = -0.75

        details = (
            f"left_eye={'YES' if left_vis else 'NO'}({left_score:.2f}) | "
            f"right_eye={'YES' if right_vis else 'NO'}({right_score:.2f}) | "
            f"eyes_count={eyes_visible}/2"
        )

        return OrientationResult(
            orientation=orientation,
            confidence=confidence,
            face_visibility=left_score + right_score,
            nose_depth_signal=0.0,
            shoulder_score=0.0,
            combined_score=combined,
            details=details,
        )

    def close(self):
        self.landmarker.close()


# ---------------------------------------------------------------------------
# Image helpers
# ---------------------------------------------------------------------------

def load_image_from_url(url: str, max_dim: int = 1280) -> Optional[np.ndarray]:
    try:
        resp = http_requests.get(url, timeout=20)
        resp.raise_for_status()
        pil = Image.open(BytesIO(resp.content)).convert("RGB")
        if max(pil.size) > max_dim:
            pil.thumbnail((max_dim, max_dim), Image.LANCZOS)
        return cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2BGR)
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Core logic
# ---------------------------------------------------------------------------

def extract_front_back(
    image_urls: List[str], detector: FrontBackDetector, detector_lock: Lock
) -> dict:
    front_candidates: list[tuple[float, float, int]] = []
    back_candidates: list[tuple[float, float, int]] = []
    urls_with_results: list[tuple[str, Optional[OrientationResult]]] = []

    for idx, url in enumerate(image_urls):
        image = load_image_from_url(url)
        if image is None:
            urls_with_results.append((url, None))
            continue

        # MediaPipe FaceLandmarker isn't guaranteed thread-safe.
        # Keep one shared instance for performance, guarded by a lock.
        with detector_lock:
            result = detector.detect(image)
        urls_with_results.append((url, result))

        if result is None:
            continue

        if result.orientation == "FRONT":
            front_candidates.append((result.confidence, result.combined_score, idx))
        elif result.orientation == "BACK":
            back_candidates.append((result.confidence, -result.combined_score, idx))
        elif result.orientation == "SIDE":
            front_candidates.append(
                (result.confidence * 0.5, result.combined_score, idx)
            )

    best_front_url, best_front_conf = None, 0.0
    best_back_url, best_back_conf = None, 0.0

    if front_candidates:
        front_candidates.sort(key=lambda x: (x[0], x[1]), reverse=True)
        best_idx = front_candidates[0][2]
        best_front_url = urls_with_results[best_idx][0]
        best_front_conf = front_candidates[0][0]

    if back_candidates:
        back_candidates.sort(key=lambda x: (x[0], x[1]), reverse=True)
        best_idx = back_candidates[0][2]
        best_back_url = urls_with_results[best_idx][0]
        best_back_conf = back_candidates[0][0]

    return {
        "front_url": best_front_url,
        "front_confidence": round(best_front_conf, 4),
        "back_url": best_back_url,
        "back_confidence": round(best_back_conf, 4),
    }


# ---------------------------------------------------------------------------
# FastAPI
# ---------------------------------------------------------------------------

@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.detector = FrontBackDetector()
    app.state.detector_lock = Lock()
    try:
        yield
    finally:
        app.state.detector.close()


app = FastAPI(title="Front/Back View API", lifespan=lifespan)


class DetectResponse(BaseModel):
    front_url: Optional[str] = None
    front_confidence: float = 0.0
    back_url: Optional[str] = None
    back_confidence: float = 0.0


class ClassifyResponse(BaseModel):
    is_front: int


@app.get("/")
def root():
    return {"status": "ok", "message": "Front/Back View Detection API"}


def parse_image_urls(payload: Any) -> List[str]:
    """
    Accepts multiple request shapes and normalizes them to List[str].
    Supported examples:
      {"image_urls": ["..."]}
      {"imageUrls": ["..."]}
      {"urls": ["..."]}
      {"image_urls": "..."}
      {"image_urls": [{"url": "..."}, {"image_url": "..."}]}
      ["...", "..."]
      "..."
    """
    if isinstance(payload, dict):
        for key in ("image_urls", "imageUrls", "urls", "images"):
            if key in payload:
                payload = payload[key]
                break
        else:
            # Allow single-url payloads at root level too:
            # {"url": "..."} / {"image_url": "..."} / {"imageUrl": "..."}
            direct = payload.get("url") or payload.get("image_url") or payload.get("imageUrl")
            if isinstance(direct, str) and direct.strip():
                payload = [direct.strip()]
            else:
                raise HTTPException(
                    status_code=400,
                    detail=(
                        "Request body must contain one of: image_urls, imageUrls, urls, images, "
                        "or a single url/image_url/imageUrl"
                    ),
                )

    if isinstance(payload, str):
        payload = [payload]

    if not isinstance(payload, list):
        raise HTTPException(
            status_code=400,
            detail="image_urls must be a URL string or a list of URL strings",
        )

    normalized: List[str] = []
    for item in payload:
        if isinstance(item, str):
            url = item.strip()
            if url:
                normalized.append(url)
            continue

        if isinstance(item, dict):
            candidate = item.get("url") or item.get("image_url") or item.get("imageUrl")
            if isinstance(candidate, str) and candidate.strip():
                normalized.append(candidate.strip())
                continue

        raise HTTPException(
            status_code=400,
            detail=(
                "Each image entry must be a URL string or an object containing "
                "'url'/'image_url'/'imageUrl'"
            ),
        )

    if not normalized:
        raise HTTPException(status_code=400, detail="image_urls must not be empty")

    return normalized


def parse_single_image_url(payload: Any) -> str:
    image_urls = parse_image_urls(payload)
    if len(image_urls) != 1:
        raise HTTPException(
            status_code=400,
            detail="This endpoint expects exactly one image URL",
        )
    return image_urls[0]


@app.post("/detect", response_model=DetectResponse)
def detect(request: Request, payload: Any = Body(...)):
    image_urls = parse_image_urls(payload)
    result = extract_front_back(
        image_urls, request.app.state.detector, request.app.state.detector_lock
    )
    return DetectResponse(**result)


@app.post("/classify", response_model=ClassifyResponse)
def classify(request: Request, payload: Any = Body(...)):
    image_url = parse_single_image_url(payload)
    image = load_image_from_url(image_url)
    if image is None:
        raise HTTPException(status_code=400, detail="Unable to load image from URL")

    with request.app.state.detector_lock:
        result = request.app.state.detector.detect(image)

    is_front = 1 if result and result.orientation == "FRONT" else 0
    return ClassifyResponse(is_front=is_front)