api.py

import io
import sys
import cv2
import numpy as np
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from pathlib import Path
from typing import Dict, Any, Optional
import json
from loguru import logger

# Configure Loguru for JSON output to stdout
logger.remove()
logger.add(sys.stdout, format="{message}", serialize=True)

# ── Compatibility shims for models trained with custom YOLO fork ──────────────
from ultralytics.nn.tasks import DetectionModel
from ultralytics.utils.loss import E2ELoss, v8DetectionLoss
class WeightedDetectionLoss(v8DetectionLoss): pass
class WeightedE2ELoss(E2ELoss): pass
class WeightedDetectionModel(DetectionModel): pass
import __main__
__main__.WeightedDetectionLoss  = WeightedDetectionLoss
__main__.WeightedE2ELoss        = WeightedE2ELoss
__main__.WeightedDetectionModel = WeightedDetectionModel
# ─────────────────────────────────────────────────────────────────────────────

# Import the services
from pipelines.ibr.service import load_models as load_ibr_models, predict_ibr_realtime, validate_ibr_step
from pipelines.flat_ribbon.service import load_models as load_fr_models, predict_fr_realtime, validate_fr_step
from pipelines.multi_tube.service import load_models as load_mt_models, predict_mt_realtime, validate_mt_step

app = FastAPI(title="HFCL OTDR Real-time API", description="High-performance in-memory inference for OTDR cables.")

# Enable CORS for all origins, methods, and headers
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

BASE_DIR = Path(__file__).resolve().parent
MODELS_DIR = BASE_DIR / "models"

# Global config (could be moved to a file)
CONFIG = {
    "ibr": {
        "RIBBON_VALUE_MAP": {"0": 5, "1": 1}
    }
}

@app.on_event("startup")
async def startup_event():
    """Load all models into memory once at startup."""
    logger.info("Initializing inference models...")
    load_ibr_models(MODELS_DIR)
    load_fr_models(MODELS_DIR)
    load_mt_models(MODELS_DIR)
    logger.info("All models loaded successfully.")

async def process_image_upload(file: UploadFile) -> np.ndarray:
    """Helper to decode uploaded image buffer into OpenCV array."""
    if not file.content_type.startswith("image/"):
        logger.warning(f"Invalid content type: {file.content_type}")
        raise HTTPException(status_code=400, detail="File must be an image.")

    try:
        contents = await file.read()
        nparr = np.frombuffer(contents, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError("Failed to decode image buffer.")
        return img
    except Exception as e:
        logger.error(f"Image processing error: {e}")
        raise HTTPException(status_code=400, detail="Invalid image encoding.")

@app.post("/api/v1/predict/ibr")
async def predict_ibr(
    file: UploadFile = File(...)
):
    logger.info(f"Received IBR prediction request: {file.filename}")
    try:
        img = await process_image_upload(file)
        # Validation
        try:
            with open(BASE_DIR / "configs/ibr/6912F.json") as f:
                val_cfg = json.load(f)
            
            # IBR requires config for markings logic
            results = predict_ibr_realtime(img, val_cfg)
            
            validate_ibr_step(val_cfg, results, file.filename)
            results["validation"] = {"status": "success"}
        except Exception as e:
            results["validation"] = {"status": "failed", "error": str(e)}

        logger.info(f"IBR Inference complete. Status: {results.get('status')}")
        return results
    except HTTPException:
        raise
    except Exception as e:
        logger.exception(f"Unexpected error during IBR inference: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/v1/predict/flat_ribbon")
async def predict_flat_ribbon(
    file: UploadFile = File(...)
):
    logger.info(f"Received Flat-Ribbon prediction request: {file.filename}")
    try:
        img = await process_image_upload(file)
        results = predict_fr_realtime(img)
        
        # Validation
        try:
            with open(BASE_DIR / "configs/flat_ribbon/config.json") as f:
                val_cfg = json.load(f)
            validate_fr_step(val_cfg, results, file.filename)
            results["validation"] = {"status": "success"}
        except Exception as e:
            results["validation"] = {"status": "failed", "error": str(e)}

        logger.info(f"Flat-Ribbon Inference complete. Status: {results.get('status')}")
        return results
    except HTTPException:
        raise
    except Exception as e:
        logger.exception(f"Unexpected error during Flat-Ribbon inference: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/v1/predict/multi_tube")
async def predict_multi_tube(
    file: UploadFile = File(...)
):
    logger.info(f"Received Multi-Tube prediction request: {file.filename}")
    try:
        img = await process_image_upload(file)
        results = predict_mt_realtime(img)
        
        # Validation
        try:
            with open(BASE_DIR / "configs/multi_tube/config.json") as f:
                val_cfg = json.load(f)
            validate_mt_step(val_cfg, results, file.filename)
            results["validation"] = {"status": "success"}
        except Exception as e:
            results["validation"] = {"status": "failed", "error": str(e)}

        logger.info(f"Multi-Tube Inference complete. Status: {results.get('status')}")
        return results
    except HTTPException:
        raise
    except Exception as e:
        logger.exception(f"Unexpected error during Multi-Tube inference: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    return {"status": "healthy"}