Spaces:
Running on Zero
Running on Zero
| 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 | |