| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | import torchvision.models as models |
| | import torch |
| | from einops import rearrange |
| |
|
| |
|
| | class ModelRes_ft(nn.Module): |
| | def __init__( |
| | self, |
| | res_base_model, |
| | out_size, |
| | imagenet_pretrain=False, |
| | linear_probe=False, |
| | use_base=True, |
| | ): |
| | super(ModelRes_ft, self).__init__() |
| | self.resnet_dict = { |
| | "resnet18": models.resnet18(pretrained=imagenet_pretrain), |
| | "resnet50": models.resnet50(pretrained=imagenet_pretrain), |
| | } |
| | resnet = self._get_res_basemodel(res_base_model) |
| | self.use_base = use_base |
| |
|
| | if not self.use_base: |
| | num_ftrs = int(resnet.fc.in_features / 2) |
| | self.res_features = nn.Sequential(*list(resnet.children())[:-3]) |
| | self.res_l1_anatomy = nn.Linear(num_ftrs, num_ftrs) |
| | self.res_l2_anatomy = nn.Linear(num_ftrs, 256) |
| | self.res_l1_pathology = nn.Linear(num_ftrs, num_ftrs) |
| | self.res_l2_pathology = nn.Linear(num_ftrs, 256) |
| |
|
| | self.mask_generator = nn.Linear(num_ftrs, num_ftrs) |
| | self.back = nn.Linear(256, num_ftrs) |
| | self.last_res = nn.Sequential(*list(resnet.children())[-3:-1]) |
| | else: |
| | self.res_features = nn.Sequential(*list(resnet.children())[:-1]) |
| | self.res_out = nn.Linear(int(resnet.fc.in_features), out_size) |
| |
|
| | def _get_res_basemodel(self, res_model_name): |
| | try: |
| | res_model = self.resnet_dict[res_model_name] |
| | print("Image feature extractor:", res_model_name) |
| | return res_model |
| | except: |
| | raise ( |
| | "Invalid model name. Check the config file and pass one of: resnet18 or resnet50" |
| | ) |
| |
|
| | def image_encoder(self, xis): |
| | |
| | """ |
| | 16 torch.Size([16, 1024, 14, 14]) |
| | torch.Size([16, 196, 1024]) |
| | torch.Size([3136, 1024]) |
| | torch.Size([16, 196, 256]) |
| | """ |
| | batch_size = xis.shape[0] |
| | res_fea = self.res_features(xis) |
| | res_fea = rearrange(res_fea, "b d n1 n2 -> b (n1 n2) d") |
| | x = rearrange(res_fea, "b n d -> (b n) d") |
| | mask = self.mask_generator(x) |
| | x_pathology = mask * x |
| | x_pathology = self.res_l1_pathology(x_pathology) |
| | x_pathology = F.relu(x_pathology) |
| |
|
| | x_pathology = self.res_l2_pathology(x_pathology) |
| |
|
| | out_emb_pathology = rearrange(x_pathology, "(b n) d -> b n d", b=batch_size) |
| | out_emb_pathology = self.back(out_emb_pathology) |
| | out_emb_pathology = rearrange(out_emb_pathology, "b (n1 n2) d -> b d n1 n2", n1=14, n2=14) |
| | out_emb_pathology = self.last_res(out_emb_pathology) |
| | out_emb_pathology = out_emb_pathology.squeeze() |
| |
|
| | return out_emb_pathology |
| |
|
| | def forward(self, img, linear_probe=False): |
| | if self.use_base: |
| | x = self.res_features(img) |
| | else: |
| | x = self.image_encoder(img) |
| |
|
| | x = x.squeeze() |
| | if linear_probe: |
| | return x |
| | else: |
| | x = self.res_out(x) |
| | return x |
| |
|
