upload code

src/Datasets.py

@@ -0,0 +1,269 @@
+import os
+import json
+import torch
+import numpy as np
+from torch.utils.data import Dataset
+from PIL import Image
+from tqdm import tqdm
+import faiss
+import torch.nn.functional as F
+from sentence_transformers import SentenceTransformer
+import torchvision.transforms as transforms
+from random import choice
+
+
+class CCDataset(Dataset):
+    def __init__(self, json_file, root_dir, vocab, transform, split, max_length, s_pretrained, device):
+        super(CCDataset, self).__init__()
+        self.vocab = vocab
+        self.split = split
+        self.max_length = max_length
+        self.device = device
+        self.transform = transform
+        assert self.split in {'train', 'val', 'test'}
+
+        s_model = SentenceTransformer(s_pretrained)
+        self.s_model = s_model.to(device)
+
+        self.root_dir = root_dir
+        self.convert = transforms.ToTensor()
+
+        with open(json_file) as f:
+            data = json.load(f)['images']
+
+        self.raw_dataset = [entry for entry in data if entry['split'] == split]
+        self.sentences = []
+        self.embeddings = []
+
+        self.images = []
+        self.captions = []
+        for record in tqdm(self.raw_dataset, desc='Tokenize ' + self.split):
+            self.sentences.extend(self.tokenize(record['sentences']))
+
+        for record in tqdm(self.raw_dataset, desc='Embeddings ' + self.split):
+            self.embeddings.extend(self.compute_embeddings(record['sentences']))
+
+        self.preprocess()
+        del self.raw_dataset
+        del self.sentences
+        del self.embeddings
+        del self.s_model
+
+    def tokenize(self, batch):
+        for elem in batch:
+            tokens = [self.vocab[x] if x in self.vocab.keys() else self.vocab['UNK'] for x in elem['tokens']]
+            if len(tokens) > self.max_length - 2:
+                continue
+
+            tokens = [self.vocab['START']] + tokens + [self.vocab['END']]
+
+            mask = [False] * len(tokens)
+
+            diff = self.max_length - len(tokens)
+            tokens += [self.vocab['PAD']] * diff
+            mask += [True] * diff  # True = pad
+
+            elem['input_ids'] = tokens
+            elem['mask'] = mask
+
+        if len(batch) < 5:
+            diff = 5 - len(batch)
+            batch += [choice(batch) for _ in range(diff)]
+
+        assert len(batch) == 5
+        return batch
+
+    def compute_embeddings(self, batch):
+        sents = [x['raw'].strip() for x in batch]
+        embs = self.s_model.encode(sents)
+        return embs
+
+    def __len__(self):
+        return len(self.captions)
+
+    def __getitem__(self, idx):
+        img_idx = idx // 5 if self.split == 'train' else idx
+        elem = self.captions[idx]
+        for k, v in self.images[img_idx].items():
+            elem[k] = v
+        return elem
+
+    def preprocess(self):
+        idx = 0
+        prev_idx = -1
+        pbar = tqdm(total=len(self.sentences), desc='Preprocessing ' + self.split)
+        while idx < len(self.sentences):
+            img_idx = idx // 5
+            assert (self.sentences[idx]['imgid'] == self.raw_dataset[img_idx]['imgid'])
+
+            input_ids = torch.tensor(self.sentences[idx]['input_ids'], dtype=torch.long)
+            mask = torch.tensor(self.sentences[idx]['mask'], dtype=torch.bool)
+            raws = [x['raw'] for x in self.raw_dataset[img_idx]['sentences']]
+            flag = -1 if self.raw_dataset[img_idx]['changeflag'] == 0 else self.raw_dataset[img_idx]['imgid']
+            flag = torch.tensor(flag, dtype=torch.long)
+            embs = torch.tensor(self.embeddings[idx]) if len(self.embeddings) > 0 else None
+
+            self.captions.append({'input_ids': input_ids, 'pad_masks': mask, 'raws': raws, 'flags': flag, 'embs': embs})
+
+            if img_idx != prev_idx:
+                before_img_path = os.path.join(self.root_dir, self.raw_dataset[img_idx]['filepath'], 'A',
+                                               self.raw_dataset[img_idx]['filename'])
+                image_before = Image.open(before_img_path)
+                after_img_path = os.path.join(self.root_dir, self.raw_dataset[img_idx]['filepath'], 'B',
+                                              self.raw_dataset[img_idx]['filename'])
+                image_after = Image.open(after_img_path)
+
+                image_before = self.transform(image_before).unsqueeze(0)
+                image_after = self.transform(image_after).unsqueeze(0)
+
+                self.images.append({'image_before': image_before, 'image_after': image_after, 'flags': flag})
+                prev_idx = img_idx
+
+            inc = 1 if self.split == 'train' else 5
+            idx += inc
+            pbar.update(inc)
+
+        pbar.close()
+
+
+class Batcher:
+    def __init__(self, dataset, batch_size, max_len, device, hd=0, model=None, shuffle=False):
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.hd = hd
+        self.max_len = max_len
+        self.device = device
+        self.model = model
+        self.index = None
+        self.visual = None
+        self.textual = None
+
+        self.ptr = 0
+        self.indices = np.arange(len(self.dataset))
+        self.shuffle = shuffle
+
+        if shuffle:
+            np.random.shuffle(self.indices)
+
+        if model and hd > 0 and self.dataset.split == 'train':
+            self.create_index()
+
+    def __iter__(self):
+        return self
+
+    def __len__(self):
+        return len(self.dataset) // self.batch_size
+
+    def __next__(self):
+        if self.ptr >= len(self.dataset):
+            self.ptr = 0
+            self.index = None
+            self.visual = None
+            self.textual = None
+
+            if self.shuffle:
+                np.random.shuffle(self.indices)
+            if self.model and self.hd > 0 and self.dataset.split == 'train':
+                self.create_index()
+
+            raise StopIteration
+
+        batched = 0
+        samples = []
+        hard_negatives = []
+        while self.ptr < len(self.dataset) and batched < self.batch_size:
+            sample = self.dataset[self.indices[self.ptr]]
+            samples.append(sample)
+
+            if self.hd > 0 and self.dataset.split == 'train':
+                hard_neg = self.mine_negatives(self.indices[self.ptr], self.hd)
+                hard_negatives.extend(hard_neg)
+
+            self.ptr += 1
+            batched += 1
+
+        return self.create_batch(samples + hard_negatives)
+
+    def get_elem(self, idx):
+        return self.dataset[idx]
+
+    @torch.no_grad()
+    def create_index(self):
+        is_training = self.model.training
+        self.model.eval()
+        self.index = faiss.IndexFlatIP(self.model.feature_dim)
+        prev_img = None
+        for idx in tqdm(range(len(self.dataset)), desc='Creating index'):
+            sample = self.dataset[idx]
+            imgs1, imgs2, = sample['image_before'], sample['image_after']
+            input_ids, mask = sample['input_ids'], sample['pad_masks']
+
+            if idx // 5 != prev_img:
+                imgs1 = imgs1.to(self.device)
+                imgs2 = imgs2.to(self.device)
+                vis_emb, _, = self.model.encoder(imgs1, imgs2)
+                self.visual = torch.cat([self.visual, vis_emb.cpu()]) if self.visual is not None else vis_emb.cpu()
+                prev_img = prev_img + 1 if prev_img is not None else 0
+
+            input_ids = input_ids.unsqueeze(0).to(self.device)
+            mask = mask.unsqueeze(0).to(self.device)
+            _, text_emb, _, _ = self.model.decoder(input_ids, None, mask, None)
+            self.textual = torch.cat([self.textual, text_emb.cpu()]) if self.textual is not None else text_emb.cpu()
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        self.visual = F.normalize(self.visual, p=2, dim=1)
+        self.textual = F.normalize(self.textual, p=2, dim=1)
+        self.index.add(self.visual)
+        if is_training:
+            self.model.train()
+
+    def mine_negatives(self, idx, n):
+        negatives = []
+        m = 4
+        label = self.dataset[idx]['flags'].item()
+
+        while len(negatives) < n and (n * m) < self.index.ntotal:
+            k = n * m
+            indeces = self.index.search(self.textual[idx].unsqueeze(0), k)[1][0]
+            indeces = [x * 5 for x in indeces]
+            negatives = [self.dataset[x] for x in indeces if self.dataset[x]['flags'].item() != label][:n]
+            m *= 2
+
+        return negatives
+
+    def create_batch(self, samples):
+        images_before = images_after = input_ids = pad_mask = labels = flags = embs = None
+        raws = []
+
+        for sample in samples:
+            img1 = sample['image_before']
+            img2 = sample['image_after']
+
+            tokens = sample['input_ids']
+            mask = sample['pad_masks']
+            flag = sample['flags']
+            emb = sample['embs']
+
+            tokens = tokens.unsqueeze(0)
+            mask = mask.unsqueeze(0)
+            flag = flag.unsqueeze(0)
+            lab = tokens.clone() * ~mask
+            lab += torch.tensor([[-100]], dtype=torch.long).repeat(1, self.max_len) * mask
+            if emb is not None:
+                emb = emb.unsqueeze(0)
+
+            images_before = torch.cat([images_before, img1]) if images_before is not None else img1
+            images_after = torch.cat([images_after, img2]) if images_after is not None else img2
+            input_ids = torch.cat([input_ids, tokens]) if input_ids is not None else tokens
+            labels = torch.cat([labels, lab]) if labels is not None else lab
+            pad_mask = torch.cat([pad_mask, mask]) if pad_mask is not None else mask
+            flags = torch.cat([flags, flag]) if flags is not None else flag
+            if emb is not None:
+                embs = torch.cat([embs, emb]) if embs is not None else emb
+
+            raws.append(sample['raws'])
+
+        return {'images_before': images_before, 'images_after': images_after, 'input_ids': input_ids,
+                'pad_mask': pad_mask, 'labels': labels, 'flags': flags, 'raws': raws, 'embs': embs}

src/Loss.py

