Spaces:
Running on Zero
Running on Zero
File size: 2,987 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 | import torch.nn as nn
import torch.nn.functional as F
from .dpt_backbone import DPTRefineNetStack
class ProgressiveOutputHead(nn.Module):
"""Progressive channel-reduction head (GroupNorm for AMP/small-batch stability)."""
def __init__(self, in_features=256, num_classes=21, intermediate_features=32, groups=8):
super().__init__()
self.num_classes = num_classes
self.output_conv1 = nn.Sequential(
nn.Conv2d(in_features, in_features // 2, 3, padding=1, bias=False),
nn.GroupNorm(num_groups=groups, num_channels=in_features // 2),
nn.ReLU(inplace=True),
)
self.output_conv2 = nn.Sequential(
nn.Conv2d(in_features // 2, intermediate_features, 3, padding=1, bias=False),
nn.GroupNorm(num_groups=max(1, min(groups, intermediate_features)), num_channels=intermediate_features),
nn.ReLU(inplace=True),
nn.Conv2d(intermediate_features, num_classes, 1),
)
self.final_activation = nn.Identity() # sigmoid/softmax applied in the loss
self._init_weights()
def _init_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
# Background-prior bias on the final logits to avoid early sigmoid overflow.
final_conv = self.output_conv2[-1]
if isinstance(final_conv, nn.Conv2d) and final_conv.bias is not None:
nn.init.constant_(final_conv.bias, -2.2 if self.num_classes == 1 else -0.5)
def forward(self, x):
return self.final_activation(self.output_conv2(self.output_conv1(x)))
class OriginalDPTSegmentationDecoder(DPTRefineNetStack):
def __init__(self, in_channels, num_classes=21, features=256, target_size=(512, 512)):
super().__init__(features=features, use_bn=False)
self.target_size = target_size
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.resize_layers = nn.ModuleList([
nn.ConvTranspose2d(out_channels[0], out_channels[0], kernel_size=4, stride=4),
nn.ConvTranspose2d(out_channels[1], out_channels[1], kernel_size=2, stride=2),
nn.Identity(),
nn.Conv2d(out_channels[3], out_channels[3], kernel_size=3, stride=2, padding=1),
])
self.output_head = ProgressiveOutputHead(
in_features=features, num_classes=num_classes, intermediate_features=32, groups=8)
def forward(self, features):
proj = [p(f) for p, f in zip(self.projects, features)]
resized = [r(f) for r, f in zip(self.resize_layers, proj)]
logits = self.output_head(self.fuse(resized))
if self.target_size is not None:
logits = F.interpolate(logits, size=self.target_size, mode='bilinear', align_corners=True)
return logits
|