Upload 13 files

build_tag.py

@@ -0,0 +1,89 @@
+class Tag(object):
+    def __init__(self):
+        self.static_tags = self.__load_static_tags()
+        self.id2tags = self.__load_id2tags()
+        self.tags2id = self.__load_tags2id()
+
+    def array2tags(self, array):
+        tags = []
+        for id in array:
+            tags.append(self.id2tags[id])
+        return tags
+
+    def tags2array(self, tags):
+        array = []
+        for tag in self.static_tags:
+            if tag in tags:
+                array.append(1)
+            else:
+                array.append(0)
+        return array
+
+    def inv_tags2array(self, array):
+        tags = []
+        for i, value in enumerate(array):
+            if value != 0:
+                tags.append(self.id2tags[i])
+        return tags
+
+    def __load_id2tags(self):
+        id2tags = {}
+        for i, tag in enumerate(self.static_tags):
+            id2tags[i] = tag
+        return id2tags
+
+    def __load_tags2id(self):
+        tags2id = {}
+        for i, tag in enumerate(self.static_tags):
+            tags2id[tag] = i
+        return tags2id
+
+    def __load_static_tags(self):
+        static_tags_name = ['cardiac monitor', 'lymphatic diseases', 'pulmonary disease', 'osteophytes', 'foreign body',
+                            'dish', 'aorta, thoracic', 'atherosclerosis', 'histoplasmosis', 'hypoventilation',
+                            'catheterization, central venous', 'pleural effusions', 'pleural effusion', 'callus',
+                            'sternotomy', 'lymph nodes', 'tortuous aorta', 'stent', 'interstitial pulmonary edema',
+                            'cholecystectomies', 'neoplasm', 'central venous catheter', 'pneumothorax',
+                            'metastatic disease', 'vena cava, superior', 'cholecystectomy', 'scoliosis',
+                            'subcutaneous emphysema', 'thoracolumbar scoliosis', 'spinal osteophytosis',
+                            'pulmonary fibroses', 'rib fractures', 'sarcoidosis', 'eventration', 'fibrosis', 'spine',
+                            'obstructive lung disease', 'pneumonitis', 'osteopenia', 'air trapping', 'demineralization',
+                            'mass lesion', 'pulmonary hypertension', 'pleural diseases', 'pleural thickening',
+                            'calcifications of the aorta', 'calcinosis', 'cystic fibrosis', 'empyema', 'catheter',
+                            'lymph', 'pericardial effusion', 'lung cancer', 'rib fracture', 'granulomatous disease',
+                            'chronic obstructive pulmonary disease', 'rib', 'clip', 'aortic ectasia', 'shoulder',
+                            'scarring', 'scleroses', 'adenopathy', 'emphysemas', 'pneumonectomy', 'infection',
+                            'aspiration', 'bilateral pleural effusion', 'bulla', 'lumbar vertebrae', 'lung neoplasms',
+                            'lymphadenopathy', 'hyperexpansion', 'ectasia', 'bronchiectasis', 'nodule', 'pneumonia',
+                            'right-sided pleural effusion', 'osteoarthritis', 'thoracic spondylosis', 'picc',
+                            'cervical fusion', 'tracheostomies', 'fusion', 'thoracic vertebrae', 'catheters',
+                            'emphysema', 'trachea', 'surgery', 'cervical spine fusion', 'hypertension, pulmonary',
+                            'pneumoperitoneum', 'scar', 'atheroscleroses', 'aortic calcifications', 'volume overload',
+                            'right upper lobe pneumonia', 'apical granuloma', 'diaphragms', 'copd', 'kyphoses',
+                            'spinal fractures', 'fracture', 'clavicle', 'focal atelectasis', 'collapse',
+                            'thoracotomies', 'congestive heart failure', 'calcified lymph nodes', 'edema',
+                            'degenerative disc diseases', 'cervical vertebrae', 'diaphragm', 'humerus', 'heart failure',
+                            'normal', 'coronary artery bypass', 'pulmonary atelectasis', 'lung diseases, interstitial',
+                            'pulmonary disease, chronic obstructive', 'opacity', 'deformity', 'chronic disease',
+                            'pleura', 'aorta', 'tuberculoses', 'hiatal hernia', 'scolioses', 'pleural fluid',
+                            'malignancy', 'kyphosis', 'bronchiectases', 'congestion', 'discoid atelectasis', 'nipple',
+                            'bronchitis', 'pulmonary artery', 'cardiomegaly', 'thoracic aorta', 'arthritic changes',
+                            'pulmonary edema', 'vascular calcification', 'sclerotic', 'central venous catheters',
+                            'catheterization', 'hydropneumothorax', 'aortic valve', 'hyperinflation', 'prostheses',
+                            'pacemaker, artificial', 'bypass grafts', 'pulmonary fibrosis', 'multiple myeloma',
+                            'postoperative period', 'cabg', 'right lower lobe pneumonia', 'granuloma',
+                            'degenerative change', 'atelectasis', 'inflammation', 'effusion', 'cicatrix',
+                            'tracheostomy', 'aortic diseases', 'sarcoidoses', 'granulomas', 'interstitial lung disease',
+                            'infiltrates', 'displaced fractures', 'chronic lung disease', 'picc line',
+                            'intubation, gastrointestinal', 'lung diseases', 'multiple pulmonary nodules',
+                            'intervertebral disc degeneration', 'pulmonary emphysema', 'spine curvature', 'fibroses',
+                            'chronic granulomatous disease', 'degenerative disease', 'atelectases', 'ribs',
+                            'pulmonary arterial hypertension', 'edemas', 'pectus excavatum', 'lung granuloma',
+                            'plate-like atelectasis', 'enlarged heart', 'hilar calcification', 'heart valve prosthesis',
+                            'tuberculosis', 'old injury', 'patchy atelectasis', 'histoplasmoses', 'exostoses',
+                            'mastectomies', 'right atrium', 'large hiatal hernia', 'hernia, hiatal', 'aortic aneurysm',
+                            'lobectomy', 'spinal fusion', 'spondylosis', 'ascending aorta', 'granulomatous infection',
+                            'fractures, bone', 'calcified granuloma', 'degenerative joint disease',
+                            'intubation, intratracheal', 'others']
+
+        return static_tags_name

