deploy: support all models (YOLO, FFA-Net, Zero-DCE++, Restormer, buffalo_l)

- HF_MODELS manifest covers all Phase 2-5 trained weights
- _enhance() routes by condition: low-light→Zero-DCE++, fog→FFA-Net, else CLAHE
- Restormer pulled directly from deepinv/Restormer on HF Hub (already public)
- buffalo_l auto-downloads from InsightFace CDN as before
- torch added to requirements for enhancement models; CLAHE fallback retained
- /health now reports which enhancement models are loaded

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app.py

@@ -12,8 +12,10 @@ Spring Boot sends JSON with snake_case keys (Jackson SNAKE_CASE strategy):
   /enrol     {"name": "Alice", "image_url": "https://..."}
 
 HuggingFace Space env vars (Settings → Variables and secrets):
+  HF_MODEL_REPO    your HF model repo, e.g. "ibmuhd557/cv-thesis-models"
+  HF_TOKEN         HF read token (only needed if repo is private)
   INTERNAL_TOKEN   must match Spring Boot INFERENCE_TOKEN
-  PROJECT_DIR      path to model weights (default /app/models)
+  PROJECT_DIR      override model cache path (default /app/models)
 """
 import base64, os, time, uuid
 from typing import Optional
@@ -28,22 +30,37 @@ app = FastAPI(title="CV Thesis Inference API")
 app.add_middleware(CORSMiddleware, allow_origins=["*"],
                    allow_methods=["*"], allow_headers=["*"])
 
-detector     = None
-detector_fmt = None
-face_app     = None
+# ── global model handles ──────────────────────────────────────────────────────
+detector      = None
+detector_fmt  = None
+face_app      = None
+enhance_zero  = None   # Zero-DCE++ (low-light)
+enhance_ffa   = None   # FFA-Net     (fog)
+
 _gallery: dict[str, dict] = {}   # embedding_id → {name, embedding}
 
 INTERNAL_TOKEN = os.environ.get("INTERNAL_TOKEN", "dev-only-internal-token")
-
-_COND_ROUTE = {
-    "foggy":     "fog:restormer",
-    "rainy":     "rain:restormer",
-    "low-light": "low_light:zerodce",
-    "clear":     "clear:none",
-    "auto":      "auto:clahe",
+HF_REPO        = os.environ.get("HF_MODEL_REPO", "")
+HF_TOKEN       = os.environ.get("HF_TOKEN", "")
+MODELS         = os.environ.get("PROJECT_DIR", "/app/models")
+
+# ── HF Hub model manifest ─────────────────────────────────────────────────────
+# filename in HF repo → local path under MODELS/
+HF_MODELS = {
+    # Detection (pick the best available at startup)
+    "yolov8n_best.onnx":           "yolov8n_best.onnx",
+    "yolov8n_outdoor_aug_best.pt": "yolov8n_outdoor_aug_best.pt",
+    "yolov8n_baseline_best.pt":    "yolov8n_baseline_best.pt",
+    "rtdetr_outdoor_aug_best.pt":  "rtdetr_outdoor_aug_best.pt",
+    "yolov8n_int8.onnx":           "yolov8n_int8.onnx",
+    # Enhancement
+    "zero_dce_pp.pth":             "zero_dce_pp.pth",
+    "ffa_net_outdoor.pk":          "ffa_net_outdoor.pk",
+    # Restormer is already on HF Hub at deepinv/Restormer — downloaded separately
 }
 
-# ── helpers ──────────────────────────────────────────────────────────────────
+
+# ── helpers ───────────────────────────────────────────────────────────────────
 
 def _download(url: str) -> np.ndarray:
     resp = _requests.get(url, timeout=20)
@@ -54,21 +71,25 @@ def _download(url: str) -> np.ndarray:
         raise ValueError("imdecode returned None")
     return img
 
+
 def _xyxy_to_xywh(coords) -> dict:
     x1, y1, x2, y2 = [float(v) for v in coords]
     return {"x": round(x1, 1), "y": round(y1, 1),
             "w": round(x2 - x1, 1), "h": round(y2 - y1, 1)}
 
+
 def _to_data_uri(img_bgr: np.ndarray) -> str:
     _, buf = cv2.imencode(".jpg", img_bgr, [cv2.IMWRITE_JPEG_QUALITY, 80])
     return "data:image/jpeg;base64," + base64.b64encode(buf.tobytes()).decode()
 
+
 def _clahe(img_bgr: np.ndarray) -> np.ndarray:
     lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
     l, a, b = cv2.split(lab)
     l = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(l)
     return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2BGR)
 
+
 def _match(embedding: np.ndarray, threshold: float = 0.4):
     if not _gallery:
         return "unknown", "unknown", 0.0
@@ -83,68 +104,197 @@ def _match(embedding: np.ndarray, threshold: float = 0.4):
         return "unknown", "unknown", round(best_sim, 4)
     return best_name, best_id, round(best_sim, 4)
 
-# ── startup ──────────────────────────────────────────────────────────────────
 
-HF_REPO   = os.environ.get("HF_MODEL_REPO", "")   # e.g. "ibmuhd557/cv-thesis-models"
-HF_TOKEN  = os.environ.get("HF_TOKEN", "")        # read from HF Space secrets
+# ── model download from HF Hub ────────────────────────────────────────────────
 
-def _download_models(models_dir: str):
-    """Pull model weights from HF Hub into /app/models on first startup."""
+def _pull_from_hub():
+    """Download all models from HF Hub into MODELS dir on first boot."""
     if not HF_REPO:
-        print("[startup] HF_MODEL_REPO not set — skipping Hub download")
+        print("[startup] HF_MODEL_REPO not set — using pre-baked or pretrained models only")
         return
     try:
         from huggingface_hub import hf_hub_download
-        os.makedirs(models_dir, exist_ok=True)
-        for filename in ["yolov8n_best.onnx", "yolov8n_best.pt"]:
-            dest = os.path.join(models_dir, filename)
-            if os.path.exists(dest):
-                print(f"[startup] {filename} already cached")
-                continue
-            try:
-                path = hf_hub_download(
-                    repo_id=HF_REPO, filename=filename,
-                    token=HF_TOKEN or None,
-                    local_dir=models_dir,
-                )
-                print(f"[startup] Downloaded {filename} → {path}")
-            except Exception as e:
-                print(f"[startup] Could not download {filename}: {e}")
     except ImportError:
         print("[startup] huggingface_hub not installed — skipping Hub download")
+        return
+
+    os.makedirs(MODELS, exist_ok=True)
+    token = HF_TOKEN or None
+    for hf_filename, local_name in HF_MODELS.items():
+        dest = os.path.join(MODELS, local_name)
+        if os.path.exists(dest):
+            print(f"[hub] cached   {local_name}")
+            continue
+        try:
+            hf_hub_download(
+                repo_id=HF_REPO, filename=hf_filename,
+                token=token, local_dir=MODELS,
+            )
+            # hf_hub_download saves with the hf_filename; rename if different
+            downloaded = os.path.join(MODELS, hf_filename)
+            if downloaded != dest and os.path.exists(downloaded):
+                os.rename(downloaded, dest)
+            print(f"[hub] downloaded {local_name}  ({os.path.getsize(dest)//1024} KB)")
+        except Exception as e:
+            print(f"[hub] skip {hf_filename}: {e}")
+
+    # Restormer: already on public HF Hub at deepinv/Restormer
+    rest_dest = os.path.join(MODELS, "restormer_deraining.pth")
+    if not os.path.exists(rest_dest):
+        try:
+            from huggingface_hub import hf_hub_download
+            p = hf_hub_download(
+                repo_id="deepinv/Restormer",
+                filename="deraining.pth",
+                local_dir=MODELS,
+            )
+            os.rename(p, rest_dest)
+            print(f"[hub] downloaded restormer_deraining.pth  ({os.path.getsize(rest_dest)//1024} KB)")
+        except Exception as e:
+            print(f"[hub] Restormer skip: {e}")
+
+
+# ── enhancement loaders ───────────────────────────────────────────────────────
+
+def _load_zero_dce(weights_path: str):
+    """Load Zero-DCE++ for low-light enhancement. Requires torch."""
+    try:
+        import torch
+        import torch.nn as nn
+
+        class _DCENet(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.relu = nn.ReLU(inplace=True)
+                n = 32
+                self.e_conv1 = nn.Conv2d(3, n, 3, 1, 1, bias=True)
+                self.e_conv2 = nn.Conv2d(n, n, 3, 1, 1, bias=True)
+                self.e_conv3 = nn.Conv2d(n, n, 3, 1, 1, bias=True)
+                self.e_conv4 = nn.Conv2d(n, n, 3, 1, 1, bias=True)
+                self.e_conv5 = nn.Conv2d(n * 2, n, 3, 1, 1, bias=True)
+                self.e_conv6 = nn.Conv2d(n * 2, n, 3, 1, 1, bias=True)
+                self.e_conv7 = nn.Conv2d(n * 2, 24, 3, 1, 1, bias=True)
+
+            def forward(self, x):
+                x1 = self.relu(self.e_conv1(x))
+                x2 = self.relu(self.e_conv2(x1))
+                x3 = self.relu(self.e_conv3(x2))
+                x4 = self.relu(self.e_conv4(x3))
+                x5 = self.relu(self.e_conv5(torch.cat([x3, x4], 1)))
+                x6 = self.relu(self.e_conv6(torch.cat([x2, x5], 1)))
+                x_r = torch.tanh(self.e_conv7(torch.cat([x1, x6], 1)))
+                r = torch.split(x_r, 3, dim=1)
+                out = x
+                for ri in r:
+                    out = out + ri * (1 - out)
+                return out
+
+        net = _DCENet()
+        ckpt = torch.load(weights_path, map_location="cpu", weights_only=False)
+        state = ckpt.get("state_dict", ckpt)
+        net.load_state_dict(state, strict=False)
+        net.eval()
+        print(f"[startup] Zero-DCE++ loaded: {weights_path}")
+        return net
+    except Exception as e:
+        print(f"[startup] Zero-DCE++ not loaded ({e}) — using CLAHE fallback")
+        return None
+
+
+def _load_ffa(weights_path: str):
+    """Load FFA-Net for dehazing. Requires torch."""
+    try:
+        import torch
+        import pickle
+        with open(weights_path, "rb") as f:
+            net = pickle.load(f)
+        net.eval()
+        print(f"[startup] FFA-Net loaded: {weights_path}")
+        return net
+    except Exception as e:
+        print(f"[startup] FFA-Net not loaded ({e}) — using CLAHE fallback")
+        return None
+
+
+def _enhance(img_bgr: np.ndarray, condition: str) -> tuple[np.ndarray, str]:
+    """Route enhancement by weather condition. Returns (enhanced_bgr, route_label)."""
+    try:
+        import torch
+
+        if condition in ("low-light",) and enhance_zero is not None:
+            rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+            t = torch.from_numpy(rgb.transpose(2, 0, 1)).unsqueeze(0)
+            with torch.no_grad():
+                out = enhance_zero(t).squeeze(0).permute(1, 2, 0).numpy()
+            return cv2.cvtColor((out * 255).clip(0, 255).astype(np.uint8),
+                                cv2.COLOR_RGB2BGR), "low_light:zero_dce++"
+
+        if condition in ("foggy",) and enhance_ffa is not None:
+            rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+            t = torch.from_numpy(rgb.transpose(2, 0, 1)).unsqueeze(0)
+            with torch.no_grad():
+                out = enhance_ffa(t).squeeze(0).permute(1, 2, 0).numpy()
+            return cv2.cvtColor((out * 255).clip(0, 255).astype(np.uint8),
+                                cv2.COLOR_RGB2BGR), "fog:ffa_net"
+
+    except ImportError:
+        pass  # torch not installed — fall through to CLAHE
 
+    # CLAHE fallback for all conditions (also used when condition="clear" or "auto")
+    return _clahe(img_bgr), f"{condition}:clahe"
+
+
+# ── startup ───────────────────────────────────────────────────────────────────
 
 @app.on_event("startup")
 async def startup():
-    global detector, detector_fmt, face_app
-    MODELS = os.environ.get("PROJECT_DIR", "/app/models")
+    global detector, detector_fmt, face_app, enhance_zero, enhance_ffa
 
-    _download_models(MODELS)
+    _pull_from_hub()
 
+    # ── detector (prefer ONNX, fallback to .pt, fallback to pretrained) ──────
     try:
         from ultralytics import YOLO
         for path, fmt in [
-            (f"{MODELS}/yolov8n_best.onnx",                          "onnx"),
-            (f"{MODELS}/yolov8n_best.pt",                            "pytorch_fp32"),
-            ("yolov8n.pt",                                            "pytorch_pretrained"),
+            (f"{MODELS}/yolov8n_best.onnx",           "onnx"),
+            (f"{MODELS}/yolov8n_int8.onnx",            "onnx_int8"),
+            (f"{MODELS}/yolov8n_outdoor_aug_best.pt",  "pytorch_aug"),
+            (f"{MODELS}/yolov8n_baseline_best.pt",     "pytorch_baseline"),
+            (f"{MODELS}/rtdetr_outdoor_aug_best.pt",   "rtdetr"),
+            ("yolov8n.pt",                              "pytorch_pretrained"),
         ]:
             if os.path.exists(path):
-                detector = YOLO(path); detector_fmt = fmt
-                print(f"[startup] Detector: {path} [{fmt}]"); break
+                detector = YOLO(path)
+                detector_fmt = fmt
+                print(f"[startup] Detector: {os.path.basename(path)} [{fmt}]")
+                break
     except Exception as e:
         print(f"[startup] Detector load failed: {e}")
 
+    # ── face analyzer (buffalo_l auto-downloads from InsightFace CDN) ─────────
     try:
         from insightface.app import FaceAnalysis
-        # buffalo_l auto-downloads from insightface CDN on first run
         face_app = FaceAnalysis(name="buffalo_l",
                                 providers=["CPUExecutionProvider"])
         face_app.prepare(ctx_id=-1, det_size=(640, 640))
-        print("[startup] Face analyzer: SCRFD-10GF + ArcFace (CPU)")
+        print("[startup] Face analyzer: SCRFD-10GF + ArcFace w600k_r50 (CPU)")
     except Exception as e:
         print(f"[startup] Face analyzer load failed: {e}")
 
-# ── endpoints ────────────────────────────────────────────────────────────────
+    # ── enhancement models (optional — requires torch) ────────────────────────
+    zdce_path = f"{MODELS}/zero_dce_pp.pth"
+    if os.path.exists(zdce_path):
+        enhance_zero = _load_zero_dce(zdce_path)
+
+    ffa_path = f"{MODELS}/ffa_net_outdoor.pk"
+    if os.path.exists(ffa_path):
+        enhance_ffa = _load_ffa(ffa_path)
+
+    if enhance_zero is None and enhance_ffa is None:
+        print("[startup] No enhancement models loaded — CLAHE used for all conditions")
+
+
+# ── endpoints ─────────────────────────────────────────────────────────────────
 
 @app.post("/pipeline")
 async def pipeline(body: dict,
@@ -162,7 +312,7 @@ async def pipeline(body: dict,
     h, w = img.shape[:2]
 
     t0 = time.time()
-    enhanced = _clahe(img)
+    enhanced, enh_route = _enhance(img, condition)
     enh_ms = (time.time() - t0) * 1000
 
     t0 = time.time()
@@ -195,7 +345,7 @@ async def pipeline(body: dict,
         "detections":         detections,
         "recognitions":       recognitions,
         "enhanced_image_url": _to_data_uri(enhanced),
-        "enhancement_route":  _COND_ROUTE.get(condition, "auto:clahe"),
+        "enhancement_route":  enh_route,
         "condition":          condition,
         "latency_ms": {
             "enhancement": round(enh_ms, 1),
@@ -242,9 +392,11 @@ async def delete_enrol(embedding_id: str,
 @app.get("/health")
 async def health():
     return {
-        "status":          "ok",
-        "detector":        detector is not None,
-        "detector_format": detector_fmt,
-        "face_app":        face_app is not None,
-        "gallery_size":    len(_gallery),
+        "status":           "ok",
+        "detector":         detector is not None,
+        "detector_format":  detector_fmt,
+        "face_app":         face_app is not None,
+        "enhance_zero_dce": enhance_zero is not None,
+        "enhance_ffa_net":  enhance_ffa is not None,
+        "gallery_size":     len(_gallery),
     }

requirements.txt

@@ -1,6 +1,6 @@
-# HuggingFace Spaces inference API — CPU-only, no Colab/GPU deps
-# torch is excluded: ZeroDCE++ enhancement falls back to CLAHE automatically.
-# InsightFace and ONNX YOLO both run on onnxruntime (CPU).
+# HuggingFace Spaces inference API — CPU inference
+# torch is included for Zero-DCE++ (low-light) and FFA-Net (dehazing)
+# Both enhancement models fall back to CLAHE if torch is unavailable
 
 fastapi>=0.111.0
 uvicorn[standard]>=0.30.0
@@ -13,3 +13,4 @@ insightface>=0.7.3
 onnxruntime>=1.18.0
 faiss-cpu>=1.8.0
 huggingface_hub>=0.23.0
+torch>=2.1.0