jkottu/image-captioning-chest-xrays
0
1import torch2import torch.nn as nn3import torchvision4import numpy as np5from torch.autograd import Variable6import torchvision.models as models7import transformers8import torchvision.transforms9 10import torchxrayvision as xrv11from transformers import ViTModel, ViTConfig12 13 14 15class VisualFeatureExtractor(nn.Module):16 def __init__(self, model_name='densenet201', pretrained=False):17 super(VisualFeatureExtractor, self).__init__()18 self.model_name = 'chexnet'19 self.pretrained = pretrained20 self.model, self.out_features, self.avg_func, self.bn, self.linear = self.__get_model()21 self.activation = nn.ReLU()22 23 def __get_model(self):24 model = None25 out_features = None26 func = None27 28 if self.model_name == 'resnet152':29 resnet = models.resnet152(pretrained=self.pretrained)30 modules = list(resnet.children())[:-2]31 model = nn.Sequential(*modules)32 out_features = resnet.fc.in_features33 func = torch.nn.AvgPool2d(kernel_size=7, stride=1, padding=0)34 35 36 elif self.model_name == 'densenet201':37 densenet = models.densenet201(pretrained=self.pretrained)38 modules = list(densenet.features)39 model = nn.Sequential(*modules)40 func = torch.nn.AvgPool2d(kernel_size=7, stride=1, padding=0)41 out_features = densenet.classifier.in_features42 43 44 elif self.model_name == 'chexnet':45 print("vit chest xray pretrained model loading")46 # Load the Vision Transformer (ViT) model configuration47 config = ViTConfig.from_pretrained('nickmuchi/vit-finetuned-chest-xray-pneumonia')48 49 # Initialize the ViT model with the specific configuration50 vit_model = ViTModel(config)51 52 # Load the state dict specifically, excluding 'classifier.bias', 'classifier.weight'53 state_dict = torch.load('model/pytorch_model.bin', map_location=torch.device('cpu'))54 state_dict = {k: v for k, v in state_dict.items() if not k.startswith('classifier')}55 vit_model.load_state_dict(state_dict, strict=False)56 57 model = vit_model58 out_features = config.hidden_size59 60 linear = nn.Linear(in_features=out_features, out_features=out_features)61 bn = nn.BatchNorm1d(num_features=out_features, momentum=0.1)62 63 return model, out_features, func, bn, linear64 65 def forward(self, images):66 """67 :param images: Input images68 :return: visual_features, avg_features69 """70 model_output = self.model(images)71 72 # Extract the pooler_output73 74 pooler_output = model_output.pooler_output75 76 # Apply the linear layer, batch normalization, and activation77 avg_features = self.activation(self.bn(self.linear(pooler_output)))78 79 return model_output.last_hidden_state, avg_features 80 81 # def forward(self, images):82 # """83 # :param images:84 # :return:85 # """86 # visual_features = self.model(images)87 88 # avg_features = self.avg_func(visual_features).squeeze()89 # # avg_features = self.activation(self.bn(self.linear(visual_features)))90 91 # return visual_features, avg_features92 93 94class MLC(nn.Module):95 def __init__(self,96 classes=210,97 sementic_features_dim=512,98 fc_in_features=2048,99 k=10,100 ):101 super(MLC, self).__init__()102 pretrained_model_name="nickmuchi/vit-finetuned-chest-xray-pneumonia"103 vit_config = ViTConfig.from_pretrained(pretrained_model_name)104 self.vit = ViTModel(vit_config)105 106 # Adjust the classifier to your number of classes107 self.classifier = nn.Linear(in_features=vit_config.hidden_size, out_features=classes)108 self.embed = nn.Embedding(classes, sementic_features_dim)109 self.k = k110 self.sigmoid = nn.Sigmoid()111 self.__init_weight()112 113 def __init_weight(self):114 nn.init.xavier_uniform_(self.classifier.weight)115 if self.classifier.bias is not None:116 self.classifier.bias.data.fill_(0)117 118 def forward(self, avg_features):119 120 121 tags = self.sigmoid(self.classifier(avg_features))122 semantic_features = self.embed(torch.topk(tags, self.k)[1])123 return tags, semantic_features124 125# class MLC(nn.Module):126# def __init__(self,127# classes=210,128# sementic_features_dim=512,129# fc_in_features=2048,130# k=10):131# super(MLC, self).__init__()132# self.classifier = nn.Linear(in_features=fc_in_features, out_features=classes)133# self.embed = nn.Embedding(classes, sementic_features_dim)134# self.k = k135# self.sigmoid = nn.Sigmoid()136# self.__init_weight()137 138# def __init_weight(self):139# # Example: Initialize weights with a different strategy140# nn.init.xavier_uniform_(self.classifier.weight)141# if self.classifier.bias is not None:142# self.classifier.bias.data.fill_(0)143 144# def forward(self, avg_features):145# tags = self.sigmoid(self.classifier(avg_features))146# semantic_features = self.embed(torch.topk(tags, self.k)[1])147# return tags, semantic_features148 149 150class CoAttention(nn.Module):151 def __init__(self,152 version='v1',153 embed_size=512,154 hidden_size=512,155 visual_size=2048,156 k=10,157 momentum=0.1):158 super(CoAttention, self).__init__()159 self.version = version160 self.W_v = nn.Linear(in_features=visual_size, out_features=visual_size)161 self.bn_v = nn.BatchNorm1d(num_features=visual_size, momentum=momentum)162 163 self.W_v_h = nn.Linear(in_features=hidden_size, out_features=visual_size)164 self.bn_v_h = nn.BatchNorm1d(num_features=visual_size, momentum=momentum)165 166 self.W_v_att = nn.Linear(in_features=visual_size, out_features=visual_size)167 self.bn_v_att = nn.BatchNorm1d(num_features=visual_size, momentum=momentum)168 169 self.W_a = nn.Linear(in_features=hidden_size, out_features=hidden_size)170 self.bn_a = nn.BatchNorm1d(num_features=k, momentum=momentum)171 172 self.W_a_h = nn.Linear(in_features=hidden_size, out_features=hidden_size)173 self.bn_a_h = nn.BatchNorm1d(num_features=1, momentum=momentum)174 175 self.W_a_att = nn.Linear(in_features=hidden_size, out_features=hidden_size)176 self.bn_a_att = nn.BatchNorm1d(num_features=k, momentum=momentum)177 178 # self.W_fc = nn.Linear(in_features=visual_size, out_features=embed_size) # for v3179 self.W_fc = nn.Linear(in_features=visual_size + hidden_size, out_features=embed_size)180 self.bn_fc = nn.BatchNorm1d(num_features=embed_size, momentum=momentum)181 182 self.tanh = nn.Tanh()183 self.softmax = nn.Softmax()184 185 self.__init_weight()186 187 def __init_weight(self):188 self.W_v.weight.data.uniform_(-0.1, 0.1)189 self.W_v.bias.data.fill_(0)190 191 self.W_v_h.weight.data.uniform_(-0.1, 0.1)192 self.W_v_h.bias.data.fill_(0)193 194 self.W_v_att.weight.data.uniform_(-0.1, 0.1)195 self.W_v_att.bias.data.fill_(0)196 197 self.W_a.weight.data.uniform_(-0.1, 0.1)198 self.W_a.bias.data.fill_(0)199 200 self.W_a_h.weight.data.uniform_(-0.1, 0.1)201 self.W_a_h.bias.data.fill_(0)202 203 self.W_a_att.weight.data.uniform_(-0.1, 0.1)204 self.W_a_att.bias.data.fill_(0)205 206 self.W_fc.weight.data.uniform_(-0.1, 0.1)207 self.W_fc.bias.data.fill_(0)208 209 def forward(self, avg_features, semantic_features, h_sent):210 if self.version == 'v1':211 return self.v1(avg_features, semantic_features, h_sent)212 elif self.version == 'v2':213 return self.v2(avg_features, semantic_features, h_sent)214 elif self.version == 'v3':215 return self.v3(avg_features, semantic_features, h_sent)216 elif self.version == 'v4':217 return self.v4(avg_features, semantic_features, h_sent)218 elif self.version == 'v5':219 return self.v5(avg_features, semantic_features, h_sent)220 221 def v1(self, avg_features, semantic_features, h_sent) -> object:222 """223 only training224 :rtype: object225 """226 W_v = self.bn_v(self.W_v(avg_features))227 W_v_h = self.bn_v_h(self.W_v_h(h_sent.squeeze(1)))228 229 alpha_v = self.softmax(self.bn_v_att(self.W_v_att(self.tanh(W_v + W_v_h))))230 v_att = torch.mul(alpha_v, avg_features)231 232 W_a_h = self.bn_a_h(self.W_a_h(h_sent))233 W_a = self.bn_a(self.W_a(semantic_features))234 alpha_a = self.softmax(self.bn_a_att(self.W_a_att(self.tanh(torch.add(W_a_h, W_a)))))235 a_att = torch.mul(alpha_a, semantic_features).sum(1)236 237 ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))238 239 return ctx, alpha_v, alpha_a240 241 def v2(self, avg_features, semantic_features, h_sent) -> object:242 """243 no bn244 :rtype: object245 """246 W_v = self.W_v(avg_features)247 W_v_h = self.W_v_h(h_sent.squeeze(1))248 249 alpha_v = self.softmax(self.W_v_att(self.tanh(W_v + W_v_h)))250 v_att = torch.mul(alpha_v, avg_features)251 252 W_a_h = self.W_a_h(h_sent)253 W_a = self.W_a(semantic_features)254 alpha_a = self.softmax(self.W_a_att(self.tanh(torch.add(W_a_h, W_a))))255 a_att = torch.mul(alpha_a, semantic_features).sum(1)256 257 ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))258 259 return ctx, alpha_v, alpha_a260 261 def v3(self, avg_features, semantic_features, h_sent) -> object:262 """263 264 :rtype: object265 """266 W_v = self.bn_v(self.W_v(avg_features))267 W_v_h = self.bn_v_h(self.W_v_h(h_sent.squeeze(1)))268 269 alpha_v = self.softmax(self.W_v_att(self.tanh(W_v + W_v_h)))270 v_att = torch.mul(alpha_v, avg_features)271 272 W_a_h = self.bn_a_h(self.W_a_h(h_sent))273 W_a = self.bn_a(self.W_a(semantic_features))274 alpha_a = self.softmax(self.W_a_att(self.tanh(torch.add(W_a_h, W_a))))275 a_att = torch.mul(alpha_a, semantic_features).sum(1)276 277 ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))278 279 return ctx, alpha_v, alpha_a280 281 def v4(self, avg_features, semantic_features, h_sent):282 W_v = self.W_v(avg_features)283 W_v_h = self.W_v_h(h_sent.squeeze(1))284 285 alpha_v = self.softmax(self.W_v_att(self.tanh(torch.add(W_v, W_v_h))))286 v_att = torch.mul(alpha_v, avg_features)287 288 W_a_h = self.W_a_h(h_sent)289 W_a = self.W_a(semantic_features)290 alpha_a = self.softmax(self.W_a_att(self.tanh(torch.add(W_a_h, W_a))))291 a_att = torch.mul(alpha_a, semantic_features).sum(1)292 293 ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))294 295 return ctx, alpha_v, alpha_a296 297 def v5(self, avg_features, semantic_features, h_sent):298 W_v = self.W_v(avg_features)299 W_v_h = self.W_v_h(h_sent.squeeze(1))300 301 alpha_v = self.softmax(self.W_v_att(self.tanh(self.bn_v(torch.add(W_v, W_v_h)))))302 v_att = torch.mul(alpha_v, avg_features)303 304 W_a_h = self.W_a_h(h_sent)305 W_a = self.W_a(semantic_features)306 alpha_a = self.softmax(self.W_a_att(self.tanh(self.bn_a(torch.add(W_a_h, W_a)))))307 a_att = torch.mul(alpha_a, semantic_features).sum(1)308 309 ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))310 311 return ctx, alpha_v, alpha_a312 313 314class SentenceLSTM(nn.Module):315 def __init__(self,316 version='v1',317 embed_size=512,318 hidden_size=512,319 num_layers=1,320 dropout=0.3,321 momentum=0.1):322 super(SentenceLSTM, self).__init__()323 self.version = version324 325 self.lstm = nn.LSTM(input_size=embed_size,326 hidden_size=hidden_size,327 num_layers=num_layers,328 dropout=dropout)329 330 self.W_t_h = nn.Linear(in_features=hidden_size,331 out_features=embed_size,332 bias=True)333 self.bn_t_h = nn.BatchNorm1d(num_features=1, momentum=momentum)334 335 self.W_t_ctx = nn.Linear(in_features=embed_size,336 out_features=embed_size,337 bias=True)338 self.bn_t_ctx = nn.BatchNorm1d(num_features=1, momentum=momentum)339 340 self.W_stop_s_1 = nn.Linear(in_features=hidden_size,341 out_features=embed_size,342 bias=True)343 self.bn_stop_s_1 = nn.BatchNorm1d(num_features=1, momentum=momentum)344 345 self.W_stop_s = nn.Linear(in_features=hidden_size,346 out_features=embed_size,347 bias=True)348 self.bn_stop_s = nn.BatchNorm1d(num_features=1, momentum=momentum)349 350 self.W_stop = nn.Linear(in_features=embed_size,351 out_features=2,352 bias=True)353 self.bn_stop = nn.BatchNorm1d(num_features=1, momentum=momentum)354 355 self.W_topic = nn.Linear(in_features=embed_size,356 out_features=embed_size,357 bias=True)358 self.bn_topic = nn.BatchNorm1d(num_features=1, momentum=momentum)359 360 self.sigmoid = nn.Sigmoid()361 self.tanh = nn.Tanh()362 self.__init_weight()363 364 def __init_weight(self):365 self.W_t_h.weight.data.uniform_(-0.1, 0.1)366 self.W_t_h.bias.data.fill_(0)367 368 self.W_t_ctx.weight.data.uniform_(-0.1, 0.1)369 self.W_t_ctx.bias.data.fill_(0)370 371 self.W_stop_s_1.weight.data.uniform_(-0.1, 0.1)372 self.W_stop_s_1.bias.data.fill_(0)373 374 self.W_stop_s.weight.data.uniform_(-0.1, 0.1)375 self.W_stop_s.bias.data.fill_(0)376 377 self.W_stop.weight.data.uniform_(-0.1, 0.1)378 self.W_stop.bias.data.fill_(0)379 380 self.W_topic.weight.data.uniform_(-0.1, 0.1)381 self.W_topic.bias.data.fill_(0)382 383 def forward(self, ctx, prev_hidden_state, states=None) -> object:384 """385 :rtype: object386 """387 if self.version == 'v1':388 return self.v1(ctx, prev_hidden_state, states)389 elif self.version == 'v2':390 return self.v2(ctx, prev_hidden_state, states)391 elif self.version == 'v3':392 return self.v3(ctx, prev_hidden_state, states)393 394 def v1(self, ctx, prev_hidden_state, states=None):395 """396 v1 (only training)397 :param ctx:398 :param prev_hidden_state:399 :param states:400 :return:401 """402 ctx = ctx.unsqueeze(1)403 hidden_state, states = self.lstm(ctx, states)404 topic = self.W_topic(self.sigmoid(self.bn_t_h(self.W_t_h(hidden_state))405 + self.bn_t_ctx(self.W_t_ctx(ctx))))406 p_stop = self.W_stop(self.sigmoid(self.bn_stop_s_1(self.W_stop_s_1(prev_hidden_state))407 + self.bn_stop_s(self.W_stop_s(hidden_state))))408 return topic, p_stop, hidden_state, states409 410 def v2(self, ctx, prev_hidden_state, states=None):411 """412 v2413 :rtype: object414 """415 ctx = ctx.unsqueeze(1)416 hidden_state, states = self.lstm(ctx, states)417 topic = self.bn_topic(self.W_topic(self.tanh(self.bn_t_h(self.W_t_h(hidden_state)418 + self.W_t_ctx(ctx)))))419 p_stop = self.bn_stop(self.W_stop(self.tanh(self.bn_stop_s(self.W_stop_s_1(prev_hidden_state)420 + self.W_stop_s(hidden_state)))))421 return topic, p_stop, hidden_state, states422 423 def v3(self, ctx, prev_hidden_state, states=None):424 """425 v3426 :rtype: object427 """428 ctx = ctx.unsqueeze(1)429 hidden_state, states = self.lstm(ctx, states)430 topic = self.W_topic(self.tanh(self.W_t_h(hidden_state) + self.W_t_ctx(ctx)))431 p_stop = self.W_stop(self.tanh(self.W_stop_s_1(prev_hidden_state) + self.W_stop_s(hidden_state)))432 return topic, p_stop, hidden_state, states433 434 435class WordLSTM(nn.Module):436 def __init__(self,437 embed_size,438 hidden_size,439 vocab_size,440 num_layers,441 n_max=50):442 super(WordLSTM, self).__init__()443 self.embed = nn.Embedding(vocab_size, embed_size)444 self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)445 self.linear = nn.Linear(hidden_size, vocab_size)446 self.__init_weights()447 self.n_max = n_max448 self.vocab_size = vocab_size449 450 def __init_weights(self):451 self.embed.weight.data.uniform_(-0.1, 0.1)452 self.linear.weight.data.uniform_(-0.1, 0.1)453 self.linear.bias.data.fill_(0)454 455 def forward(self, topic_vec, captions):456 embeddings = self.embed(captions)457 embeddings = torch.cat((topic_vec, embeddings), 1)458 hidden, _ = self.lstm(embeddings)459 outputs = self.linear(hidden[:, -1, :])460 return outputs461 462 def sample(self, features, start_tokens):463 sampled_ids = np.zeros((np.shape(features)[0], self.n_max))464 sampled_ids[:, 0] = start_tokens.view(-1, )465 predicted = start_tokens466 embeddings = features467 embeddings = embeddings468 469 for i in range(1, self.n_max):470 predicted = self.embed(predicted)471 embeddings = torch.cat([embeddings, predicted], dim=1)472 hidden_states, _ = self.lstm(embeddings)473 hidden_states = hidden_states[:, -1, :]474 outputs = self.linear(hidden_states)475 predicted = torch.max(outputs, 1)[1]476 sampled_ids[:, i] = predicted477 predicted = predicted.unsqueeze(1)478 return sampled_ids479 480 481if __name__ == '__main__':482 import torchvision.transforms as transforms483 484 import warnings485 warnings.filterwarnings("ignore")486#487 extractor = VisualFeatureExtractor(model_name='resnet152')488 mlc = MLC(fc_in_features=extractor.out_features)489 co_att = CoAttention(visual_size=extractor.out_features)490 sent_lstm = SentenceLSTM()491 word_lstm = WordLSTM(embed_size=512, hidden_size=512, vocab_size=100, num_layers=1)492 493 images = torch.randn((4, 3, 224, 224))494 captions = torch.ones((4, 10)).long()495 hidden_state = torch.randn((4, 1, 512))496 497 # # image_file = '../data/images/CXR2814_IM-1239-1001.png'498# # # images = Image.open(image_file).convert('RGB')499# # # captions = torch.ones((1, 10)).long()500# # # hidden_state = torch.randn((10, 512))501# #502# norm = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))503#504# transform = transforms.Compose([505# transforms.Resize(256),506# transforms.TenCrop(224),507# transforms.Lambda(lambda crops: torch.stack([norm(transforms.ToTensor()(crop)) for crop in crops])),508# ])509 510# images = transform(images)511# images.unsqueeze_(0)512#513# # bs, ncrops, c, h, w = images.size()514# # images = images.view(-1, c, h, w)515#516 print("images:{}".format(images.shape))517 print("captions:{}".format(captions.shape))518 print("hidden_states:{}".format(hidden_state.shape))519 520 visual_features, avg_features = extractor.forward(images)521 522 print("visual_features:{}".format(visual_features.shape))523 print("avg features:{}".format(avg_features.shape))524 525 tags, semantic_features = mlc.forward(avg_features)526 527 print("tags:{}".format(tags.shape))528 print("semantic_features:{}".format(semantic_features.shape))529 530 ctx, alpht_v, alpht_a = co_att.forward(avg_features, semantic_features, hidden_state)531 532 print("ctx:{}".format(ctx.shape))533 print("alpht_v:{}".format(alpht_v.shape))534 print("alpht_a:{}".format(alpht_a.shape))535 536 topic, p_stop, hidden_state, states = sent_lstm.forward(ctx, hidden_state)537 # p_stop_avg = p_stop.view(bs, ncrops, -1).mean(1)538 539 print("Topic:{}".format(topic.shape))540 print("P_STOP:{}".format(p_stop.shape))541 # print("P_stop_avg:{}".format(p_stop_avg.shape))542 543 words = word_lstm.forward(topic, captions)544 print("words:{}".format(words.shape))545 546 cam = torch.mul(visual_features, alpht_v.view(alpht_v.shape[0], alpht_v.shape[1], 1, 1)).sum(1)547 cam.squeeze_()548 cam = cam.cpu().data.numpy()549 for i in range(cam.shape[0]):550 heatmap = cam[i]551 heatmap = heatmap / np.max(heatmap)552 print(heatmap.shape)553