@@ -0,0 +1,84 @@
+from pytorch_metric_learning.distances import CosineSimilarity
+import torch
+
+
+class InfoNCELoss():
+    def __init__(self, device, k, temperature=0.07, threshold=1.0, fna=False):
+        super(InfoNCELoss, self).__init__()
+        self.device = device
+        self.similarity = CosineSimilarity()
+        self.k = k
+        self.temperature = temperature
+        self.threshold = threshold
+        self.fna = fna
+
+    def __call__(self, x, y, labels, sts):
+        false_negatives = self.detect_false_negative(sts)
+        indices_tuple = self.get_all_pairs_indices(labels, false_negatives)
+
+        mat = self.similarity(x, y)
+        a1, p, a2, n = indices_tuple
+        pos_pair, neg_pair = [], []
+        if len(a1) > 0:
+            pos_pair = mat[a1, p]
+        if len(a2) > 0:
+            neg_pair = mat[a2, n]
+
+        if len(neg_pair) > 0 and self.k > -1:
+            paired = list(zip(neg_pair.tolist(), a2.tolist(), n.tolist()))
+            selected = sorted(paired, key=lambda x: x[0], reverse=True)[:self.k]
+            _, x, y = list(zip(*selected))
+            x = torch.tensor(x).to(a2.device)
+            y = torch.tensor(y).to(n.device)
+
+            neg_pair = mat[x, y]
+            indices_tuple = (a1, p, x, y)
+
+        return self._compute_loss(pos_pair, neg_pair, indices_tuple), len(pos_pair)
+
+    def detect_false_negative(self, embs):
+        mat = torch.matmul(embs, torch.t(embs))
+        return torch.where(mat >= self.threshold)
+
+    def _compute_loss(self, pos_pairs, neg_pairs, indices_tuple):
+        a1, p, a2, _ = indices_tuple
+
+        if len(a1) > 0 and len(a2) > 0:
+            dtype = neg_pairs.dtype
+
+            if not self.similarity.is_inverted:
+                pos_pairs = -pos_pairs
+                neg_pairs = -neg_pairs
+
+            pos_pairs = pos_pairs.unsqueeze(1) / self.temperature
+            neg_pairs = neg_pairs / self.temperature
+            n_per_p = a2.unsqueeze(0) == a1.unsqueeze(1)
+            neg_pairs = neg_pairs * n_per_p
+            neg_pairs[n_per_p == 0] = torch.finfo(dtype).min
+
+            max_val = torch.max(
+                pos_pairs, torch.max(neg_pairs, dim=1, keepdim=True)[0]
+            ).detach()
+            numerator = torch.exp(pos_pairs - max_val).squeeze(1)
+            denominator = torch.sum(torch.exp(neg_pairs - max_val), dim=1) + numerator
+            log_exp = torch.log((numerator / denominator) + torch.finfo(dtype).tiny)
+            return torch.mean(-log_exp)
+
+        return 0
+
+    def get_all_pairs_indices(self, labels, false_negatives):
+        labels1 = labels.unsqueeze(1)
+        labels2 = labels.unsqueeze(0)
+        matches = (labels1 == labels2).byte()
+        diffs = matches ^ 1
+
+        diffs[false_negatives[0], false_negatives[1]] = 0  # FNE
+        if self.fna:
+            matches[false_negatives[0], false_negatives[1]] = 1  # FNA
+
+        diffs.fill_diagonal_(0)
+        matches.fill_diagonal_(1)
+
+        a1_idx, p_idx = torch.where(matches)
+        a2_idx, n_idx = torch.where(diffs)
+        return a1_idx, p_idx, a2_idx, n_idx

src/eval.py

@@ -0,0 +1,288 @@
+import argparse
+import random
+import os
+import matplotlib.pyplot as plt
+
+import numpy as np
+import torch
+import json
+import faiss
+from tqdm import tqdm
+import torch.nn.functional as F
+import torchvision.transforms as T
+
+import open_clip
+
+from Datasets import CCDataset, Batcher
+from model import ICCModel
+from utils import get_vocabulary, unormalize, get_eval_score
+
+AT_K = sorted([1, 3, 5, 10], reverse=True)
+
+
+def captioning(args, config, model, data_loader, vocab, device):
+    scores, results = inference(config, model, data_loader, vocab, device, return_results=True)
+
+    with open(os.path.join(args.output_path, 'caption.txt'), 'w') as out:
+        for t in scores.items():
+            out.write(str(t) + '\n')
+
+    scores, _ = inference(config, model, data_loader, vocab, device, sub=True, return_results=False)
+
+    with open(os.path.join(args.output_path, 'caption_sub.txt'), 'w') as out:
+        for t in scores.items():
+            out.write(str(t) + '\n')
+
+    return results
+
+
+def retrieve(args, config, model, src_loader, device):
+    scores_p, scores_r, scores_rr = search(config, model, src_loader, device)
+    with open(os.path.join(args.output_path, 'retrieve.txt'), 'w') as out:
+        for k in AT_K:
+            out.write('P@{0} {1:.4f}\n'.format(k, scores_p[k]))
+            out.write('R@{0} {1:.4f}\n'.format(k, scores_r[k]))
+            out.write('MRR@{0} {1:.4f}\n'.format(k, scores_rr[k]))
+            out.write('\n')
+
+    scores_p, scores_r, scores_rr = search(config, model, src_loader, device, sub=True)
+    with open(os.path.join(args.output_path, 'retrieve_sub.txt'), 'w') as out:
+        for k in AT_K:
+            out.write('P@{0} {1:.4f}\n'.format(k, scores_p[k]))
+            out.write('R@{0} {1:.4f}\n'.format(k, scores_r[k]))
+            out.write('MRR@{0} {1:.4f}\n'.format(k, scores_rr[k]))
+            out.write('\n')
+
+
+@torch.no_grad()
+def search(config, model, src_loader, device, sub=False):
+    model.eval()
+
+    visual = None
+    textual = None
+    flags = []
+    embs = None
+    index = faiss.IndexFlatIP(config['d_model'])
+
+    batcher = src_loader
+
+    for batch in tqdm(batcher, desc='Indexing'):
+        imgs1, imgs2, = batch['images_before'], batch['images_after']
+        imgs1 = imgs1.to(device)
+        imgs2 = imgs2.to(device)
+        flag = batch['flags']
+        emb = batch['embs']
+        if sub and flag[0] == -1:
+            continue
+
+        flags.append(flag)
+        embs = torch.cat([embs, emb]) if embs is not None else emb
+
+        vis_emb, _ = model.encoder(imgs1, imgs2)
+        visual = torch.cat([visual, vis_emb.cpu()]) if visual is not None else vis_emb.cpu()
+
+        input_ids, mask = batch['input_ids'], batch['pad_mask']
+        input_ids = input_ids.to(device)
+        mask = mask.to(device)
+
+        _, text_emb, _, _ = model.decoder(input_ids, None, mask, None)
+        textual = torch.cat([textual, text_emb.cpu()]) if textual is not None else text_emb.cpu()
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    embs = embs.to(device)
+    sims = torch.matmul(embs, torch.t(embs))
+
+    visual = F.normalize(visual, p=2, dim=1)
+    textual = F.normalize(textual, p=2, dim=1)
+
+    index.add(visual)
+
+    scores_p = {k: [] for k in AT_K}
+    scores_r = {k: [] for k in AT_K}
+    scores_rr = {k: [] for k in AT_K}
+
+    for i in tqdm(range(textual.shape[0]), desc='Ranking'):
+        indices = None
+        query = textual[i]
+        query_lab = flags[i]
+
+        relevants = set(
+            [x for x in range(len(textual)) if flags[x] == query_lab or sims[i][x] >= config['s-threshold']])
+
+        for k in AT_K:
+            p = 0
+            r = 0
+            rr = 0
+
+            if indices is None:
+                indices = index.search(query.unsqueeze(0), k)[1][0]
+            else:
+                indices = indices[:k]
+
+            for rank, idx in enumerate(indices):
+                if idx in relevants:
+                    if p == 0:
+                        rr = 1 / (rank + 1)
+                    p += 1
+                    r += 1
+
+            scores_p[k].append(p / len(indices))
+            scores_r[k].append(r / len(relevants))
+            scores_rr[k].append(rr)
+
+    for k in AT_K:
+        scores_p[k] = sum(scores_p[k]) / len(scores_p[k])
+        scores_r[k] = sum(scores_r[k]) / len(scores_r[k])
+        scores_rr[k] = sum(scores_rr[k]) / len(scores_rr[k])
+
+    return scores_p, scores_r, scores_rr
+
+
+@torch.no_grad()
+def inference(config, model, data_loader, vocab, device, sub=False, return_results=False):
+    results = []
+    references = []
+    hypotheses = []
+    inverse_vocab = {v: k for k, v in vocab.items()}
+
+    model.eval()
+    for batch in tqdm(data_loader, desc='Inference'):
+        img1 = batch['images_before'][0].unsqueeze(0).to(device)
+        img2 = batch['images_after'][0].unsqueeze(0).to(device)
+        raws = batch['raws']
+        flags = batch['flags']
+        if sub and flags[0] == -1:
+            continue
+
+        references.append(raws[0])
+
+        input_ids = torch.tensor([[vocab['START']]], dtype=torch.long, device=device)
+        _, vis_toks = model.encoder(img1, img2)
+
+        for _ in range(config['max_len']):
+            _, _, lm_logits, weights = model.decoder(input_ids, None, None, vis_toks)
+
+            next_item = lm_logits[0][-1].topk(1)[1]
+            input_ids = torch.cat([input_ids, next_item.reshape(1, -1)], dim=1)
+            if next_item.item() == vocab['END']:
+                break
+
+        words = [inverse_vocab[x] for x in input_ids[0].cpu().tolist()]
+        sentence = ' '.join(words[1:-1]).strip()
+        hypotheses.append([sentence])
+
+        if return_results:
+            results.append(
+                (img1.cpu(), img2.cpu(), weights.detach().cpu(), vis_toks.detach().cpu(), sentence))
+
+    score_dict = get_eval_score(references, hypotheses)
+    return score_dict, results
+
+
+def plot(args, feat_size, results):
+    fig_idx = 0
+    for img1, img2, weights, diff, sentence in tqdm(results, desc='Plot'):
+        img1 = unormalize(img1)
+        img1 = img1[0].permute(1, 2, 0)  # h,w,c
+        img2 = unormalize(img2)
+        img2 = img2[0].permute(1, 2, 0)  # h,w,c
+
+        transform = T.Resize(size=(img1.size(0), img1.size(1)))
+        weights = weights[0].reshape(-1, feat_size, feat_size)
+        weights = transform(weights).permute(1, 2, 0)  # h,w,d
+        weights = torch.sum(weights, 2) / weights.shape[2]
+        after = img2  # h,w,c
+
+        feature_map = diff[:, 0, :].reshape(-1, feat_size, feat_size)  # e,h,w
+        feature_map = transform(feature_map).permute(1, 2, 0)  # h,w,c
+        feature_map = torch.sum(feature_map, 2) / feature_map.shape[2]  # h, w
+
+        fig, ax = plt.subplots(2, 2, figsize=(6, 8))
+        fig.tight_layout()
+        ax[0, 0].imshow(img1)
+        ax[0, 0].set_title("Before")
+        ax[0, 0].axis('off')
+        ax[0, 1].imshow(img2)
+        ax[0, 1].set_title("After")
+        ax[0, 1].axis('off')
+
+        ax[1, 0].set_title("Img diff")
+        ax[1, 0].imshow(feature_map)
+        ax[1, 0].axis('off')
+
+        ax[1, 1].set_title("Att weights")
+        ax[1, 1].imshow(after, interpolation='nearest')
+        ax[1, 1].imshow(weights, interpolation='bilinear', alpha=0.5)
+        ax[1, 1].axis('off')
+
+        fig.text(.1, .05, sentence, wrap=True)
+
+        with open(os.path.join(args.output_path, str(fig_idx) + '.png'), 'wb') as f:
+            plt.savefig(f)
+            plt.close()
+            fig_idx += 1
+
+
+def run(args, config):
+    print('Initializing...')
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    random.seed(args.seed)
+    torch.backends.cudnn.deterministic = True
+
+    device = torch.device('cpu')
+    if torch.cuda.is_available():
+        device = torch.device('cuda')
+
+    if os.path.exists(args.vocab):
+        with open(args.vocab, 'r') as infile:
+            vocab = json.load(infile)
+    else:
+        vocab = get_vocabulary(args.annotation_json, args.vocab)
+
+    clip, _, preprocess = open_clip.create_model_and_transforms(config['backbone'])
+
+    model = ICCModel(device, clip, config['backbone'], config['d_model'],
+                     len(vocab), config['max_len'], config['num_heads'], config['h_dim'], config['a_dim'],
+                     config['encoder_layers'], config['decoder_layers'], config['dropout'],
+                     learnable=config['learnable'], fine_tune=config['fine_tune'],
+                     tie_embeddings=config['tie_embeddings'], prenorm=config['prenorm'])
+
+    model.load_state_dict(torch.load(args.model, map_location=device))
+    model = model.to(device)
+    del clip
+
+    print('Loading...')
+    test_set = CCDataset(args.annotation_json, args.image_dir, vocab, preprocess, 'test', config['max_len'],
+                         config['s-transformers'], device)
+    test_loader = Batcher(test_set, 1, config['max_len'], device)
+
+    print('Final evaluation...')
+    results = captioning(args, config, model, test_loader, vocab, device)
+    retrieve(args, config, model, test_loader, device)
+    plot(args, model.encoder.encoder.feat_size, results)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model', type=str, default='../input/model_best.pt')
+    parser.add_argument('--annotation_json', type=str, default='../input/Levir_CC/LevirCCcaptions.json')
+    parser.add_argument('--image_dir', type=str, default='../input/Levir_CC/images/')
+    parser.add_argument('--vocab', type=str, default='../input/levir_vocab.json')
+
+    parser.add_argument('--config', type=str, default='../config.json')
+    parser.add_argument('--output_path', type=str, default='../output/')
+    parser.add_argument('--seed', type=int, default=42)
+
+    args = parser.parse_args()
+
+    with open(args.config, 'r') as config_file:
+        config = json.load(config_file)
+
+    run(args, config)
+
+
+if __name__ == '__main__':
+    main()