dataset.py

@@ -0,0 +1,151 @@
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+import os
+import json
+from utils.build_vocab import Vocabulary, JsonReader
+import numpy as np
+from torchvision import transforms
+import pickle
+
+
+class ChestXrayDataSet(Dataset):
+    def __init__(self,
+                 image_dir,
+                 caption_json,
+                 file_list,
+                 vocabulary,
+                 s_max=10,
+                 n_max=50,
+                 transforms=None):
+        self.image_dir = image_dir
+        self.caption = JsonReader(caption_json)
+        self.file_names, self.labels = self.__load_label_list(file_list)
+        self.vocab = vocabulary
+        self.transform = transforms
+        self.s_max = s_max
+        self.n_max = n_max
+
+    def __load_label_list(self, file_list):
+        labels = []
+        filename_list = []
+        with open(file_list, 'r') as f:
+            for line in f:
+                items = line.split()
+                image_name = items[0]
+                label = items[1:]
+                label = [int(i) for i in label]
+                image_name = '{}.png'.format(image_name)
+                filename_list.append(image_name)
+                labels.append(label)
+        return filename_list, labels
+
+    def __getitem__(self, index):
+        image_name = self.file_names[index]
+        image = Image.open(os.path.join(self.image_dir, image_name)).convert('RGB')
+        label = self.labels[index]
+        if self.transform is not None:
+            image = self.transform(image)
+        try:
+            text = self.caption[image_name]
+        except Exception as err:
+            text = 'normal. '
+
+        target = list()
+        max_word_num = 0
+        for i, sentence in enumerate(text.split('. ')):
+            if i >= self.s_max:
+                break
+            sentence = sentence.split()
+            if len(sentence) == 0 or len(sentence) == 1 or len(sentence) > self.n_max:
+                continue
+            tokens = list()
+            tokens.append(self.vocab('<start>'))
+            tokens.extend([self.vocab(token) for token in sentence])
+            tokens.append(self.vocab('<end>'))
+            if max_word_num < len(tokens):
+                max_word_num = len(tokens)
+            target.append(tokens)
+        sentence_num = len(target)
+        return image, image_name, list(label / np.sum(label)), target, sentence_num, max_word_num
+
+    def __len__(self):
+        return len(self.file_names)
+
+
+def collate_fn(data):
+    images, image_id, label, captions, sentence_num, max_word_num = zip(*data)
+    images = torch.stack(images, 0)
+
+    max_sentence_num = max(sentence_num)
+    max_word_num = max(max_word_num)
+
+    targets = np.zeros((len(captions), max_sentence_num + 1, max_word_num))
+    prob = np.zeros((len(captions), max_sentence_num + 1))
+
+    for i, caption in enumerate(captions):
+        for j, sentence in enumerate(caption):
+            targets[i, j, :len(sentence)] = sentence[:]
+            prob[i][j] = len(sentence) > 0
+
+    return images, image_id, torch.Tensor(label), targets, prob
+
+
+def get_loader(image_dir,
+               caption_json,
+               file_list,
+               vocabulary,
+               transform,
+               batch_size,
+               s_max=10,
+               n_max=50,
+               shuffle=False):
+    dataset = ChestXrayDataSet(image_dir=image_dir,
+                               caption_json=caption_json,
+                               file_list=file_list,
+                               vocabulary=vocabulary,
+                               s_max=s_max,
+                               n_max=n_max,
+                               transforms=transform)
+    data_loader = torch.utils.data.DataLoader(dataset=dataset,
+                                              batch_size=batch_size,
+                                              shuffle=shuffle,
+                                              collate_fn=collate_fn)
+    return data_loader
+
+
+if __name__ == '__main__':
+    vocab_path = '../data/vocab.pkl'
+    image_dir = '../data/images'
+    caption_json = '../data/debugging_captions.json'
+    file_list = '../data/debugging.txt'
+    batch_size = 6
+    resize = 256
+    crop_size = 224
+
+    transform = transforms.Compose([
+        transforms.Resize(resize),
+        transforms.RandomCrop(crop_size),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406),
+                             (0.229, 0.224, 0.225))])
+
+    with open(vocab_path, 'rb') as f:
+        vocab = pickle.load(f)
+
+    data_loader = get_loader(image_dir=image_dir,
+                             caption_json=caption_json,
+                             file_list=file_list,
+                             vocabulary=vocab,
+                             transform=transform,
+                             batch_size=batch_size,
+                             shuffle=False)
+
+    for i, (image, image_id, label, target, prob) in enumerate(data_loader):
+        print(image.shape)
+        print(image_id)
+        print(label)
+        print(target)
+        print(prob)
+        break

extractor.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2cec084672668d0a2d9c4e2451f1f9c71069fff1ad0bb09759953625b4ec731c
+size 348017030

loss.py

