Spaces:
Running
Running
File size: 4,269 Bytes
917aae9 | 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 109 110 111 112 113 114 115 116 117 | 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 |