src/eval_func/bleu/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2015 Xinlei Chen, Hao Fang, Tsung-Yi Lin, and Ramakrishna Vedantam
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

src/eval_func/bleu/__init__.py

@@ -0,0 +1 @@
+__author__ = 'tylin'

src/eval_func/bleu/bleu.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python
+# 
+# File Name : bleu.py
+#
+# Description : Wrapper for BLEU scorer.
+#
+# Creation Date : 06-01-2015
+# Last Modified : Thu 19 Mar 2015 09:13:28 PM PDT
+# Authors : Hao Fang <hfang@uw.edu> and Tsung-Yi Lin <tl483@cornell.edu>
+
+from .bleu_scorer import BleuScorer
+
+
+class Bleu:
+    def __init__(self, n=4):
+        # default compute Blue score up to 4
+        self._n = n
+        self._hypo_for_image = {}
+        self.ref_for_image = {}
+
+    def compute_score(self, gts, res):
+
+        bleu_scorer = BleuScorer(n=self._n)
+        for i in range(len(res)):
+            hypo = res[i]
+            ref = gts[i]
+
+            # Sanity check.
+            assert(type(hypo) is list)
+            assert(len(hypo) == 1)
+            assert(type(ref) is list)
+            assert(len(ref) >= 1)
+
+            bleu_scorer += (hypo[0], ref)
+
+        #score, scores = bleu_scorer.compute_score(option='shortest')
+        score, scores = bleu_scorer.compute_score(option='closest', verbose=1)
+        #score, scores = bleu_scorer.compute_score(option='average', verbose=1)
+
+        # return (bleu, bleu_info)
+        return score, scores
+
+    def method(self):
+        return "Bleu"

src/eval_func/bleu/bleu_scorer.py

