Add/update predictor.py

predictor.py

@@ -5,12 +5,15 @@ Unified CHM inference for WEO-SAS/chm-meta (v1) and WEO-SAS/chm-meta-v2 (v2).
 
 Both versions expose the same interface — only the model directory changes:
 
-    predictor = CHMPredictor("WEO-SAS/chm-meta")      # v1: SSL ViT-H + DPT
-    predictor = CHMPredictor("WEO-SAS/chm-meta-v2")   # v2: DINOv3 ViT-L + DPT
+    predictor = CHMPredictor("./chm-meta")      # v1: SSL ViT-H + DPT
+    predictor = CHMPredictor("./chm-meta-v2")   # v2: DINOv3 ViT-L + DPT
 
     chm = predictor.predict(image)              # (3,H,W) float32 → (H,W) metres
     predictor.predict_tif("in.tif", "out.tif") # full GeoTIFF pipeline
 
+When called from chm_pt.py the pre-built model is injected via model= so that
+weights are not loaded twice.
+
 Requirements
 ------------
 Both versions: torch, numpy, rasterio, Pillow
@@ -38,11 +41,19 @@ class CHMPredictor:
 
     Parameters
     ----------
-    model_dir : local path to a downloaded WEO-SAS CHM model repo
-    device    : torch device (auto-detected if None)
+    model_dir  : local path to a downloaded WEO-SAS CHM model repo
+    device     : torch device (auto-detected if None)
+    model      : pre-built model; bypasses weights loading (used by chm_pt.py)
+    processor  : pre-built HF processor for v2; bypasses processor loading
     """
 
-    def __init__(self, model_dir: str, device: Optional[torch.device] = None):
+    def __init__(
+        self,
+        model_dir:  str,
+        device:     Optional[torch.device] = None,
+        model       = None,
+        processor   = None,
+    ):
         model_dir = Path(model_dir)
         with open(model_dir / "predictor_config.json") as f:
             cfg = json.load(f)
@@ -59,17 +70,21 @@ class CHMPredictor:
             "cuda" if torch.cuda.is_available() else "cpu"
         )
 
-        weights_path = model_dir / cfg["weights_file"]
-
-        if self.model_version == "v1":
-            self._load_v1(model_dir, weights_path)
+        if model is not None:
+            # Pre-built model injected by chm_pt.py — skip weight loading
+            self.model     = model.to(self.device)
+            self.processor = processor
+        elif self.model_version == "v1":
+            self._load_v1(model_dir, model_dir / cfg["weights_file"])
         elif self.model_version == "v2":
             self._load_v2(model_dir)
         else:
             raise ValueError(f"Unknown model_version '{self.model_version}' in predictor_config.json")
 
+        self.model.eval()
+
     # ------------------------------------------------------------------
-    # Model loading
+    # Model loading (only used when model= is not injected)
     # ------------------------------------------------------------------
 
     def _load_v1(self, model_dir: Path, weights_path: Path) -> None:
@@ -78,7 +93,7 @@ class CHMPredictor:
 
         self.model     = SSLModule(ssl_path=str(weights_path), local_path=str(weights_path))
         self.processor = None
-        self.model.to(self.device).eval()
+        self.model.to(self.device)
 
     def _load_v2(self, model_dir: Path) -> None:
         try:
@@ -91,7 +106,7 @@ class CHMPredictor:
 
         self.model     = CHMv2ForDepthEstimation.from_pretrained(str(model_dir))
         self.processor = CHMv2ImageProcessorFast.from_pretrained(str(model_dir))
-        self.model.to(self.device).eval()
+        self.model.to(self.device)
 
     # ------------------------------------------------------------------
     # Per-tile inference
@@ -109,12 +124,11 @@ class CHMPredictor:
         """tile: (3, patch_size, patch_size) float32 in [0, 1] → (patch_size, patch_size)"""
         from PIL import Image  # noqa: PLC0415
 
-        # Processor expects uint8 HWC; it applies its own rescale + normalisation
-        arr_hwc  = (tile * 255).clip(0, 255).astype(np.uint8).transpose(1, 2, 0)
-        pil_img  = Image.fromarray(arr_hwc)
-        H, W     = pil_img.height, pil_img.width
-        inputs   = self.processor(images=pil_img, return_tensors="pt")
-        inputs   = {k: v.to(self.device) for k, v in inputs.items()}
+        arr_hwc = (tile * 255).clip(0, 255).astype(np.uint8).transpose(1, 2, 0)
+        pil_img = Image.fromarray(arr_hwc)
+        H, W    = pil_img.height, pil_img.width
+        inputs  = self.processor(images=pil_img, return_tensors="pt")
+        inputs  = {k: v.to(self.device) for k, v in inputs.items()}
         with torch.no_grad():
             outputs = self.model(**inputs)
         depth = self.processor.post_process_depth_estimation(
@@ -149,11 +163,9 @@ class CHMPredictor:
         if image.ndim != 3 or image.shape[0] != 3:
             raise ValueError(f"Expected (3, H, W), got {image.shape}")
 
-        _, H, W  = image.shape
-        ps       = self.patch_size
-        st       = self.stride
+        _, H, W = image.shape
+        ps, st  = self.patch_size, self.stride
 
-        # Image fits in a single tile
         if H <= ps and W <= ps:
             pad = np.zeros((3, ps, ps), dtype=np.float32)
             pad[:, :H, :W] = image
@@ -173,7 +185,7 @@ class CHMPredictor:
                 tile = np.zeros((3, ps, ps), dtype=np.float32)
                 tile[:, :th, :tw] = image[:, y:y2, x:x2]
 
-                pred = self._infer_tile(tile)        # (ps, ps)
+                pred = self._infer_tile(tile)
                 output[y:y2, x:x2] += pred[:th, :tw]
                 count [y:y2, x:x2] += 1.0
 
@@ -207,14 +219,13 @@ class CHMPredictor:
             arr     = src.read([b + 1 for b in bands]).astype(np.float32)
             profile = src.profile.copy()
 
-        # Percentile normalise to [0, 1] per band
         for b in range(arr.shape[0]):
             vmin = float(np.nanpercentile(arr[b], 1))
             vmax = float(np.nanpercentile(arr[b], 99))
             arr[b] = np.clip((arr[b] - vmin) / max(vmax - vmin, 1e-6), 0.0, 1.0)
 
         print(f"CHM inference  model={self.model_version}  input={arr.shape}  {input_path}")
-        chm = self.predict(arr)    # (H, W)
+        chm = self.predict(arr)
         print(f"Output shape {chm.shape}  range [{chm.min():.2f}, {chm.max():.2f}] m")
 
         out_profile = profile.copy()