@@ -0,0 +1,78 @@
+import torch
+import numpy as np
+from torch.nn.modules import loss
+
+
+class WARPLoss(loss.Module):
+    def __init__(self, num_labels=204):
+        super(WARPLoss, self).__init__()
+        self.rank_weights = [1.0 / 1]
+        for i in range(1, num_labels):
+            self.rank_weights.append(self.rank_weights[i - 1] + (1.0 / i + 1))
+
+    def forward(self, input, target) -> object:
+        """
+
+        :rtype:
+        :param input: Deep features tensor Variable of size batch x n_attrs.
+        :param target: Ground truth tensor Variable of size batch x n_attrs.
+        :return:
+        """
+        batch_size = target.size()[0]
+        n_labels = target.size()[1]
+        max_num_trials = n_labels - 1
+        loss = 0.0
+
+        for i in range(batch_size):
+
+            for j in range(n_labels):
+                if target[i, j] == 1:
+
+                    neg_labels_idx = np.array([idx for idx, v in enumerate(target[i, :]) if v == 0])
+                    neg_idx = np.random.choice(neg_labels_idx, replace=False)
+                    sample_score_margin = 1 - input[i, j] + input[i, neg_idx]
+                    num_trials = 0
+
+                    while sample_score_margin < 0 and num_trials < max_num_trials:
+                        neg_idx = np.random.choice(neg_labels_idx, replace=False)
+                        num_trials += 1
+                        sample_score_margin = 1 - input[i, j] + input[i, neg_idx]
+
+                    r_j = np.floor(max_num_trials / num_trials)
+                    weight = self.rank_weights[r_j]
+
+                    for k in range(n_labels):
+                        if target[i, k] == 0:
+                            score_margin = 1 - input[i, j] + input[i, k]
+                            loss += (weight * torch.clamp(score_margin, min=0.0))
+        return loss
+
+
+class MultiLabelSoftmaxRegressionLoss(loss.Module):
+    def __init__(self):
+        super(MultiLabelSoftmaxRegressionLoss, self).__init__()
+
+    def forward(self, input, target) -> object:
+        return -1 * torch.sum(input * target)
+
+
+class LossFactory(object):
+    def __init__(self, type, num_labels=156):
+        self.type = type
+        if type == 'BCE':
+            # self.activation_func = torch.nn.Sigmoid()
+            self.loss = torch.nn.BCELoss()
+        elif type == 'CE':
+            self.loss = torch.nn.CrossEntropyLoss()
+        elif type == 'WARP':
+            self.activation_func = torch.nn.Softmax()
+            self.loss = WARPLoss(num_labels=num_labels)
+        elif type == 'MSR':
+            self.activation_func = torch.nn.LogSoftmax()
+            self.loss = MultiLabelSoftmaxRegressionLoss()
+
+    def compute_loss(self, output, target):
+        # output = self.activation_func(output)
+        # if self.type == 'NLL' or self.type == 'WARP' or self.type == 'MSR':
+        #     target /= torch.sum(target, 1).view(-1, 1)
+        return self.loss(output, target)

mlc.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b099a01f553ee36e1b4ecedae8b19f1decd6fb35d1aebd3ed42ae9fa6bf61080
+size 346714524

models.py