@@ -0,0 +1,263 @@
+#!/usr/bin/env python
+
+# bleu_scorer.py
+# David Chiang <chiang@isi.edu>
+
+# Copyright (c) 2004-2006 University of Maryland. All rights
+# reserved. Do not redistribute without permission from the
+# author. Not for commercial use.
+
+# Modified by: 
+# Hao Fang <hfang@uw.edu>
+# Tsung-Yi Lin <tl483@cornell.edu>
+
+'''Provides:
+cook_refs(refs, n=4): Transform a list of reference sentences as strings into a form usable by cook_test().
+cook_test(test, refs, n=4): Transform a test sentence as a string (together with the cooked reference sentences) into a form usable by score_cooked().
+'''
+
+import copy
+import sys, math, re
+from collections import defaultdict
+
+def precook(s, n=4, out=False):
+    """Takes a string as input and returns an object that can be given to
+    either cook_refs or cook_test. This is optional: cook_refs and cook_test
+    can take string arguments as well."""
+    words = s.split()
+    counts = defaultdict(int)
+    for k in range(1,n+1):
+        for i in range(len(words)-k+1):
+            ngram = tuple(words[i:i+k])
+            counts[ngram] += 1
+    return (len(words), counts)
+
+def cook_refs(refs, eff=None, n=4): ## lhuang: oracle will call with "average"
+    '''Takes a list of reference sentences for a single segment
+    and returns an object that encapsulates everything that BLEU
+    needs to know about them.'''
+
+    reflen = []
+    maxcounts = {}
+    for ref in refs:
+        rl, counts = precook(ref, n)
+        reflen.append(rl)
+        for (ngram,count) in counts.items():
+            maxcounts[ngram] = max(maxcounts.get(ngram,0), count)
+
+    # Calculate effective reference sentence length.
+    if eff == "shortest":
+        reflen = min(reflen)
+    elif eff == "average":
+        reflen = float(sum(reflen))/len(reflen)
+
+    ## lhuang: N.B.: leave reflen computaiton to the very end!!
+    
+    ## lhuang: N.B.: in case of "closest", keep a list of reflens!! (bad design)
+
+    return (reflen, maxcounts)
+
+def cook_test(test, xxx_todo_changeme, eff=None, n=4):
+    '''Takes a test sentence and returns an object that
+    encapsulates everything that BLEU needs to know about it.'''
+    (reflen, refmaxcounts) = xxx_todo_changeme
+    testlen, counts = precook(test, n, True)
+
+    result = {}
+
+    # Calculate effective reference sentence length.
+    
+    if eff == "closest":
+        result["reflen"] = min((abs(l-testlen), l) for l in reflen)[1]
+    else: ## i.e., "average" or "shortest" or None
+        result["reflen"] = reflen
+
+    result["testlen"] = testlen
+
+    result["guess"] = [max(0,testlen-k+1) for k in range(1,n+1)]
+
+    result['correct'] = [0]*n
+    for (ngram, count) in counts.items():
+        result["correct"][len(ngram)-1] += min(refmaxcounts.get(ngram,0), count)
+
+    return result
+
+class BleuScorer(object):
+    """Bleu scorer.
+    """
+
+    __slots__ = "n", "crefs", "ctest", "_score", "_ratio", "_testlen", "_reflen", "special_reflen"
+    # special_reflen is used in oracle (proportional effective ref len for a node).
+
+    def copy(self):
+        ''' copy the refs.'''
+        new = BleuScorer(n=self.n)
+        new.ctest = copy.copy(self.ctest)
+        new.crefs = copy.copy(self.crefs)
+        new._score = None
+        return new
+
+    def __init__(self, test=None, refs=None, n=4, special_reflen=None):
+        ''' singular instance '''
+
+        self.n = n
+        self.crefs = []
+        self.ctest = []
+        self.cook_append(test, refs)
+        self.special_reflen = special_reflen
+
+    def cook_append(self, test, refs):
+        '''called by constructor and __iadd__ to avoid creating new instances.'''
+        
+        if refs is not None:
+            self.crefs.append(cook_refs(refs))
+            if test is not None:
+                cooked_test = cook_test(test, self.crefs[-1])
+                self.ctest.append(cooked_test) ## N.B.: -1
+            else:
+                self.ctest.append(None) # lens of crefs and ctest have to match
+
+        self._score = None ## need to recompute
+
+    def ratio(self, option=None):
+        self.compute_score(option=option)
+        return self._ratio
+
+    def score_ratio(self, option=None):
+        '''return (bleu, len_ratio) pair'''
+        return (self.fscore(option=option), self.ratio(option=option))
+
+    def score_ratio_str(self, option=None):
+        return "%.4f (%.2f)" % self.score_ratio(option)
+
+    def reflen(self, option=None):
+        self.compute_score(option=option)
+        return self._reflen
+
+    def testlen(self, option=None):
+        self.compute_score(option=option)
+        return self._testlen        
+
+    def retest(self, new_test):
+        if type(new_test) is str:
+            new_test = [new_test]
+        assert len(new_test) == len(self.crefs), new_test
+        self.ctest = []
+        for t, rs in zip(new_test, self.crefs):
+            self.ctest.append(cook_test(t, rs))
+        self._score = None
+
+        return self
+
+    def rescore(self, new_test):
+        ''' replace test(s) with new test(s), and returns the new score.'''
+        
+        return self.retest(new_test).compute_score()
+
+    def size(self):
+        assert len(self.crefs) == len(self.ctest), "refs/test mismatch! %d<>%d" % (len(self.crefs), len(self.ctest))
+        return len(self.crefs)
+
+    def __iadd__(self, other):
+        '''add an instance (e.g., from another sentence).'''
+
+        if type(other) is tuple:
+            ## avoid creating new BleuScorer instances
+            self.cook_append(other[0], other[1])
+        else:
+            assert self.compatible(other), "incompatible BLEUs."
+            self.ctest.extend(other.ctest)
+            self.crefs.extend(other.crefs)
+            self._score = None ## need to recompute
+
+        return self        
+
+    def compatible(self, other):
+        return isinstance(other, BleuScorer) and self.n == other.n
+
+    def single_reflen(self, option="average"):
+        return self._single_reflen(self.crefs[0][0], option)
+
+    def _single_reflen(self, reflens, option=None, testlen=None):
+        
+        if option == "shortest":
+            reflen = min(reflens)
+        elif option == "average":
+            reflen = float(sum(reflens))/len(reflens)
+        elif option == "closest":
+            reflen = min((abs(l-testlen), l) for l in reflens)[1]
+        else:
+            assert False, "unsupported reflen option %s" % option
+
+        return reflen
+
+    def recompute_score(self, option=None, verbose=0):
+        self._score = None
+        return self.compute_score(option, verbose)
+        
+    def compute_score(self, option=None, verbose=0):
+        n = self.n
+        small = 1e-9
+        tiny = 1e-15 ## so that if guess is 0 still return 0
+        bleu_list = [[] for _ in range(n)]
+
+        if self._score is not None:
+            return self._score
+
+        if option is None:
+            option = "average" if len(self.crefs) == 1 else "closest"
+
+        self._testlen = 0
+        self._reflen = 0
+        totalcomps = {'testlen':0, 'reflen':0, 'guess':[0]*n, 'correct':[0]*n}
+
+        # for each sentence
+        for comps in self.ctest:            
+            testlen = comps['testlen']
+            self._testlen += testlen
+
+            if self.special_reflen is None: ## need computation
+                reflen = self._single_reflen(comps['reflen'], option, testlen)
+            else:
+                reflen = self.special_reflen
+
+            self._reflen += reflen
+                
+            for key in ['guess','correct']:
+                for k in range(n):
+                    totalcomps[key][k] += comps[key][k]
+
+            # append per image bleu score
+            bleu = 1.
+            for k in range(n):
+                bleu *= (float(comps['correct'][k]) + tiny) \
+                        /(float(comps['guess'][k]) + small) 
+                bleu_list[k].append(bleu ** (1./(k+1)))
+            ratio = (testlen + tiny) / (reflen + small) ## N.B.: avoid zero division
+            if ratio < 1:
+                for k in range(n):
+                    bleu_list[k][-1] *= math.exp(1 - 1/ratio)
+
+            # if verbose > 1:
+            #     print(comps, reflen)
+
+        totalcomps['reflen'] = self._reflen
+        totalcomps['testlen'] = self._testlen
+
+        bleus = []
+        bleu = 1.
+        for k in range(n):
+            bleu *= float(totalcomps['correct'][k] + tiny) \
+                    / (totalcomps['guess'][k] + small)
+            bleus.append(bleu ** (1./(k+1)))
+        ratio = (self._testlen + tiny) / (self._reflen + small) ## N.B.: avoid zero division
+        if ratio < 1:
+            for k in range(n):
+                bleus[k] *= math.exp(1 - 1/ratio)
+
+        # if verbose > 0:
+        #     print(totalcomps)
+        #     print("ratio:", ratio)
+
+        self._score = bleus
+        return self._score, bleu_list

src/eval_func/cider/__init__.py

@@ -0,0 +1 @@
+__author__ = 'tylin'

src/eval_func/cider/cider.py

@@ -0,0 +1,51 @@
+# Filename: cider.py
+#
+# Description: Describes the class to compute the CIDEr (Consensus-Based Image Description Evaluation) Metric 
+#               by Vedantam, Zitnick, and Parikh (http://arxiv.org/abs/1411.5726)
+#
+# Creation Date: Sun Feb  8 14:16:54 2015
+#
+# Authors: Ramakrishna Vedantam <vrama91@vt.edu> and Tsung-Yi Lin <tl483@cornell.edu>
+
+from eval_func.cider.cider_scorer import CiderScorer
+import pdb
+
+class Cider:
+    """
+    Main Class to compute the CIDEr metric 
+
+    """
+    def __init__(self, test=None, refs=None, n=4, sigma=6.0):
+        # set cider to sum over 1 to 4-grams
+        self._n = n
+        # set the standard deviation parameter for gaussian penalty
+        self._sigma = sigma
+
+    def compute_score(self, gts, res):
+        """
+        Main function to compute CIDEr score
+        :param  hypo_for_image (dict) : dictionary with key <image> and value <tokenized hypothesis / candidate sentence>
+                ref_for_image (dict)  : dictionary with key <image> and value <tokenized reference sentence>
+        :return: cider (float) : computed CIDEr score for the corpus 
+        """
+
+        cider_scorer = CiderScorer(n=self._n, sigma=self._sigma)
+
+        for i in range(len(res)):
+            hypo = res[i]
+            ref = gts[i]
+
+            # Sanity check.
+            assert(type(hypo) is list)
+            assert(len(hypo) == 1)
+            assert(type(ref) is list)
+            assert(len(ref) > 0)
+
+            cider_scorer += (hypo[0], ref)
+
+        (score, scores) = cider_scorer.compute_score()
+
+        return score, scores
+
+    def method(self):
+        return "CIDEr"

