Spaces:
Sleeping
Sleeping
| 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) | |