import torch import torch.nn as nn import torch.nn.functional as F class BasicBlock(nn.Module): expansion = 1 def __init__(self, in_planes, planes, stride=1): super(BasicBlock, self).__init__() self.conv1 = nn.Conv2d( in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False ) self.bn1 = nn.BatchNorm2d(planes) self.conv2 = nn.Conv2d( planes, planes, kernel_size=3, stride=1, padding=1, bias=False ) self.bn2 = nn.BatchNorm2d(planes) self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( nn.Conv2d( in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False, ), nn.BatchNorm2d(self.expansion * planes), ) def forward(self, x): out = torch.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) out += self.shortcut(x) out = torch.relu(out) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, in_planes, planes, stride=1): super(Bottleneck, self).__init__() self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False) self.bn1 = nn.BatchNorm2d(planes) self.conv2 = nn.Conv2d( planes, planes, kernel_size=3, stride=stride, padding=1, bias=False ) self.bn2 = nn.BatchNorm2d(planes) self.conv3 = nn.Conv2d( planes, self.expansion * planes, kernel_size=1, bias=False ) self.bn3 = nn.BatchNorm2d(self.expansion * planes) self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( nn.Conv2d( in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False, ), nn.BatchNorm2d(self.expansion * planes), ) def forward(self, x): out = torch.relu(self.bn1(self.conv1(x))) out = torch.relu(self.bn2(self.conv2(out))) out = self.bn3(self.conv3(out)) out += self.shortcut(x) out = torch.relu(out) return out class ResNet(nn.Module): def __init__(self, block, num_blocks, num_classes=1000): super(ResNet, self).__init__() self.in_planes = 64 self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = nn.BatchNorm2d(64) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1) self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2) self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2) self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512 * block.expansion, num_classes) def _make_layer(self, block, planes, num_blocks, stride): strides = [stride] + [1] * (num_blocks - 1) layers = [] for stride in strides: layers.append(block(self.in_planes, planes, stride)) self.in_planes = planes * block.expansion return nn.Sequential(*layers) def forward(self, x): out = torch.relu(self.bn1(self.conv1(x))) out = self.maxpool(out) out = self.layer1(out) out = self.layer2(out) out = self.layer3(out) out = self.layer4(out) out = self.avgpool(out) out = torch.flatten(out, 1) out = self.fc(out) return out def ResNet18(num_classes=1000): return ResNet(BasicBlock, [2, 2, 2, 2], num_classes) def ResNet34(num_classes=1000): return ResNet(BasicBlock, [3, 4, 6, 3], num_classes) def ResNet50(num_classes=1000): return ResNet(Bottleneck, [3, 4, 6, 3], num_classes) def ResNet101(num_classes=1000): return ResNet(Bottleneck, [3, 4, 23, 3], num_classes) def ResNet152(num_classes=1000): return ResNet(Bottleneck, [3, 8, 36, 3], num_classes) class ClassifierHead(nn.Module): def __init__(self, in_features, num_classes): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) self.max_pool = nn.AdaptiveMaxPool2d((1, 1)) self.classifier = nn.Sequential( nn.Linear(in_features * 2, 1024), nn.BatchNorm1d(1024), nn.ReLU(), nn.Dropout(0.5), nn.Linear(1024, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(0.3), nn.Linear(512, num_classes), ) def forward(self, x): avg_pooled = self.avg_pool(x).flatten(1) max_pooled = self.max_pool(x).flatten(1) features = torch.cat([avg_pooled, max_pooled], dim=1) return self.classifier(features) class ResNetUNet(ResNet): def __init__(self, block, num_blocks, num_classes=1000): super().__init__(block, num_blocks, num_classes) # Calculate encoder channel sizes self.enc_channels = [ 64, 64 * block.expansion, 128 * block.expansion, 256 * block.expansion, 512 * block.expansion, ] # Replace t_max_avg_pooling with standard avgpool in_features = 512 * block.expansion self.classifier_head = ClassifierHead(in_features, num_classes) # Decoder layers remain the same self.decoder5 = nn.Sequential( nn.Conv2d(2048 + 1024, 1024, 3, padding=1), nn.BatchNorm2d(1024), nn.ReLU(inplace=True), nn.Conv2d(1024, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), ) self.decoder4 = nn.Sequential( nn.Conv2d(512 + 512, 512, 3, padding=1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.Conv2d(512, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), ) self.decoder3 = nn.Sequential( nn.Conv2d(256 + 256, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Conv2d(256, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), ) self.decoder2 = nn.Sequential( nn.Conv2d(128 + 64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(128, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), ) self.final_conv = nn.Sequential( nn.Conv2d(64, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(32, 1, 1), nn.Sigmoid(), ) def forward(self, x): input_size = x.shape[-2:] # Encoder path x = torch.relu(self.bn1(self.conv1(x))) e1 = self.maxpool(x) e2 = self.layer1(e1) e3 = self.layer2(e2) e4 = self.layer3(e3) e5 = self.layer4(e4) # Get segmentation first e4_resized = F.interpolate( e4, size=e5.shape[-2:], mode="bilinear", align_corners=True ) d5 = self.decoder5(torch.cat([e5, e4_resized], dim=1)) e3_resized = F.interpolate( e3, size=d5.shape[-2:], mode="bilinear", align_corners=True ) d4 = self.decoder4(torch.cat([d5, e3_resized], dim=1)) e2_resized = F.interpolate( e2, size=d4.shape[-2:], mode="bilinear", align_corners=True ) d3 = self.decoder3(torch.cat([d4, e2_resized], dim=1)) e1_resized = F.interpolate( e1, size=d3.shape[-2:], mode="bilinear", align_corners=True ) d2 = self.decoder2(torch.cat([d3, e1_resized], dim=1)) seg_out = self.final_conv(d2) seg_out = F.interpolate( seg_out, size=input_size, mode="bilinear", align_corners=True ) # Use segmentation to mask features before classification # Upsample segmentation mask to match feature size attention_mask = F.interpolate( seg_out, size=e5.shape[2:], mode="bilinear", align_corners=True ) # Apply attention mask to features attended_features = e5 * (0.25 + attention_mask) # Use new classifier head cls_out = self.classifier_head(attended_features) return cls_out, seg_out # Helper functions without K and T parameters def ResNet18UNet(num_classes=1000): return ResNetUNet(BasicBlock, [2, 2, 2, 2], num_classes) def ResNet34UNet(num_classes=1000): return ResNetUNet(BasicBlock, [3, 4, 6, 3], num_classes) def ResNet50UNet(num_classes=1000): return ResNetUNet(Bottleneck, [3, 4, 6, 3], num_classes) def ResNet101UNet(num_classes=1000): return ResNetUNet(Bottleneck, [3, 4, 23, 3], num_classes) def ResNet152UNet(num_classes=1000): return ResNetUNet(Bottleneck, [3, 8, 36, 3], num_classes)