| """Standalone MobileNetV3-FPN string segmentation model.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
|
|
|
|
| CHECKPOINT_FORMATS = {"yoyo_string_semantic_transfer_v1", "yoyo_string_semantic_unet_v1"} |
|
|
|
|
| def _group_count(channels: int) -> int: |
| for groups in (8, 4, 2, 1): |
| if channels % groups == 0: |
| return groups |
| return 1 |
|
|
|
|
| class _FPNRefine(nn.Module): |
| def __init__(self, channels: int): |
| super().__init__() |
| self.layers = nn.Sequential( |
| nn.Conv2d(channels, channels, 3, padding=1, groups=channels, bias=False), |
| nn.GroupNorm(_group_count(channels), channels), |
| nn.SiLU(inplace=True), |
| nn.Conv2d(channels, channels, 1, bias=False), |
| nn.GroupNorm(_group_count(channels), channels), |
| nn.SiLU(inplace=True), |
| ) |
|
|
| def forward(self, value: torch.Tensor) -> torch.Tensor: |
| return self.layers(value) |
|
|
|
|
| class MobileNetV3FPNString(nn.Module): |
| """MobileNetV3-Large encoder with a high-resolution FPN decoder.""" |
|
|
| _FEATURE_INDICES = (1, 3, 6, 12, 16) |
| _FEATURE_CHANNELS = (16, 24, 40, 112, 960) |
|
|
| def __init__(self, decoder_channels: int = 32): |
| super().__init__() |
| from torchvision.models import mobilenet_v3_large |
|
|
| self.encoder = mobilenet_v3_large(weights=None).features |
| self.lateral = nn.ModuleList( |
| nn.Sequential( |
| nn.Conv2d(input_channels, decoder_channels, 1, bias=False), |
| nn.GroupNorm(_group_count(decoder_channels), decoder_channels), |
| ) |
| for input_channels in self._FEATURE_CHANNELS |
| ) |
| self.refine = nn.ModuleList(_FPNRefine(decoder_channels) for _ in range(4)) |
| self.classifier = nn.Sequential( |
| _FPNRefine(decoder_channels), |
| nn.Conv2d(decoder_channels, 1, 1), |
| ) |
|
|
| def forward(self, value: torch.Tensor) -> torch.Tensor: |
| input_size = value.shape[-2:] |
| features = [] |
| selected = set(self._FEATURE_INDICES) |
| for index, layer in enumerate(self.encoder): |
| value = layer(value) |
| if index in selected: |
| features.append(value) |
| pyramid = self.lateral[-1](features[-1]) |
| for level in range(len(features) - 2, -1, -1): |
| pyramid = F.interpolate(pyramid, size=features[level].shape[-2:], mode="bilinear", align_corners=False) |
| pyramid = self.refine[level](pyramid + self.lateral[level](features[level])) |
| logits = self.classifier(pyramid) |
| return F.interpolate(logits, size=input_size, mode="bilinear", align_corners=False) |
|
|
|
|
| def load_model(checkpoint: str | Path, device: str | torch.device = "cpu") -> tuple[nn.Module, dict[str, Any]]: |
| """Load a StrandSeg-Lite checkpoint without importing the training repo.""" |
| target = torch.device(device) |
| payload = torch.load(Path(checkpoint), map_location=target, weights_only=True) |
| if not isinstance(payload, dict) or payload.get("format") not in CHECKPOINT_FORMATS: |
| raise ValueError(f"Unsupported StrandSeg-Lite checkpoint: {checkpoint}") |
| config = dict(payload.get("model_config") or {}) |
| architecture = str(config.get("architecture", "mobilenet_v3_fpn")).lower() |
| if architecture != "mobilenet_v3_fpn": |
| raise ValueError(f"This package supports mobilenet_v3_fpn, got {architecture!r}") |
| model = MobileNetV3FPNString(decoder_channels=max(16, int(config.get("base_channels", 16)) * 2)) |
| model.load_state_dict(payload["state_dict"]) |
| model.to(target).eval() |
| return model, payload |
|
|
