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="/")