Upload dpt_head.py with huggingface_hub

dpt_head.py

@@ -0,0 +1,702 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""DPT (Dense Prediction Transformer) depth head in PyTorch.
+
+Ported from the Scenic/Flax implementation at:
+  research/vision/scene_understanding/imsight/modules/dpt.py
+  scenic/projects/dense_features/models/decoders.py
+
+Architecture:
+  ReassembleBlocks → 4×Conv3x3 → 4×FeatureFusionBlock → project → DepthHead
+"""
+
+import io
+import os
+import urllib.request
+import zipfile
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+# ── Building blocks ─────────────────────────────────────────────────────────
+
+
+class PreActResidualConvUnit(nn.Module):
+    """Pre-activation residual convolution unit."""
+
+    def __init__(self, features: int):
+        super().__init__()
+        self.conv1 = nn.Conv2d(features, features, 3, padding=1, bias=False)
+        self.conv2 = nn.Conv2d(features, features, 3, padding=1, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = F.relu(x)
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        return x + residual
+
+
+class FeatureFusionBlock(nn.Module):
+    """Fuses features with optional residual input, then upsamples 2×."""
+
+    def __init__(self, features: int, has_residual: bool = False,
+                 expand: bool = False):
+        super().__init__()
+        self.has_residual = has_residual
+        if has_residual:
+            self.residual_unit = PreActResidualConvUnit(features)
+        self.main_unit = PreActResidualConvUnit(features)
+        out_features = features // 2 if expand else features
+        self.out_conv = nn.Conv2d(features, out_features, 1, bias=True)
+
+    def forward(self, x: torch.Tensor,
+                residual: torch.Tensor = None) -> torch.Tensor:
+        if self.has_residual and residual is not None:
+            if residual.shape != x.shape:
+                residual = F.interpolate(
+                    residual, size=x.shape[2:], mode="bilinear",
+                    align_corners=False)
+            residual = self.residual_unit(residual)
+            x = x + residual
+        x = self.main_unit(x)
+        # Upsample 2× with align_corners=True (matches Scenic reference)
+        x = F.interpolate(x, scale_factor=2, mode="bilinear",
+                          align_corners=True)
+        x = self.out_conv(x)
+        return x
+
+
+class ReassembleBlocks(nn.Module):
+    """Projects and resizes intermediate ViT features to different scales."""
+
+    def __init__(self, input_embed_dim: int = 1024,
+                 out_channels: tuple = (128, 256, 512, 1024),
+                 readout_type: str = "project"):
+        super().__init__()
+        self.readout_type = readout_type
+
+        # 1×1 conv to project to per-level channels
+        self.out_projections = nn.ModuleList([
+            nn.Conv2d(input_embed_dim, ch, 1) for ch in out_channels
+        ])
+
+        # Spatial resize layers: 4× up, 2× up, identity, 2× down
+        self.resize_layers = nn.ModuleList([
+            nn.ConvTranspose2d(out_channels[0], out_channels[0],
+                               kernel_size=4, stride=4, padding=0),
+            nn.ConvTranspose2d(out_channels[1], out_channels[1],
+                               kernel_size=2, stride=2, padding=0),
+            nn.Identity(),
+            nn.Conv2d(out_channels[3], out_channels[3], 3, stride=2,
+                      padding=1),
+        ])
+
+        # Readout projection (concatenate cls_token with patch features)
+        if readout_type == "project":
+            self.readout_projects = nn.ModuleList([
+                nn.Linear(2 * input_embed_dim, input_embed_dim)
+                for _ in out_channels
+            ])
+
+    def forward(self, features):
+        """Process list of (cls_token, spatial_features) tuples.
+
+        Args:
+            features: list of (cls_token [B,D], patch_feats [B,D,H,W])
+
+        Returns:
+            list of tensors at different scales.
+        """
+        out = []
+        for i, (cls_token, x) in enumerate(features):
+            B, D, H, W = x.shape
+
+            if self.readout_type == "project":
+                # Flatten spatial → (B, HW, D)
+                x_flat = x.flatten(2).transpose(1, 2)
+                # Expand cls_token → (B, HW, D)
+                readout = cls_token.unsqueeze(1).expand(-1, x_flat.shape[1], -1)
+                # Concat + project + GELU
+                x_cat = torch.cat([x_flat, readout], dim=-1)
+                x_proj = F.gelu(self.readout_projects[i](x_cat))
+                # Reshape back to spatial
+                x = x_proj.transpose(1, 2).reshape(B, D, H, W)
+
+            # 1×1 projection
+            x = self.out_projections[i](x)
+            # Spatial resize
+            x = self.resize_layers[i](x)
+            out.append(x)
+        return out
+
+
+class DPTDepthHead(nn.Module):
+    """Full DPT head + depth classification decoder.
+
+    Takes 4 intermediate ViT features and produces a depth map.
+    """
+
+    def __init__(self, input_embed_dim: int = 1024,
+                 channels: int = 256,
+                 post_process_channels: tuple = (128, 256, 512, 1024),
+                 readout_type: str = "project",
+                 num_depth_bins: int = 256,
+                 min_depth: float = 1e-3,
+                 max_depth: float = 10.0):
+        super().__init__()
+        self.num_depth_bins = num_depth_bins
+        self.min_depth = min_depth
+        self.max_depth = max_depth
+
+        # Reassemble: project + resize
+        self.reassemble = ReassembleBlocks(
+            input_embed_dim=input_embed_dim,
+            out_channels=post_process_channels,
+            readout_type=readout_type,
+        )
+
+        # 3×3 convs to map each level to `channels`
+        self.convs = nn.ModuleList([
+            nn.Conv2d(ch, channels, 3, padding=1, bias=False)
+            for ch in post_process_channels
+        ])
+
+        # Fusion blocks: first has no residual, rest have residual
+        self.fusion_blocks = nn.ModuleList([
+            FeatureFusionBlock(channels, has_residual=False),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+        ])
+
+        # Final projection
+        self.project = nn.Conv2d(channels, channels, 3, padding=1, bias=True)
+
+        # Depth classification head (Dense layer)
+        self.depth_head = nn.Linear(channels, num_depth_bins)
+
+    def forward(self, intermediate_features, image_size=None):
+        """Run DPT depth prediction.
+
+        Args:
+            intermediate_features: list of 4 (cls_token, patch_feats) tuples
+            image_size: (H, W) to resize output to, or None
+
+        Returns:
+            depth map tensor (B, 1, H, W)
+        """
+        # Reassemble
+        x = self.reassemble(intermediate_features)
+        # 3×3 conv per level
+        x = [self.convs[i](feat) for i, feat in enumerate(x)]
+
+        # Fuse bottom-up: start from deepest (x[-1])
+        out = self.fusion_blocks[0](x[-1])
+        for i in range(1, 4):
+            out = self.fusion_blocks[i](out, residual=x[-(i + 1)])
+
+        # Project
+        out = self.project(out)
+        out = F.relu(out)
+
+        # Depth classification
+        # out: (B, C, H, W) → (B, H, W, C)
+        out = out.permute(0, 2, 3, 1)
+        out = self.depth_head(out)  # (B, H, W, num_bins)
+
+        # Classification-based depth prediction
+        bin_centers = torch.linspace(
+            self.min_depth, self.max_depth, self.num_depth_bins,
+            device=out.device)
+        out = F.relu(out) + self.min_depth
+        out_norm = out / out.sum(dim=-1, keepdim=True)
+        depth = torch.einsum("bhwn,n->bhw", out_norm, bin_centers)
+        depth = depth.unsqueeze(1)  # (B, 1, H, W)
+
+        # Resize to original image size
+        if image_size is not None:
+            depth = F.interpolate(depth, size=image_size, mode="bilinear",
+                                  align_corners=False)
+
+        return depth
+
+
+class DPTNormalsHead(nn.Module):
+    """Full DPT head + surface normals decoder.
+
+    Takes 4 intermediate ViT features and produces a normal map.
+    """
+
+    def __init__(self, input_embed_dim: int = 1024,
+                 channels: int = 256,
+                 post_process_channels: tuple = (128, 256, 512, 1024),
+                 readout_type: str = "project"):
+        super().__init__()
+
+        # Reassemble: project + resize
+        self.reassemble = ReassembleBlocks(
+            input_embed_dim=input_embed_dim,
+            out_channels=post_process_channels,
+            readout_type=readout_type,
+        )
+
+        # 3×3 convs to map each level to `channels`
+        self.convs = nn.ModuleList([
+            nn.Conv2d(ch, channels, 3, padding=1, bias=False)
+            for ch in post_process_channels
+        ])
+
+        # Fusion blocks: first has no residual, rest have residual
+        self.fusion_blocks = nn.ModuleList([
+            FeatureFusionBlock(channels, has_residual=False),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+        ])
+
+        # Final projection
+        self.project = nn.Conv2d(channels, channels, 3, padding=1, bias=True)
+
+        # Normals head (Dense layer)
+        self.normals_head = nn.Linear(channels, 3)
+
+    def forward(self, intermediate_features, image_size=None):
+        """Run DPT normals prediction.
+
+        Args:
+            intermediate_features: list of 4 (cls_token, patch_feats) tuples
+            image_size: (H, W) to resize output to, or None
+
+        Returns:
+            normal map tensor (B, 3, H, W)
+        """
+        # Reassemble
+        x = self.reassemble(intermediate_features)
+        # 3×3 conv per level
+        x = [self.convs[i](feat) for i, feat in enumerate(x)]
+
+        # Fuse bottom-up: start from deepest (x[-1])
+        out = self.fusion_blocks[0](x[-1])
+        for i in range(1, 4):
+            out = self.fusion_blocks[i](out, residual=x[-(i + 1)])
+
+        # Project
+        out = self.project(out)
+        
+        # Normals head
+        # out: (B, C, H, W) → (B, H, W, C)
+        out = out.permute(0, 2, 3, 1)
+        out = self.normals_head(out)  # (B, H, W, 3)
+
+        # Normalize to unit length
+        out = F.normalize(out, p=2, dim=-1)
+
+        # Resize to original image size
+        if image_size is not None:
+            # PyTorch interpolate expects (B, C, H, W)
+            out = out.permute(0, 3, 1, 2)
+            out = F.interpolate(out, size=image_size, mode="bilinear",
+                                  align_corners=False)
+        else:
+            out = out.permute(0, 3, 1, 2)
+
+        return out
+
+
+class DPTSegmentationHead(nn.Module):
+    """Full DPT head + segmentation decoder.
+
+    Takes 4 intermediate ViT features and produces a segmentation map.
+    """
+
+    def __init__(self, input_embed_dim: int = 1024,
+                 channels: int = 256,
+                 post_process_channels: tuple = (128, 256, 512, 1024),
+                 readout_type: str = "project",
+                 num_classes: int = 150):
+        super().__init__()
+
+        # Reassemble: project + resize
+        self.reassemble = ReassembleBlocks(
+            input_embed_dim=input_embed_dim,
+            out_channels=post_process_channels,
+            readout_type=readout_type,
+        )
+
+        # 3×3 convs to map each level to `channels`
+        self.convs = nn.ModuleList([
+            nn.Conv2d(ch, channels, 3, padding=1, bias=False)
+            for ch in post_process_channels
+        ])
+
+        # Fusion blocks: first has no residual, rest have residual
+        self.fusion_blocks = nn.ModuleList([
+            FeatureFusionBlock(channels, has_residual=False),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+            FeatureFusionBlock(channels, has_residual=True),
+        ])
+
+        # Final projection
+        self.project = nn.Conv2d(channels, channels, 3, padding=1, bias=True)
+
+        # Segmentation head (Dense layer)
+        self.segmentation_head = nn.Linear(channels, num_classes)
+
+    def forward(self, intermediate_features, image_size=None):
+        """Run DPT segmentation prediction.
+
+        Args:
+            intermediate_features: list of 4 (cls_token, patch_feats) tuples
+            image_size: (H, W) to resize output to, or None
+
+        Returns:
+            segmentation map tensor (B, num_classes, H, W)
+        """
+        # Reassemble
+        x = self.reassemble(intermediate_features)
+        # 3×3 conv per level
+        x = [self.convs[i](feat) for i, feat in enumerate(x)]
+
+        # Fuse bottom-up: start from deepest (x[-1])
+        out = self.fusion_blocks[0](x[-1])
+        for i in range(1, 4):
+            out = self.fusion_blocks[i](out, residual=x[-(i + 1)])
+
+        # Project
+        out = self.project(out)
+        
+        # Segmentation head
+        # out: (B, C, H, W) → (B, H, W, C)
+        out = out.permute(0, 2, 3, 1)
+        out = self.segmentation_head(out)  # (B, H, W, num_classes)
+
+        # Resize to original image size
+        if image_size is not None:
+            # PyTorch interpolate expects (B, C, H, W)
+            out = out.permute(0, 3, 1, 2)
+            out = F.interpolate(out, size=image_size, mode="bilinear",
+                                  align_corners=False)
+        else:
+            out = out.permute(0, 3, 1, 2)
+
+        return out
+
+
+# ── Weight loading from Scenic/Flax checkpoint ─────────────────────────────
+
+
+def _load_npy_from_zip(zf, name):
+    """Load a single .npy array from a zipfile."""
+    with zf.open(name) as f:
+        return np.load(io.BytesIO(f.read()))
+
+
+def _conv_kernel_flax_to_torch(w):
+    """Convert Flax conv kernel (H,W,Cin,Cout) → PyTorch (Cout,Cin,H,W)."""
+    return torch.from_numpy(w.transpose(3, 2, 0, 1).copy())
+
+
+def _conv_transpose_kernel_flax_to_torch(w):
+    """Convert Flax ConvTranspose kernel (H,W,Cin,Cout) → PyTorch (Cin,Cout,H,W)."""
+    return torch.from_numpy(w.transpose(2, 3, 0, 1).copy())
+
+
+def _linear_kernel_flax_to_torch(w):
+    """Convert Flax Dense kernel (in,out) → PyTorch Linear (out,in)."""
+    return torch.from_numpy(w.T.copy())
+
+
+def _bias(w):
+    return torch.from_numpy(w.copy())
+
+
+def load_dpt_weights(model: DPTDepthHead, zip_path: str):
+    """Load Scenic/Flax DPT weights from a zip/npz file into PyTorch model."""
+    zf = zipfile.ZipFile(zip_path, "r")
+    npy = lambda name: _load_npy_from_zip(zf, name)
+    sd = {}
+    prefix = "decoder/dpt/"
+
+    # --- ReassembleBlocks ---
+    for i in range(4):
+        # out_projections (Conv2d 1×1)
+        sd[f"reassemble.out_projections.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/kernel.npy"))
+        sd[f"reassemble.out_projections.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/bias.npy"))
+
+        # readout_projects (Linear)
+        sd[f"reassemble.readout_projects.{i}.weight"] = _linear_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/kernel.npy"))
+        sd[f"reassemble.readout_projects.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/bias.npy"))
+
+    # resize_layers: 0=ConvTranspose, 1=ConvTranspose, 2=Identity, 3=Conv
+    sd["reassemble.resize_layers.0.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/kernel.npy"))
+    sd["reassemble.resize_layers.0.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/bias.npy"))
+    sd["reassemble.resize_layers.1.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/kernel.npy"))
+    sd["reassemble.resize_layers.1.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/bias.npy"))
+    # resize_layers_2 = Identity (no weights)
+    sd["reassemble.resize_layers.3.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/kernel.npy"))
+    sd["reassemble.resize_layers.3.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/bias.npy"))
+
+    # --- Convs (3×3, no bias) ---
+    for i in range(4):
+        sd[f"convs.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}convs_{i}/kernel.npy"))
+
+    # --- Fusion blocks ---
+    for i in range(4):
+        fb = f"{prefix}fusion_blocks_{i}/"
+        if i == 0:
+            # No residual unit, only 1 PreActResidualConvUnit
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+        else:
+            # Residual unit (index 0) + main unit (index 1)
+            sd[f"fusion_blocks.{i}.residual_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.residual_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv2/kernel.npy"))
+
+        # out_conv (Conv2d 1×1)
+        sd[f"fusion_blocks.{i}.out_conv.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{fb}Conv_0/kernel.npy"))
+        sd[f"fusion_blocks.{i}.out_conv.bias"] = _bias(
+            npy(f"{fb}Conv_0/bias.npy"))
+
+    # --- Project ---
+    sd["project.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}project/kernel.npy"))
+    sd["project.bias"] = _bias(
+        npy(f"{prefix}project/bias.npy"))
+
+    # --- Depth classification head ---
+    sd["depth_head.weight"] = _linear_kernel_flax_to_torch(
+        npy("decoder/pixel_depth_classif/kernel.npy"))
+    sd["depth_head.bias"] = _bias(
+        npy("decoder/pixel_depth_classif/bias.npy"))
+
+    zf.close()
+
+    # Load into model
+    missing, unexpected = model.load_state_dict(sd, strict=True)
+    if missing:
+        print(f"WARNING: Missing keys: {missing}")
+    if unexpected:
+        print(f"WARNING: Unexpected keys: {unexpected}")
+    print(f"Loaded DPT depth head weights ({len(sd)} tensors)")
+    return model
+
+
+def load_normals_weights(model: DPTNormalsHead, zip_path: str):
+    """Load Scenic/Flax DPT weights from a zip/npz file into PyTorch model."""
+    zf = zipfile.ZipFile(zip_path, "r")
+    npy = lambda name: _load_npy_from_zip(zf, name)
+    sd = {}
+    prefix = "decoder/dpt/"
+
+    # --- ReassembleBlocks ---
+    for i in range(4):
+        # out_projections (Conv2d 1×1)
+        sd[f"reassemble.out_projections.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/kernel.npy"))
+        sd[f"reassemble.out_projections.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/bias.npy"))
+
+        # readout_projects (Linear)
+        sd[f"reassemble.readout_projects.{i}.weight"] = _linear_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/kernel.npy"))
+        sd[f"reassemble.readout_projects.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/bias.npy"))
+
+    # resize_layers: 0=ConvTranspose, 1=ConvTranspose, 2=Identity, 3=Conv
+    sd["reassemble.resize_layers.0.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/kernel.npy"))
+    sd["reassemble.resize_layers.0.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/bias.npy"))
+    sd["reassemble.resize_layers.1.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/kernel.npy"))
+    sd["reassemble.resize_layers.1.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/bias.npy"))
+    # resize_layers_2 = Identity (no weights)
+    sd["reassemble.resize_layers.3.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/kernel.npy"))
+    sd["reassemble.resize_layers.3.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/bias.npy"))
+
+    # --- Convs (3×3, no bias) ---
+    for i in range(4):
+        sd[f"convs.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}convs_{i}/kernel.npy"))
+
+    # --- Fusion blocks ---
+    for i in range(4):
+        fb = f"{prefix}fusion_blocks_{i}/"
+        if i == 0:
+            # No residual unit, only 1 PreActResidualConvUnit
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+        else:
+            # Residual unit (index 0) + main unit (index 1)
+            sd[f"fusion_blocks.{i}.residual_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.residual_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv2/kernel.npy"))
+
+        # out_conv (Conv2d 1×1)
+        sd[f"fusion_blocks.{i}.out_conv.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{fb}Conv_0/kernel.npy"))
+        sd[f"fusion_blocks.{i}.out_conv.bias"] = _bias(
+            npy(f"{fb}Conv_0/bias.npy"))
+
+    # --- Project ---
+    sd["project.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}project/kernel.npy"))
+    sd["project.bias"] = _bias(
+        npy(f"{prefix}project/bias.npy"))
+
+    # --- Normals head ---
+    sd["normals_head.weight"] = _linear_kernel_flax_to_torch(
+        npy("decoder/pixel_normals/kernel.npy"))
+    sd["normals_head.bias"] = _bias(
+        npy("decoder/pixel_normals/bias.npy"))
+
+    zf.close()
+
+    # Load into model
+    missing, unexpected = model.load_state_dict(sd, strict=True)
+    if missing:
+        print(f"WARNING: Missing keys: {missing}")
+    if unexpected:
+        print(f"WARNING: Unexpected keys: {unexpected}")
+    print(f"Loaded DPT normals head weights ({len(sd)} tensors)")
+    return model
+
+
+def load_segmentation_weights(model: DPTSegmentationHead, zip_path: str):
+    """Load Scenic/Flax DPT weights from a zip/npz file into PyTorch model."""
+    zf = zipfile.ZipFile(zip_path, "r")
+    npy = lambda name: _load_npy_from_zip(zf, name)
+    sd = {}
+    prefix = "decoder/dpt/"
+
+    # --- ReassembleBlocks ---
+    for i in range(4):
+        # out_projections (Conv2d 1×1)
+        sd[f"reassemble.out_projections.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/kernel.npy"))
+        sd[f"reassemble.out_projections.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/out_projection_{i}/bias.npy"))
+
+        # readout_projects (Linear)
+        sd[f"reassemble.readout_projects.{i}.weight"] = _linear_kernel_flax_to_torch(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/kernel.npy"))
+        sd[f"reassemble.readout_projects.{i}.bias"] = _bias(
+            npy(f"{prefix}reassemble_blocks/readout_projects_{i}/bias.npy"))
+
+    # resize_layers: 0=ConvTranspose, 1=ConvTranspose, 2=Identity, 3=Conv
+    sd["reassemble.resize_layers.0.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/kernel.npy"))
+    sd["reassemble.resize_layers.0.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_0/bias.npy"))
+    sd["reassemble.resize_layers.1.weight"] = _conv_transpose_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/kernel.npy"))
+    sd["reassemble.resize_layers.1.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_1/bias.npy"))
+    # resize_layers_2 = Identity (no weights)
+    sd["reassemble.resize_layers.3.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/kernel.npy"))
+    sd["reassemble.resize_layers.3.bias"] = _bias(
+        npy(f"{prefix}reassemble_blocks/resize_layers_3/bias.npy"))
+
+    # --- Convs (3×3, no bias) ---
+    for i in range(4):
+        sd[f"convs.{i}.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{prefix}convs_{i}/kernel.npy"))
+
+    # --- Fusion blocks ---
+    for i in range(4):
+        fb = f"{prefix}fusion_blocks_{i}/"
+        if i == 0:
+            # No residual unit, only 1 PreActResidualConvUnit
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+        else:
+            # Residual unit (index 0) + main unit (index 1)
+            sd[f"fusion_blocks.{i}.residual_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.residual_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_0/conv2/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv1.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv1/kernel.npy"))
+            sd[f"fusion_blocks.{i}.main_unit.conv2.weight"] = _conv_kernel_flax_to_torch(
+                npy(f"{fb}PreActResidualConvUnit_1/conv2/kernel.npy"))
+
+        # out_conv (Conv2d 1×1)
+        sd[f"fusion_blocks.{i}.out_conv.weight"] = _conv_kernel_flax_to_torch(
+            npy(f"{fb}Conv_0/kernel.npy"))
+        sd[f"fusion_blocks.{i}.out_conv.bias"] = _bias(
+            npy(f"{fb}Conv_0/bias.npy"))
+
+    # --- Project ---
+    sd["project.weight"] = _conv_kernel_flax_to_torch(
+        npy(f"{prefix}project/kernel.npy"))
+    sd["project.bias"] = _bias(
+        npy(f"{prefix}project/bias.npy"))
+
+    # --- Segmentation head ---
+    sd["segmentation_head.weight"] = _linear_kernel_flax_to_torch(
+        npy("decoder/pixel_segmentation/kernel.npy"))
+    sd["segmentation_head.bias"] = _bias(
+        npy("decoder/pixel_segmentation/bias.npy"))
+
+    zf.close()
+
+    # Load into model
+    missing, unexpected = model.load_state_dict(sd, strict=True)
+    if missing:
+        print(f"WARNING: Missing keys: {missing}")
+    if unexpected:
+        print(f"WARNING: Unexpected keys: {unexpected}")
+    print(f"Loaded DPT segmentation head weights ({len(sd)} tensors)")
+    return model