Spaces:
Running on Zero
Running on Zero
File size: 5,320 Bytes
2267636 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """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)
|