Base app.

.gitattributes

@@ -32,3 +32,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/saiapr_tc-12.zip filter=lfs diff=lfs merge=lfs -text
+cache/20211220_191132_refclef_32_512_resnet50_8_6_8_0.1_0.1_0.1_0_0.0001_0.0_12_4_90_1_0_0_0/best.ckpt filter=lfs diff=lfs merge=lfs -text
+data/val-sim_metric.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,10 +1,10 @@
 ---
-title: RECModel
-emoji: 💩
-colorFrom: purple
-colorTo: pink
+title: ProbingREC
+emoji: 👁
+colorFrom: blue
+colorTo: gray
 sdk: gradio
-sdk_version: 3.9
+sdk_version: 3.4
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,119 @@
+from models import IntuitionKillingMachine
+from transforms import undo_box_transforms_batch, ToTensor, Normalize, SquarePad, Resize, NormalizeBoxCoords
+from torchvision.transforms import Compose
+from encoders import get_tokenizer
+from PIL import Image, ImageDraw
+from zipfile import ZipFile
+from copy import copy
+import gradio as gr
+import pandas as pd
+import torch
+
+def parse_model_args(model_path):
+    _, _, dataset, max_length, input_size, backbone, num_heads, num_layers, num_conv, _, _, mu, mask_pooling = model_path.split('_')[:13]
+    return {
+        'dataset': dataset,
+        'max_length': int(max_length),
+        'input_size': int(input_size),
+        'backbone': backbone,
+        'num_heads': int(num_heads),
+        'num_layers': int(num_layers),
+        'num_conv': int(num_conv),
+        'mu': float(mu),
+        'mask_pooling': bool(mask_pooling == '1')
+    }
+
+
+class Prober:
+    def __init__(self,
+                 df_path=None,
+                 dataset_path=None,
+                 model_checkpoint=None):
+        params = parse_model_args(model_checkpoint)
+        mean = [0.485, 0.456, 0.406]
+        sdev = [0.229, 0.224, 0.225]
+        self.tokenizer = get_tokenizer()
+        self.df = pd.read_json(df_path)[['sample_idx', 'bbox', 'file_path', 'sent']]
+        self.df.loc[:, "image_id"] = self.df.file_path.apply(lambda x: int(x.split('/')[-1][:-4]))
+        self.df.file_path = self.df.file_path.apply(lambda x: x.replace('refer/data/images/', ''))
+        self.model = IntuitionKillingMachine(
+            backbone=params['backbone'],
+            pretrained=True,
+            num_heads=params['num_heads'],
+            num_layers=params['num_layers'],
+            num_conv=params['num_conv'],
+            segmentation_head=bool(params['mu'] > 0.0),
+            mask_pooling=params['mask_pooling']
+        ) 
+        self.load_model(model_checkpoint)
+        self.transform = Compose([
+            ToTensor(),
+            Normalize(mean, sdev),
+            SquarePad(),
+            Resize(size=(params['input_size'], params['input_size'])),
+            NormalizeBoxCoords(),
+        ])
+        self.max_length = 30
+        self.zipfile = ZipFile(dataset_path, 'r')
+
+    def load_model(self, model_checkpoint):
+        checkpoint = torch.load(
+            model_checkpoint, map_location=lambda storage, loc: storage
+        )
+
+        # strip 'model.' from pl checkpoint
+        state_dict = {
+            k[len('model.'):]: v
+            for k, v in checkpoint['state_dict'].items()
+        }
+
+        missing, _ = self.model.load_state_dict(state_dict, strict=False)
+
+        # ensure the only missing keys are those of the segmentation head only
+        assert [k for k in missing if 'segm' not in k] == []
+
+        self.model = self.model.eval()
+
+
+    @torch.no_grad()
+    def probe(self, idx, re, search_by_sample_id: bool= True):
+        if search_by_sample_id:
+            img_path, target, = self.df.loc[idx][['file_path','bbox']].values
+        else: 
+            img_path, target = self.df[self.df.image_id == idx][['file_path','bbox']].values[0]
+        img = Image.open(self.zipfile.open(img_path)).convert('RGB')
+        W0, H0 = img.size
+        sample = {
+            'image': img,
+            'image_size': (H0, W0),  # image original size
+            'bbox': torch.tensor([copy(target)]),
+            'bbox_raw': torch.tensor([copy(target)]),
+            'mask': torch.ones((1, H0, W0), dtype=torch.float32),  # visibiity mask
+            'mask_bbox': None,  # target bbox mask
+        } 
+        sample = self.transform(sample)
+        tok = self.tokenizer(re,
+                             max_length=30,
+                             return_tensors='pt',
+                             truncation=True)
+        inn = {'image': torch.stack([sample['image']]),
+               'mask': torch.stack([sample['mask']]),
+               'tok': tok}
+        output = undo_box_transforms_batch(self.model(inn)[0],
+                                           [sample['tr_param']]).numpy().tolist()[0]
+        img1 = ImageDraw.Draw(img)
+        #img1.rectangle(target, outline ="#0000FF00", width=3)
+        img1.rectangle(output, outline ="#00FF0000", width=3)
+        return img
+
+
+prober = Prober(
+    df_path = 'data/val-sim_metric.json',
+    dataset_path = "data/saiapr_tc-12.zip",
+    model_checkpoint= "cache/20211220_191132_refclef_32_512_resnet50_8_6_8_0.1_0.1_0.1_0_0.0001_0.0_12_4_90_1_0_0_0/best.ckpt"
+)
+
+demo = gr.Interface(fn=prober.probe, inputs=["number", "text", "checkbox"], outputs="image")
+
+demo.queue(concurrency_count=10)
+demo.launch(debug=True)

backbones.py

@@ -0,0 +1,82 @@
+import torch
+
+import torch.nn as nn
+
+from torchvision.ops.misc import FrozenBatchNorm2d
+
+from torchvision.models import resnet, detection, segmentation
+
+import timm
+
+
+# https://detectron2.readthedocs.io/en/latest/modules/layers.html#detectron2.layers.FrozenBatchNorm2d.convert_frozen_batchnorm
+@torch.no_grad()
+def convert_frozen_batchnorm(module):
+    bn_module = (
+        nn.modules.batchnorm.BatchNorm2d,
+        nn.modules.batchnorm.SyncBatchNorm
+    )
+    res = module
+    if isinstance(module, bn_module):
+        res = FrozenBatchNorm2d(module.num_features)
+        if module.affine:
+            res.weight.data = module.weight.data.clone().detach()
+            res.bias.data = module.bias.data.clone().detach()
+        res.running_mean.data = module.running_mean.data
+        res.running_var.data = module.running_var.data
+        res.eps = module.eps
+    else:
+        for name, child in module.named_children():
+            new_child = convert_frozen_batchnorm(child)
+            if new_child is not child:
+                res.add_module(name, new_child)
+    return res
+
+
+def get_backbone(backbone, pretrained=True):
+    if backbone in ('resnet18', 'resnet34', 'resnet50', 'resnet101'):
+        # pretrained on ImageNet for classification
+        model = resnet.__dict__[backbone](
+            pretrained=pretrained, norm_layer=FrozenBatchNorm2d
+        )
+    elif backbone == 'resnet50d':
+        # pretrained on COCO for detection
+        model = convert_frozen_batchnorm(
+            detection.fasterrcnn_resnet50_fpn(pretrained=pretrained).backbone.body
+        )
+    elif backbone == 'resnet50s':
+        # pretrained on COCO for segmentation
+        model = convert_frozen_batchnorm(
+            segmentation.deeplabv3_resnet50(pretrained=pretrained).backbone
+        )
+    elif backbone == 'resnet101s':
+        # pretrained on COCO for segmentation
+        model = convert_frozen_batchnorm(
+            segmentation.deeplabv3_resnet101(pretrained=pretrained).backbone
+        )
+
+    elif backbone in ('cspdarknet53', 'efficientnet-b0', 'efficientnet-b3'):
+        # model = convert_frozen_batchnorm(
+        #     timm.create_model(
+        #         backbone.replace('-', '_'),
+        #         pretrained=True,
+        #         features_only=True,
+        #         #out_indices=(1, 2, 3, 4)
+        #     )
+        # )
+        model = convert_frozen_batchnorm(
+            timm.create_model(
+                backbone.replace('-', '_'),
+                pretrained=pretrained,
+                num_classes=0,
+                global_pool=''
+            )
+        )
+
+    else:
+        raise RuntimeError(f'{backbone} is not a valid backbone')
+
+    # empty cache (dealloc modules other than the backbone)
+    torch.cuda.empty_cache()
+
+    return model

cache/20211220_191132_refclef_32_512_resnet50_8_6_8_0.1_0.1_0.1_0_0.0001_0.0_12_4_90_1_0_0_0/best.ckpt

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

datasets.py

