Spaces:
Running
Running
File size: 2,693 Bytes
f5288df |
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 |
import torch.nn as nn
from monai.networks.nets import resnet101, resnet50, resnet18, ViT
import torch
## resnet50 architecture, FC layers converted to I
class ResNet50_3D(nn.Module):
def __init__(self):
super(ResNet50_3D, self).__init__()
resnet = resnet50(pretrained=False) # assuming you're not using a pretrained model
resnet.conv1 = nn.Conv3d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
hidden_dim = resnet.fc.in_features
self.backbone = resnet
self.backbone.fc = nn.Identity()
def forward(self, x):
x = self.backbone(x)
return x
class Classifier(nn.Module):
""" Classifier class with FC layer and single output neuron """
def __init__(self, d_model, hidden_dim=1024, num_classes=1):
super(Classifier, self).__init__()
self.fc = nn.Linear(d_model, num_classes)
def forward(self, x):
x = self.fc(x)
return x
class Backbone(nn.Module):
""" ResNet 3D Backbone"""
def __init__(self):
super(Backbone, self).__init__()
resnet = resnet50(pretrained=False) # assuming you're not using a pretrained model
resnet.conv1 = nn.Conv3d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
hidden_dim = resnet.fc.in_features
self.backbone = resnet
self.backbone.fc = nn.Identity()
def forward(self, x):
x = self.backbone(x)
return x
class SingleScanModel(nn.Module):
""" End to end model with backbone and classifier"""
def __init__(self, backbone, classifier):
super(SingleScanModel, self).__init__()
self.backbone = backbone
self.classifier = classifier
self.dropout = nn.Dropout(p=0.2)
def forward(self, x):
x = self.backbone(x)
x = self.dropout(x)
x = self.classifier(x)
return x
class SingleScanModelBP(nn.Module):
""" End to end model with backbone and classifier that takes 2 input scans at once"""
def __init__(self, backbone, classifier):
super(SingleScanModelBP, self).__init__()
self.backbone = backbone
self.classifier = classifier
self.dropout = nn.Dropout(p=0.2)
self.bilinear_pooling = nn.Bilinear(in1_features=2048, in2_features=2048, out_features=512)
def forward(self, x):
x = [self.backbone(scan) for scan in x.split(1, dim=1)]
features = torch.stack(x, dim=1).squeeze(2)
merged_features = torch.mean(features, dim=1) # Shape: (batch_size, feature_dim)
merged_features = self.dropout(merged_features)
output = self.classifier(merged_features)
return output |