src/eval_func/cider/cider_scorer.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python
+# Tsung-Yi Lin <tl483@cornell.edu>
+# Ramakrishna Vedantam <vrama91@vt.edu>
+
+import copy
+from collections import defaultdict
+import numpy as np
+import pdb
+import math
+
+def precook(s, n=4, out=False):
+    """
+    Takes a string as input and returns an object that can be given to
+    either cook_refs or cook_test. This is optional: cook_refs and cook_test
+    can take string arguments as well.
+    :param s: string : sentence to be converted into ngrams
+    :param n: int    : number of ngrams for which representation is calculated
+    :return: term frequency vector for occuring ngrams
+    """
+    words = s.split()
+    counts = defaultdict(int)
+    for k in range(1,n+1):
+        for i in range(len(words)-k+1):
+            ngram = tuple(words[i:i+k])
+            counts[ngram] += 1
+    return counts
+
+def cook_refs(refs, n=4): ## lhuang: oracle will call with "average"
+    '''Takes a list of reference sentences for a single segment
+    and returns an object that encapsulates everything that BLEU
+    needs to know about them.
+    :param refs: list of string : reference sentences for some image
+    :param n: int : number of ngrams for which (ngram) representation is calculated
+    :return: result (list of dict)
+    '''
+    return [precook(ref, n) for ref in refs]
+
+def cook_test(test, n=4):
+    '''Takes a test sentence and returns an object that
+    encapsulates everything that BLEU needs to know about it.
+    :param test: list of string : hypothesis sentence for some image
+    :param n: int : number of ngrams for which (ngram) representation is calculated
+    :return: result (dict)
+    '''
+    return precook(test, n, True)
+
+class CiderScorer(object):
+    """CIDEr scorer.
+    """
+
+    def copy(self):
+        ''' copy the refs.'''
+        new = CiderScorer(n=self.n)
+        new.ctest = copy.copy(self.ctest)
+        new.crefs = copy.copy(self.crefs)
+        return new
+
+    def __init__(self, test=None, refs=None, n=4, sigma=6.0):
+        ''' singular instance '''
+        self.n = n
+        self.sigma = sigma
+        self.crefs = []
+        self.ctest = []
+        self.document_frequency = defaultdict(float)
+        self.cook_append(test, refs)
+        self.ref_len = None
+
+    def cook_append(self, test, refs):
+        '''called by constructor and __iadd__ to avoid creating new instances.'''
+
+        if refs is not None:
+            self.crefs.append(cook_refs(refs))
+            if test is not None:
+                self.ctest.append(cook_test(test)) ## N.B.: -1
+            else:
+                self.ctest.append(None) # lens of crefs and ctest have to match
+
+    def size(self):
+        assert len(self.crefs) == len(self.ctest), "refs/test mismatch! %d<>%d" % (len(self.crefs), len(self.ctest))
+        return len(self.crefs)
+
+    def __iadd__(self, other):
+        '''add an instance (e.g., from another sentence).'''
+
+        if type(other) is tuple:
+            ## avoid creating new CiderScorer instances
+            self.cook_append(other[0], other[1])
+        else:
+            self.ctest.extend(other.ctest)
+            self.crefs.extend(other.crefs)
+
+        return self
+    def compute_doc_freq(self):
+        '''
+        Compute term frequency for reference data.
+        This will be used to compute idf (inverse document frequency later)
+        The term frequency is stored in the object
+        :return: None
+        '''
+        for refs in self.crefs:
+            # refs, k ref captions of one image
+            for ngram in set([ngram for ref in refs for (ngram,count) in ref.items()]):
+                self.document_frequency[ngram] += 1
+            # maxcounts[ngram] = max(maxcounts.get(ngram,0), count)
+
+    def compute_cider(self):
+        def counts2vec(cnts):
+            """
+            Function maps counts of ngram to vector of tfidf weights.
+            The function returns vec, an array of dictionary that store mapping of n-gram and tf-idf weights.
+            The n-th entry of array denotes length of n-grams.
+            :param cnts:
+            :return: vec (array of dict), norm (array of float), length (int)
+            """
+            vec = [defaultdict(float) for _ in range(self.n)]
+            length = 0
+            norm = [0.0 for _ in range(self.n)]
+            for (ngram, term_freq) in cnts.items():
+                # give word count 1 if it doesn't appear in reference corpus
+                df = np.log(max(1.0, self.document_frequency[ngram]))
+                # ngram index
+                n = len(ngram)-1
+                # tf (term_freq) * idf (precomputed idf) for n-grams
+                vec[n][ngram] = float(term_freq)*(self.ref_len - df)
+                # compute norm for the vector.  the norm will be used for computing similarity
+                norm[n] += pow(vec[n][ngram], 2)
+
+                if n == 1:
+                    length += term_freq
+            norm = [np.sqrt(n) for n in norm]
+            return vec, norm, length
+
+        def sim(vec_hyp, vec_ref, norm_hyp, norm_ref, length_hyp, length_ref):
+            '''
+            Compute the cosine similarity of two vectors.
+            :param vec_hyp: array of dictionary for vector corresponding to hypothesis
+            :param vec_ref: array of dictionary for vector corresponding to reference
+            :param norm_hyp: array of float for vector corresponding to hypothesis
+            :param norm_ref: array of float for vector corresponding to reference
+            :param length_hyp: int containing length of hypothesis
+            :param length_ref: int containing length of reference
+            :return: array of score for each n-grams cosine similarity
+            '''
+            delta = float(length_hyp - length_ref)
+            # measure consine similarity
+            val = np.array([0.0 for _ in range(self.n)])
+            for n in range(self.n):
+                # ngram
+                for (ngram,count) in vec_hyp[n].items():
+                    # vrama91 : added clipping
+                    val[n] += min(vec_hyp[n][ngram], vec_ref[n][ngram]) * vec_ref[n][ngram]
+
+                if (norm_hyp[n] != 0) and (norm_ref[n] != 0):
+                    val[n] /= (norm_hyp[n]*norm_ref[n])
+
+                assert(not math.isnan(val[n]))
+                # vrama91: added a length based gaussian penalty
+                val[n] *= np.e**(-(delta**2)/(2*self.sigma**2))
+            return val
+
+        # compute log reference length
+        self.ref_len = np.log(float(len(self.crefs)))
+        if len(self.crefs) == 1:
+            self.ref_len = 1
+        scores = []
+        for test, refs in zip(self.ctest, self.crefs):
+            # compute vector for test captions
+            vec, norm, length = counts2vec(test)
+            # compute vector for ref captions
+            score = np.array([0.0 for _ in range(self.n)])
+            for ref in refs:
+                vec_ref, norm_ref, length_ref = counts2vec(ref)
+                score += sim(vec, vec_ref, norm, norm_ref, length, length_ref)
+            # change by vrama91 - mean of ngram scores, instead of sum
+            score_avg = np.mean(score)
+            # divide by number of references
+            score_avg /= len(refs)
+            # multiply score by 10
+            score_avg *= 10.0
+            # append score of an image to the score list
+            scores.append(score_avg)
+        return scores
+
+    def compute_score(self, option=None, verbose=0):
+        # compute idf
+        self.compute_doc_freq()
+        # assert to check document frequency
+        assert(len(self.ctest) >= max(self.document_frequency.values()))
+        # compute cider score
+        score = self.compute_cider()
+        # debug
+        # print score
+        return np.mean(np.array(score)), np.array(score)

src/eval_func/rouge/__init__.py

@@ -0,0 +1 @@
+__author__ = 'vrama91'

src/eval_func/rouge/rouge.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python
+
+#
+
+# File Name : rouge.py
+
+#
+
+# Description : Computes ROUGE-L metric as described by Lin and Hovey (2004)
+
+#
+
+# Creation Date : 2015-01-07 06:03
+
+# Author : Ramakrishna Vedantam <vrama91@vt.edu>
+
+
+import numpy as np
+
+import pdb
+
+
+def my_lcs(string, sub):
+    """
+
+    Calculates longest common subsequence for a pair of tokenized strings
+
+    :param string : list of str : tokens from a string split using whitespace
+
+    :param sub : list of str : shorter string, also split using whitespace
+
+    :returns: length (list of int): length of the longest common subsequence between the two strings
+
+
+
+    Note: my_lcs only gives length of the longest common subsequence, not the actual LCS
+
+    """
+
+    if (len(string) < len(sub)):
+        sub, string = string, sub
+
+    lengths = [[0 for i in range(0, len(sub) + 1)] for j in range(0, len(string) + 1)]
+
+    for j in range(1, len(sub) + 1):
+
+        for i in range(1, len(string) + 1):
+
+            if (string[i - 1] == sub[j - 1]):
+
+                lengths[i][j] = lengths[i - 1][j - 1] + 1
+
+            else:
+
+                lengths[i][j] = max(lengths[i - 1][j], lengths[i][j - 1])
+
+    return lengths[len(string)][len(sub)]
+
+
+class Rouge():
+    '''
+
+    Class for computing ROUGE-L score for a set of candidate sentences for the MS COCO test set
+
+
+
+    '''
+
+    def __init__(self):
+
+        # vrama91: updated the value below based on discussion with Hovey
+
+        self.beta = 1.2
+
+    def calc_score(self, candidate, refs):
+
+        """
+
+        Compute ROUGE-L score given one candidate and references for an image
+
+        :param candidate: str : candidate sentence to be evaluated
+
+        :param refs: list of str : COCO reference sentences for the particular image to be evaluated
+
+        :returns score: int (ROUGE-L score for the candidate evaluated against references)
+
+        """
+
+
+        assert (len(candidate) == 1)
+
+        assert (len(refs) > 0)
+
+        prec = []
+
+        rec = []
+
+        # split into tokens
+
+        token_c = candidate[0].split(" ")
+
+        for reference in refs:
+            # split into tokens
+            hh =1
+
+            token_r = reference.split(" ")
+
+            # compute the longest common subsequence
+
+            lcs = my_lcs(token_r, token_c)
+
+            prec.append(lcs / float(len(token_c)))
+
+            rec.append(lcs / float(len(token_r)))
+
+        prec_max = max(prec)
+
+        rec_max = max(rec)
+
+        if (prec_max != 0 and rec_max != 0):
+
+            score = ((1 + self.beta ** 2) * prec_max * rec_max) / float(rec_max + self.beta ** 2 * prec_max)
+
+        else:
+
+            score = 0.0
+
+        return score
+
+    def compute_score(self, references, hypotheses):
+
+        """
+
+        Computes Rouge-L score given a set of reference and candidate sentences for the dataset
+
+        Invoked by evaluate_captions.py
+
+        :param hypo_for_image: dict : candidate / test sentences with "image name" key and "tokenized sentences" as values
+
+        :param ref_for_image: dict : reference MS-COCO sentences with "image name" key and "tokenized sentences" as values
+
+        :returns: average_score: float (mean ROUGE-L score computed by averaging scores for all the images)
+
+        """
+
+        # assert (gts.keys() == res.keys())
+        #
+        # imgIds = gts.keys()
+
+        score = []
+
+        for i in range(len(hypotheses)):
+            hypo = hypotheses[i]
+            ref = references[i]
+
+            score.append(self.calc_score(hypo, ref))
+
+            # Sanity check.
+
+            assert (type(hypo) is list)
+
+            assert (len(hypo) == 1)
+
+            assert (type(ref) is list)
+
+            assert (len(ref) > 0)
+
+        average_score = np.mean(np.array(score))
+
+        return average_score, np.array(score)
+
+    def method(self):
+
+        return "Rouge"

src/model.py

