app.py

import cv2
import numpy as np
import gradio as gr

from src.inference import load_model, infer, process_with_logits, crop
from src.detection import load_detector, detect


from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

# import requests

DETECTOR_MODEL = "models/detector_quantized.onnx"
LIVENESS_MODEL = "models/best_model_quantized.onnx"

# Load once at import time (so it's fast per request)
face_detector = load_detector(DETECTOR_MODEL, (320, 320))
liveness_session, input_name = load_model(LIVENESS_MODEL)

def spoof_default_result():
    return {
        "is_real": False,
        "status": "spoof",
        "logit_diff": 0.00,
        "real_logit": 0.00,
        "spoof_logit": 0.00,
        "confidence": 0.00,
        "p_real": 0.00,
    }

def resize_for_detection(image_bgr, max_side=1280):
    h, w = image_bgr.shape[:2]
    if max(h, w) <= max_side:
        return image_bgr, 1.0
    scale = max_side / float(max(h, w))
    new_w, new_h = int(w * scale), int(h * scale)
    resized = cv2.resize(image_bgr, (new_w, new_h), interpolation=cv2.INTER_AREA)
    return resized, scale

def scale_bbox_to_original(bbox, inv_scale):
    return {
        "x": bbox["x"] * inv_scale,
        "y": bbox["y"] * inv_scale,
        "width": bbox["width"] * inv_scale,
        "height": bbox["height"] * inv_scale,
    }

FIXED_THRESHOLD = 0.4
FIXED_MARGIN = 5
FIXED_BBOX_EXPANSION = 1.5

# def predict_liveness(image_in, threshold=FIXED_THRESHOLD, margin=FIXED_MARGIN, bbox_expansion_factor=FIXED_BBOX_EXPANSION):
#     # image_bgr is a numpy array from gr.Image(type="numpy")

#     # image_in is RGB from Gradio [web:122]
#     image_bgr = cv2.cvtColor(image_in, cv2.COLOR_RGB2BGR)
#     if image_bgr is None:
#         return {"faces": [], "result": spoof_default_result(), "error": "No image provided"}

#     if face_detector is None or liveness_session is None:
#         return {"faces": [], "result": spoof_default_result(), "error": "Models not loaded"}

#     # threshold -> logit threshold (same as demo.py)
#     p = max(1e-6, min(1 - 1e-6, float(threshold)))
#     logit_threshold = float(np.log(p / (1 - p)))

#     # detect on resized
#     det_bgr, scale = resize_for_detection(image_bgr, max_side=1280)
#     det_rgb = cv2.cvtColor(det_bgr, cv2.COLOR_BGR2RGB)
#     faces = detect(det_rgb, face_detector, margin=int(margin))

#     if not faces:
#         return {"faces": [], "result": spoof_default_result()}

#     inv_scale = 1.0 / scale if scale != 0 else 1.0
#     image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

#     face_crops = []
#     meta = []
#     for face in faces:
#         det_conf = face.get("confidence", None)
#         det_conf = round(float(det_conf), 2) if det_conf is not None else None

#         bbox_orig = scale_bbox_to_original(face["bbox"], inv_scale)
#         x, y, w, h = bbox_orig["x"], bbox_orig["y"], bbox_orig["width"], bbox_orig["height"]

#         try:
#             face_crop = crop(image_rgb, (x, y, x + w, y + h), float(bbox_expansion_factor))
#         except Exception:
#             continue

#         face_crops.append(face_crop)
#         meta.append({
#             "bbox": {"x": int(x), "y": int(y), "width": int(w), "height": int(h)},
#             "det_confidence": det_conf
#         })

#     if not face_crops:
#         return {"faces": [], "result": spoof_default_result()}

#     predictions = infer(face_crops, liveness_session, input_name, model_img_size=128)
#     if not predictions:
#         return {"faces": [], "result": spoof_default_result()}

#     out_faces = []
#     for m, pred in zip(meta, predictions):
#         r = process_with_logits(pred, logit_threshold)  # already rounded (your change)
#         out_faces.append({**m, "liveness": r})

#     best = max(out_faces, key=lambda d: (d["det_confidence"] or 0.0))
#     return best["liveness"]


def predict_liveness(image_in, threshold=FIXED_THRESHOLD, margin=FIXED_MARGIN, bbox_expansion_factor=FIXED_BBOX_EXPANSION):
    """Return only the liveness result dict under all conditions."""
    
    # Convert RGB to BGR
    if image_in is None:
        return spoof_default_result()

    image_bgr = cv2.cvtColor(image_in, cv2.COLOR_RGB2BGR)
    if image_bgr is None:
        return spoof_default_result()

    # Ensure models are loaded
    if face_detector is None or liveness_session is None:
        return spoof_default_result()

    # Convert threshold to logit
    p = max(1e-6, min(1 - 1e-6, float(threshold)))
    logit_threshold = float(np.log(p / (1 - p)))

    # Detect faces
    det_bgr, scale = resize_for_detection(image_bgr, max_side=1280)
    det_rgb = cv2.cvtColor(det_bgr, cv2.COLOR_BGR2RGB)
    faces = detect(det_rgb, face_detector, margin=int(margin))

    if not faces:
        return spoof_default_result()

    inv_scale = 1.0 / scale if scale != 0 else 1.0
    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

    face_crops = []
    meta = []

    for face in faces:
        det_conf = face.get("confidence", None)
        det_conf = round(float(det_conf), 2) if det_conf is not None else None

        bbox_orig = scale_bbox_to_original(face["bbox"], inv_scale)
        x, y, w, h = bbox_orig["x"], bbox_orig["y"], bbox_orig["width"], bbox_orig["height"]

        try:
            face_crop = crop(image_rgb, (x, y, x + w, y + h), float(bbox_expansion_factor))
        except Exception:
            continue

        face_crops.append(face_crop)
        meta.append({
            "bbox": {"x": int(x), "y": int(y), "width": int(w), "height": int(h)},
            "det_confidence": det_conf
        })

    if not face_crops:
        return spoof_default_result()

    predictions = infer(face_crops, liveness_session, input_name, model_img_size=128)
    if not predictions:
        return spoof_default_result()

    out_faces = []
    for m, pred in zip(meta, predictions):
        r = process_with_logits(pred, logit_threshold)
        out_faces.append({**m, "liveness": r})

    # Return only the liveness dict of the best face
    best = max(out_faces, key=lambda d: (d["det_confidence"] or 0.0))
    return best["liveness"]

demo = gr.Interface(
    fn=predict_liveness,
    inputs=[
        gr.Image(type="numpy", label="Upload image", height=400, width=500),
        # gr.Slider(0.01, 0.99, value=0.5, step=0.01, label="Threshold"),
        # gr.Slider(0, 50, value=5, step=1, label="Margin"),
        # gr.Slider(1.0, 2.0, value=1.5, step=0.1, label="BBox expansion"),
    ],
    outputs=gr.JSON(label="Result JSON"),
    title="SOL9X: Face Liveness Detection",
)



if __name__ == "__main__":
    demo.launch()