@@ -0,0 +1,552 @@
+import torch
+import torch.nn as nn
+import torchvision
+import numpy as np
+from torch.autograd import Variable
+import torchvision.models as models
+import transformers
+import torchvision.transforms
+
+import torchxrayvision as xrv
+from transformers import ViTModel, ViTConfig
+
+
+
+class VisualFeatureExtractor(nn.Module):
+    def __init__(self, model_name='densenet201', pretrained=False):
+        super(VisualFeatureExtractor, self).__init__()
+        self.model_name = 'chexnet'
+        self.pretrained = pretrained
+        self.model, self.out_features, self.avg_func, self.bn, self.linear = self.__get_model()
+        self.activation = nn.ReLU()
+
+    def __get_model(self):
+        model = None
+        out_features = None
+        func = None
+        
+        if self.model_name == 'resnet152':
+            resnet = models.resnet152(pretrained=self.pretrained)
+            modules = list(resnet.children())[:-2]
+            model = nn.Sequential(*modules)
+            out_features = resnet.fc.in_features
+            func = torch.nn.AvgPool2d(kernel_size=7, stride=1, padding=0)
+        
+        
+        elif self.model_name == 'densenet201':
+            densenet = models.densenet201(pretrained=self.pretrained)
+            modules = list(densenet.features)
+            model = nn.Sequential(*modules)
+            func = torch.nn.AvgPool2d(kernel_size=7, stride=1, padding=0)
+            out_features = densenet.classifier.in_features
+
+
+        elif self.model_name == 'chexnet':
+            print("vit chest xray pretrained model loading")
+            # Load the Vision Transformer (ViT) model configuration
+            config = ViTConfig.from_pretrained('nickmuchi/vit-finetuned-chest-xray-pneumonia')
+
+            # Initialize the ViT model with the specific configuration
+            vit_model = ViTModel(config)
+
+            # Load the state dict specifically, excluding 'classifier.bias', 'classifier.weight'
+            state_dict = torch.load('vit-coatten/utils/pytorch_model.bin', map_location=torch.device('cpu'))
+            state_dict = {k: v for k, v in state_dict.items() if not k.startswith('classifier')}
+            vit_model.load_state_dict(state_dict, strict=False)
+
+            model = vit_model
+            out_features = config.hidden_size
+
+        linear = nn.Linear(in_features=out_features, out_features=out_features)
+        bn = nn.BatchNorm1d(num_features=out_features, momentum=0.1)
+
+        return model, out_features, func, bn, linear
+
+    def forward(self, images):
+        """
+        :param images: Input images
+        :return: visual_features, avg_features
+        """
+        model_output = self.model(images)
+
+        # Extract the pooler_output
+
+        pooler_output = model_output.pooler_output
+
+        # Apply the linear layer, batch normalization, and activation
+        avg_features = self.activation(self.bn(self.linear(pooler_output)))
+
+        return model_output.last_hidden_state, avg_features     
+
+    # def forward(self, images):
+    #     """
+    #     :param images:
+    #     :return:
+    #     """
+    #     visual_features = self.model(images)
+        
+    #     avg_features = self.avg_func(visual_features).squeeze()
+    #     # avg_features = self.activation(self.bn(self.linear(visual_features)))
+
+    #     return visual_features, avg_features
+
+
+class MLC(nn.Module):
+    def __init__(self,
+                 classes=210,
+                 sementic_features_dim=512,
+                 fc_in_features=2048,
+                 k=10,
+                 ):
+        super(MLC, self).__init__()
+        pretrained_model_name="nickmuchi/vit-finetuned-chest-xray-pneumonia"
+        vit_config = ViTConfig.from_pretrained(pretrained_model_name)
+        self.vit = ViTModel(vit_config)
+
+        # Adjust the classifier to your number of classes
+        self.classifier = nn.Linear(in_features=vit_config.hidden_size, out_features=classes)
+        self.embed = nn.Embedding(classes, sementic_features_dim)
+        self.k = k
+        self.sigmoid = nn.Sigmoid()
+        self.__init_weight()
+
+    def __init_weight(self):
+        nn.init.xavier_uniform_(self.classifier.weight)
+        if self.classifier.bias is not None:
+            self.classifier.bias.data.fill_(0)
+
+    def forward(self, avg_features):
+
+
+        tags = self.sigmoid(self.classifier(avg_features))
+        semantic_features = self.embed(torch.topk(tags, self.k)[1])
+        return tags, semantic_features
+    
+# class MLC(nn.Module):
+#     def __init__(self,
+#                  classes=210,
+#                  sementic_features_dim=512,
+#                  fc_in_features=2048,
+#                  k=10):
+#         super(MLC, self).__init__()
+#         self.classifier = nn.Linear(in_features=fc_in_features, out_features=classes)
+#         self.embed = nn.Embedding(classes, sementic_features_dim)
+#         self.k = k
+#         self.sigmoid = nn.Sigmoid()
+#         self.__init_weight()
+
+#     def __init_weight(self):
+#         # Example: Initialize weights with a different strategy
+#         nn.init.xavier_uniform_(self.classifier.weight)
+#         if self.classifier.bias is not None:
+#             self.classifier.bias.data.fill_(0)
+
+#     def forward(self, avg_features):
+#         tags = self.sigmoid(self.classifier(avg_features))
+#         semantic_features = self.embed(torch.topk(tags, self.k)[1])
+#         return tags, semantic_features
+
+
+class CoAttention(nn.Module):
+    def __init__(self,
+                 version='v1',
+                 embed_size=512,
+                 hidden_size=512,
+                 visual_size=2048,
+                 k=10,
+                 momentum=0.1):
+        super(CoAttention, self).__init__()
+        self.version = version
+        self.W_v = nn.Linear(in_features=visual_size, out_features=visual_size)
+        self.bn_v = nn.BatchNorm1d(num_features=visual_size, momentum=momentum)
+
+        self.W_v_h = nn.Linear(in_features=hidden_size, out_features=visual_size)
+        self.bn_v_h = nn.BatchNorm1d(num_features=visual_size, momentum=momentum)
+
+        self.W_v_att = nn.Linear(in_features=visual_size, out_features=visual_size)
+        self.bn_v_att = nn.BatchNorm1d(num_features=visual_size, momentum=momentum)
+
+        self.W_a = nn.Linear(in_features=hidden_size, out_features=hidden_size)
+        self.bn_a = nn.BatchNorm1d(num_features=k, momentum=momentum)
+
+        self.W_a_h = nn.Linear(in_features=hidden_size, out_features=hidden_size)
+        self.bn_a_h = nn.BatchNorm1d(num_features=1, momentum=momentum)
+
+        self.W_a_att = nn.Linear(in_features=hidden_size, out_features=hidden_size)
+        self.bn_a_att = nn.BatchNorm1d(num_features=k, momentum=momentum)
+
+        # self.W_fc = nn.Linear(in_features=visual_size, out_features=embed_size)  # for v3
+        self.W_fc = nn.Linear(in_features=visual_size + hidden_size, out_features=embed_size)
+        self.bn_fc = nn.BatchNorm1d(num_features=embed_size, momentum=momentum)
+
+        self.tanh = nn.Tanh()
+        self.softmax = nn.Softmax()
+
+        self.__init_weight()
+
+    def __init_weight(self):
+        self.W_v.weight.data.uniform_(-0.1, 0.1)
+        self.W_v.bias.data.fill_(0)
+
+        self.W_v_h.weight.data.uniform_(-0.1, 0.1)
+        self.W_v_h.bias.data.fill_(0)
+
+        self.W_v_att.weight.data.uniform_(-0.1, 0.1)
+        self.W_v_att.bias.data.fill_(0)
+
+        self.W_a.weight.data.uniform_(-0.1, 0.1)
+        self.W_a.bias.data.fill_(0)
+
+        self.W_a_h.weight.data.uniform_(-0.1, 0.1)
+        self.W_a_h.bias.data.fill_(0)
+
+        self.W_a_att.weight.data.uniform_(-0.1, 0.1)
+        self.W_a_att.bias.data.fill_(0)
+
+        self.W_fc.weight.data.uniform_(-0.1, 0.1)
+        self.W_fc.bias.data.fill_(0)
+
+    def forward(self, avg_features, semantic_features, h_sent):
+        if self.version == 'v1':
+            return self.v1(avg_features, semantic_features, h_sent)
+        elif self.version == 'v2':
+            return self.v2(avg_features, semantic_features, h_sent)
+        elif self.version == 'v3':
+            return self.v3(avg_features, semantic_features, h_sent)
+        elif self.version == 'v4':
+            return self.v4(avg_features, semantic_features, h_sent)
+        elif self.version == 'v5':
+            return self.v5(avg_features, semantic_features, h_sent)
+
+    def v1(self, avg_features, semantic_features, h_sent) -> object:
+        """
+        only training
+        :rtype: object
+        """
+        W_v = self.bn_v(self.W_v(avg_features))
+        W_v_h = self.bn_v_h(self.W_v_h(h_sent.squeeze(1)))
+
+        alpha_v = self.softmax(self.bn_v_att(self.W_v_att(self.tanh(W_v + W_v_h))))
+        v_att = torch.mul(alpha_v, avg_features)
+
+        W_a_h = self.bn_a_h(self.W_a_h(h_sent))
+        W_a = self.bn_a(self.W_a(semantic_features))
+        alpha_a = self.softmax(self.bn_a_att(self.W_a_att(self.tanh(torch.add(W_a_h, W_a)))))
+        a_att = torch.mul(alpha_a, semantic_features).sum(1)
+
+        ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))
+
+        return ctx, alpha_v, alpha_a
+
+    def v2(self, avg_features, semantic_features, h_sent) -> object:
+        """
+        no bn
+        :rtype: object
+        """
+        W_v = self.W_v(avg_features)
+        W_v_h = self.W_v_h(h_sent.squeeze(1))
+
+        alpha_v = self.softmax(self.W_v_att(self.tanh(W_v + W_v_h)))
+        v_att = torch.mul(alpha_v, avg_features)
+
+        W_a_h = self.W_a_h(h_sent)
+        W_a = self.W_a(semantic_features)
+        alpha_a = self.softmax(self.W_a_att(self.tanh(torch.add(W_a_h, W_a))))
+        a_att = torch.mul(alpha_a, semantic_features).sum(1)
+
+        ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))
+
+        return ctx, alpha_v, alpha_a
+
+    def v3(self, avg_features, semantic_features, h_sent) -> object:
+        """
+
+        :rtype: object
+        """
+        W_v = self.bn_v(self.W_v(avg_features))
+        W_v_h = self.bn_v_h(self.W_v_h(h_sent.squeeze(1)))
+
+        alpha_v = self.softmax(self.W_v_att(self.tanh(W_v + W_v_h)))
+        v_att = torch.mul(alpha_v, avg_features)
+
+        W_a_h = self.bn_a_h(self.W_a_h(h_sent))
+        W_a = self.bn_a(self.W_a(semantic_features))
+        alpha_a = self.softmax(self.W_a_att(self.tanh(torch.add(W_a_h, W_a))))
+        a_att = torch.mul(alpha_a, semantic_features).sum(1)
+
+        ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))
+
+        return ctx, alpha_v, alpha_a
+
+    def v4(self, avg_features, semantic_features, h_sent):
+        W_v = self.W_v(avg_features)
+        W_v_h = self.W_v_h(h_sent.squeeze(1))
+
+        alpha_v = self.softmax(self.W_v_att(self.tanh(torch.add(W_v, W_v_h))))
+        v_att = torch.mul(alpha_v, avg_features)
+
+        W_a_h = self.W_a_h(h_sent)
+        W_a = self.W_a(semantic_features)
+        alpha_a = self.softmax(self.W_a_att(self.tanh(torch.add(W_a_h, W_a))))
+        a_att = torch.mul(alpha_a, semantic_features).sum(1)
+
+        ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))
+
+        return ctx, alpha_v, alpha_a
+
+    def v5(self, avg_features, semantic_features, h_sent):
+        W_v = self.W_v(avg_features)
+        W_v_h = self.W_v_h(h_sent.squeeze(1))
+
+        alpha_v = self.softmax(self.W_v_att(self.tanh(self.bn_v(torch.add(W_v, W_v_h)))))
+        v_att = torch.mul(alpha_v, avg_features)
+
+        W_a_h = self.W_a_h(h_sent)
+        W_a = self.W_a(semantic_features)
+        alpha_a = self.softmax(self.W_a_att(self.tanh(self.bn_a(torch.add(W_a_h, W_a)))))
+        a_att = torch.mul(alpha_a, semantic_features).sum(1)
+
+        ctx = self.W_fc(torch.cat([v_att, a_att], dim=1))
+
+        return ctx, alpha_v, alpha_a
+
+
+class SentenceLSTM(nn.Module):
+    def __init__(self,
+                 version='v1',
+                 embed_size=512,
+                 hidden_size=512,
+                 num_layers=1,
+                 dropout=0.3,
+                 momentum=0.1):
+        super(SentenceLSTM, self).__init__()
+        self.version = version
+
+        self.lstm = nn.LSTM(input_size=embed_size,
+                            hidden_size=hidden_size,
+                            num_layers=num_layers,
+                            dropout=dropout)
+
+        self.W_t_h = nn.Linear(in_features=hidden_size,
+                               out_features=embed_size,
+                               bias=True)
+        self.bn_t_h = nn.BatchNorm1d(num_features=1, momentum=momentum)
+
+        self.W_t_ctx = nn.Linear(in_features=embed_size,
+                                 out_features=embed_size,
+                                 bias=True)
+        self.bn_t_ctx = nn.BatchNorm1d(num_features=1, momentum=momentum)
+
+        self.W_stop_s_1 = nn.Linear(in_features=hidden_size,
+                                    out_features=embed_size,
+                                    bias=True)
+        self.bn_stop_s_1 = nn.BatchNorm1d(num_features=1, momentum=momentum)
+
+        self.W_stop_s = nn.Linear(in_features=hidden_size,
+                                  out_features=embed_size,
+                                  bias=True)
+        self.bn_stop_s = nn.BatchNorm1d(num_features=1, momentum=momentum)
+
+        self.W_stop = nn.Linear(in_features=embed_size,
+                                out_features=2,
+                                bias=True)
+        self.bn_stop = nn.BatchNorm1d(num_features=1, momentum=momentum)
+
+        self.W_topic = nn.Linear(in_features=embed_size,
+                                 out_features=embed_size,
+                                 bias=True)
+        self.bn_topic = nn.BatchNorm1d(num_features=1, momentum=momentum)
+
+        self.sigmoid = nn.Sigmoid()
+        self.tanh = nn.Tanh()
+        self.__init_weight()
+
+    def __init_weight(self):
+        self.W_t_h.weight.data.uniform_(-0.1, 0.1)
+        self.W_t_h.bias.data.fill_(0)
+
+        self.W_t_ctx.weight.data.uniform_(-0.1, 0.1)
+        self.W_t_ctx.bias.data.fill_(0)
+
+        self.W_stop_s_1.weight.data.uniform_(-0.1, 0.1)
+        self.W_stop_s_1.bias.data.fill_(0)
+
+        self.W_stop_s.weight.data.uniform_(-0.1, 0.1)
+        self.W_stop_s.bias.data.fill_(0)
+
+        self.W_stop.weight.data.uniform_(-0.1, 0.1)
+        self.W_stop.bias.data.fill_(0)
+
+        self.W_topic.weight.data.uniform_(-0.1, 0.1)
+        self.W_topic.bias.data.fill_(0)
+
+    def forward(self, ctx, prev_hidden_state, states=None) -> object:
+        """
+        :rtype: object
+        """
+        if self.version == 'v1':
+            return self.v1(ctx, prev_hidden_state, states)
+        elif self.version == 'v2':
+            return self.v2(ctx, prev_hidden_state, states)
+        elif self.version == 'v3':
+            return self.v3(ctx, prev_hidden_state, states)
+
+    def v1(self, ctx, prev_hidden_state, states=None):
+        """
+        v1 (only training)
+        :param ctx:
+        :param prev_hidden_state:
+        :param states:
+        :return:
+        """
+        ctx = ctx.unsqueeze(1)
+        hidden_state, states = self.lstm(ctx, states)
+        topic = self.W_topic(self.sigmoid(self.bn_t_h(self.W_t_h(hidden_state))
+                                          + self.bn_t_ctx(self.W_t_ctx(ctx))))
+        p_stop = self.W_stop(self.sigmoid(self.bn_stop_s_1(self.W_stop_s_1(prev_hidden_state))
+                                          + self.bn_stop_s(self.W_stop_s(hidden_state))))
+        return topic, p_stop, hidden_state, states
+
+    def v2(self, ctx, prev_hidden_state, states=None):
+        """
+        v2
+        :rtype: object
+        """
+        ctx = ctx.unsqueeze(1)
+        hidden_state, states = self.lstm(ctx, states)
+        topic = self.bn_topic(self.W_topic(self.tanh(self.bn_t_h(self.W_t_h(hidden_state)
+                                                                 + self.W_t_ctx(ctx)))))
+        p_stop = self.bn_stop(self.W_stop(self.tanh(self.bn_stop_s(self.W_stop_s_1(prev_hidden_state)
+                                                                   + self.W_stop_s(hidden_state)))))
+        return topic, p_stop, hidden_state, states
+
+    def v3(self, ctx, prev_hidden_state, states=None):
+        """
+        v3
+        :rtype: object
+        """
+        ctx = ctx.unsqueeze(1)
+        hidden_state, states = self.lstm(ctx, states)
+        topic = self.W_topic(self.tanh(self.W_t_h(hidden_state) + self.W_t_ctx(ctx)))
+        p_stop = self.W_stop(self.tanh(self.W_stop_s_1(prev_hidden_state) + self.W_stop_s(hidden_state)))
+        return topic, p_stop, hidden_state, states
+
+
+class WordLSTM(nn.Module):
+    def __init__(self,
+                 embed_size,
+                 hidden_size,
+                 vocab_size,
+                 num_layers,
+                 n_max=50):
+        super(WordLSTM, self).__init__()
+        self.embed = nn.Embedding(vocab_size, embed_size)
+        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
+        self.linear = nn.Linear(hidden_size, vocab_size)
+        self.__init_weights()
+        self.n_max = n_max
+        self.vocab_size = vocab_size
+
+    def __init_weights(self):
+        self.embed.weight.data.uniform_(-0.1, 0.1)
+        self.linear.weight.data.uniform_(-0.1, 0.1)
+        self.linear.bias.data.fill_(0)
+
+    def forward(self, topic_vec, captions):
+        embeddings = self.embed(captions)
+        embeddings = torch.cat((topic_vec, embeddings), 1)
+        hidden, _ = self.lstm(embeddings)
+        outputs = self.linear(hidden[:, -1, :])
+        return outputs
+
+    def sample(self, features, start_tokens):
+        sampled_ids = np.zeros((np.shape(features)[0], self.n_max))
+        sampled_ids[:, 0] = start_tokens.view(-1, )
+        predicted = start_tokens
+        embeddings = features
+        embeddings = embeddings
+
+        for i in range(1, self.n_max):
+            predicted = self.embed(predicted)
+            embeddings = torch.cat([embeddings, predicted], dim=1)
+            hidden_states, _ = self.lstm(embeddings)
+            hidden_states = hidden_states[:, -1, :]
+            outputs = self.linear(hidden_states)
+            predicted = torch.max(outputs, 1)[1]
+            sampled_ids[:, i] = predicted
+            predicted = predicted.unsqueeze(1)
+        return sampled_ids
+
+
+if __name__ == '__main__':
+    import torchvision.transforms as transforms
+
+    import warnings
+    warnings.filterwarnings("ignore")
+#
+    extractor = VisualFeatureExtractor(model_name='resnet152')
+    mlc = MLC(fc_in_features=extractor.out_features)
+    co_att = CoAttention(visual_size=extractor.out_features)
+    sent_lstm = SentenceLSTM()
+    word_lstm = WordLSTM(embed_size=512, hidden_size=512, vocab_size=100, num_layers=1)
+
+    images = torch.randn((4, 3, 224, 224))
+    captions = torch.ones((4, 10)).long()
+    hidden_state = torch.randn((4, 1, 512))
+
+    # # image_file = '../data/images/CXR2814_IM-1239-1001.png'
+#     # # images = Image.open(image_file).convert('RGB')
+#     # # captions = torch.ones((1, 10)).long()
+#     # # hidden_state = torch.randn((10, 512))
+# #
+# norm = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+#
+# transform = transforms.Compose([
+#     transforms.Resize(256),
+#     transforms.TenCrop(224),
+#     transforms.Lambda(lambda crops: torch.stack([norm(transforms.ToTensor()(crop)) for crop in crops])),
+# ])
+
+# images = transform(images)
+# images.unsqueeze_(0)
+#
+# # bs, ncrops, c, h, w = images.size()
+# # images = images.view(-1, c, h, w)
+#
+    print("images:{}".format(images.shape))
+    print("captions:{}".format(captions.shape))
+    print("hidden_states:{}".format(hidden_state.shape))
+
+    visual_features, avg_features = extractor.forward(images)
+
+    print("visual_features:{}".format(visual_features.shape))
+    print("avg features:{}".format(avg_features.shape))
+
+    tags, semantic_features = mlc.forward(avg_features)
+
+    print("tags:{}".format(tags.shape))
+    print("semantic_features:{}".format(semantic_features.shape))
+
+    ctx, alpht_v, alpht_a = co_att.forward(avg_features, semantic_features, hidden_state)
+
+    print("ctx:{}".format(ctx.shape))
+    print("alpht_v:{}".format(alpht_v.shape))
+    print("alpht_a:{}".format(alpht_a.shape))
+
+    topic, p_stop, hidden_state, states = sent_lstm.forward(ctx, hidden_state)
+    # p_stop_avg = p_stop.view(bs, ncrops, -1).mean(1)
+
+    print("Topic:{}".format(topic.shape))
+    print("P_STOP:{}".format(p_stop.shape))
+    # print("P_stop_avg:{}".format(p_stop_avg.shape))
+
+    words = word_lstm.forward(topic, captions)
+    print("words:{}".format(words.shape))
+
+    cam = torch.mul(visual_features, alpht_v.view(alpht_v.shape[0], alpht_v.shape[1], 1, 1)).sum(1)
+    cam.squeeze_()
+    cam = cam.cpu().data.numpy()
+    for i in range(cam.shape[0]):
+        heatmap = cam[i]
+        heatmap = heatmap / np.max(heatmap)
+        print(heatmap.shape)

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4b052d087f705ba06b16aa03c01dfdf37f36f0f8ab7b136cda1524bba8ab09d
+size 343280753

