app.py

from __future__ import annotations

import ast
import io
import logging
import os
import re
from pathlib import Path
from typing import Any, Dict, List, Tuple

import cv2
import gradio as gr
import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("doclayout-yolo-server")

MODEL_PATH = Path(
    os.getenv("DOCLAYOUT_MODEL_PATH", "./models/doclayout_yolo_docstructbench_imgsz1024.onnx")
)

# Fallback download if model doesn't exist locally
if not MODEL_PATH.exists():
    fallback_dir = Path("./models")
    fallback_dir.mkdir(parents=True, exist_ok=True)
    MODEL_PATH = fallback_dir / "doclayout_yolo_docstructbench_imgsz1024.onnx"
    if not MODEL_PATH.exists():
        import urllib.request

        url = "https://huggingface.co/rabbit1906/Yolo/resolve/main/doclayout_yolo_docstructbench_imgsz1024.onnx"
        logger.info(f"Downloading model from {url} to {MODEL_PATH}...")
        urllib.request.urlretrieve(url, str(MODEL_PATH))
        logger.info("Download complete.")

# Inspect available execution providers (CUDA GPU or CPU fallback)
available_providers = ort.get_available_providers()
logger.info(f"Available ONNX providers: {available_providers}")
providers = []
if "CUDAExecutionProvider" in available_providers:
    providers.append("CUDAExecutionProvider")
providers.append("CPUExecutionProvider")

sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

# Load ONNX model directly with graceful fallback
try:
    session = ort.InferenceSession(str(MODEL_PATH), sess_options=sess_options, providers=providers)
except Exception as e:
    logger.warning(f"Failed to initialize with {providers}: {e}. Retrying with CPUExecutionProvider...")
    session = ort.InferenceSession(str(MODEL_PATH), sess_options=sess_options, providers=["CPUExecutionProvider"])

metadata = session.get_modelmeta().custom_metadata_map
STRIDE = ast.literal_eval(metadata.get("stride", "32"))
NAMES = ast.literal_eval(metadata.get("names", "{}"))

active_providers = session.get_providers()
is_gpu_active = any("cuda" in p.lower() or "dml" in p.lower() for p in active_providers)
logger.info(f"Model initialized successfully! Active Providers: {active_providers} (GPU: {is_gpu_active})")


def resize_and_pad(image: np.ndarray, target_size: int = 1024, stride: int = 32) -> Tuple[np.ndarray, float, Tuple[int, int]]:
    h, w = image.shape[:2]
    r = min(target_size / h, target_size / w)
    new_h, new_w = int(round(h * r)), int(round(w * r))
    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

    pad_w = (target_size - new_w) % stride
    pad_h = (target_size - new_h) % stride
    dw, dh = pad_w / 2, pad_h / 2
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))

    padded = cv2.copyMakeBorder(
        resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)
    )
    return padded, r, (left, top)


def run_yolo_inference(img_bgr: np.ndarray, conf_threshold: float = 0.25) -> List[Dict[str, Any]]:
    orig_h, orig_w = img_bgr.shape[:2]
    padded, ratio, (pad_x, pad_y) = resize_and_pad(img_bgr, target_size=1024, stride=STRIDE)

    # Convert BGR to RGB, transpose to CHW, normalize to 0..1
    blob = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB).transpose(2, 0, 1)
    blob = np.expand_dims(blob, axis=0).astype(np.float32) / 255.0

    input_name = session.get_inputs()[0].name
    outputs = session.run(None, {input_name: blob})
    preds = outputs[0]  # Shape: (1, num_boxes, 6) -> [x0, y0, x1, y1, conf, class_id]
    if len(preds.shape) == 3:
        preds = preds[0]

    results = []
    for row in preds:
        conf = float(row[4])
        if conf < conf_threshold:
            continue
        cls_id = int(row[5])
        class_name = NAMES.get(cls_id, str(cls_id))

        # Scale coordinates back to original image dimensions
        x0 = max(0.0, (float(row[0]) - pad_x) / ratio)
        y0 = max(0.0, (float(row[1]) - pad_y) / ratio)
        x1 = min(float(orig_w), (float(row[2]) - pad_x) / ratio)
        y1 = min(float(orig_h), (float(row[3]) - pad_y) / ratio)

        results.append({
            "box": [round(x0, 2), round(y0, 2), round(x1, 2), round(y1, 2)],
            "label": class_name,
            "class_id": cls_id,
            "confidence": round(conf, 4),
        })

    return results


