fix(auth): gate cv_dev bypass; feat(inference): ONNX serving; data: channel-aware loader

Co-authored-by: Olufemi Taiwo <117665354+femi23@users.noreply.github.com>

Author: obielin

requirements.txt

@@ -49,6 +49,9 @@ optuna>=3.1.0
 # Explainability & Governance
 shap>=0.42.0
 
+# Optimised inference (serve exported ONNX models)
+onnxruntime>=1.16.0
+
 # Testing and Development
 pytest>=7.0.0
 pytest-cov>=3.0.0

src/climatevision/api/auth.py

@@ -10,6 +10,7 @@
 import hashlib
 import hmac
 import logging
+import os
 import secrets
 from datetime import datetime
 from typing import Optional
@@ -77,13 +78,17 @@ def validate_key(self, api_key: str) -> Optional[dict]:
         if not api_key or not api_key.startswith("cv_"):
             return None
 
-        # Development bypass — allow cv_dev for local testing
+        # Development bypass — only when explicitly enabled. Off by default so
+        # production is secure even if CLIMATEVISION_ALLOW_DEV_KEY is unset.
         if api_key == "cv_dev":
-            return {
-                "id": 0,
-                "name": "Development",
-                "demo": True,
-            }
+            if os.environ.get("CLIMATEVISION_ALLOW_DEV_KEY", "0") == "1":
+                return {
+                    "id": 0,
+                    "name": "Development",
+                    "demo": True,
+                }
+            logger.warning("cv_dev key rejected: dev bypass disabled in this environment.")
+            return None
 
         # Check cache first
         key_hash = self.hash_key(api_key)

src/climatevision/data/dataset.py

@@ -82,11 +82,13 @@ def __init__(
         transform: Optional[Callable] = None,
         normalizer: Optional[Callable] = None,
         image_size: int = 256,
+        n_channels: int = 4,
     ):
         self.root = Path(root)
         self.transform = transform
         self.normalizer = normalizer
         self.image_size = image_size
+        self.n_channels = n_channels
 
         image_dir = self.root / "images"
         mask_dir  = self.root / "masks"
@@ -119,13 +121,13 @@ def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
         image = _load_tif(img_path)   # (C, H, W) float32
         mask  = _load_mask(mask_path) # (H, W)    uint8
 
-        # Ensure 4 bands (pad with zeros if fewer)
+        # Ensure exactly n_channels bands (pad with zeros if fewer, truncate if more)
         c, h, w = image.shape
-        if c < 4:
-            pad = np.zeros((4 - c, h, w), dtype=np.float32)
+        if c < self.n_channels:
+            pad = np.zeros((self.n_channels - c, h, w), dtype=np.float32)
             image = np.concatenate([image, pad], axis=0)
-        elif c > 4:
-            image = image[:4]
+        elif c > self.n_channels:
+            image = image[:self.n_channels]
 
         # Ensure spatial size — pad if smaller, random crop via transform
         if h < self.image_size or w < self.image_size:
