Skip to content

xwhy.models.TorchvisionSegmentation

Bases: BaseSegmentation

Segmentation backend for standard torchvision models.

Supports dynamic loading of models like DeepLabV3+, FCN, and LRASPP.

Source code in src/xwhy/models/segmentation/torchvision_models.py 
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
class TorchvisionSegmentation(BaseSegmentation):
    """Segmentation backend for standard torchvision models.

    Supports dynamic loading of models like DeepLabV3+, FCN, and LRASPP.
    """

    # Map model names to their respective initialization functions and default weights
    _MODEL_REGISTRY: ClassVar[Mapping[str, tuple[Callable[..., Any], Any]]] = {
        "deeplabv3_resnet101": (
            deeplabv3_resnet101,
            DeepLabV3_ResNet101_Weights.DEFAULT,
        ),
        "deeplabv3_resnet50": (deeplabv3_resnet50, DeepLabV3_ResNet50_Weights.DEFAULT),
        "deeplabv3_mobilenet_v3_large": (
            deeplabv3_mobilenet_v3_large,
            DeepLabV3_MobileNet_V3_Large_Weights.DEFAULT,
        ),
        "fcn_resnet50": (fcn_resnet50, FCN_ResNet50_Weights.DEFAULT),
        "lraspp_mobilenet_v3_large": (
            lraspp_mobilenet_v3_large,
            LRASPP_MobileNet_V3_Large_Weights.DEFAULT,
        ),
    }

    def __init__(
        self,
        *,
        settings: Settings,
        model_name: str = "deeplabv3_resnet101",
        seed: int = 42,
        device: torch.device | str | None = None,
        **kwargs: Any,  # noqa: ANN401
    ) -> None:
        """Initialize the Torchvision segmentation backend.

        Args:
            settings: Global application settings for cache directories.
            model_name: The torchvision segmentation model identifier.
            seed: Random seed for reproducible inference.
            device: Target computation device.
            **kwargs: Additional arbitrary keyword arguments.

        """
        self._settings = settings
        self._model_name = model_name
        self._seed = seed

        if self._model_name not in self._MODEL_REGISTRY:
            raise ValueError(
                f"Unsupported model '{self._model_name}'. "
                f"Available models: {list(self._MODEL_REGISTRY.keys())}"
            )

        self._rng = np.random.default_rng(self._seed)

        if device is None:
            self._device = torch.device(
                "cuda" if torch.cuda.is_available() else "cpu",
            )
        else:
            self._device = torch.device(device) if isinstance(device, str) else device

        # Torchvision resources
        self._weights: Any | None = None
        self._model: Any | None = None
        self._preprocess: Any | None = None
        self._class_names: list[str] = []

    @property
    def model(self) -> Any:  # noqa: ANN401
        """Read-only property to access the underlying raw segmentation model.

        Raises:
            RuntimeError: If the model has not been loaded yet.

        Returns:
            The loaded torchvision segmentation model.

        """
        if self._model is None:
            raise RuntimeError(
                f"Model '{self._model_name}' is not loaded. Call .load() first."
            )
        return self._model

    @property
    def preprocess_fn(self) -> Callable[..., Any] | None:
        """Read-only property to access the preprocessing transform function.

        Returns:
            The torchvision transform function configured for the segmentation model.

        """
        return self._preprocess

    @property
    def class_names(self) -> list[str]:
        """Read-only property to access the segmentation semantic class names.

        Returns:
            A list of class names supported by the loaded model.

        """
        if not self._class_names:
            logger.warning("Model not loaded yet. Loading now to fetch class names.")
            self.load()
        return self._class_names

    def __call__(self, inputs: torch.Tensor) -> torch.Tensor:
        """Execute the forward pass of the segmentation model.

        Args:
            inputs: A PyTorch tensor containing the preprocessed images.
                Expected shape is typically (B, C, H, W).

        Raises:
            RuntimeError: If the model has not been loaded yet.

        Returns:
            A PyTorch tensor containing the segmentation logits or masks
            of shape (B, num_classes, H, W).

        """
        if self._model is None:
            _, model = self.load()
        else:
            model = self._model

        inputs = inputs.to(self._device)

        with torch.no_grad():
            outputs = model(inputs)

            # Torchvision segmentation models return an OrderedDict.
            # The main output is stored in the "out" key.
            if isinstance(outputs, dict) and "out" in outputs:
                logits = outputs["out"]
            else:
                logits = outputs

        return logits  # type: ignore[no-any-return]

    def _set_seed(self) -> None:
        """Set random seeds for reproducibility."""
        logger.debug(f"Setting seeds to {self._seed} for reproducibility...")
        self._rng = np.random.default_rng(self._seed)
        torch.manual_seed(self._seed)

        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(self._seed)
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False

    def _get_cache_dir(self) -> Path:
        """Retrieve the caching directory from application settings."""
        cache_dir = getattr(self._settings, "segmentation_cache_dir", None)
        if not cache_dir:
            cache_dir = Path.home() / ".cache" / "xwhy" / "segmentation"
        return cache_dir

    def load(self) -> tuple[Any, Any]:
        """Load the specified model and preprocessing transforms into memory.

        Returns:
            A tuple containing the initialized (preprocess_transforms, model).

        """
        if self._model is not None and self._preprocess is not None:
            return self._preprocess, self._model

        self._set_seed()
        cache_dir = self._get_cache_dir()
        cache_dir.mkdir(parents=True, exist_ok=True)

        # Point torch cache to our internal directory to prevent hub subfolder mismatch
        os.environ["TORCH_HOME"] = str(cache_dir)

        logger.info(f"Loading {self._model_name} segmentation model...")

        # Dynamically fetch the model builder and weights from the registry
        model_builder, self._weights = self._MODEL_REGISTRY[self._model_name]

        # Load model and explicitly set to eval mode
        self._model = model_builder(weights=self._weights).to(self._device)
        self._model.eval()

        # Extract the correct preprocessing pipeline and classes
        self._preprocess = self._weights.transforms()

        # Meta dictionary safely fallback to empty list if categories are absent
        self._class_names = getattr(self._weights, "meta", {}).get("categories", [])
        logger.info(
            f"Segmentation model classes loaded: {len(self._class_names)} classes."
        )

        return self._preprocess, self._model

    def predict(self, inputs: torch.Tensor) -> torch.Tensor:
        """Run segmentation on preprocessed tensor inputs.

        Args:
            inputs: A preprocessed image tensor of shape (B, C, H, W).

        Returns:
            A tensor of logits/masks (B, num_classes, H, W).

        """
        return self.__call__(inputs)