@@ -0,0 +1,282 @@
+import os
+
+import json
+
+import random
+
+import torch
+
+import ijson
+
+import numpy as np
+
+from PIL import Image
+
+from torchvision.transforms import ToTensor
+
+from torchvision.ops import box_convert, clip_boxes_to_image
+
+from re_classifier import REClassifier
+
+from utils import progressbar
+
+
+def collate_fn(batch):
+    image = torch.stack([s['image'] for s in batch], dim=0)
+
+    image_size = torch.FloatTensor([s['image_size'] for s in batch])
+
+    # bbox = torch.stack([s['bbox'] for s in batch], dim=0)
+    bbox = torch.cat([s['bbox'] for s in batch], dim=0)
+
+    # bbox_raw = torch.stack([s['bbox_raw'] for s in batch], dim=0)
+    bbox_raw = torch.cat([s['bbox_raw'] for s in batch], dim=0)
+
+    expr = [s['expr'] for s in batch]
+
+    tok = None
+    if batch[0]['tok'] is not None:
+        tok = {
+            'input_ids': torch.cat([s['tok']['input_ids'] for s in batch], dim=0),
+            'attention_mask': torch.cat([s['tok']['attention_mask'] for s in batch], dim=0)
+        }
+
+        # dynamic batching
+        max_length = max([s['tok']['length'] for s in batch])
+        tok = {
+            'input_ids': tok['input_ids'][:, :max_length],
+            'attention_mask': tok['attention_mask'][:, :max_length],
+        }
+
+    mask = None
+    if batch[0]['mask'] is not None:
+        mask = torch.stack([s['mask'] for s in batch], dim=0)
+
+    mask_bbox = None
+    if batch[0]['mask_bbox'] is not None:
+        mask_bbox = torch.stack([s['mask_bbox'] for s in batch], dim=0)
+
+    tr_param = [s['tr_param'] for s in batch]
+
+    return {
+        'image': image,
+        'image_size': image_size,
+        'bbox': bbox,
+        'bbox_raw': bbox_raw,
+        'expr': expr,
+        'tok': tok,
+        'tr_param': tr_param,
+        'mask': mask,
+        'mask_bbox': mask_bbox,
+    }
+
+
+class RECDataset(torch.utils.data.Dataset):
+    def __init__(self, transform=None, tokenizer=None, max_length=32, with_mask_bbox=False):
+        super().__init__()
+        self.samples = []  # list of samples: [(file_name, expresion, bbox)]
+        self.transform = transform
+        self.tokenizer = tokenizer
+        self.max_length = int(max_length)
+        self.with_mask_bbox = bool(with_mask_bbox)
+
+    def tokenize(self, inp, max_length):
+        return self.tokenizer(
+            inp,
+            return_tensors='pt',
+            padding='max_length',
+            return_token_type_ids=False,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            truncation=True,
+            max_length=max_length
+        )
+
+    def print_stats(self):
+        print(f'{len(self.samples)} samples')
+        lens = [len(expr.split()) for _, expr, _ in self.samples]
+        print('expression lengths stats: '
+              f'min={np.min(lens):.1f}, '
+              f'mean={np.mean(lens):.1f}, '
+              f'median={np.median(lens):.1f}, '
+              f'max={np.max(lens):.1f}, '
+              f'99.9P={np.percentile(lens, 99.9):.1f}'
+        )
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        file_name, expr, bbox = self.samples[idx]
+
+        if not os.path.exists(file_name):
+            raise IOError(f'{file_name} not found')
+        img = Image.open(file_name).convert('RGB')
+
+        # if isinstance(expr, (list, tuple)):
+        #     expr = random.choice(expr)
+
+        # image size as read from disk (PIL)
+        W0, H0 = img.size
+
+        # # ensure box coordinates fall inside the image
+        # bbox = clip_boxes_to_image(bbox, (H0, W0))
+        # assert torch.all(bbox[:, (0, 1)] <= bbox[:, (2, 3)])  # xyxy format
+
+        sample = {
+            'image': img,
+            'image_size': (H0, W0),  # image original size
+            'bbox': bbox.clone(),  # box transformations are inplace ops
+            'bbox_raw': bbox.clone(),  # raw boxes w/o any transformation (in pixels)
+            'expr': expr,
+            'tok': None,
+            'mask': torch.ones((1, H0, W0), dtype=torch.float32),  # visibiity mask
+            'mask_bbox': None,  # target bbox mask
+        }
+
+        # apply transforms
+        if self.transform is None:
+            sample['image'] = ToTensor()(sample['image'])
+        else:
+            sample = self.transform(sample)
+
+        # tokenize after the transformations (just in case there where a left<>right substitution)
+        if self.tokenizer is not None:
+            sample['tok'] = self.tokenize(sample['expr'], self.max_length)
+            sample['tok']['length'] = sample['tok']['attention_mask'].sum(1).item()
+
+        # bbox segmentation mask
+        if self.with_mask_bbox:
+            # image size after transforms
+            _, H, W = sample['image'].size()
+
+            # transformed bbox in pixels
+            bbox = sample['bbox'].clone()
+            bbox[:, (0, 2)] *= W
+            bbox[:, (1, 3)] *= H
+            bbox = clip_boxes_to_image((bbox + 0.5).long(), (H, W))
+
+            # output mask
+            sample['mask_bbox'] = torch.zeros((1, H, W), dtype=torch.float32)
+            for x1, y1, x2, y2 in bbox.tolist():
+                sample['mask_bbox'][:, y1:y2+1, x1:x2+1] = 1.0
+
+        return sample
+
+
+class RegionDescriptionsVisualGnome(RECDataset):
+    def __init__(self, data_root, transform=None, tokenizer=None,
+                 max_length=32, with_mask_bbox=False):
+        super().__init__(transform=transform, tokenizer=tokenizer,
+                         max_length=max_length, with_mask_bbox=with_mask_bbox)
+
+
+        # if available, read COCO IDs from the val, testA and testB splits from
+        # the RefCOCO dataset
+        try:
+            with open('./refcoco_valtest_ids.txt', 'r') as fh:
+                refcoco_ids = [int(lin.strip()) for lin in fh.readlines()]
+        except:
+            refcoco_ids = []
+
+        def path_from_url(fname):
+            return os.path.join(data_root, fname[fname.index('VG_100K'):])
+
+        with open(os.path.join(data_root, 'image_data.json'), 'r') as f:
+            image_data = {
+                data['image_id']: path_from_url(data['url'])
+                for data in json.load(f)
+                if data['coco_id'] is None or data['coco_id'] not in refcoco_ids
+            }
+        print(f'{len(image_data)} images')
+
+        self.samples = []
+
+        with open(os.path.join(data_root, 'region_descriptions.json'), 'r') as f:
+            for record in progressbar(ijson.items(f, 'item.regions.item'), desc='loading data'):
+                if record['image_id'] not in image_data:
+                    continue
+                file_name = image_data[record['image_id']]
+
+                expr = record['phrase']
+
+                bbox = [record['x'], record['y'], record['width'], record['height']]
+                bbox = torch.atleast_2d(torch.FloatTensor(bbox))
+                bbox = box_convert(bbox, 'xywh', 'xyxy')  # xyxy
+
+                self.samples.append((file_name, expr, bbox))
+
+        self.print_stats()
+
+
+class ReferDataset(RECDataset):
+    def __init__(self, data_root, dataset, split_by, split, transform=None,
+                 tokenizer=None, max_length=32, with_mask_bbox=False):
+        super().__init__(transform=transform, tokenizer=tokenizer,
+                         max_length=max_length, with_mask_bbox=with_mask_bbox)
+
+        # https://github.com/lichengunc/refer
+        try:
+            import sys
+            sys.path.append('refer')
+            from refer import REFER
+        except:
+            raise RuntimeError('create a symlink to valid refer compilation '
+                               '(see https://github.com/lichengunc/refer)')
+
+        refer = REFER(data_root, dataset, split_by)
+        ref_ids = sorted(refer.getRefIds(split=split))
+
+        self.samples = []
+
+        for rid in progressbar(ref_ids, desc='loading data'):
+            ref = refer.Refs[rid]
+            ann = refer.refToAnn[rid]
+
+            file_name = refer.Imgs[ref['image_id']]['file_name']
+            if dataset == 'refclef':
+                file_name = os.path.join(
+                    'refer', 'data', 'images', 'saiapr_tc-12', file_name
+                )
+            else:
+                coco_set = file_name.split('_')[1]
+                file_name = os.path.join(
+                    'refer', 'data', 'images', 'mscoco', coco_set, file_name
+                )
+
+            bbox = ann['bbox']
+            bbox = torch.atleast_2d(torch.FloatTensor(bbox))
+            bbox = box_convert(bbox, 'xywh', 'xyxy')  # xyxy
+
+            sentences = [s['sent'] for s in ref['sentences']]
+            if 'train' in split:  # remove repeated expresions
+                sentences = list(set(sentences))
+            sentences = sorted(sentences)
+
+            self.samples += [(file_name, expr, bbox) for expr in sentences]
+
+        self.print_stats()
+
+
+class RefCLEF(ReferDataset):
+    def __init__(self, *args, **kwargs):
+        assert args[0] in ('train', 'val', 'test')
+        super().__init__('refer/data', 'refclef', 'berkeley', *args, **kwargs)
+
+
+class RefCOCO(ReferDataset):
+    def __init__(self, *args, **kwargs):
+        assert args[0] in ('train', 'val', 'trainval', 'testA', 'testB')
+        super().__init__('refer/data', 'refcoco', 'unc', *args, **kwargs)
+
+
+class RefCOCOp(ReferDataset):
+    def __init__(self, *args, **kwargs):
+        assert args[0] in ('train', 'val', 'trainval', 'testA', 'testB')
+        super().__init__('refer/data', 'refcoco+', 'unc', *args, **kwargs)
+
+
+class RefCOCOg(ReferDataset):
+    def __init__(self, *args, **kwargs):
+        assert args[0] in ('train', 'val', 'test')
+        super().__init__('refer/data', 'refcocog', 'umd', *args, **kwargs)