@@ -217,6 +219,7 @@ def create_dataloaders(
     normalizer: Optional[Callable] = None,
     pin_memory: bool = True,
     use_weighted_sampler: bool = True,
+    n_channels: int = 4,
 ) -> dict[str, DataLoader]:
     """
     Build train / val / test DataLoaders from a data directory.
@@ -252,6 +255,7 @@ def create_dataloaders(
             transform=transform,
             normalizer=normalizer,
             image_size=image_size,
+            n_channels=n_channels,
         )
 
         sampler = None

src/climatevision/inference/pipeline.py

@@ -67,7 +67,60 @@ def _find_best_checkpoint(analysis_type: str) -> Optional[Path]:
     return candidates[0] if candidates else None
 
 
-def _load_model(analysis_type: str = "deforestation") -> tuple[UNet, torch.device]:
+def _find_onnx(analysis_type: str) -> Optional[Path]:
+    """
+    Search for a trained ONNX model for an analysis type.
+    Priority: config.yaml ``onnx_weights`` > newest models/<analysis_type>_*/model.onnx
+    > models/unet_<analysis_type>.onnx
+    """
+    model_cfg = get_model_config(analysis_type)
+    config_path = model_cfg.get("onnx_weights")
+    if config_path:
+        p = _PROJECT_ROOT / config_path
+        if p.exists():
+            return p
+
+    typed = sorted(
+        _MODELS_DIR.glob(f"{analysis_type}_*/model.onnx"),
+        key=lambda p: p.stat().st_mtime,
+        reverse=True,
+    )
+    if typed:
+        return typed[0]
+
+    direct = _MODELS_DIR / f"unet_{analysis_type}.onnx"
+    return direct if direct.exists() else None
+
+
+class _ONNXSegmenter:
+    """
+    Drop-in replacement for :class:`UNet` in :func:`run_inference`.
+
+    Mimics the minimal torch.nn.Module surface that run_inference relies on
+    (``n_channels``, ``n_classes``, ``.eval()``, ``.to()``, and being callable
+    on a ``(N, C, H, W)`` tensor returning a logits tensor) so the rest of the
+    pipeline needs no changes whether the backing model is PyTorch or ONNX.
+    """
+
+    def __init__(self, session: Any, n_channels: int, n_classes: int, input_name: str):
+        self._session = session
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self._input_name = input_name
+
+    def eval(self) -> "_ONNXSegmenter":
+        return self
+
+    def to(self, _device: Any) -> "_ONNXSegmenter":  # device handled by the ORT provider
+        return self
+
+    def __call__(self, tensor: "torch.Tensor") -> "torch.Tensor":
+        arr = tensor.detach().cpu().numpy().astype(np.float32)
+        logits = self._session.run(None, {self._input_name: arr})[0]
+        return torch.from_numpy(np.ascontiguousarray(logits))
+
+
+def _load_model(analysis_type: str = "deforestation") -> tuple[Any, torch.device]:
     """Load (or return cached) U-Net model configured for the analysis type."""
     if analysis_type in _model_cache:
         return _model_cache[analysis_type]
@@ -104,6 +157,17 @@ def _load_model(analysis_type: str = "deforestation") -> tuple[UNet, torch.devic
             checkpoint.get("val_iou", 0.0),
         )
     else:
+        onnx_path = _find_onnx(analysis_type)
+        onnx_model = _try_load_onnx(onnx_path, n_channels, n_classes) if onnx_path else None
+        if onnx_model is not None:
+            logger.info(
+                "Loaded %s model from ONNX %s (%d channels, %d classes).",
+                analysis_type, onnx_path, n_channels, n_classes,
+            )
+            onnx_model.eval()
+            _model_cache[analysis_type] = (onnx_model, device)
+            return onnx_model, device
+
         logger.warning(
             "No trained model found for %s under %s — using untrained weights (demo).",
             analysis_type,
@@ -117,6 +181,28 @@ def _load_model(analysis_type: str = "deforestation") -> tuple[UNet, torch.devic
     return model, device
 
 
+def _try_load_onnx(
+    onnx_path: Path, n_channels: int, n_classes: int
+) -> Optional["_ONNXSegmenter"]:
+    """Build an ONNX inference session, or return None if onnxruntime is unavailable."""
+    try:
+        import onnxruntime as ort
+    except ImportError:
+        logger.warning(
+            "Found ONNX model at %s but onnxruntime is not installed; "
+            "add 'onnxruntime' to requirements to serve it. Falling back.",
+            onnx_path,
+        )
+        return None
+
+    providers = ["CPUExecutionProvider"]
+    if "CUDAExecutionProvider" in ort.get_available_providers():
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    session = ort.InferenceSession(str(onnx_path), providers=providers)
+    input_name = session.get_inputs()[0].name
+    return _ONNXSegmenter(session, n_channels, n_classes, input_name)
+
+
 # ---------------------------------------------------------------------------
 # Sentinel-2 normalisation statistics (matches preprocessing.py)
 # Band order: [Red, Green, Blue, NIR]