model property

Read-only property to access the underlying raw segmentation model.

Raises:

Type Description
RuntimeError

If the model has not been loaded yet.

Returns:

Type Description
Any

The loaded torchvision segmentation model.

preprocess_fn property

Read-only property to access the preprocessing transform function.

Returns:

Type Description
Callable[..., Any] | None

The torchvision transform function configured for the segmentation model.

class_names property

Read-only property to access the segmentation semantic class names.

Returns:

Type Description
list[str]

A list of class names supported by the loaded model.

__init__(*, settings, model_name='deeplabv3_resnet101', seed=42, device=None, **kwargs)

Initialize the Torchvision segmentation backend.

Parameters:

Name Type Description Default
settings Settings

Global application settings for cache directories.

 required
model_name str

The torchvision segmentation model identifier.

 'deeplabv3_resnet101'
seed int

Random seed for reproducible inference.

 42
device device | str | None

Target computation device.

 None
**kwargs Any

Additional arbitrary keyword arguments.

 {}
Source code in src/xwhy/models/segmentation/torchvision_models.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def __init__(
    self,
    *,
    settings: Settings,
    model_name: str = "deeplabv3_resnet101",
    seed: int = 42,
    device: torch.device | str | None = None,
    **kwargs: Any,  # noqa: ANN401
) -> None:
    """Initialize the Torchvision segmentation backend.

    Args:
        settings: Global application settings for cache directories.
        model_name: The torchvision segmentation model identifier.
        seed: Random seed for reproducible inference.
        device: Target computation device.
        **kwargs: Additional arbitrary keyword arguments.

    """
    self._settings = settings
    self._model_name = model_name
    self._seed = seed

    if self._model_name not in self._MODEL_REGISTRY:
        raise ValueError(
            f"Unsupported model '{self._model_name}'. "
            f"Available models: {list(self._MODEL_REGISTRY.keys())}"
        )

    self._rng = np.random.default_rng(self._seed)

    if device is None:
        self._device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu",
        )
    else:
        self._device = torch.device(device) if isinstance(device, str) else device

    # Torchvision resources
    self._weights: Any | None = None
    self._model: Any | None = None
    self._preprocess: Any | None = None
    self._class_names: list[str] = []

