Spaces:
Running on Zero
Running on Zero
| """DPT-style depth decoders that consume spatial feature maps. | |
| Both heads take a list of 4 spatial feature maps (FLUX + DINO + concepts) with | |
| possibly different channel counts, run the shared DPT RefineNet cascade, and | |
| upsample to a depth map: | |
| - DPTHeadSpatial : single 2x upsample in the head (bilinear resize to GT after). | |
| - DPTHeadHighRes : 4x learned 2x upsampling (16x) with optional image skip connections. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from .dpt_backbone import DPTRefineNetStack | |
| class DPTHeadSpatial(DPTRefineNetStack): | |
| def __init__(self, in_channels=[3840, 3840, 3840, 3840], features=256, | |
| num_classes=1, use_bn=False, head_features=32): | |
| super().__init__(features=features, use_bn=use_bn) | |
| self.head_features = head_features | |
| self.num_classes = num_classes | |
| out_channels = [256, 512, 1024, 1024] | |
| self.projects = nn.ModuleList([nn.Conv2d(in_ch, oc, 1) for in_ch, oc in zip(in_channels, out_channels)]) | |
| self.pre_fuse = nn.ModuleList([ | |
| nn.Sequential(nn.GroupNorm(1, oc), nn.Conv2d(oc, oc, 1), nn.ReLU(inplace=True)) | |
| for oc in out_channels]) | |
| self.head = nn.Sequential( | |
| nn.Conv2d(features, features // 2, 3, padding=1), | |
| nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True), | |
| nn.Conv2d(features // 2, head_features, 3, padding=1), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(head_features, num_classes, 1), | |
| ) | |
| def forward(self, features): | |
| resized = [pf(proj(f)) for f, proj, pf in zip(features, self.projects, self.pre_fuse)] | |
| return self.head(self.fuse(resized)) | |
| class DPTHeadHighRes(DPTRefineNetStack): | |
| """High-resolution variant: progressive 4x (2x) learned upsampling = 16x total, | |
| with optional skip connections from the original image.""" | |
| def __init__(self, in_channels=[3840, 3840, 3840, 3840], features=256, | |
| num_classes=1, use_bn=False, head_features=64, use_skip_connections=True): | |
| super().__init__(features=features, use_bn=use_bn) | |
| self.head_features = head_features | |
| self.num_classes = num_classes | |
| self.use_skip_connections = use_skip_connections | |
| out_channels = [256, 512, 1024, 1024] | |
| self.projects = nn.ModuleList([nn.Conv2d(in_ch, oc, 1) for in_ch, oc in zip(in_channels, out_channels)]) | |
| self.pre_fuse = nn.ModuleList([ | |
| nn.Sequential(nn.GroupNorm(1, oc), nn.Conv2d(oc, oc, 1), nn.ReLU(inplace=True)) | |
| for oc in out_channels]) | |
| if use_skip_connections: | |
| def skip_enc(stride, out_ch): | |
| return nn.Sequential( | |
| nn.Conv2d(3, out_ch, 3, stride=stride, padding=1), nn.ReLU(inplace=True), | |
| nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.ReLU(inplace=True)) | |
| self.skip_enc_2x = skip_enc(2, 16) | |
| self.skip_enc_4x = skip_enc(4, 16) | |
| self.skip_enc_8x = skip_enc(8, 16) | |
| self.skip_enc_full = skip_enc(1, 8) | |
| def up_block(c_in, c_out): | |
| return nn.Sequential( | |
| nn.Conv2d(c_in, c_out, 3, padding=1), nn.ReLU(inplace=True), | |
| nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True), | |
| nn.Conv2d(c_out, c_out, 3, padding=1), nn.ReLU(inplace=True)) | |
| skip = 16 if use_skip_connections else 0 | |
| self.up1 = up_block(features, features // 2) | |
| self.fuse1 = nn.Conv2d(features // 2 + skip, features // 2, 1) if use_skip_connections else None | |
| self.up2 = up_block(features // 2, features // 4) | |
| self.fuse2 = nn.Conv2d(features // 4 + skip, features // 4, 1) if use_skip_connections else None | |
| self.up3 = up_block(features // 4, head_features) | |
| self.fuse3 = nn.Conv2d(head_features + skip, head_features, 1) if use_skip_connections else None | |
| self.up4 = up_block(head_features, head_features // 2) | |
| self.fuse4 = nn.Conv2d(head_features // 2 + (8 if use_skip_connections else 0), head_features // 2, 1) if use_skip_connections else None | |
| self.output = nn.Conv2d(head_features // 2, num_classes, 3, padding=1) | |
| def _apply_skip(self, x, fuse, skip): | |
| if not self.use_skip_connections or skip is None: | |
| return x | |
| if skip.shape[-2:] != x.shape[-2:]: | |
| skip = F.interpolate(skip, size=x.shape[-2:], mode='bilinear', align_corners=True) | |
| return fuse(torch.cat([x, skip], dim=1)) | |
| def forward(self, features, image=None): | |
| if self.use_skip_connections and image is not None: | |
| skip_8x, skip_4x = self.skip_enc_8x(image), self.skip_enc_4x(image) | |
| skip_2x, skip_full = self.skip_enc_2x(image), self.skip_enc_full(image) | |
| else: | |
| skip_8x = skip_4x = skip_2x = skip_full = None | |
| resized = [pf(proj(f)) for f, proj, pf in zip(features, self.projects, self.pre_fuse)] | |
| path_1 = self.fuse(resized, keep_layer1_size=True) # keep patch resolution | |
| x = self._apply_skip(self.up1(path_1), self.fuse1, skip_8x) | |
| x = self._apply_skip(self.up2(x), self.fuse2, skip_4x) | |
| x = self._apply_skip(self.up3(x), self.fuse3, skip_2x) | |
| x = self._apply_skip(self.up4(x), self.fuse4, skip_full) | |
| return self.output(x) | |