# ── FastAPI Application ──────────────────────────────────────────────
fastapi_app = FastAPI(title="DocLayout-YOLO GPU Service", version="1.0.0")


@fastapi_app.get("/health")
def health():
    return {
        "status": "ok",
        "service": "DocLayout-YOLO",
        "providers": active_providers,
        "gpu_accelerated": is_gpu_active,
        "model_file": str(MODEL_PATH.name),
        "classes": NAMES,
    }


@fastapi_app.post("/predict")
async def predict_api(file: UploadFile = File(...), conf: float = 0.25):
    """RPC inference endpoint compatible with DocuFlow Translate-V2."""
    contents = await file.read()
    nparr = np.frombuffer(contents, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img is None:
        raise HTTPException(status_code=400, detail="Invalid image file")

    boxes = run_yolo_inference(img, conf_threshold=conf)
    return JSONResponse(content={
        "status": "success",
        "gpu": is_gpu_active,
        "width": img.shape[1],
        "height": img.shape[0],
        "total_detections": len(boxes),
        "boxes": boxes,
    })


# ── Gradio Web UI ────────────────────────────────────────────────────
CLASS_COLORS = {
    "title": (255, 0, 0),       # Red
    "text": (0, 128, 255),      # Orange/Blue
    "table": (0, 200, 0),       # Green
    "figure": (200, 0, 200),    # Magenta
    "formula": (0, 255, 255),   # Cyan
    "caption": (128, 128, 0),   # Yellow
}


def gradio_predict(image: Image.Image, conf_threshold: float):
    if image is None:
        return None, "Vui lòng tải ảnh lên!"
    img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    detections = run_yolo_inference(img_bgr, conf_threshold=conf_threshold)

    # Draw bounding boxes
    annotated = img_bgr.copy()
    for d in detections:
        x0, y0, x1, y1 = [int(v) for v in d["box"]]
        label = d["label"]
        conf = d["confidence"]
        color = CLASS_COLORS.get(label.lower(), (0, 255, 0))

        cv2.rectangle(annotated, (x0, y0), (x1, y1), color, 2)
        tag = f"{label} {conf:.2f}"
        (w, h), _ = cv2.getTextSize(tag, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(annotated, (x0, max(0, y0 - 18)), (x0 + w, y0), color, -1)
        cv2.putText(
            annotated, tag, (x0, max(12, y0 - 4)),
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA
        )

    res_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
    summary_text = (
        f"✅ Dò tìm thấy {len(detections)} khối cấu trúc.\n"
        f"⚡ Tăng tốc phần cứng: {'NVIDIA GPU (CUDA)' if is_gpu_active else 'CPU'}\n\n"
        + "\n".join([f"- [{d['label'].upper()}] độ tin cậy {d['confidence']*100:.1f}% tại tọa độ {d['box']}" for d in detections[:15]])
    )
    return res_rgb, summary_text


demo = gr.Interface(
    fn=gradio_predict,
    inputs=[
        gr.Image(type="pil", label="Tải ảnh trang PDF (PNG / JPG)"),
        gr.Slider(minimum=0.1, maximum=0.9, value=0.25, step=0.05, label="Ngưỡng tin cậy (Confidence Threshold)"),
    ],
    outputs=[
        gr.Image(type="numpy", label="Ảnh nhận diện Bố cục (Layout Segmentation)"),
        gr.Textbox(label="Chi tiết phát hiện & Trạng thái phần cứng", lines=10),
    ],
    title="🚀 DocLayout-YOLO Inference API & Visualizer",
    description="Dịch vụ AI phân tích bố cục tài liệu chuyên nghiệp (Document Layout Analysis) cho DocuFlow Translate-V2.",
)

# Mount Gradio into FastAPI root
app = gr.mount_gradio_app(fastapi_app, demo, path="/")