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)