@@ -0,0 +1,450 @@
+import torch
+from torch import nn
+from einops import rearrange
+import math
+from torch import Tensor
+import torchvision.models as models
+
+from torch.nn import functional as F
+
+
+class ICCModel(nn.Module):
+    def __init__(self, device, pretrained, backbone, d_model, vocab_size, max_len,
+                 num_heads, h_dim, a_dim, encoder_layers, decoder_layers, dropout,
+                 learnable=False, fine_tune=True, tie_embeddings=True, prenorm=False):
+        super(ICCModel, self).__init__()
+
+        self.feature_dim = d_model
+        visual = pretrained.visual if pretrained else None
+        self.encoder = ImagesEncoder(device, visual, backbone, d_model, num_heads, h_dim, a_dim, dropout,
+                                     encoder_layers, fine_tune)
+
+        self.decoder = Decoder(device, d_model, vocab_size, max_len, num_heads,
+                               decoder_layers, dropout,
+                               learnable=learnable, tie_embeddings=tie_embeddings, prenorm=prenorm)
+
+    def forward(self, img1, img2, input_ids, labels, attention_mask):
+        vis_emb, vis_toks = self.encoder(img1, img2)
+        cap_loss, text_emb, lm_logits, weights = self.decoder(input_ids, labels, attention_mask, vis_toks)
+        return cap_loss, vis_emb, text_emb, vis_toks, lm_logits, weights
+
+
+class ImagesEncoder(nn.Module):
+    def __init__(self, device, pretrained, backbone, d_model, num_heads, h_dim, a_dim, dropout,
+                 encoder_layers, fine_tune):
+        super(ImagesEncoder, self).__init__()
+        self.encoder = Encoder(pretrained, backbone, d_model, fine_tune)
+        self.encoder_trans = AttentiveEncoder(device, encoder_layers,
+                                              [self.encoder.feat_size, self.encoder.feat_size, d_model], num_heads,
+                                              hidden_dim=h_dim, attention_dim=a_dim, dropout=dropout)
+
+        self.cos = torch.nn.CosineSimilarity(dim=1)
+        self.Conv1 = nn.Conv2d(d_model * 2, d_model, kernel_size=1)
+        self.LN = resblock(d_model, d_model)
+
+        self.att_pool = nn.MultiheadAttention(d_model, num_heads)
+        self.att_pool_norm = nn.LayerNorm(d_model)
+        self.img_queries = nn.Parameter(torch.randn(1, d_model))
+
+    def forward(self, img1, img2):
+        feat1 = self.encoder(img1)
+        feat2 = self.encoder(img2)
+        x1, x2 = self.encoder_trans(feat1, feat2)  # batch_size, channel, enc_image_size, enc_image_size
+
+        x_sam = self.cos(x1, x2)
+        x = torch.cat([x1, x2], dim=1) + x_sam.unsqueeze(1)  # batch_size, 2channel, enc_image_size, enc_image_size
+        x = self.LN(self.Conv1(x))
+        batch, channel = x.size(0), x.size(1)
+        x = x.view(batch, channel, -1).permute(2, 0, 1)  # h*w, batch, dim
+
+        img_queries = self.img_queries.unsqueeze(1).repeat(1, x.shape[1], 1)  # L,N,E
+        img_emb = self.att_pool(img_queries, x, x, need_weights=False)[0]
+        img_emb = self.att_pool_norm(img_emb)  # 1, batch, d_model
+
+        cls = img_emb[0]
+        return cls, x
+
+
+class Encoder(nn.Module):
+    def __init__(self, pretrained, backbone, d_model, fine_tune):
+        super(Encoder, self).__init__()
+        self.backbone = backbone
+
+        if 'rn' in backbone.lower():
+            modules = list(pretrained.children())[:-1]
+            self.net = nn.Sequential(*modules)
+            self.feat_dim = 2048
+            self.feat_size = 7
+        elif 'b-32' in backbone.lower():
+            self.net = pretrained
+            self.net.output_tokens = True
+            self.feat_dim = 768
+            self.feat_size = 7
+        elif 'l-14' in backbone.lower():
+            self.net = pretrained
+            self.net.output_tokens = True
+            self.feat_dim = 1024
+            self.feat_size = 16
+        elif backbone == 'resnet50':
+            net = models.resnet50(pretrained=True)
+            modules = list(net.children())[:-2]
+            self.net = nn.Sequential(*modules)
+            self.feat_dim = 2048
+            self.feat_size = 8
+        elif backbone == 'resnet101':
+            net = models.resnet101(pretrained=True)
+            modules = list(net.children())[:-2]
+            self.net = nn.Sequential(*modules)
+            self.feat_dim = 2048
+            self.feat_size = 8
+
+        self.proj = None
+        if self.feat_dim != d_model:
+            self.proj = nn.Conv2d(self.feat_dim, d_model, kernel_size=1)
+
+        self.fine_tune(fine_tune)
+
+    def forward(self, image):
+        feat = self.net(image)  # batch, feat_dim, feat_size, feat_size
+        if 'vit' in self.backbone.lower():
+            feat = feat[1].reshape(-1, self.feat_size, self.feat_size, self.feat_dim).permute(0, 3, 1, 2)
+
+        if self.proj:
+            feat = self.proj(feat)
+
+        return feat
+
+    def fine_tune(self, fine_tune=True):
+        for p in self.net.parameters():
+            p.requires_grad = False
+
+        if 'resnet' in self.backbone:
+            to_finetune = list(self.net.children())[-5:]
+        elif 'vit' in self.backbone.lower():
+            to_finetune = list(self.net.children())[-2:]  # only transformer layers
+        else:
+            to_finetune = list(self.net.children())[-3:]  # only fine-tune convolutional blocks 2 through 4
+
+        for c in to_finetune:
+            for p in c.parameters():
+                p.requires_grad = fine_tune
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout=0.):
+        super(FeedForward, self).__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class MultiHeadAtt(nn.Module):
+    def __init__(self, dim_q, dim_kv, attention_dim, heads=8, dropout=0.):
+        super(MultiHeadAtt, self).__init__()
+        project_out = not (heads == 1 and attention_dim == dim_kv)
+        self.heads = heads
+        self.scale = (attention_dim // self.heads) ** -0.5
+
+        self.to_q = nn.Linear(dim_q, attention_dim, bias=False)
+        self.to_k = nn.Linear(dim_kv, attention_dim, bias=False)
+        self.to_v = nn.Linear(dim_kv, attention_dim, bias=False)
+        self.attend = nn.Softmax(dim=-1)
+        self.dropout = nn.Dropout(dropout)
+        self.to_out = nn.Sequential(
+            nn.Linear(attention_dim, dim_q),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x1, x2, x3):
+        q = self.to_q(x1)
+        k = self.to_k(x2)
+        v = self.to_k(x3)
+        q = rearrange(q, 'b n (h d) -> b h n d', h=self.heads)
+        k = rearrange(k, 'b n (h d) -> b h n d', h=self.heads)
+        v = rearrange(v, 'b n (h d) -> b h n d', h=self.heads)
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.dropout(self.attend(dots))
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)  # (b,n,dim)
+
+
+class Transformer(nn.Module):
+    def __init__(self, dim_q, dim_kv, heads, attention_dim, hidden_dim, dropout=0., norm_first=False):
+        super(Transformer, self).__init__()
+        self.norm_first = norm_first
+        self.att = MultiHeadAtt(dim_q, dim_kv, attention_dim, heads=heads, dropout=dropout)
+        self.feedforward = FeedForward(dim_q, hidden_dim, dropout=dropout)
+        self.norm1 = nn.LayerNorm(dim_q)
+        self.norm2 = nn.LayerNorm(dim_q)
+
+    def forward(self, x1, x2, x3):
+        if self.norm_first:
+            x = self.att(self.norm1(x1), self.norm1(x2), self.norm1(x3)) + x1
+            x = self.feedforward(self.norm2(x)) + x
+        else:
+            x = self.norm1(self.att(x1, x2, x3) + x1)
+            x = self.norm2(self.feedforward(x) + x)
+
+        return x
+
+
+class AttentiveEncoder(nn.Module):
+    def __init__(self, device, n_layers, feature_size, heads, hidden_dim=512, attention_dim=512, dropout=0.):
+        super(AttentiveEncoder, self).__init__()
+        h_feat, w_feat, channels = feature_size
+
+        self.device = device
+        self.h_embedding = nn.Embedding(h_feat, int(channels / 2))
+        self.w_embedding = nn.Embedding(w_feat, int(channels / 2))
+        self.selftrans = nn.ModuleList([])
+        for i in range(n_layers):
+            self.selftrans.append(nn.ModuleList([
+                Transformer(channels, channels, heads, attention_dim, hidden_dim, dropout, norm_first=False),
+                Transformer(channels * 2, channels * 2, heads, attention_dim, hidden_dim, dropout, norm_first=False),
+            ]))
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, img1, img2):
+        batch, c, h, w = img1.shape
+        pos_h = torch.arange(h).to(self.device)
+        pos_w = torch.arange(w).to(self.device)
+        embed_h = self.w_embedding(pos_h)
+        embed_w = self.h_embedding(pos_w)
+        pos_embedding = torch.cat([embed_w.unsqueeze(0).repeat(h, 1, 1),
+                                   embed_h.unsqueeze(1).repeat(1, w, 1)],
+                                  dim=-1)
+        pos_embedding = pos_embedding.permute(2, 0, 1).unsqueeze(0).repeat(batch, 1, 1, 1)
+        img1 = img1 + pos_embedding
+        img2 = img2 + pos_embedding
+        img1 = img1.view(batch, c, -1).transpose(-1, 1)  # batch, hw, c
+        img2 = img2.view(batch, c, -1).transpose(-1, 1)
+        img_sa1, img_sa2 = img1, img2
+
+        for (l, m) in self.selftrans:
+            img_sa1 = l(img_sa1, img_sa1, img_sa1) + img_sa1
+            img_sa2 = l(img_sa2, img_sa2, img_sa2) + img_sa2
+            img = torch.cat([img_sa1, img_sa2], dim=-1)
+            img = m(img, img, img)
+            img_sa1 = img[:, :, :c] + img1
+            img_sa2 = img[:, :, c:] + img2
+
+        img1 = img_sa1.reshape(batch, h, w, c).transpose(-1, 1)
+        img2 = img_sa2.reshape(batch, h, w, c).transpose(-1, 1)
+
+        return img1, img2
+
+
+class resblock(nn.Module):
+    def __init__(self, inchannel, outchannel, stride=1, shortcut=None):
+        super(resblock, self).__init__()
+        self.left = nn.Sequential(
+            nn.Conv2d(inchannel, int(outchannel / 2), kernel_size=1),
+            # nn.LayerNorm(int(outchannel/2),dim=1),
+            nn.BatchNorm2d(int(outchannel / 2)),
+            nn.ReLU(),
+            nn.Conv2d(int(outchannel / 2), int(outchannel / 2), kernel_size=3, stride=1, padding=1),
+            # nn.LayerNorm(int(outchannel/2),dim=1),
+            nn.BatchNorm2d(int(outchannel / 2)),
+            nn.ReLU(),
+            nn.Conv2d(int(outchannel / 2), outchannel, kernel_size=1),
+            # nn.LayerNorm(int(outchannel / 1),dim=1)
+            nn.BatchNorm2d(outchannel)
+        )
+        self.right = shortcut
+
+    def forward(self, x):
+        out = self.left(x)
+        residual = x
+        out = out + residual
+        return F.relu(out)
+
+
+class Decoder(nn.Module):
+    def __init__(self, device, h_dim, vocab_size, max_len, n_head, n_layers, dropout=0.10,
+                 learnable=False, tie_embeddings=True, prenorm=False):
+
+        super(Decoder, self).__init__()
+
+        self.embed_dim = h_dim
+        self.vocab_size = vocab_size
+        self.dropout = dropout
+        self.device = device
+
+        self.tokens_embed = nn.Embedding(vocab_size, self.embed_dim)
+        self.position_encoding = PositionalEncoding(self.embed_dim, dropout=dropout, max_len=max_len,
+                                                    device=device, learnable=learnable)
+
+        self.uni_decoder = nn.ModuleList(
+            [DecoderLayer(h_dim, h_dim, n_head, dim_feedforward=h_dim * 4, dropout=self.dropout, prenorm=prenorm,
+                          crossattention=False) for _ in range(n_layers)])
+
+        self.cross_decoder = nn.ModuleList(
+            [DecoderLayer(h_dim, h_dim, n_head, dim_feedforward=h_dim * 4, dropout=self.dropout, prenorm=prenorm,
+                          crossattention=True) for _ in range(n_layers)])
+
+        self.lm_head = nn.Linear(h_dim, vocab_size, bias=False)
+        if tie_embeddings:
+            self.tokens_embed.weight = self.lm_head.weight
+        self.dropout = nn.Dropout(p=self.dropout)
+        self.init_weights()
+        self.loss_fn = nn.CrossEntropyLoss()
+
+    def init_weights(self):
+        self.tokens_embed.weight.data.uniform_(-0.1, 0.1)
+        self.lm_head.weight.data.uniform_(-0.1, 0.1)
+
+    def forward(self, input_ids=None, labels=None, pad_mask=None, img_emb=None):
+        att_weights = None
+        mask = torch.tril(torch.ones(input_ids.shape[1], input_ids.shape[1]))
+        mask = ~mask.bool()
+        mask = mask.to(self.device)
+
+        inputs_embeds = self.tokens_embed(input_ids)
+        inputs_embeds = self.position_encoding(inputs_embeds)  # batch, seq, e_dim
+        inputs_embeds = inputs_embeds.permute(1, 0, 2)  # seq, batch, e_dim
+
+        # seq, batch, emb_dim
+        out = inputs_embeds
+        for block in self.uni_decoder:
+            out, _ = block(out, None, tgt_mask=mask, tgt_key_padding_mask=pad_mask)
+
+        if pad_mask is not None:  # not inference
+            cls = []
+            for i in range(pad_mask.shape[0]):
+                end = pad_mask[i].shape[0] - pad_mask[i].count_nonzero()
+                cls.append(out[end - 1, i, :])
+
+            cls = torch.stack(cls)  # batch, emb_dim
+        else:
+            cls = None
+
+        if img_emb is None:
+            return None, cls, None, None
+
+        for block in self.cross_decoder:
+            out, att_weights = block(out, img_emb, tgt_mask=mask, tgt_key_padding_mask=pad_mask)
+
+        lm_logits = self.lm_head(self.dropout(out))  # seq, batch, voc_dim
+        lm_logits = lm_logits.permute(1, 0, 2)  # batch, seq, voc_dim
+
+        if labels is not None:  # not inference
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_logits = shift_logits.view(-1, shift_logits.size(-1))
+            shift_labels = labels[..., 1:].contiguous()
+            shift_labels = shift_labels.view(-1)
+            loss = self.loss_fn(shift_logits, shift_labels)
+        else:
+            loss = None
+
+        return loss, cls, lm_logits, att_weights
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout, max_len, device, learnable=False):
+        super(PositionalEncoding, self).__init__()
+        self.learnable = learnable
+        self.max_len = max_len
+        self.device = device
+        self.dropout = nn.Dropout(p=dropout)
+
+        if not learnable:
+            pe = torch.zeros(max_len, d_model)
+            position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+            div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+            pe[:, 0::2] = torch.sin(position * div_term)
+            pe[:, 1::2] = torch.cos(position * div_term)
+            pe = pe.unsqueeze(0)
+            self.register_buffer('pe', pe)
+        else:
+            self.pos_emb = nn.Embedding(max_len, int(d_model))
+
+    def forward(self, x):
+        if self.learnable:
+            position_ids = torch.arange(x.size(1), dtype=torch.long).to(self.device)
+            position_ids = position_ids.unsqueeze(0).view(-1, x.size(1))  # batch, seq
+            x = x + self.pos_emb(position_ids)
+        else:
+            x = x + self.pe[:, :x.size(1), :]
+        return self.dropout(x)
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, d_model, img_dim, nhead, dim_feedforward=2048, dropout=0.1, layer_norm_eps=1e-5,
+                 prenorm=False, crossattention=False):
+        super(DecoderLayer, self).__init__()
+
+        self.prenorm = prenorm
+        self.crossattention = crossattention
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+
+        if crossattention:
+            self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, kdim=img_dim, vdim=img_dim)
+            self.mha_dropout = nn.Dropout(dropout)
+            self.mha_norm = nn.LayerNorm(d_model, eps=layer_norm_eps)
+
+        self.ff_linear1 = nn.Linear(d_model, dim_feedforward)
+        self.ff_dropout = nn.Dropout(dropout)
+        self.ff_linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.sa_norm = nn.LayerNorm(d_model, eps=layer_norm_eps)
+        self.ff_norm = nn.LayerNorm(d_model, eps=layer_norm_eps)
+
+        self.sa_dropout = nn.Dropout(dropout)
+        self.ff_dropout = nn.Dropout(dropout)
+
+        self.activation = nn.GELU()
+
+    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask=None,
+                memory_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None):
+        att_weight = None
+        x = tgt
+
+        if self.prenorm:
+            x = x + self._sa_block(self.sa_norm(x), tgt_mask, tgt_key_padding_mask)
+            if self.crossattention:
+                enc_att, att_weight = self._mha_block(self.mha_norm(x), memory, memory_mask, memory_key_padding_mask)
+                x = x + enc_att
+            x = x + self._ff_block(self.ff_norm(x))
+        else:
+            x = self.sa_norm(x + self._sa_block(x, tgt_mask, tgt_key_padding_mask))
+            if self.crossattention:
+                enc_att, att_weight = self._mha_block(x, memory, memory_mask, memory_key_padding_mask)
+                x = self.mha_norm(x + enc_att)
+
+            x = self.ff_norm(x + self._ff_block(x))
+        return x, att_weight
+
+    def _sa_block(self, x, attn_mask, key_padding_mask):
+        x = self.self_attn(x, x, x,  # L,N,E
+                           attn_mask=attn_mask,  # L, S
+                           key_padding_mask=key_padding_mask,  # N, S
+                           is_causal=True,
+                           need_weights=False)[0]
+        return self.sa_dropout(x)
+
+    def _mha_block(self, x, mem, attn_mask, key_padding_mask):
+        x, att_weight = self.cross_attn(x, mem, mem,
+                                        attn_mask=attn_mask,
+                                        key_padding_mask=key_padding_mask,
+                                        is_causal=False,
+                                        need_weights=True)
+        return self.mha_dropout(x), att_weight
+
+    def _ff_block(self, x):
+        x = self.ff_linear2(self.ff_dropout(self.activation(self.ff_linear1(x))))
+        return self.ff_dropout(x)