sentence.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33aa4fa7339a3958307299a75d72cd1d6d0d144cadb101cfbe5c40af205741bc
+size 22081920

tester.py

@@ -0,0 +1,283 @@
+
+
+import time
+import pickle
+import torch
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+from torch.autograd import Variable
+from PIL import Image
+import cv2
+
+from utils.models import *
+from utils.dataset import *
+from utils.loss import *
+from utils.build_tag import *
+
+
+class CaptionSampler(object):
+    def __init__(self):
+        # Default configuration values
+        self.args = {
+            "model_dir": "/Users/jkottu/Desktop/image-captioning-chest-xrays/vit-coatten",
+            "image_dir": "./data/images",
+            "caption_json": "data/new_data/captions.json",
+            "vocab_path": "/Users/jkottu/Desktop/image-captioning-chest-xrays/vit-coatten/vocab.pkl",
+            "file_lists": "data/new_data/test_data.txt",
+            "load_model_path": "train_best_loss.pth.tar",
+            "resize": 224,
+            "cam_size": 224,
+            "generate_dir": "cam",
+            "result_path": "results",
+            "result_name": "debug",
+            "momentum": 0.1,
+            "visual_model_name": "densenet201", 
+            "pretrained": False,
+            "classes": 210,
+            "sementic_features_dim": 512,
+            "k": 10, 
+            "attention_version": "v4",
+            "embed_size": 512,
+            "hidden_size": 512,
+            "sent_version": "v1",
+            "sentence_num_layers": 2,
+            "dropout": 0.1,
+            "word_num_layers": 1,
+            "s_max": 10,
+            "n_max": 30,
+            "batch_size": 8,
+            "lambda_tag": 10000,
+            "lambda_stop": 10,
+            "lambda_word": 1,
+            "cuda": False  # Keep CUDA disabled by default
+        }
+
+        self.vocab = self.__init_vocab()
+        self.tagger = self.__init_tagger()
+        self.transform = self.__init_transform()
+        self.model_state_dict = self.__load_mode_state_dict()
+
+        self.extractor = self.__init_visual_extractor()
+        self.mlc = self.__init_mlc()
+        self.co_attention = self.__init_co_attention()
+        self.sentence_model = self.__init_sentence_model()
+        self.word_model = self.__init_word_word()
+
+        self.ce_criterion = self._init_ce_criterion()
+        self.mse_criterion = self._init_mse_criterion()
+
+    @staticmethod
+    def _init_ce_criterion():
+        return nn.CrossEntropyLoss(size_average=False, reduce=False)
+
+    @staticmethod
+    def _init_mse_criterion():
+        return nn.MSELoss()
+
+
+    def sample(self, image_file):
+        self.extractor.eval()
+        self.mlc.eval()
+        self.co_attention.eval()
+        self.sentence_model.eval()
+        self.word_model.eval()
+        
+        # imageData = Image.open(image_file).convert('RGB')
+        imageData = self.transform(imageData)
+        imageData = imageData.unsqueeze_(0)
+
+        print(imageData.shape)
+    
+        image = self.__to_var(imageData, requires_grad=False)
+    
+        visual_features, avg_features = self.extractor.forward(image)
+    
+        tags, semantic_features = self.mlc(avg_features)
+        sentence_states = None
+        prev_hidden_states = self.__to_var(torch.zeros(image.shape[0], 1, self.args["hidden_size"]))
+    
+        pred_sentences = []
+    
+        for i in range(self.args["s_max"]):
+            ctx, alpha_v, alpha_a = self.co_attention.forward(avg_features, semantic_features, prev_hidden_states)
+            topic, p_stop, hidden_state, sentence_states = self.sentence_model.forward(ctx,
+                                                                                       prev_hidden_states,
+                                                                                       sentence_states)
+            p_stop = p_stop.squeeze(1)
+            p_stop = torch.max(p_stop, 1)[1].unsqueeze(1)
+
+            start_tokens = np.zeros((topic.shape[0], 1))
+            start_tokens[:, 0] = self.vocab('<start>')
+            start_tokens = self.__to_var(torch.Tensor(start_tokens).long(), requires_grad=False)
+
+            sampled_ids = self.word_model.sample(topic, start_tokens)
+            prev_hidden_states = hidden_state
+
+            sampled_ids = sampled_ids * p_stop.numpy()
+
+            
+            pred_sentences.append(self.__vec2sent(sampled_ids[0]))
+    
+        return pred_sentences
+
+
+    def __init_cam_path(self, image_file):
+        generate_dir = os.path.join(self.args["model_dir"], self.args["generate_dir"])
+        if not os.path.exists(generate_dir):
+            os.makedirs(generate_dir)
+
+        image_dir = os.path.join(generate_dir, image_file)
+
+        if not os.path.exists(image_dir):
+            os.makedirs(image_dir)
+        return image_dir
+
+    def __save_json(self, result):
+        result_path = os.path.join(self.args["model_dir"], self.args["result_path"])
+        if not os.path.exists(result_path):
+            os.makedirs(result_path)
+        with open(os.path.join(result_path, '{}.json'.format(self.args["result_name"])), 'w') as f:
+            json.dump(result, f)
+
+    def __load_mode_state_dict(self):
+        try:
+            model_state_dict = torch.load(os.path.join(self.args["model_dir"], self.args["load_model_path"]), map_location=torch.device('cpu'))
+            print("[Load Model-{} Succeed!]".format(self.args["load_model_path"]))
+            print("Load From Epoch {}".format(model_state_dict['epoch']))
+            return model_state_dict
+        except Exception as err:
+            print("[Load Model Failed] {}".format(err))
+            raise err
+
+    def __init_tagger(self):
+        return Tag()
+
+    def __vec2sent(self, array):
+        sampled_caption = []
+        for word_id in array:
+            word = self.vocab.get_word_by_id(word_id)
+            if word == '<start>':
+                continue
+            if word == '<end>' or word == '<pad>':
+                break
+            sampled_caption.append(word)
+        return ' '.join(sampled_caption)
+
+    def __init_vocab(self):
+        with open(self.args["vocab_path"], 'rb') as f:
+            vocab = pickle.load(f)
+        return vocab
+
+    def __init_data_loader(self, file_list):
+        data_loader = get_loader(image_dir=self.args.image_dir,
+                                 caption_json=self.args.caption_json,
+                                 file_list=file_list,
+                                 vocabulary=self.vocab,
+                                 transform=self.transform,
+                                 batch_size=self.args.batch_size,
+                                 s_max=self.args.s_max,
+                                 n_max=self.args.n_max,
+                                 shuffle=False)
+        return data_loader
+
+    def __init_transform(self):
+        transform = transforms.Compose([
+            transforms.Resize((self.args["resize"], self.args["resize"])),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406),
+                                 (0.229, 0.224, 0.225))])
+        return transform
+
+    def __to_var(self, x, requires_grad=True):
+        if self.args["cuda"]:
+            x = x.cuda()
+        return Variable(x, requires_grad=requires_grad)
+
+    def __init_visual_extractor(self):
+        model = VisualFeatureExtractor(model_name=self.args["visual_model_name"],
+                                       pretrained=self.args["pretrained"])
+
+        if self.model_state_dict is not None:
+            print("Visual Extractor Loaded!")
+            model.load_state_dict(self.model_state_dict['extractor'])
+
+        if self.args["cuda"]:
+            model = model.cuda()
+
+        return model
+
+    def __init_mlc(self):
+        model = MLC(classes=self.args["classes"],
+                    sementic_features_dim=self.args["sementic_features_dim"],
+                    fc_in_features=self.extractor.out_features,
+                    k=self.args["k"])
+
+        if self.model_state_dict is not None:
+            print("MLC Loaded!")
+            model.load_state_dict(self.model_state_dict['mlc'])
+
+        if self.args["cuda"]:
+            model = model.cuda()
+
+        return model
+
+    def __init_co_attention(self):
+        model = CoAttention(version=self.args["attention_version"],
+                            embed_size=self.args["embed_size"],
+                            hidden_size=self.args["hidden_size"],
+                            visual_size=self.extractor.out_features,
+                            k=self.args["k"],
+                            momentum=self.args["momentum"])
+
+        if self.model_state_dict is not None:
+            print("Co-Attention Loaded!")
+            model.load_state_dict(self.model_state_dict['co_attention'])
+
+        if self.args["cuda"]:
+            model = model.cuda()
+
+        return model
+
+    def __init_sentence_model(self):
+        model = SentenceLSTM(version=self.args["sent_version"],
+                             embed_size=self.args["embed_size"],
+                             hidden_size=self.args["hidden_size"],
+                             num_layers=self.args["sentence_num_layers"],
+                             dropout=self.args["dropout"],
+                             momentum=self.args["momentum"])
+
+        if self.model_state_dict is not None:
+            print("Sentence Model Loaded!")
+            model.load_state_dict(self.model_state_dict['sentence_model'])
+
+        if self.args["cuda"]:
+            model = model.cuda()
+
+        return model
+
+    def __init_word_word(self):
+        model = WordLSTM(vocab_size=len(self.vocab),
+                         embed_size=self.args["embed_size"],
+                         hidden_size=self.args["hidden_size"],
+                         num_layers=self.args["word_num_layers"],
+                         n_max=self.args["n_max"])
+
+        if self.model_state_dict is not None:
+            print("Word Model Loaded!")
+            model.load_state_dict(self.model_state_dict['word_model'])
+
+        if self.args["cuda"]:
+            model = model.cuda()
+
+        return model
+
+
+    
+def main(image):
+    sampler = CaptionSampler()
+    # image = 'sample_images/CXR195_IM-0618-1001.png'
+    caption  = sampler.sample(image)
+    print(caption[0])
+
+    return caption[0]
+

train_best_loss.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16673cee2882ab65a5f2e4fb23cb1b2d25cd484713ece7cdc910d791b5b79a59
+size 1535115128

val_best_loss.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74099c8e559e56a355ed8e9d8d7b1408ad98be7abc07f21122e5e57e93cf6dfc
+size 1535112572

vocab.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b87d71ea39483d3f3af078210a08d3b685f92ee679aa3d6030904844587d4d8
+size 31925

word.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41a13a1ba5b085ab7d3848cb42dc7cbbc866d7f00fad4e8a838c1369e511b949
+size 13216728