embeddings.py

@@ -0,0 +1,182 @@
+import math
+
+import torch
+
+from torch import nn
+
+
+# adapted from https://pytorch.org/tutorials/beginner/transformer_tutorial.html
+class PositionEmbedding1D(nn.Module):
+    def __init__(self, embedding_dim, dropout=0.1, max_len=128):
+        super().__init__()
+
+        # self.dropout = nn.Dropout(p=dropout)
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, embedding_dim, 2) * (-math.log(10000.0) / embedding_dim))
+        pe = torch.zeros(max_len, embedding_dim)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)  # .transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        # # x: Tensor, shape [batch_size, seq_len, embedding_dim]
+        # x = x + self.pe[:, :x.size(1)]
+        # return self.dropout(x)
+        N, T, _ = x.size()
+        return self.pe[:, :T].repeat(N, 1, 1)
+
+
+class LearnedPositionEmbedding1D(nn.Module):
+    def __init__(self, embedding_dim, max_len=128):
+        super().__init__()
+        self.pe = nn.Parameter(torch.Tensor(1, max_len, embedding_dim))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.xavier_normal_(self.pe)
+
+    def forward(self, x):
+        N, T, _ = x.size()
+        return self.pe[:, :T].repeat(N, 1, 1)
+
+
+# https://huggingface.co/transformers/_modules/transformers/models/detr/modeling_detr.html
+class PositionEmbedding2D(nn.Module):
+    def __init__(self, embedding_dim, temperature=10000, normalize=False,
+                 scale=None):
+        super().__init__()
+        assert embedding_dim % 2 == 0
+        self.half_embedding_dim = embedding_dim // 2
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, pixel_values, pixel_mask):
+        assert pixel_mask is not None, "No pixel mask provided"
+        if pixel_mask.dim() == 4 and pixel_mask.size(1) == 1:
+            pixel_mask = pixel_mask.squeeze(1)
+        y_embed = pixel_mask.cumsum(1, dtype=torch.float32)
+        x_embed = pixel_mask.cumsum(2, dtype=torch.float32)
+        if self.normalize:
+            y_embed = y_embed / (y_embed[:, -1:, :] + 1e-6) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + 1e-6) * self.scale
+
+        dim_t = torch.arange(self.half_embedding_dim, dtype=torch.float32, device=pixel_values.device)
+        dim_t = self.temperature ** (2 * torch.divide(dim_t, 2, rounding_mode='floor') / self.half_embedding_dim)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack((
+            pos_x[:, :, :, 0::2].sin(),
+            pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((
+            pos_y[:, :, :, 0::2].sin(),
+            pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        return pos
+
+
+# https://huggingface.co/transformers/_modules/transformers/models/detr/modeling_detr.html
+class LearnedPositionEmbedding2D(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+        assert embedding_dim % 2 == 0, 'embedding dimensionality must be even'
+        self.rows_embeddings = nn.Embedding(50, embedding_dim//2)
+        self.cols_embeddings = nn.Embedding(50, embedding_dim//2)
+
+    def forward(self, pixel_values, pixel_mask=None):
+        h, w = pixel_values.shape[-2:]
+        i = torch.arange(w, device=pixel_values.device)
+        j = torch.arange(h, device=pixel_values.device)
+        x_emb = self.cols_embeddings(i)
+        y_emb = self.rows_embeddings(j)
+        pos = torch.cat([x_emb.unsqueeze(0).repeat(h, 1, 1), y_emb.unsqueeze(1).repeat(1, w, 1)], dim=-1)
+        pos = pos.permute(2, 0, 1)
+        pos = pos.unsqueeze(0)
+        pos = pos.repeat(pixel_values.shape[0], 1, 1, 1)
+        return pos
+
+
+class Box8PositionEmbedding2D(nn.Module):
+    def __init__(self, embedding_dim, with_projection=True):
+        super().__init__()
+
+        self.proj = None
+        if with_projection:
+            self.proj = nn.Linear(8, embedding_dim)
+            nn.init.xavier_normal_(self.proj.weight)
+            nn.init.zeros_(self.proj.bias)
+
+    def forward(self, fmap, fmap_mask=None):
+        N, _, H, W = fmap.size()
+
+        y1, x1 = torch.meshgrid(
+            torch.arange(H, device=fmap.device, dtype=torch.float)/H,
+            torch.arange(W, device=fmap.device, dtype=torch.float)/W
+        )
+        y2, x2 = x1+1.0/W, y1+1.0/H
+        ww, hh = x2-x1, y2-y1
+        # x1, y1 = 2*x1-1, 2*y1-1
+        # x2, y2 = 2*x2-1, 2*y2-1
+        xc, yc = x1+0.5/W, y1+0.5/H
+
+        pos = torch.stack([x1, y1, x2, y2, xc, yc, ww, hh], dim=-1)
+        if self.proj is not None:
+            pos = self.proj(pos)
+        pos = pos.permute(2, 0, 1)
+        pos = pos.unsqueeze(0).repeat(N, 1, 1, 1)
+        return pos
+
+    def encode_boxes(self, boxes):
+        x1, y1, x2, y2 = boxes.unbind(-1)
+        ww, hh = x2-x1, y2-y1
+        xc, yc = x1+0.5*ww, y1+0.5*hh
+        pos = torch.stack([x1, y1, x2, y2, xc, yc, ww, hh], dim=-1)
+        if self.proj is not None:
+            pos = self.proj(pos)
+        return pos
+
+
+class RelativePositionEmbedding2D(nn.Module):
+    def __init__(self, embedding_dim, spatial_bins=(16, 16), with_projection=True):
+        super().__init__()
+
+        assert isinstance(spatial_bins, (list, tuple)) and len(spatial_bins) == 2
+        self.spatial_bins = spatial_bins
+
+        self.proj = None
+        if with_projection:
+            self.proj = nn.Linear(2*spatial_bins[0]*spatial_bins[1], embedding_dim)
+            nn.init.xavier_normal_(self.proj.weight)
+            nn.init.zeros_(self.proj.bias)
+
+    def forward(self, fmap, fmap_mask=None):
+        N, _, H, W = fmap.size()
+
+        BH, BW = self.spatial_bins
+        yc, xc = torch.meshgrid(
+            0.5/BH + torch.arange(BH, device=fmap.device, dtype=torch.float)/BH,
+            0.5/BW + torch.arange(BW, device=fmap.device, dtype=torch.float)/BW
+        )
+
+        pos = torch.stack([xc, yc], dim=-1).view(-1, 1, 2)
+        pos = (pos - pos.transpose(0, 1)).reshape(BH, BW, -1)  # relative positions
+
+        if self.proj is not None:
+            pos = self.proj(pos)
+
+        pos = pos.permute(2, 0, 1)
+        pos = pos.unsqueeze(0)
+
+        if H != BH or W != BW:
+            pos = nn.functional.interpolate(pos, (H, W), mode='nearest')
+
+        pos = pos.repeat(N, 1, 1, 1)
+
+        return pos

encoders.py

@@ -0,0 +1,414 @@
+import os
+
+from collections import OrderedDict
+
+import torch
+
+import transformers
+
+import torch.nn.functional as F
+
+from torch import nn
+
+from torchvision.models import detection
+
+from backbones import get_backbone
+
+from embeddings import Box8PositionEmbedding2D
+
+EPS = 1e-5
+
+TRANSFORMER_MODEL = 'bert-base-uncased'
+# TRANSFORMER_MODEL = 'distilroberta-base'
+
+
+def get_tokenizer(cache=None):
+    if cache is None:
+        return transformers.BertTokenizer.from_pretrained(TRANSFORMER_MODEL)
+
+    model_path = os.path.join(cache, TRANSFORMER_MODEL)
+    os.makedirs(model_path, exist_ok=True)
+
+    if os.path.exists(os.path.join(model_path, 'config.json')):
+        return transformers.BertTokenizer.from_pretrained(model_path)
+
+    tokenizer = transformers.BertTokenizer.from_pretrained(TRANSFORMER_MODEL)
+    tokenizer.save_pretrained(model_path)
+
+    return tokenizer
+
+
+def weight_init(m):
+    if isinstance(m, nn.Conv2d):
+        nn.init.xavier_normal_(m.weight, gain=nn.init.calculate_gain('relu'))
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.Linear):
+        nn.init.xavier_normal_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.Embedding):
+        nn.init.xavier_normal_(m.weight)
+
+
+class ImageEncoder(nn.Module):
+    def __init__(self, backbone='resnet50', out_channels=256, pretrained=True,
+                 freeze_pretrained=False, with_pos=True):
+        super().__init__()
+
+        model = get_backbone(backbone, pretrained)
+
+        if pretrained and freeze_pretrained:
+            for p in model.parameters():
+                p.requires_grad = False
+
+        if 'resnet' in backbone:
+            self.backbone = detection.backbone_utils.IntermediateLayerGetter(
+                model, return_layers=OrderedDict({'layer4': 'output'})
+            )
+            channels = 512 if backbone in ('resnet18', 'resnet34') else 2048
+
+        elif backbone in ('cspdarknet53', 'efficientnet-b0', 'efficientnet-b3'):
+            output_layer_name = list(model.named_children())[-1][0]
+            self.backbone = detection.backbone_utils.IntermediateLayerGetter(
+                model, return_layers=OrderedDict({output_layer_name: 'output'})
+            )
+            channels = {
+                'cspdarknet53': 1024,
+                'efficientnet-b0': 1280,
+                'efficientnet-b3': 1536
+            }[backbone]
+
+        else:
+            raise RuntimeError('not a valid backbone')
+
+        in_channels = channels+8 if with_pos else channels
+
+        self.proj = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, (1, 1), 1, bias=False),
+            nn.GroupNorm(1, out_channels, eps=EPS),
+            # nn.ReLU(inplace=True),
+        )
+        self.proj.apply(weight_init)
+
+        self.pos_emb = None
+        if with_pos:
+            self.pos_emb = Box8PositionEmbedding2D(with_projection=False)
+
+        self.out_channels = out_channels
+
+    def forward(self, img, mask=None):
+        x = self.backbone(img)['output']
+        if self.pos_emb is not None:
+            x = torch.cat([x, self.pos_emb(x)], dim=1)
+        x = self.proj(x)  # NxDxHxW
+
+        x_mask = None
+        if mask is not None:
+            _, _, H, W = x.size()
+            x_mask = F.interpolate(mask, (H, W), mode='bilinear')
+            x_mask = (x_mask > 0.5).long()
+
+        return x, x_mask
+
+
+class FPNImageEncoder(nn.Module):
+    def __init__(self,
+                 backbone='resnet50', out_channels=256, pretrained=True,
+                 freeze_pretrained=False, with_pos=True):
+        super().__init__()
+
+        model = get_backbone(backbone, pretrained)
+
+        if pretrained and freeze_pretrained:
+            for p in model.parameters():
+                p.requires_grad = False
+
+        if 'resnet' in backbone:
+            if backbone in ('resnet18', 'resnet34'):
+                in_channels_list = [64, 128, 256, 512]
+            else:
+                in_channels_list = [256, 512, 1024, 2048]
+            return_layers = OrderedDict({
+                'layer1': '0', 'layer2': '1', 'layer3': '2', 'layer4': '3'
+            })
+
+        # elif backbone == 'cspdarknet53':
+        #     in_channels_list = [128, 256, 512, 1024]
+        #     return_layers = OrderedDict({
+        #         '1':'0', '2':'1', '3':'2', '4':'3'
+        #     })
+
+        else:
+            raise RuntimeError('not a valid backbone')
+
+        self.backbone = model
+
+        self.fpn = detection.backbone_utils.BackboneWithFPN(
+            backbone=self.backbone,
+            return_layers=return_layers,
+            in_channels_list=in_channels_list,
+            out_channels=out_channels
+        )
+
+        self.fpn.fpn.extra_blocks = None   # removes the 'pool' layer added by default
+
+        self.out_channels = out_channels
+
+        in_channels = int(out_channels + float(with_pos) * 8)
+
+        self.proj = nn.ModuleDict({
+            level: nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, (1, 1), 1, bias=False),
+                nn.GroupNorm(1, out_channels, eps=EPS),
+                # nn.ReLU(inplace=True),
+            ) for level in return_layers.values()
+        })
+        self.proj.apply(weight_init)
+
+        self.pos_emb = None
+        if with_pos:
+            self.pos_emb = Box8PositionEmbedding2D(with_projection=False)
+
+    def forward(self, x, mask=None):
+        x = self.fpn(x)
+
+        # smallest feature map (eg. 16x16 for an input of 512x512 pixels)
+        _, _, H, W = list(x.values())[-1].size()
+
+        x_out = None
+        for level, fmap in x.items():
+            # fmap = torch.relu(fmap)  # FPN blocks end in a conv2d, w/o activ.
+            if self.pos_emb is not None:
+                fmap = torch.cat([fmap, self.pos_emb(fmap)], dim=1)  # +Pos
+            fmap = self.proj[level](fmap)   # Conv+BN+ReLU
+            fmap = F.interpolate(fmap, (H, W), mode='nearest')  # to a smaller size
+            if x_out is None:
+                x_out = fmap
+            else:
+                x_out += fmap
+
+        x_mask = None
+        if mask is not None:
+            x_mask = F.interpolate(mask, (H, W), mode='bilinear')
+            x_mask = (x_mask > 0.5).long()
+
+        return x_out, x_mask
+
+
+class TransformerImageEncoder(nn.Module):
+    def __init__(self,
+                 backbone='resnet50', out_channels=256, pretrained=True,
+                 freeze_pretrained=False, num_heads=8, num_layers=6,
+                 dropout_p=0.1):
+        super().__init__()
+
+        model = get_backbone(backbone, pretrained)
+
+        if pretrained and freeze_pretrained:
+            for p in model.parameters():
+                p.requires_grad = False
+
+        if 'resnet' in backbone:
+            self.backbone = detection.backbone_utils.IntermediateLayerGetter(
+                model, return_layers=OrderedDict({'layer4': 'output'})
+            )
+            channels = 512 if backbone in ('resnet18', 'resnet34') else 2048
+
+        elif backbone in ('cspdarknet53', 'efficientnet-b0', 'efficientnet-b3'):
+            output_layer_name = list(model.named_children())[-1][0]
+            self.backbone = detection.backbone_utils.IntermediateLayerGetter(
+                model, return_layers=OrderedDict({output_layer_name: 'output'})
+            )
+            channels = {
+                'cspdarknet53': 1024,
+                'efficientnet-b0': 1280,
+                'efficientnet-b3': 1536
+            }[backbone]
+
+        else:
+            raise RuntimeError('not a valid backbone')
+
+        self.proj = nn.Sequential(
+            nn.Conv2d(channels, out_channels, (1, 1), 1, bias=False),
+            nn.GroupNorm(1, out_channels, eps=EPS),
+            # nn.ReLU(inplace=True),
+        )
+        self.proj.apply(weight_init)
+
+        from transformers_pos import (
+            TransformerEncoder,
+            TransformerEncoderLayer,
+        )
+
+        self.encoder = TransformerEncoder(
+            TransformerEncoderLayer(
+                d_model=out_channels,
+                nhead=num_heads,
+                dropout=dropout_p,
+                batch_first=True
+            ),
+            num_layers=num_layers
+        )
+
+        self.pos_emb = Box8PositionEmbedding2D(embedding_dim=out_channels)
+
+        self.out_channels = out_channels
+
+    def flatten(self, x):
+        N, _, H, W = x.size()
+        x = x.to(memory_format=torch.channels_last)
+        x = x.permute(0, 2, 3, 1).view(N, H*W, -1)  # NxHWxD
+        return x
+
+    def forward(self, img, mask=None):
+        x = self.backbone(img)['output']
+        x = self.proj(x)  # NxDxHxW
+
+        N, _, H, W = x.size()
+
+        pos = self.pos_emb(x)  # NxDxHxW
+        pos = self.flatten(pos)  # NxRxD
+
+        x = self.flatten(x)  # NxRxD
+
+        # visibility mask
+        x_mask = None
+        if mask is not None:
+            x_mask = F.interpolate(mask, (H, W), mode='bilinear')
+            x_mask = (x_mask > 0.5).long()
+
+        if mask is None:
+            x = self.encoder(x, pos=pos)  # NxRxD
+        else:
+            mask = self.flatten(x_mask).squeeze(-1)
+            x = self.encoder(x, src_key_padding_mask=(mask==0), pos=pos)  # NxRxD
+
+        x = x.permute(0, 2, 1).view(N, -1, H, W)  # NxDxHxW
+
+        return x, x_mask
+
+
+class LanguageEncoder(nn.Module):
+    def __init__(self, out_features=256, dropout_p=0.2,
+                 freeze_pretrained=False, global_pooling=True):
+        super().__init__()
+        self.language_model = transformers.AutoModel.from_pretrained(
+            TRANSFORMER_MODEL
+        )
+
+        if freeze_pretrained:
+            for p in self.language_model.parameters():
+                p.requires_grad = False
+
+        self.out_features = out_features
+
+        self.proj = nn.Sequential(
+            nn.Linear(768, out_features),
+            nn.LayerNorm(out_features, eps=1e-5),
+            # nn.ReLU(inplace=True),
+            # nn.Dropout(dropout_p),
+        )
+        self.proj.apply(weight_init)
+
+        self.global_pooling = bool(global_pooling)
+
+    def forward(self, z):
+        res = self.language_model(
+            input_ids=z['input_ids'],
+            position_ids=None,
+            attention_mask=z['attention_mask']
+        )
+
+        if self.global_pooling:
+            z, z_mask = self.proj(res.pooler_output), None
+        else:
+            z, z_mask = self.proj(res.last_hidden_state), z['attention_mask']
+
+        return z, z_mask
+
+
+class RNNLanguageEncoder(nn.Module):
+    def __init__(self,
+                 model_type='gru', hidden_size=1024, num_layers=2,
+                 out_features=256, dropout_p=0.2, global_pooling=True):
+        super().__init__()
+        self.embeddings = transformers.AutoModel.from_pretrained(
+            TRANSFORMER_MODEL
+        ).embeddings.word_embeddings
+        self.embeddings.weight.requires_grad = True
+
+        # self.dropout_emb = nn.Dropout(0.5)
+        self.dropout_emb = nn.Dropout(dropout_p)
+
+        assert model_type in ('gru', 'lstm')
+        self.rnn = (nn.GRU if model_type == 'gru' else nn.LSTM)(
+            input_size=self.embeddings.weight.size(1),
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            dropout=dropout_p,
+            batch_first=True,
+            bidirectional=True
+        )
+
+        self.proj = nn.Sequential(
+            nn.Linear(2*hidden_size, out_features),
+            nn.LayerNorm(out_features, eps=1e-5),
+            # nn.ReLU(inplace=True),
+            # nn.Dropout(dropout_p),
+        )
+        self.proj.apply(weight_init)
+
+        self.out_features = out_features
+
+        self.global_pooling = bool(global_pooling)
+        assert global_pooling  # only w/ global pooling
+
+    def forward(self, z):
+        z_mask = z['attention_mask']
+
+        z = self.dropout_emb(self.embeddings(z['input_ids']))
+        z, h_n = self.rnn(z, None)
+
+        if isinstance(self.rnn, nn.LSTM):
+            h_n = h_n[0]
+
+        # hidden states as (num_layers, num_directions, batch, hidden_size)
+        h_n = h_n.view(self.rnn.num_layers, 2, z.size(0), self.rnn.hidden_size)
+
+        # last hidden states
+        h_n = h_n[-1].permute(1, 0, 2).reshape(z.size(0), -1)
+        h_n = self.proj(h_n)
+        return h_n, z_mask
+
+
+class SimpleEncoder(nn.Module):
+    def __init__(self, out_features=256, dropout_p=0.1, global_pooling=True):
+        super().__init__()
+        self.embeddings = transformers.AutoModel.from_pretrained(
+            TRANSFORMER_MODEL
+        ).embeddings.word_embeddings
+        self.embeddings.weight.requires_grad = True
+
+        # self.dropout_emb = nn.Dropout(0.5)
+        self.dropout_emb = nn.Dropout(dropout_p)
+
+        self.proj = nn.Sequential(
+            nn.Linear(768, out_features),
+            nn.LayerNorm(out_features, eps=1e-5),
+            # nn.ReLU(inplace=True),
+            # nn.Dropout(dropout_p),
+        )
+        self.proj.apply(weight_init)
+
+        self.out_features = out_features
+
+        self.global_pooling = bool(global_pooling)
+        assert not self.global_pooling  # only w/o global pooling
+
+    def forward(self, z):
+        z_mask = z['attention_mask']
+        z = self.embeddings(z['input_ids'])
+        z = self.proj(self.dropout_emb(z))
+        # z[:, 0] = torch.mean(z[:, 1:], 1)
+        return z, z_mask

models.py

@@ -0,0 +1,412 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torchvision.ops import box_convert
+import embeddings as emb
+import encoders as enc
+from encoders import weight_init
+
+def conv3x3(in_channels, out_channels, num_groups=0):
+    return nn.Sequential(
+        # Conv2d w/o bias since BatchNorm2d/GroupNorm already accounts for it (affine=True)
+        nn.Conv2d(in_channels, out_channels, (3, 3), 1, 1, bias=False),
+        nn.BatchNorm2d(out_channels) if num_groups < 1 else nn.GroupNorm(num_groups, out_channels),
+        nn.ReLU(inplace=True),
+    )
+
+
+class IntuitionKillingMachine(nn.Module):
+    def __init__(self,
+                 backbone='resnet50', pretrained=True, embedding_size=256,
+                 num_heads=8, num_layers=6, num_conv=4, dropout_p=0.1,
+                 segmentation_head=True, mask_pooling=True):
+        super().__init__()
+
+        if backbone.endswith('+tr'):
+            self.vis_enc = enc.TransformerImageEncoder(
+                backbone=backbone.rstrip('+tr'),
+                out_channels=embedding_size,
+                pretrained=pretrained,
+            )
+
+        elif backbone.endswith('+fpn'):
+            self.vis_enc = enc.FPNImageEncoder(
+                backbone=backbone.rstrip('+fpn'),
+                out_channels=embedding_size,
+                pretrained=pretrained,
+                with_pos=False
+            )
+        else:
+            self.vis_enc = enc.ImageEncoder(
+                backbone=backbone,
+                out_channels=embedding_size,
+                pretrained=pretrained,
+                with_pos=False
+            )
+
+        # freeze ResNet stem
+        if 'resnet' in backbone:
+            self.vis_enc.backbone.conv1.requires_grad = False
+            self.vis_enc.backbone.conv1.eval()
+
+        self.vis_pos_emb = emb.LearnedPositionEmbedding2D(
+            embedding_dim=embedding_size
+        )
+
+        self.lan_enc = enc.LanguageEncoder(
+            out_features=embedding_size,
+            global_pooling=False,
+            dropout_p=dropout_p
+        )
+
+        self.lan_pos_emb = emb.LearnedPositionEmbedding1D(
+            embedding_dim=embedding_size
+        )
+
+        from transformers_pos import (
+            XTransformerEncoder,
+            TransformerEncoder,
+            TransformerEncoderLayer,
+        )
+
+        self.encoder = TransformerEncoder(
+            TransformerEncoderLayer(
+                d_model=embedding_size,
+                nhead=num_heads,
+                dropout=dropout_p,
+                batch_first=True
+            ),
+            num_layers=num_layers
+        )
+
+        # ---
+        # CONV PRE-HEAD (NECK?)
+
+        if num_conv > 0:
+            self.pre_head = nn.Sequential(*[
+                conv3x3(embedding_size, embedding_size) for _ in range(num_conv)
+            ])
+            self.pre_head.apply(weight_init)
+        else:
+            self.pre_head = nn.Identity()
+
+        # ---
+        # OUTPUT HEADS
+
+        # box prediction
+        self.head = nn.Sequential(
+            nn.Linear(embedding_size, 4, bias=True),
+            nn.Sigmoid()
+        )
+        self.head.apply(weight_init)
+
+        # box segmentation mask
+        self.segm_head = None
+        if segmentation_head:
+            self.segm_head = nn.Sequential(
+                nn.Conv2d(embedding_size, 1, (3, 3), 1, 1, bias=True),
+                #nn.Sigmoid()
+            )
+            self.segm_head.apply(weight_init)
+
+        # ---
+
+        self.mask_pooling = bool(mask_pooling)
+
+        if self.mask_pooling and self.segm_head is None:
+            raise RuntimeError('mask pooling w/o a segmentation head does not makes sense')
+
+        self.embedding_size = embedding_size
+
+    # def slow_param_ids(self, **kwargs):
+    #     return []
+
+    def slow_param_ids(self, slow_visual_backbone=True, slow_language_backbone=True):
+        ids = []
+
+        if slow_visual_backbone:
+            ids += [id(p) for p in self.vis_enc.backbone.parameters()]
+            if hasattr(self.vis_enc, 'encoder'):  # +tr
+                ids += [id(p) for p in self.vis_enc.encoder.parameters()]
+
+        if slow_language_backbone:
+            if isinstance(self.lan_enc, enc.LanguageEncoder):
+                ids += [id(p) for p in self.lan_enc.language_model.parameters()]
+            else:
+                ids += [id(p) for p in self.lan_enc.embeddings.parameters()]
+
+        return ids
+
+    def flatten(self, x):
+        N, D, H, W = x.size()
+        x = x.to(memory_format=torch.channels_last)
+        x = x.permute(0, 2, 3, 1).view(N, H*W, D)
+        return x  # NxHWxD
+
+    def unflatten(self, x, size):
+        N, R, D = x.size()
+        H, W = size
+        assert R == H*W, 'wrong tensor size'
+        x = x.permute(0, 2, 1).to(memory_format=torch.contiguous_format)
+        x = x.view(N, D, H, W)
+        return x  # NxDxHxW
+
+    def forward(self, input):
+        img, mask, tok = input['image'], input['mask'], input['tok']
+
+        # ---
+        # VISUAL EMBEDDINGS
+
+        x, x_mask = self.vis_enc(img, mask)   # NxDxHxW, NxHxW
+        x_pos = self.vis_pos_emb(x, x_mask)
+
+        N, D, H, W = x.size()  # save dims before flatten
+
+        x = self.flatten(x)  # NxRxD
+        x_mask = self.flatten(x_mask).squeeze(-1)  # NxR
+        x_pos = self.flatten(x_pos)   # NxRxD
+
+        # ---
+        # LANGUAGE EMBEDDINGS
+
+        z, z_mask = self.lan_enc(tok)   # NxTxD, NxT
+        z_pos = self.lan_pos_emb(z)  # NxTxD
+
+        # ---
+        # V+L TRANSFORMER
+
+        # [...visual...]+[[CLS]...language tokens...[SEP]]
+        xz = torch.cat([x, z], dim=1)
+        xz_mask = torch.cat([x_mask, z_mask], dim=1)
+        xz_pos = torch.cat([x_pos, z_pos], dim=1)
+
+        xz = self.encoder(xz, src_key_padding_mask=(xz_mask==0), pos=xz_pos)  #, size=(H,W))
+
+        # restore spatiality of visual embeddings after cross-modal encoding
+        xz_vis = xz[:, :H*W, ...]
+        xz_vis = self.unflatten(xz_vis, (H, W))
+
+        x_mask = self.unflatten(x_mask.unsqueeze(-1), (H, W))
+
+        # ---
+
+        # convolutional pre-head
+        xz_vis = self.pre_head(xz_vis)
+
+        # ---
+
+        # segmentation head w/ (opt.) pooling
+        segm_mask, pooled_feat = None, None
+        if self.segm_head is not None:
+            segm_mask = torch.sigmoid(self.segm_head(xz_vis)) * x_mask
+            if self.mask_pooling:  # box mask guided pooling
+                pooled_feat = (segm_mask * xz_vis).sum((2, 3)) / segm_mask.sum((2, 3))
+            segm_mask = F.interpolate(segm_mask, img.size()[2:], mode='bilinear', align_corners=True)
+
+        # if not mask_pooling, do the pooling using all visual feats (equiv. to a uniform mask)
+        if pooled_feat is None:
+            pooled_feat = (x_mask * xz_vis).sum((2, 3)) / x_mask.sum((2, 3))
+
+        # bbox prediction
+        pred = self.head(pooled_feat)
+        pred = box_convert(pred, 'cxcywh', 'xyxy')
+
+        return pred, segm_mask
+
+class HeadlessMachine(nn.Module):
+    def __init__(self,
+                 backbone='resnet50', pretrained=True, embedding_size=256,
+                 num_heads=8, num_layers=6, num_conv=4, dropout_p=0.1,
+                 segmentation_head=True, mask_pooling=True):
+        super().__init__()
+
+        if backbone.endswith('+tr'):
+            self.vis_enc = enc.TransformerImageEncoder(
+                backbone=backbone.rstrip('+tr'),
+                out_channels=embedding_size,
+                pretrained=pretrained,
+            )
+
+        elif backbone.endswith('+fpn'):
+            self.vis_enc = enc.FPNImageEncoder(
+                backbone=backbone.rstrip('+fpn'),
+                out_channels=embedding_size,
+                pretrained=pretrained,
+                with_pos=False
+            )
+        else:
+            self.vis_enc = enc.ImageEncoder(
+                backbone=backbone,
+                out_channels=embedding_size,
+                pretrained=pretrained,
+                with_pos=False
+            )
+
+        # freeze ResNet stem
+        if 'resnet' in backbone:
+            self.vis_enc.backbone.conv1.requires_grad = False
+            self.vis_enc.backbone.conv1.eval()
+
+        self.vis_pos_emb = emb.LearnedPositionEmbedding2D(
+            embedding_dim=embedding_size
+        )
+
+        self.lan_enc = enc.LanguageEncoder(
+            out_features=embedding_size,
+            global_pooling=False,
+            dropout_p=dropout_p
+        )
+
+        self.lan_pos_emb = emb.LearnedPositionEmbedding1D(
+            embedding_dim=embedding_size
+        )
+
+        from transformers_pos import (
+            XTransformerEncoder,
+            TransformerEncoder,
+            TransformerEncoderLayer,
+        )
+
+        self.encoder = TransformerEncoder(
+            TransformerEncoderLayer(
+                d_model=embedding_size,
+                nhead=num_heads,
+                dropout=dropout_p,
+                batch_first=True
+            ),
+            num_layers=num_layers
+        )
+
+        # ---
+        # CONV PRE-HEAD (NECK?)
+
+        if num_conv > 0:
+            self.pre_head = nn.Sequential(*[
+                conv3x3(embedding_size, embedding_size) for _ in range(num_conv)
+            ])
+            self.pre_head.apply(weight_init)
+        else:
+            self.pre_head = nn.Identity()
+
+        # ---
+        # OUTPUT HEADS
+
+        # box prediction
+        self.head = nn.Sequential(
+            nn.Linear(embedding_size, 4, bias=True),
+            nn.Sigmoid()
+        )
+        self.head.apply(weight_init)
+
+        # box segmentation mask
+        self.segm_head = None
+        if segmentation_head:
+            self.segm_head = nn.Sequential(
+                nn.Conv2d(embedding_size, 1, (3, 3), 1, 1, bias=True),
+                #nn.Sigmoid()
+            )
+            self.segm_head.apply(weight_init)
+
+        # ---
+
+        self.mask_pooling = bool(mask_pooling)
+
+        if self.mask_pooling and self.segm_head is None:
+            raise RuntimeError('mask pooling w/o a segmentation head does not makes sense')
+
+        self.embedding_size = embedding_size
+
+    # def slow_param_ids(self, **kwargs):
+    #     return []
+
+    def slow_param_ids(self, slow_visual_backbone=True, slow_language_backbone=True):
+        ids = []
+
+        if slow_visual_backbone:
+            ids += [id(p) for p in self.vis_enc.backbone.parameters()]
+            if hasattr(self.vis_enc, 'encoder'):  # +tr
+                ids += [id(p) for p in self.vis_enc.encoder.parameters()]
+
+        if slow_language_backbone:
+            if isinstance(self.lan_enc, enc.LanguageEncoder):
+                ids += [id(p) for p in self.lan_enc.language_model.parameters()]
+            else:
+                ids += [id(p) for p in self.lan_enc.embeddings.parameters()]
+
+        return ids
+
+    def flatten(self, x):
+        N, D, H, W = x.size()
+        x = x.to(memory_format=torch.channels_last)
+        x = x.permute(0, 2, 3, 1).view(N, H*W, D)
+        return x  # NxHWxD
+
+    def unflatten(self, x, size):
+        N, R, D = x.size()
+        H, W = size
+        assert R == H*W, 'wrong tensor size'
+        x = x.permute(0, 2, 1).to(memory_format=torch.contiguous_format)
+        x = x.view(N, D, H, W)
+        return x  # NxDxHxW
+
+    def forward(self, input):
+        img, mask, tok = input['image'], input['mask'], input['tok']
+
+        # ---
+        # VISUAL EMBEDDINGS
+
+        x, x_mask = self.vis_enc(img, mask)   # NxDxHxW, NxHxW
+        x_pos = self.vis_pos_emb(x, x_mask)
+
+        N, D, H, W = x.size()  # save dims before flatten
+
+        x = self.flatten(x)  # NxRxD
+        x_mask = self.flatten(x_mask).squeeze(-1)  # NxR
+        x_pos = self.flatten(x_pos)   # NxRxD
+
+        # ---
+        # LANGUAGE EMBEDDINGS
+
+        z, z_mask = self.lan_enc(tok)   # NxTxD, NxT
+        z_pos = self.lan_pos_emb(z)  # NxTxD
+
+        # ---
+        # V+L TRANSFORMER
+
+        # [...visual...]+[[CLS]...language tokens...[SEP]]
+        xz = torch.cat([x, z], dim=1)
+        xz_mask = torch.cat([x_mask, z_mask], dim=1)
+        xz_pos = torch.cat([x_pos, z_pos], dim=1)
+
+        xz = self.encoder(xz, src_key_padding_mask=(xz_mask==0), pos=xz_pos)  #, size=(H,W))
+
+        # restore spatiality of visual embeddings after cross-modal encoding
+        xz_vis = xz[:, :H*W, ...]
+        xz_vis = self.unflatten(xz_vis, (H, W))
+
+        x_mask = self.unflatten(x_mask.unsqueeze(-1), (H, W))
+
+        # ---
+
+        # convolutional pre-head
+        xz_vis = self.pre_head(xz_vis)
+
+        # ---
+
+        # segmentation head w/ (opt.) pooling
+        segm_mask, pooled_feat = None, None
+        if self.segm_head is not None:
+            segm_mask = torch.sigmoid(self.segm_head(xz_vis)) * x_mask
+            if self.mask_pooling:  # box mask guided pooling
+                pooled_feat = (segm_mask * xz_vis).sum((2, 3)) / segm_mask.sum((2, 3))
+            segm_mask = F.interpolate(segm_mask, img.size()[2:], mode='bilinear', align_corners=True)
+
+        # if not mask_pooling, do the pooling using all visual feats (equiv. to a uniform mask)
+        if pooled_feat is None:
+            pooled_feat = (x_mask * xz_vis).sum((2, 3)) / x_mask.sum((2, 3))
+
+        # bbox prediction
+        pred = self.head(pooled_feat)
+        pred = box_convert(pred, 'cxcywh', 'xyxy')
+
+        return pred, segm_mask

requirements.txt

@@ -0,0 +1,5 @@
+Pillow==9.1.0
+timm==0.6.7
+torch==1.9.0
+torchvision==0.10.0
+transformers==4.12.3

testing_loading.py

@@ -0,0 +1,97 @@
+from models import IntuitionKillingMachine
+from transforms import undo_box_transforms_batch, ToTensor, Normalize, SquarePad, Resize, NormalizeBoxCoords
+from torchvision.transforms import Compose
+from encoders import get_tokenizer
+from PIL import Image, ImageDraw
+from zipfile import ZipFile
+from copy import copy
+import pandas as pd
+import torch
+
+def parse_model_args(model_path):
+    _, _, dataset, max_length, input_size, backbone, num_heads, num_layers, num_conv, _, _, mu, mask_pooling = model_path.split('_')[:13]
+    return {
+        'dataset': dataset,
+        'max_length': int(max_length),
+        'input_size': int(input_size),
+        'backbone': backbone,
+        'num_heads': int(num_heads),
+        'num_layers': int(num_layers),
+        'num_conv': int(num_conv),
+        'mu': float(mu),
+        'mask_pooling': bool(mask_pooling == '1')
+    }
+
+
+class Prober:
+    def __init__(self,
+                 df_path=None,
+                 dataset_path=None,
+                 model_checkpoint=None):
+        params = parse_model_args(model_checkpoint)
+        mean = [0.485, 0.456, 0.406]
+        sdev = [0.229, 0.224, 0.225]
+        self.tokenizer = get_tokenizer()
+        self.df = pd.read_json(df_path)[['sample_idx', 'bbox', 'file_path', 'sent']]
+        self.df.loc[:, "image_id"] = self.df.file_path.apply(lambda x: int(x.split('/')[-1][:-4]))
+        self.df.file_path = self.df.file_path.apply(lambda x: x.replace('refer/data/images/', ''))
+        self.model = IntuitionKillingMachine(
+            backbone=params['backbone'],
+            pretrained=True,
+            num_heads=params['num_heads'],
+            num_layers=params['num_layers'],
+            num_conv=params['num_conv'],
+            segmentation_head=bool(params['mu'] > 0.0),
+            mask_pooling=params['mask_pooling']
+        ) 
+        self.transform = Compose([
+            ToTensor(),
+            Normalize(mean, sdev),
+            SquarePad(),
+            Resize(size=(params['input_size'], params['input_size'])),
+            NormalizeBoxCoords(),
+        ])
+        self.max_length = 30
+        self.zipfile = ZipFile(dataset_path, 'r')
+
+    @torch.no_grad()
+    def probe(self, idx, re, search_by_sample_id: bool= True):
+        if search_by_sample_id:
+            img_path, target, = self.df.loc[idx][['file_path','bbox']].values
+        else: 
+            img_path, target = self.df[self.df.image_id == idx][['file_path','bbox']].values[0]
+        img = Image.open(self.zipfile.open(img_path)).convert('RGB')
+        W0, H0 = img.size
+        sample = {
+            'image': img,
+            'image_size': (H0, W0),  # image original size
+            'bbox': torch.tensor([copy(target)]),
+            'bbox_raw': torch.tensor([copy(target)]),
+            'mask': torch.ones((1, H0, W0), dtype=torch.float32),  # visibiity mask
+            'mask_bbox': None,  # target bbox mask
+        } 
+        print('inn bbox: ', sample['bbox'])
+        sample = self.transform(sample)
+        tok = self.tokenizer(re,
+                             max_length=30,
+                             return_tensors='pt',
+                             truncation=True)
+        inn = {'image': torch.stack([sample['image']]),
+               'mask': torch.stack([sample['mask']]),
+               'bbox': torch.stack([sample['bbox']]),
+               'tok': tok}
+        output = undo_box_transforms_batch(self.model(inn)[0],
+                                           [sample['tr_param']]).numpy().tolist()[0]
+        img1 = ImageDraw.Draw(img)
+        #img1.rectangle(target, outline ="#0000FF00", width=3)
+        img1.rectangle(output, outline ="#00FF0000", width=3)
+        return img
+
+if __name__ == "__main__":
+    prober = Prober(
+        df_path = 'data/val-sim_metric.json',
+        dataset_path = "data/saiapr_tc-12.zip",
+        model_checkpoint= "cache/20211220_191132_refclef_32_512_resnet50_8_6_8_0.1_0.1_0.1_0_0.0001_0.0_12_4_90_1_0_0_0/best.ckpt"
+    )
+    prober.probe(0, "tree")
+    print("Done")

transformers_pos.py

@@ -0,0 +1,198 @@
+import copy
+from typing import Optional, Any
+
+import torch
+
+from torch import Tensor
+from torch import nn
+from torch.nn import functional as F
+
+
+def conv3x3(in_channels, out_channels, num_groups=0):
+    return nn.Sequential(
+        # Conv2d w/o bias since BatchNorm2d/GroupNorm already accounts for it (affine=True)
+        nn.Conv2d(in_channels, out_channels, (3, 3), 1, 1, bias=False),
+        nn.BatchNorm2d(out_channels) if num_groups < 1 else nn.GroupNorm(num_groups, out_channels),
+        nn.ReLU(inplace=True),
+    )
+
+
+class XTransformerEncoder(nn.Module):
+    __constants__ = ['norm']
+    def __init__(self, encoder_layer, num_layers, num_conv=2, norm=None):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+        d_model = encoder_layer.linear1.in_features
+        self.conv = nn.ModuleList([
+            nn.Sequential(*[
+                conv3x3(d_model, d_model) for _ in range(num_conv)
+            ]) for _ in range(num_layers)
+        ])
+
+    def flatten(self, x):
+        N, D, H, W = x.size()
+        x = x.to(memory_format=torch.channels_last)
+        x = x.permute(0, 2, 3, 1).view(N, H*W, D)
+        return x  # NxHWxD
+
+    def unflatten(self, x, size):
+        N, R, D = x.size()
+        H, W = size
+        assert R == H*W, 'wrong tensor size'
+        x = x.permute(0, 2, 1).to(memory_format=torch.contiguous_format)
+        x = x.view(N, D, H, W)
+        return x  # NxDxHxW
+
+    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, size=None) -> Tensor:
+        output = src
+
+        for i, mod in enumerate(self.layers):
+            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos)
+
+            vis = self.unflatten(output[:, :size[0]*size[1]], size)
+            vis = self.flatten(self.conv[i](vis))
+
+            output = torch.cat([vis, output[:, size[0]*size[1]:]], dim=1)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerEncoder(nn.Module):
+    r"""TransformerEncoder is a stack of N encoder layers
+
+    Args:
+        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+
+    Examples::
+        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = transformer_encoder(src)
+    """
+    __constants__ = ['norm']
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super(TransformerEncoder, self).__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        output = src
+
+        for mod in self.layers:
+            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerEncoderLayer(nn.Module):
+    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
+    This standard encoder layer is based on the paper "Attention Is All You Need".
+    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
+    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
+    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
+    in a different way during application.
+
+    Args:
+        d_model: the number of expected features in the input (required).
+        nhead: the number of heads in the multiheadattention models (required).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        activation: the activation function of intermediate layer, relu or gelu (default=relu).
+        layer_norm_eps: the eps value in layer normalization components (default=1e-5).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False``.
+
+    Examples::
+        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = encoder_layer(src)
+
+    Alternatively, when ``batch_first`` is ``True``:
+        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
+        >>> src = torch.rand(32, 10, 512)
+        >>> out = encoder_layer(src)
+    """
+    __constants__ = ['batch_first']
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu",
+                 layer_norm_eps=1e-5, batch_first=False,
+                 device=None, dtype=None) -> None:
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super(TransformerEncoderLayer, self).__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
+                                               **factory_kwargs)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward, **factory_kwargs)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model, **factory_kwargs)
+
+        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+
+    def __setstate__(self, state):
+        if 'activation' not in state:
+            state['activation'] = F.relu
+        super(TransformerEncoderLayer, self).__setstate__(state)
+
+    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None) -> Tensor:
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+
+        q = k = src if pos is None else src + pos
+
+        src2 = self.self_attn(q, k, src, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))

transforms.py

@@ -0,0 +1,276 @@
+import torch
+
+from torchvision import transforms
+
+from torchvision.transforms import Compose
+
+from PIL import Image
+
+
+class ToTensor(transforms.ToTensor):
+    def __call__(self, input):
+        if not isinstance(input, dict):
+            return super().__call__(input)
+        assert 'image' in input
+        input['image'] = super().__call__(input['image'])
+        return input
+
+
+class Normalize(transforms.Normalize):
+    def __call__(self, input):
+        if not isinstance(input, dict):
+            return super().__call__(input)
+        assert 'image' in input
+        input['image'] = super().__call__(input['image'])
+        return input
+
+
+class NormalizeBoxCoords(transforms.ToTensor):
+    def __call__(self, input):
+        if not isinstance(input, dict):
+            return super().__call__(input)
+        assert 'image' in input and 'bbox' in input
+        _, H, W = input['image'].size()
+        input['bbox'][:, (0, 2)] /= W
+        input['bbox'][:, (1, 3)] /= H
+
+        if 'tr_param' not in input:
+            input['tr_param'] = []
+        input['tr_param'].append({'normalize_box_coords': (H, W)})
+
+        return input
+
+
+class SquarePad(torch.nn.Module):
+    def __call__(self, input):
+        if isinstance(input, Image.Image):
+            raise NotImplementedError('put the SquarePad transform after ToTensor')
+
+        assert 'image' in input
+        _, h, w = input['image'].size()
+
+        max_wh = max(w, h)
+        xp = int(0.5 * (max_wh - w))
+        yp = int(0.5 * (max_wh - h))
+        padding = (xp, yp, (max_wh-xp)-w, (max_wh-yp)-h)
+
+        input['image'] = transforms.functional.pad(
+            input['image'], padding, fill=0, padding_mode='constant'
+        )
+        # input['image'] = transforms.functional.pad(
+        #     input['image'], padding, padding_mode='edge'
+        # )
+
+        if 'mask' in input:
+            input['mask'] = transforms.functional.pad(
+                input['mask'], padding, fill=0, padding_mode='constant'
+            )
+
+        if 'bbox' in input:
+            input['bbox'][:, (0, 2)] += xp
+            input['bbox'][:, (1, 3)] += yp
+
+        if 'tr_param' not in input:
+            input['tr_param'] = []
+        input['tr_param'].append({'square_pad': padding})
+
+        return input
+
+
+class Resize(transforms.Resize):
+    def __call__(self, input):
+        if not isinstance(input, dict):
+            return super().__call__(input)
+
+        assert 'image' in input
+
+        if not torch.is_tensor(input['image']):
+            raise NotImplementedError('put the Resize transform after ToTensor')
+
+        _, img_h, img_w = input['image'].size()
+
+        if isinstance(self.size, int):
+            dst_h = self.size if img_h < img_w else int(self.size * img_h / img_w)
+            dst_w = self.size if img_w < img_h else int(self.size * img_w / img_h)
+        else:
+            dst_h, dst_w = self.size
+
+        input['image'] = super().__call__(input['image'])
+
+        if 'mask' in input:
+            input['mask'] = super().__call__(input['mask'])
+
+        sx, sy = dst_w / img_w, dst_h / img_h
+
+        if 'bbox' in input:
+            input['bbox'][:, (0, 2)] *= sx
+            input['bbox'][:, (1, 3)] *= sy
+
+        if 'tr_param' not in input:
+            input['tr_param'] = []
+        input['tr_param'].append({'resize': (sx, sy)})
+
+        return input
+
+
+class RandomHorizontalFlip(transforms.RandomHorizontalFlip):
+    def __call__(self, input):
+        if not isinstance(input, dict):
+            return super().__call__(input)
+
+        assert 'image' in input
+
+        if not torch.is_tensor(input['image']):
+            raise NotImplementedError('use Resize after ToTensor')
+
+        result = super().__call__(input['image'])
+        if result is input['image']:  # not flipped
+            return input
+        input['image'] = result
+
+        if 'mask' in input:
+            input['mask'] = torch.flip(input['mask'], dims=(-1,))
+
+        img_w = input['image'].size(2)
+
+        if 'bbox' in input:
+            input['bbox'][:, (0, 2)] = img_w - input['bbox'][:, (2, 0)]
+
+        if 'expr' in input:
+            input['expr'] = input['expr'].replace('left', '<LEFT>').replace('right', 'left').replace('<LEFT>', 'right')
+
+        return input
+
+
+class RandomAffine(transforms.RandomAffine):
+    def get_params(self, *args, **kwargs):
+        self.params = super().get_params(*args, **kwargs)
+        return self.params
+
+    def __call__(self, input):
+        if not isinstance(input, dict):
+            return super().__call__(input)
+
+        assert 'image' in input
+
+        if not torch.is_tensor(input['image']):
+            raise NotImplementedError('put the Resize transform after ToTensor')
+
+        #self.fill = input['image'].mean((1,2))  # set fill value to the mean pixel value
+        result = super().__call__(input['image'])
+        if result is input['image']:  # not transformed
+            return input
+        input['image'] = result
+
+        _, img_h, img_w = input['image'].size()
+
+        angle, translate, scale, shear = self.params
+        center = (img_w * 0.5, img_h * 0.5)
+        matrix = transforms.functional._get_inverse_affine_matrix(center, angle, translate, scale, shear)
+        matrix = torch.FloatTensor([matrix[:3], matrix[3:], [0, 0, 1]])
+        matrix = torch.linalg.inv(matrix)
+
+        if 'mask' in input:
+            input['mask'] = transforms.functional.affine(
+                input['mask'], *self.params, self.interpolation, self.fill
+            )
+
+        if 'bbox' in input:
+            for i, (x1, y1, x2, y2) in enumerate(input['bbox']):
+                pt = matrix @ torch.FloatTensor([
+                    [x1, y1, 1],
+                    [x2, y1, 1],
+                    [x2, y2, 1],
+                    [x1, y2, 1]
+                ]).T
+                x_min, y_min, _ = pt.min(dim=1).values
+                x_max, y_max, _ = pt.max(dim=1).values
+                input['bbox'][i, :] = torch.FloatTensor([x_min, y_min, x_max, y_max])
+
+        # if 'tr_param' not in input:
+        #     input['tr_param'] = []
+        # input['tr_param'].append({'random_affine': matrix[:2, :].tolist()})
+
+        return input
+
+
+class ColorJitter(transforms.ColorJitter):
+    def __call__(self, input):
+        if not isinstance(input, dict):
+            return super().__call__(input)
+        assert 'image' in input
+        input['image'] = super().__call__(input['image'])
+        return input
+
+
+def get_transform(split, input_size=512):
+    mean = [0.485, 0.456, 0.406]
+    sdev = [0.229, 0.224, 0.225]
+
+    if split in ('train', 'trainval'):
+        transform = Compose([
+            # ColorJitter(brightness=0.5, saturation=0.5),  # before normalization
+            ToTensor(),
+            Normalize(mean, sdev),  # first normalize so that the mean is ~0
+            SquarePad(),  # zero pad (approx mean pixel value)
+            Resize(size=(input_size, input_size)),
+            # RandomHorizontalFlip(p=0.5),
+            RandomAffine(degrees=5, translate=(0.1, 0.1), scale=(0.9, 1.1)),
+            NormalizeBoxCoords(),
+        ])
+    elif split in ('val', 'test', 'testA', 'testB', 'testC'):
+        transform = Compose([
+            ToTensor(),
+            Normalize(mean, sdev),
+            SquarePad(),
+            Resize(size=(input_size, input_size)),
+            NormalizeBoxCoords(),
+        ])
+    elif split in ('visu',):
+        transform = Compose([
+            ToTensor(),
+            SquarePad(),
+            Resize(size=(input_size, input_size)),
+            NormalizeBoxCoords(),
+        ])
+    else:
+        raise ValueError(f'\'{split}\' is not a valid data split')
+
+    return transform
+
+
+def denormalize(img):
+    mean = [0.485, 0.456, 0.406]
+    sdev = [0.229, 0.224, 0.225]
+    return Normalize(
+        mean=[-m/s for m, s in zip(mean, sdev)], std=[1./s for s in sdev]
+    )(img)
+
+
+def undo_box_transforms(bbox, tr_param):
+    # undo validation mode transformations
+    bbox = bbox.clone()
+    for tr in tr_param[::-1]:
+        if 'resize' in tr:
+            sx, sy = tr['resize']
+            bbox[:, (0, 2)] /= sx
+            bbox[:, (1, 3)] /= sy
+        elif 'square_pad' in tr:
+            px, py, _, _ = tr['square_pad']
+            bbox[:, (0, 2)] -= px
+            bbox[:, (1, 3)] -= py
+        elif 'normalize_box_coords' in tr:
+            img_h, img_w = tr['normalize_box_coords']
+            bbox[:, (0, 2)] *= img_w
+            bbox[:, (1, 3)] *= img_h
+        else:
+            continue
+    return bbox
+
+
+def undo_box_transforms_batch(bbox, tr_param):
+    output = []
+    for i in range(bbox.size(0)):
+        bb = undo_box_transforms(torch.atleast_2d(bbox[i]), tr_param[i])
+        output.append(bb)
+    return torch.cat(output, dim=0)