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) | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class SAM(nn.Module): | |
| def __init__(self, bias=False): | |
| super(SAM, self).__init__() | |
| self.bias = bias | |
| self.conv = nn.Conv2d( | |
| in_channels=2, | |
| out_channels=1, | |
| kernel_size=7, | |
| stride=1, | |
| padding=3, | |
| dilation=1, | |
| bias=self.bias, | |
| ) | |
| def forward(self, x): | |
| max = torch.max(x, 1)[0].unsqueeze(1) | |
| avg = torch.mean(x, 1).unsqueeze(1) | |
| concat = torch.cat((max, avg), dim=1) | |
| output = self.conv(concat) | |
| output = F.sigmoid(output) * x | |
| return output | |
| class CAM(nn.Module): | |
| def __init__(self, channels, r): | |
| super(CAM, self).__init__() | |
| self.channels = channels | |
| self.r = r | |
| self.linear = nn.Sequential( | |
| nn.Linear( | |
| in_features=self.channels, | |
| out_features=self.channels // self.r, | |
| bias=True, | |
| ), | |
| nn.ReLU(inplace=True), | |
| nn.Linear( | |
| in_features=self.channels // self.r, | |
| out_features=self.channels, | |
| bias=True, | |
| ), | |
| ) | |
| def forward(self, x): | |
| max = F.adaptive_max_pool2d(x, output_size=1) | |
| avg = F.adaptive_avg_pool2d(x, output_size=1) | |
| b, c, _, _ = x.size() | |
| linear_max = self.linear(max.view(b, c)).view(b, c, 1, 1) | |
| linear_avg = self.linear(avg.view(b, c)).view(b, c, 1, 1) | |
| output = linear_max + linear_avg | |
| output = F.sigmoid(output) * x | |
| return output | |
| class CBAM(nn.Module): | |
| def __init__(self, channels, r): | |
| super(CBAM, self).__init__() | |
| self.channels = channels | |
| self.r = r | |
| self.sam = SAM(bias=False) | |
| self.cam = CAM(channels=self.channels, r=self.r) | |
| def forward(self, x): | |
| output = self.cam(x) | |
| output = self.sam(output) | |
| return output + x | |
| 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) | |
| # Get the expansion factor | |
| expansion = block.expansion | |
| # Calculate encoder channel sizes | |
| self.enc_channels = [ | |
| 64, | |
| 64 * block.expansion, | |
| 128 * block.expansion, | |
| 256 * block.expansion, | |
| 512 * block.expansion, | |
| ] | |
| in_features = 512 * block.expansion | |
| self.classifier_head = ClassifierHead(in_features, num_classes) | |
| self.cbam = CBAM(channels=512 * block.expansion, r=16) | |
| # Calculate encoder channel sizes | |
| self.decoder5 = nn.Sequential( | |
| nn.Conv2d((512 * expansion) + (256 * expansion), 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.decoder4 = nn.Sequential( | |
| nn.Conv2d(256 + (128 * expansion), 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.decoder3 = nn.Sequential( | |
| nn.Conv2d(128 + (64 * expansion), 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.decoder2 = nn.Sequential( | |
| nn.Conv2d(64 + 64, 64, 3, padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(64, 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 | |
| ) | |
| attended_features = self.cbam(e5) | |
| # 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 = attended_features * (0.25 + attention_mask) | |
| cls_out = self.classifier_head(attended_features) | |
| return cls_out, seg_out | |
| 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) |