Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torchvision.models as models | |
| from transformers import SegformerForSemanticSegmentation | |
| # ========================================== | |
| # 1. CBAM Components (unchanged) | |
| # ========================================== | |
| class ChannelAttention(nn.Module): | |
| def __init__(self, in_planes, ratio=4): | |
| super(ChannelAttention, self).__init__() | |
| self.avg_pool = nn.AdaptiveAvgPool2d(1) | |
| self.max_pool = nn.AdaptiveMaxPool2d(1) | |
| self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False) | |
| self.relu1 = nn.ReLU() | |
| self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False) | |
| self.sigmoid = nn.Sigmoid() | |
| def forward(self, x): | |
| avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x)))) | |
| max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x)))) | |
| out = avg_out + max_out | |
| return self.sigmoid(out) | |
| class SpatialAttention(nn.Module): | |
| def __init__(self, kernel_size=7): | |
| super(SpatialAttention, self).__init__() | |
| assert kernel_size in (3, 7), 'kernel size must be 3 or 7' | |
| padding = 3 if kernel_size == 7 else 1 | |
| self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) | |
| self.sigmoid = nn.Sigmoid() | |
| def forward(self, x): | |
| avg_out = torch.mean(x, dim=1, keepdim=True) | |
| max_out, _ = torch.max(x, dim=1, keepdim=True) | |
| x_cat = torch.cat([avg_out, max_out], dim=1) | |
| out = self.conv1(x_cat) | |
| return self.sigmoid(out) | |
| class CBAM(nn.Module): | |
| def __init__(self, channels=19, reduction=4, spatial_kernel_size=7): | |
| super(CBAM, self).__init__() | |
| self.ca = ChannelAttention(channels, ratio=reduction) | |
| self.sa = SpatialAttention(kernel_size=spatial_kernel_size) | |
| def forward(self, x): | |
| x = x * self.ca(x) | |
| x = x * self.sa(x) | |
| return x | |
| # ========================================== | |
| # 2. Main Model → MODIFIED: ONLY O3 (1 output) | |
| # ========================================== | |
| class PollutionDifferenceModel(nn.Module): | |
| def __init__(self, num_classes=19, pollution_dims=1): # ✅ 5 → 1 | |
| super(PollutionDifferenceModel, self).__init__() | |
| self.backbone = SegformerForSemanticSegmentation.from_pretrained( | |
| "nvidia/segformer-b5-finetuned-cityscapes-1024-1024", | |
| use_safetensors=True | |
| ) | |
| for param in self.backbone.parameters(): | |
| param.requires_grad = False | |
| self.cbam = CBAM(channels=num_classes) | |
| self.convnext = models.convnext_tiny(weights=models.ConvNeXt_Tiny_Weights.DEFAULT) | |
| for param in self.convnext.parameters(): | |
| param.requires_grad = True | |
| original_stem = self.convnext.features[0][0] | |
| self.convnext.features[0][0] = nn.Conv2d( | |
| in_channels=num_classes, | |
| out_channels=original_stem.out_channels, | |
| kernel_size=original_stem.kernel_size, | |
| stride=original_stem.stride, | |
| padding=original_stem.padding, | |
| bias=(original_stem.bias is not None) | |
| ) | |
| nn.init.kaiming_normal_(self.convnext.features[0][0].weight, mode='fan_out', nonlinearity='relu') | |
| self.convnext.classifier[2] = nn.Identity() | |
| convnext_out_dim = 768 | |
| # ✅ Final output = 1 (O3 only) | |
| self.mlp_decoder = nn.Sequential( | |
| nn.Linear(convnext_out_dim, 256), | |
| nn.GELU(), | |
| nn.Dropout(0.3), | |
| nn.Linear(256, 64), | |
| nn.GELU(), | |
| nn.Dropout(0.3), | |
| nn.Linear(64, pollution_dims) | |
| ) | |
| def get_semantic_map(self, x): | |
| outputs = self.backbone(pixel_values=x) | |
| seg_logits = outputs.logits | |
| probs = torch.nn.functional.softmax(seg_logits, dim=1) | |
| return probs | |
| def forward(self, imgA, imgB): | |
| mapA = self.get_semantic_map(imgA) | |
| mapB = self.get_semantic_map(imgB) | |
| diff_map = mapA - mapB | |
| attended_diff_map = self.cbam(diff_map) | |
| z_diff = self.convnext(attended_diff_map) | |
| pred_pollution_delta = self.mlp_decoder(z_diff) | |
| return pred_pollution_delta |