Upload folder using huggingface_hub

README.md

@@ -7,46 +7,46 @@ tags:
 pipeline_tag: image-segmentation
 ---
 
-# MetSeg — Macro-structure Segmentation (DPT)
+# Lung structures Segmentation (DPT)
 
-Pathology macro-structure segmentation. DPT backbone with a custom HF wrapper
-bundled via `trust_remote_code`.
+Pathology segmentation for lung structures (blood vessels and airways). 
+- Encoder (freezed): H-optimus-0 ViT backbone (pretrained on histopathology data).
+- Decoder (trained): custom DPT head with multi-scale feature fusion.
 
 ## Usage
 
+The model expects a normalized `(B, 3, H, W)` float tensor as `pixel_values`.
+Use ImageNet mean/std — same stats applied at training time (matches the
+H-optimus-0 backbone's expected input distribution).
+
+Input image: 224x224 @ 1.5 MPP
+
 ```python
-from PIL import Image
+import numpy as np
 import torch
+from PIL import Image
+from torchvision.transforms import ToTensor, Normalize, Resize, Compose
 from transformers import AutoModel
 
 model = AutoModel.from_pretrained("RendeiroLab/MetPredict-lung-structure-segmentation", trust_remote_code=True).eval()
-
-# Preprocess: PIL / ndarray / tensor — single image or list/batch.
-# Applies ImageNet normalization (mean/std stored in config.json).
-img = Image.open("tile.png")
-pixel_values = model.preprocess(img)            # → (1, 3, H, W) on model device
-
-with torch.no_grad():
+device = next(model.parameters()).device
+
+transform = Compose([
+    ToTensor(),
+    Resize((224, 224)),
+    Normalize(
+        mean=[0.485, 0.456, 0.406], 
+        std=[0.229, 0.224, 0.225]
+    ),
+])
+
+img = Image.open("tile.png").convert("RGB")
+x = transform(img)
+pixel_values = x.unsqueeze(0).to(device)
+
+with torch.inference_mode():
     out = model(pixel_values)
-logits = out.logits                              # (1, n_classes, H, W)
-pred = logits.argmax(dim=1)                      # (1, H, W)
+logits = out.logits                                    # (1, n_classes, H, W)
+pred = logits.argmax(dim=1)                            # (1, H, W)
 ```
 
-`preprocess` accepts:
-- `PIL.Image` (RGB) — single or list
-- `numpy.ndarray` shape `(H, W, 3)` or `(B, H, W, 3)`, uint8 or float
-- `torch.Tensor` shape `(3, H, W)` or `(B, 3, H, W)`, uint8 or float
-
-No resize is applied. H and W must be divisible by the backbone patch size
-(typically 14 or 16). Tile or pad upstream as needed.
-
-If you already have a normalized tensor, you can call the model directly with
-`pixel_values=...`.
-
-## Alternative: portable `torch.export`
-
-```python
-import torch
-m = torch.export.load("model.pt2").module()
-y = m(torch.randn(1, 3, 224, 224))
-```

config.json

@@ -23,16 +23,6 @@
   "decoder_readout": "cat",
   "dtype": "float32",
   "encoder_depth": 4,
-  "image_mean": [
-    0.485,
-    0.456,
-    0.406
-  ],
-  "image_std": [
-    0.229,
-    0.224,
-    0.225
-  ],
   "in_channels": 3,
   "model_type": "metpredict_dpt",
   "n_classes": 3,

hf_model.py

@@ -7,9 +7,8 @@ classes are reconstructable in a clean env.
 """
 from __future__ import annotations
 
-from typing import Any, Literal, Optional, Sequence, Union, cast
+from typing import Any, Literal, Optional, cast
 
-import numpy as np
 import torch
 import torch.nn.functional as F
 from transformers import (
@@ -22,11 +21,6 @@ from transformers.modeling_outputs import SemanticSegmenterOutput
 
 from .dpt import DPT
 
-# ImageNet stats — matches the `A.Normalize()` default used at training time
-# (correct for Virchow2 / H-optimus-0 backbones).
-_IMAGENET_MEAN = (0.485, 0.456, 0.406)
-_IMAGENET_STD = (0.229, 0.224, 0.225)
-
 
 class DPTConfig(PretrainedConfig):
     model_type = "metpredict_dpt"
@@ -42,8 +36,6 @@ class DPTConfig(PretrainedConfig):
         decoder_readout: str = "cat",
         activation: Optional[str] = None,
         in_channels: int = 3,
-        image_mean: Sequence[float] = _IMAGENET_MEAN,
-        image_std: Sequence[float] = _IMAGENET_STD,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -56,8 +48,6 @@ class DPTConfig(PretrainedConfig):
         self.decoder_readout = decoder_readout
         self.activation = activation
         self.in_channels = in_channels
-        self.image_mean = list(image_mean)
-        self.image_std = list(image_std)
         # `auto_map` makes the repo loadable as AutoModel without local imports.
         self.auto_map = {
             "AutoConfig": "hf_model.DPTConfig",
@@ -69,6 +59,7 @@ class DPTForSegmentation(PreTrainedModel):
     config_class = DPTConfig
     base_model_prefix = "dpt"
     main_input_name = "pixel_values"
+    all_tied_weights_keys: dict = {}  # To be compatible with transformers 4.x and 5.x
 
     def __init__(self, config: DPTConfig):
         super().__init__(config)
@@ -86,82 +77,8 @@ class DPTForSegmentation(PreTrainedModel):
             activation=config.activation,
         )
         self.dpt = DPT(**dpt_kwargs)
-        # Normalize stats kept as buffers so they move with `.to(device)` and
-        # show up in `state_dict` for inspection but are excluded from grads.
-        self.register_buffer(
-            "image_mean",
-            torch.tensor(config.image_mean, dtype=torch.float32).view(1, -1, 1, 1),
-            persistent=False,
-        )
-        self.register_buffer(
-            "image_std",
-            torch.tensor(config.image_std, dtype=torch.float32).view(1, -1, 1, 1),
-            persistent=False,
-        )
         # Skip post_init weight init — DPT.initialize() already ran inside DPT.__init__.
 
-    @torch.no_grad()
-    def preprocess(
-        self,
-        images: Union[
-            "PIL.Image.Image",  # noqa: F821
-            np.ndarray,
-            torch.Tensor,
-            list,
-        ],
-    ) -> torch.Tensor:
-        """Convert raw inputs into a model-ready ``pixel_values`` tensor.
-
-        Accepts a single image or a batch, in any of:
-        - PIL.Image (RGB)
-        - numpy.ndarray, shape (H, W, 3) or (B, H, W, 3), dtype uint8 or float
-        - torch.Tensor, shape (3, H, W) or (B, 3, H, W), dtype uint8 or float
-
-        Returns a float tensor of shape ``(B, 3, H, W)`` normalized with the
-        ImageNet stats used during training, on the model's current device.
-
-        Notes:
-        - uint8 inputs are scaled to [0, 1] before normalization.
-        - No resize is applied: the DPT is fully convolutional but the H and W
-          must be divisible by the backbone's patch size (typically 14 or 16).
-          Tile/pad upstream as needed.
-        """
-        if isinstance(images, list):
-            tensors = [self._to_chw_float(x) for x in images]
-            batch = torch.stack(tensors, dim=0)
-        else:
-            t = self._to_chw_float(images)
-            batch = t if t.ndim == 4 else t.unsqueeze(0)
-        batch = batch.to(self.image_mean.device, dtype=torch.float32)
-        return (batch - self.image_mean) / self.image_std
-
-    @staticmethod
-    def _to_chw_float(x: Any) -> torch.Tensor:
-        """Convert a single image-like input to a CHW float tensor in [0, 1]."""
-        # Lazy import: PIL is optional at inference time.
-        try:
-            from PIL import Image as _PILImage
-        except ImportError:  # pragma: no cover
-            _PILImage = None  # type: ignore[assignment]
-        if _PILImage is not None and isinstance(x, _PILImage.Image):
-            arr = np.asarray(x.convert("RGB"))  # HWC uint8
-            t = torch.from_numpy(arr).permute(2, 0, 1).contiguous()
-        elif isinstance(x, np.ndarray):
-            t = torch.from_numpy(x)
-            if t.ndim == 3 and t.shape[-1] in (1, 3):
-                t = t.permute(2, 0, 1).contiguous()
-            elif t.ndim == 4 and t.shape[-1] in (1, 3):
-                t = t.permute(0, 3, 1, 2).contiguous()
-        elif isinstance(x, torch.Tensor):
-            t = x
-        else:
-            raise TypeError(f"Unsupported image type: {type(x).__name__}")
-        if t.dtype == torch.uint8:
-            t = t.float() / 255.0
-        else:
-            t = t.float()
-        return t
-
     def forward(
         self,
         pixel_values: torch.Tensor,