src/train.py

@@ -0,0 +1,198 @@
+import argparse
+import random
+import os
+from datetime import datetime
+
+import numpy as np
+import torch
+import json
+from torch.optim import AdamW
+from torchvision.transforms import v2
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+from transformers import get_constant_schedule_with_warmup
+
+from Datasets import CCDataset, Batcher
+from model import ICCModel
+from utils import get_vocabulary
+from Loss import InfoNCELoss
+from eval import captioning, retrieve, plot
+from huggingface_hub import hf_hub_download
+import open_clip
+
+
+def train(args, model, train_loader, valid_loader, device, infonce, optim, scheduler, writer):
+    step = 0
+    best_score = float("inf")
+    best_model = None
+
+    for epoch in range(args.epochs):
+        model.train()
+
+        for batch in tqdm(train_loader, desc='Epoch ' + str(epoch)):
+            imgs1 = batch['images_before'].to(device)
+            imgs2 = batch['images_after'].to(device)
+            toks = batch['input_ids'].to(device)
+            labs = batch['labels'].to(device)
+            flags = batch['flags'].to(device)
+            attention_mask = batch['pad_mask'].to(device)
+            embs = batch['embs'].to(device)
+
+            cap_loss, vis_emb, text_emb, _, _, _ = model(imgs1, imgs2, toks, labs, attention_mask)
+
+            con_loss, num_pos = infonce(vis_emb, text_emb, flags, embs)
+            loss = cap_loss + args.lamb * con_loss
+            loss.backward()
+
+            if args.max_grad_norm:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+            grad = torch.norm(torch.stack(
+                [torch.norm(p.grad.detach()).to(device) for p in model.parameters() if p.grad is not None]))
+
+            optim.step()
+            scheduler.step()
+            optim.zero_grad()
+
+            writer.add_scalar('train_loss', loss.item(), step)
+            writer.add_scalar('grad', grad, step)
+            writer.add_scalar('lr', scheduler.get_last_lr()[0], step)
+
+            step += 1
+
+        torch.save(model.state_dict(), args.output_path + 'model_{}.pt'.format(step))
+
+        model.eval()
+        with torch.no_grad():
+            eval_losses = torch.empty(0)
+            for batch in tqdm(valid_loader, desc='Validation ' + str(epoch)):
+                imgs1 = batch['images_before'].to(device)
+                imgs2 = batch['images_after'].to(device)
+                toks = batch['input_ids'].to(device)
+                labs = batch['labels'].to(device)
+                flags = batch['flags'].to(device)
+                attention_mask = batch['pad_mask'].to(device)
+                embs = batch['embs'].to(device)
+
+                cap_loss, vis_emb, text_emb, _, _, _ = model(imgs1, imgs2, toks, labs, attention_mask)
+
+                con_loss, _ = infonce(vis_emb, text_emb, flags, embs)
+                loss = cap_loss + args.lamb * con_loss
+                eval_losses = torch.cat([eval_losses, loss.cpu().unsqueeze(0)])
+
+            eval_score = torch.mean(eval_losses)
+            writer.add_scalar('eval_score', eval_score, step)
+
+        is_best = eval_score < best_score
+        best_score = min(eval_score, best_score)
+        if is_best:
+            best_model = step
+
+    if best_model is not None:
+        state_dict = torch.load(os.path.join(args.output_path + 'model_{}.pt'.format(best_model)), map_location=device)
+        torch.save(state_dict, args.output_path + 'model_best.pt')
+
+
+def run(args, config):
+    print('Initializing...')
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    random.seed(args.seed)
+    torch.backends.cudnn.deterministic = True
+
+    device = torch.device('cpu')
+    if torch.cuda.is_available():
+        device = torch.device('cuda')
+
+    dt_str = datetime.now().strftime("%d-%m-%Y-%H-%M-%S")
+    writer_path = args.output_path + dt_str
+    writer = SummaryWriter(writer_path)
+
+    if os.path.exists(args.vocab):
+        with open(args.vocab, 'r') as infile:
+            vocab = json.load(infile)
+    else:
+        vocab = get_vocabulary(args.annotation_json, args.vocab)
+
+    clip = None
+    preprocess = v2.Compose([
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+        v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
+    if 'resnet' not in config['backbone']:
+        checkpoint_path = hf_hub_download("chendelong/RemoteCLIP",
+                                          f"RemoteCLIP-{config['backbone']}.pt",
+                                          cache_dir=args.pretrained)
+
+        clip, _, preprocess = open_clip.create_model_and_transforms(config['backbone'])
+        ckpt = torch.load(checkpoint_path, map_location="cpu")
+        clip.load_state_dict(ckpt)
+
+    model = ICCModel(device, clip, config['backbone'], config['d_model'],
+                     len(vocab), config['max_len'], config['num_heads'], config['h_dim'], config['a_dim'],
+                     config['encoder_layers'], config['decoder_layers'], config['dropout'],
+                     learnable=config['learnable'], fine_tune=config['fine_tune'],
+                     tie_embeddings=config['tie_embeddings'], prenorm=config['prenorm'])
+    model = model.to(device)
+    del clip
+
+    print('Loading...')
+    training_set = CCDataset(args.annotation_json, args.image_dir, vocab, preprocess, 'train', config['max_len'],
+                             config['s-transformers'], device)
+    valid_set = CCDataset(args.annotation_json, args.image_dir, vocab, preprocess, 'val', config['max_len'],
+                          config['s-transformers'], device)
+    test_set = CCDataset(args.annotation_json, args.image_dir, vocab, preprocess, 'test', config['max_len'],
+                         config['s-transformers'], device)
+
+    train_loader = Batcher(training_set, args.batch_size, config['max_len'], device, args.hd, model=model, shuffle=True)
+    valid_loader = Batcher(valid_set, args.batch_size, config['max_len'], device)
+    test_loader = Batcher(test_set, 1, config['max_len'], device)
+
+    print('Training...')
+    infonce = InfoNCELoss(device, k=args.k, temperature=args.temperature, threshold=config['s-threshold'],
+                          fna=config['fna'])
+    optim = AdamW([x for x in model.parameters() if x.requires_grad], lr=args.learning_rate, eps=args.adam_epsilon)
+    scheduler = get_constant_schedule_with_warmup(optim,
+                                                  num_warmup_steps=args.warmup_steps * len(train_loader) * args.epochs)
+    train(args, model, train_loader, valid_loader, device, infonce, optim, scheduler, writer)
+
+    print('Final evaluation...')
+    model.load_state_dict(torch.load(os.path.join(args.output_path, 'model_best.pt'), map_location=device))
+    results = captioning(args, config, model, test_loader, vocab, device)
+    retrieve(args, config, model, test_loader, device)
+    plot(args, model.encoder.encoder.feat_size, results)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--annotation_json', type=str, default='../input/Levir_CC/LevirCCcaptions.json')
+    parser.add_argument('--image_dir', type=str, default='../input/Levir_CC/images/')
+    parser.add_argument('--vocab', type=str, default='../input/levir_vocab.json')
+    parser.add_argument('--pretrained', type=str, default='../../input/checkpoints')
+    parser.add_argument('--config', type=str, default='../config.json')
+    parser.add_argument('--output_path', type=str, default='../output/')
+
+    parser.add_argument('--epochs', type=int, default=50)
+    parser.add_argument('--batch_size', type=int, default=4)
+    parser.add_argument('--k', type=int, default=-1)
+    parser.add_argument('--hd', type=int, default=-1)
+    parser.add_argument('--learning_rate', type=float, default=1e-4)
+    parser.add_argument('--warmup_steps', type=float, default=0.025)
+    parser.add_argument('--lr_decay', type=float, default=0.7)
+    parser.add_argument('--adam_epsilon', type=float, default=1e-8)
+    parser.add_argument('--max_grad_norm', type=float, default=None)
+    parser.add_argument('--temperature', type=float, default=0.01)
+    parser.add_argument('--lamb', type=float, default=0.5)
+    parser.add_argument('--seed', type=int, default=42)
+
+    args = parser.parse_args()
+
+    with open(args.config, 'r') as config_file:
+        config = json.load(config_file)
+
+    run(args, config)
+
+
+if __name__ == '__main__':
+    main()

src/utils.py

@@ -0,0 +1,115 @@
+import json
+import os
+
+import pandas as pd
+import torch
+import glob
+
+from eval_func.bleu.bleu import Bleu
+from eval_func.rouge.rouge import Rouge
+from eval_func.cider.cider import Cider
+
+
+def get_eval_score(ref, hypo):
+    scorers = [
+        (Bleu(4), ["Bleu_1", "Bleu_2", "Bleu_3", "Bleu_4"]),
+        (Rouge(), "ROUGE_L"),
+        (Cider(), "CIDEr")
+    ]
+
+    score = []
+    method = []
+    for scorer, method_i in scorers:
+        score_i, scores_i = scorer.compute_score(ref, hypo)
+        score.extend(score_i) if isinstance(score_i, list) else score.append(score_i)
+        method.extend(method_i) if isinstance(method_i, list) else method.append(method_i)
+
+    score_dict = dict(zip(method, score))
+    return score_dict
+
+
+def get_vocabulary(in_path, out_file):
+    if 'levir' in in_path.lower():
+        return get_levir_vocabulary(in_path, out_file)
+    elif 'dubai' in in_path.lower():
+        return get_dubai_vocabulary(in_path, out_file)
+    elif 'clevr' in in_path.lower():
+        return get_clevr_vocabulary(in_path, out_file)
+
+
+def get_levir_vocabulary(in_path, out_file):
+    with open(in_path) as fin:
+        data = json.load(fin)['images']
+
+    sents = [y for x in data for y in x['sentences']]
+    tokens = [y for x in sents for y in x['tokens']]
+    occurencies = pd.Series(tokens).value_counts()
+    selected = occurencies[occurencies > 5]
+    vocab = {w: i + 4 for i, w in enumerate(selected.index)}
+    vocab['PAD'] = 0
+    vocab['START'] = 1
+    vocab['UNK'] = 2
+    vocab['END'] = 3
+
+    with open(out_file, 'w') as fout:
+        json.dump(vocab, fout)
+
+    return vocab
+
+
+def get_dubai_vocabulary(in_path, out_file):
+    data = []
+    for path in glob.glob(in_path + '/*.json'):
+        with open(path) as fin:
+            data.extend(json.load(fin)['images'])
+
+    sents = [y for x in data for y in x['sentences']]
+    tokens = [y for x in sents for y in x['tokens']]
+    selected = pd.Series(tokens).value_counts()
+    vocab = {w: i + 4 for i, w in enumerate(selected.index)}
+    vocab['PAD'] = 0
+    vocab['START'] = 1
+    vocab['UNK'] = 2
+    vocab['END'] = 3
+
+    with open(out_file, 'w') as fout:
+        json.dump(vocab, fout)
+
+    return vocab
+
+
+def get_clevr_vocabulary(in_path, out_file):
+    sents = []
+    with open(os.path.join(in_path, 'change_captions.json'), 'r', encoding='utf-8') as fin:
+        data = json.load(fin)
+        sents += [y for x in data for y in data[x]]
+
+    with open(os.path.join(in_path, 'no_change_captions.json'), 'r', encoding='utf-8') as fin:
+        data = json.load(fin)
+        sents += [y for x in data for y in data[x]]
+
+    tokens = [y for x in sents for y in x.split(' ')]
+    occurencies = pd.Series(tokens).value_counts()
+    vocab = {w: i + 4 for i, w in enumerate(occurencies.index)}
+    vocab['PAD'] = 0
+    vocab['START'] = 1
+    vocab['UNK'] = 2
+    vocab['END'] = 3
+
+    with open(out_file, 'w') as fout:
+        json.dump(vocab, fout)
+
+    return vocab
+
+
+def unormalize(tensor, mean=None, std=None):
+    if mean is not None and std is not None:
+        for t, m, s in zip(tensor, mean, std):
+            t.mul_(s).add_(m)
+        return torch.clip(tensor, min=0, max=1)
+
+    b, c, h, w = tensor.shape
+    tensor = tensor.view(b, -1)
+    tensor -= tensor.min(1, keepdim=True)[0]
+    tensor /= tensor.max(1, keepdim=True)[0]
+    return tensor.view(b, c, h, w)