__call__(inputs)

Execute the forward pass of the segmentation model.

Parameters:

Name Type Description Default
inputs Tensor

A PyTorch tensor containing the preprocessed images. Expected shape is typically (B, C, H, W).

 required

Raises:

Type Description
RuntimeError

If the model has not been loaded yet.

Returns:

Type Description
Tensor

A PyTorch tensor containing the segmentation logits or masks
Tensor

of shape (B, num_classes, H, W).
Source code in src/xwhy/models/segmentation/torchvision_models.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def __call__(self, inputs: torch.Tensor) -> torch.Tensor:
    """Execute the forward pass of the segmentation model.

    Args:
        inputs: A PyTorch tensor containing the preprocessed images.
            Expected shape is typically (B, C, H, W).

    Raises:
        RuntimeError: If the model has not been loaded yet.

    Returns:
        A PyTorch tensor containing the segmentation logits or masks
        of shape (B, num_classes, H, W).

    """
    if self._model is None:
        _, model = self.load()
    else:
        model = self._model

    inputs = inputs.to(self._device)

    with torch.no_grad():
        outputs = model(inputs)

        # Torchvision segmentation models return an OrderedDict.
        # The main output is stored in the "out" key.
        if isinstance(outputs, dict) and "out" in outputs:
            logits = outputs["out"]
        else:
            logits = outputs

    return logits  # type: ignore[no-any-return]

load()

Load the specified model and preprocessing transforms into memory.

Returns:

Type Description
tuple[Any, Any]

A tuple containing the initialized (preprocess_transforms, model).
Source code in src/xwhy/models/segmentation/torchvision_models.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def load(self) -> tuple[Any, Any]:
    """Load the specified model and preprocessing transforms into memory.

    Returns:
        A tuple containing the initialized (preprocess_transforms, model).

    """
    if self._model is not None and self._preprocess is not None:
        return self._preprocess, self._model

    self._set_seed()
    cache_dir = self._get_cache_dir()
    cache_dir.mkdir(parents=True, exist_ok=True)

    # Point torch cache to our internal directory to prevent hub subfolder mismatch
    os.environ["TORCH_HOME"] = str(cache_dir)

    logger.info(f"Loading {self._model_name} segmentation model...")

    # Dynamically fetch the model builder and weights from the registry
    model_builder, self._weights = self._MODEL_REGISTRY[self._model_name]

    # Load model and explicitly set to eval mode
    self._model = model_builder(weights=self._weights).to(self._device)
    self._model.eval()

    # Extract the correct preprocessing pipeline and classes
    self._preprocess = self._weights.transforms()

    # Meta dictionary safely fallback to empty list if categories are absent
    self._class_names = getattr(self._weights, "meta", {}).get("categories", [])
    logger.info(
        f"Segmentation model classes loaded: {len(self._class_names)} classes."
    )

    return self._preprocess, self._model

predict(inputs)

Run segmentation on preprocessed tensor inputs.

Parameters:

Name Type Description Default
inputs Tensor

A preprocessed image tensor of shape (B, C, H, W).

 required

Returns:

Type Description
Tensor

A tensor of logits/masks (B, num_classes, H, W).
Source code in src/xwhy/models/segmentation/torchvision_models.py 
227
228
229
230
231
232
233
234
235
236
237
def predict(self, inputs: torch.Tensor) -> torch.Tensor:
    """Run segmentation on preprocessed tensor inputs.

    Args:
        inputs: A preprocessed image tensor of shape (B, C, H, W).

    Returns:
        A tensor of logits/masks (B, num_classes, H, W).

    """
    return self.__call__(inputs)