Add PyTorch DPT decoder module with shared head and task-specific decoders

[bingyic]

Summary

Add pytorch/decoders.py implementing the DPT dense prediction decoder hierarchy in PyTorch, mirroring the Scenic/Flax decoder structure.

Architecture

The module follows a shared-head pattern where a common DPT backbone is subclassed for each task:

• DPTHead: Shared backbone performing feature reassembly at multiple scales (ReassembleBlocks) and bottom-up fusion (FeatureFusionBlocks).
• Decoder: Base class wrapping a DPTHead with a common interface.
• SegmentationDecoder(Decoder): Per-pixel class logits (num_classes output channels).
• DepthDecoder(Decoder): Monocular depth prediction (1-channel output).
• NormalsDecoder(Decoder): Surface normals prediction (3-channel output).

Weight Loading

load_decoder_weights() loads pre-converted .npz checkpoints (produced by scripts/convert_segmentation_checkpoint.py) with backward-compatible key remapping from the legacy flat format.

Usage

from tips.pytorch.decoders import SegmentationDecoder, load_decoder_weights

decoder = SegmentationDecoder(num_classes=150, input_embed_dim=1024)
load_decoder_weights(decoder, 'path/to/checkpoint.npz')
logits = decoder(intermediate_features, image_size=(480, 640))

[koertchen]

please remove this internal detail

[koertchen]

nit. remove "(Refactored as Thin Wrappers)"

[koertchen]

Remove mention of our scenic/flax internals.

[koertchen]

also here

[koertchen]

DepthDecoder actually is classification based, not regression:

class DepthDecoder(Decoder):
"""Decoder for monocular depth prediction using classification bins."""
def init(
self, min_depth: float = 0.001, max_depth: float = 10.0, **kwargs
) -> None:
super().init(out_channels=256, **kwargs)
self.min_depth = min_depth
self.max_depth = max_depth
self.register_buffer(
"bin_centers", torch.linspace(min_depth, max_depth, 256)
)
def forward(
self,
intermediate_features: List[Tuple[torch.Tensor, torch.Tensor]],
image_size: Optional[Tuple[int, int]] = None,
) -> torch.Tensor:
logits = super().forward(intermediate_features, image_size=image_size)
# Apply ReLU and shift
logits = torch.relu(logits) + self.min_depth
# Normalize to probabilities along the channel dimension
probs = logits / torch.sum(logits, dim=1, keepdim=True)
# Compute expectation: sum(prob * bin_center)
depth_map = torch.einsum(
"bchw,c->bhw", probs, self.bin_centers.to(logits.device)
)
return depth_map.unsqueeze(1)