| """Compact CNN-encoder + Lite-DPT depth model. |
| |
| The high-resolution stages are pure convolutional; attention, if present in the |
| encoder, is confined to the stride-32 bottleneck. The decoder is a lightweight |
| top-down FPN with depthwise-separable residual blocks and an ASPP bottleneck. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Any |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| try: |
| import timm |
| except ImportError as exc: |
| raise ImportError("install timm: uv add timm") from exc |
|
|
|
|
| @dataclass |
| class CompactDepthOutput: |
| predicted_depth: torch.Tensor |
| aux: dict[str, torch.Tensor] | None = None |
|
|
|
|
| class ConvNormAct(nn.Sequential): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| kernel_size: int = 3, |
| *, |
| stride: int = 1, |
| groups: int = 1, |
| dilation: int = 1, |
| activation: bool = True, |
| ) -> None: |
| padding = dilation * (kernel_size // 2) |
| layers: list[nn.Module] = [ |
| nn.Conv2d( |
| in_channels, |
| out_channels, |
| kernel_size, |
| stride=stride, |
| padding=padding, |
| groups=groups, |
| dilation=dilation, |
| bias=False, |
| ), |
| nn.BatchNorm2d(out_channels), |
| ] |
| if activation: |
| layers.append(nn.ReLU(inplace=True)) |
| super().__init__(*layers) |
|
|
|
|
| class DWResidualBlock(nn.Module): |
| """Inverted-bottleneck depthwise block: 1x1 expand, 3x3 depthwise, 1x1 project.""" |
|
|
| def __init__(self, channels: int, expansion: int = 2) -> None: |
| super().__init__() |
| hidden = channels * expansion |
| self.expand = ConvNormAct(channels, hidden, 1) |
| self.depthwise = ConvNormAct(hidden, hidden, 3, groups=hidden) |
| self.project = nn.Sequential( |
| nn.Conv2d(hidden, channels, 1, bias=False), |
| nn.BatchNorm2d(channels), |
| ) |
| self.act = nn.ReLU(inplace=True) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.act(x + self.project(self.depthwise(self.expand(x)))) |
|
|
|
|
| class ASPP(nn.Module): |
| """Lightweight ASPP with a reduced channel width for NPU/CPU friendliness.""" |
|
|
| def __init__(self, channels: int, hidden: int = 64, rates: tuple[int, ...] = (2, 4, 6)) -> None: |
| super().__init__() |
| self.reduce = ConvNormAct(channels, hidden, 1) |
| self.branches = nn.ModuleList( |
| [ConvNormAct(hidden, hidden, 3, dilation=rate) for rate in rates] |
| ) |
| self.global_branch = nn.Sequential( |
| nn.AdaptiveAvgPool2d(1), |
| nn.Conv2d(hidden, hidden, 1, bias=False), |
| nn.ReLU(inplace=True), |
| ) |
| self.project = ConvNormAct(hidden * (len(rates) + 2), channels, 1) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.reduce(x) |
| size = x.shape[-2:] |
| features = [x] |
| features.extend(branch(x) for branch in self.branches) |
| global_feature = self.global_branch(x) |
| features.append(F.interpolate(global_feature, size=size, mode="bilinear", align_corners=False)) |
| return self.project(torch.cat(features, dim=1)) |
|
|
|
|
| class EfficientViTEncoder(nn.Module): |
| """EfficientViT with explicit stride-4/8/16 taps plus the stride-32 stage.""" |
|
|
| def __init__(self, model_name: str = "efficientvit_m2.r224_in1k", pretrained: bool = True) -> None: |
| super().__init__() |
| self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0) |
| self.out_channels = (32, 64, 128, 192) |
|
|
| def forward(self, x: torch.Tensor) -> list[torch.Tensor]: |
| pe = self.model.patch_embed |
| x = pe.conv1(x) |
| x = pe.relu1(x) |
| x = pe.conv2(x) |
| x = pe.relu2(x) |
| f4 = x |
| x = pe.conv3(x) |
| x = pe.relu3(x) |
| f8 = x |
| x = pe.conv4(x) |
| f16 = x |
| x = self.model.stages[0](x) |
| x = self.model.stages[1](x) |
| f32 = x |
| return [f4, f8, f16, f32] |
|
|
|
|
| class TimmFeaturesEncoder(nn.Module): |
| """Wrapper around a timm features_only model.""" |
|
|
| def __init__( |
| self, |
| model_name: str, |
| *, |
| out_indices: tuple[int, ...], |
| pretrained: bool = True, |
| ) -> None: |
| super().__init__() |
| self.model = timm.create_model( |
| model_name, |
| pretrained=pretrained, |
| features_only=True, |
| out_indices=out_indices, |
| ) |
| info = self.model.feature_info |
| self.out_channels = tuple(info[i]["num_chs"] for i in out_indices) |
|
|
| def forward(self, x: torch.Tensor) -> list[torch.Tensor]: |
| return list(self.model(x)) |
|
|
|
|
| class LiteDPT(nn.Module): |
| """Top-down fusion decoder with depthwise residual refinement.""" |
|
|
| def __init__( |
| self, |
| in_channels: tuple[int, int, int, int], |
| channels: tuple[int, int, int, int] = (80, 112, 144, 176), |
| *, |
| blocks_per_level: int = 2, |
| aspp_hidden: int = 64, |
| deep_supervision: bool = False, |
| ) -> None: |
| super().__init__() |
| c4, c8, c16, c32 = channels |
| self.channels = channels |
| self.deep_supervision = deep_supervision |
| self.lateral4 = ConvNormAct(in_channels[0], c4, 1) |
| self.lateral8 = ConvNormAct(in_channels[1], c8, 1) |
| self.lateral16 = ConvNormAct(in_channels[2], c16, 1) |
| self.lateral32 = ConvNormAct(in_channels[3], c32, 1) |
|
|
| self.aspp = ASPP(c32, hidden=aspp_hidden) |
|
|
| self.fuse32_to16 = ConvNormAct(c32, c16, 1) |
| self.fuse16_to8 = ConvNormAct(c16, c8, 1) |
| self.fuse8_to4 = ConvNormAct(c8, c4, 1) |
|
|
| self.refine32 = nn.Sequential(*[DWResidualBlock(c32) for _ in range(blocks_per_level)]) |
| self.refine16 = nn.Sequential(*[DWResidualBlock(c16) for _ in range(blocks_per_level)]) |
| self.refine8 = nn.Sequential(*[DWResidualBlock(c8) for _ in range(blocks_per_level)]) |
| self.refine4 = nn.Sequential(*[DWResidualBlock(c4) for _ in range(blocks_per_level)]) |
|
|
| self.head = nn.Sequential( |
| ConvNormAct(c4, c4, 3, groups=c4), |
| nn.Conv2d(c4, 1, 1), |
| ) |
| if deep_supervision: |
| self.aux_heads = nn.ModuleDict( |
| { |
| "16": nn.Conv2d(c16, 1, 1), |
| "8": nn.Conv2d(c8, 1, 1), |
| } |
| ) |
| else: |
| self.aux_heads = None |
|
|
| def forward( |
| self, features: list[torch.Tensor], *, return_aux: bool = False |
| ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: |
| f4, f8, f16, f32 = features |
| x4 = self.lateral4(f4) |
| x8 = self.lateral8(f8) |
| x16 = self.lateral16(f16) |
| x32 = self.aspp(self.lateral32(f32)) |
| x32 = self.refine32(x32) |
|
|
| x16 = x16 + F.interpolate( |
| self.fuse32_to16(x32), size=x16.shape[-2:], mode="bilinear", align_corners=False |
| ) |
| x16 = self.refine16(x16) |
| aux16 = self.aux_heads["16"](x16) if self.aux_heads is not None else None |
|
|
| x8 = x8 + F.interpolate( |
| self.fuse16_to8(x16), size=x8.shape[-2:], mode="bilinear", align_corners=False |
| ) |
| x8 = self.refine8(x8) |
| aux8 = self.aux_heads["8"](x8) if self.aux_heads is not None else None |
|
|
| x4 = x4 + F.interpolate( |
| self.fuse8_to4(x8), size=x4.shape[-2:], mode="bilinear", align_corners=False |
| ) |
| x4 = self.refine4(x4) |
|
|
| depth = F.softplus(self.head(x4)) |
| if return_aux and aux16 is not None and aux8 is not None: |
| return depth, { |
| "16": F.softplus(aux16).squeeze(1), |
| "8": F.softplus(aux8).squeeze(1), |
| } |
| return depth |
|
|
|
|
| class CompactDepthModel(nn.Module): |
| """Encoder + Lite-DPT depth model with padded, full-resolution output.""" |
|
|
| ENCODERS = { |
| "efficientvit_m1": ("efficientvit_m1.r224_in1k", "efficientvit"), |
| "efficientvit_m2": ("efficientvit_m2.r224_in1k", "efficientvit"), |
| "fastvit_t8": ("fastvit_t8.apple_in1k", "timm"), |
| "mobilenetv4_conv_s": ("mobilenetv4_conv_small.e2400_r224_in1k", "timm"), |
| "repvit_m0_9": ("repvit_m0_9.dist_300e_in1k", "timm"), |
| } |
| TIMM_OUT_INDICES = { |
| "fastvit_t8": (0, 1, 2, 3), |
| "mobilenetv4_conv_s": (1, 2, 3, 4), |
| "repvit_m0_9": (0, 1, 2, 3), |
| } |
|
|
| def __init__( |
| self, |
| encoder: str = "efficientvit_m2", |
| *, |
| pretrained: bool = True, |
| channels: tuple[int, int, int, int] = (80, 112, 144, 176), |
| blocks_per_level: int = 2, |
| aspp_hidden: int = 64, |
| deep_supervision: bool = False, |
| ) -> None: |
| super().__init__() |
| if encoder not in self.ENCODERS: |
| raise ValueError(f"unknown encoder {encoder!r}; choices: {sorted(self.ENCODERS)}") |
| model_name, kind = self.ENCODERS[encoder] |
| if kind == "efficientvit": |
| self.encoder = EfficientViTEncoder(model_name, pretrained=pretrained) |
| else: |
| self.encoder = TimmFeaturesEncoder( |
| model_name, |
| out_indices=self.TIMM_OUT_INDICES[encoder], |
| pretrained=pretrained, |
| ) |
| self.encoder_name = encoder |
| self.decoder = LiteDPT( |
| self.encoder.out_channels, |
| channels=channels, |
| blocks_per_level=blocks_per_level, |
| aspp_hidden=aspp_hidden, |
| deep_supervision=deep_supervision, |
| ) |
| self.channels = channels |
| self.blocks_per_level = blocks_per_level |
| self.aspp_hidden = aspp_hidden |
| self.deep_supervision = deep_supervision |
|
|
| def forward( |
| self, pixel_values: torch.Tensor, *, return_aux: bool = False |
| ) -> CompactDepthOutput: |
| height, width = pixel_values.shape[-2:] |
| pad_h = (-height) % 32 |
| pad_w = (-width) % 32 |
| x = pixel_values |
| if pad_h or pad_w: |
| x = F.pad(x, (0, pad_w, 0, pad_h), mode="replicate") |
| features = self.encoder(x) |
| decoder_out = self.decoder(features, return_aux=return_aux) |
| if isinstance(decoder_out, tuple): |
| depth, aux = decoder_out |
| else: |
| depth, aux = decoder_out, None |
| depth = F.interpolate(depth, size=(height, width), mode="bilinear", align_corners=False) |
| return CompactDepthOutput(predicted_depth=depth.squeeze(1), aux=aux) |
|
|
| def config(self) -> dict[str, Any]: |
| return { |
| "encoder": self.encoder_name, |
| "channels": list(self.channels), |
| "blocks_per_level": self.blocks_per_level, |
| "aspp_hidden": self.aspp_hidden, |
| "deep_supervision": self.deep_supervision, |
| } |
|
|
|
|
| def count_parameters(model: nn.Module) -> int: |
| return sum(p.numel() for p in model.parameters()) |
|
|
|
|
| def save_compact(model: CompactDepthModel, path, **extra: Any) -> None: |
| torch.save( |
| { |
| "state_dict": model.state_dict(), |
| "config": model.config(), |
| "params": count_parameters(model), |
| **extra, |
| }, |
| path, |
| ) |
|
|
|
|
| def load_compact( |
| path, *, map_location="cpu", deep_supervision: bool | None = None |
| ) -> CompactDepthModel: |
| checkpoint = torch.load(path, map_location=map_location, weights_only=False) |
| config = checkpoint["config"] |
| use_deep_supervision = ( |
| config.get("deep_supervision", False) |
| if deep_supervision is None |
| else deep_supervision |
| ) |
| model = CompactDepthModel( |
| config["encoder"], |
| pretrained=False, |
| channels=tuple(config["channels"]), |
| blocks_per_level=config.get("blocks_per_level", 2), |
| aspp_hidden=config.get("aspp_hidden", 64), |
| deep_supervision=use_deep_supervision, |
| ) |
| try: |
| model.load_state_dict(checkpoint["state_dict"]) |
| except RuntimeError: |
| |
| |
| model.load_state_dict(checkpoint["state_dict"], strict=False) |
| return model |
|
|
