Upload 84 files

2024.09.27/config.yaml

@@ -0,0 +1,211 @@
+##################
+# Trainer settings
+##################
+
+
+TASK: UniCLTask
+
+NAME: 'Example Eval Configuration'
+SAVE_TIMER_LOG: true
+
+# TUTORIAL STEP 1: CHOOSE SAVE DIR
+SAVE_DIR: ''
+LOG_EVERY: 10
+LOGLEVEL_OVERRIDE: INFO
+LOG_GPU_MEM: true
+RESUME: False
+RESET_DATA_LOADER: false
+
+FP16: true
+ZERO_STAGE: 0
+DEEPSPEED: false
+# ZERO_STAGE: 1
+AMP: PYTORCH
+# USE_APEX_DDP: false
+# USE_APEX_AMP: false
+# USE_HIT: false
+
+FIND_UNUSED_PARAMETERS: false
+
+SAVE_PER_OPTIM_STEPS: 500
+EVAL_PER_OPTIM_STEPS: 250
+EVAL_AT_START: False
+# SAVE_PER_UPDATE_NUM: -1
+# EVAL_PER_UPDATE_NUM: 0 # 0: do evaluation when saving checkpoint, -1: don't do evaluation
+
+NO_AUTO_LR_SCALING: true
+GRAD_CLIPPING: 1.0 #0.07
+
+SET_SAMPLER_EPOCH: true
+
+DONT_LOAD_MODEL: true
+
+user_dir: "./MainzVision" # lower case due to it is used in mainz as such
+
+##################
+# Task settings
+##################
+
+
+
+VERBOSE: true
+WORKERS: 6
+PIN_MEMORY: true
+IMAGE_ENCODER:
+  NAME: davit_v1
+  NUM_CLASSES: 0
+  #IMAGE_SIZE: [384, 384]
+  IMAGE_SIZE: [480, 480]
+  LOAD_PRETRAINED: true
+  PRETRAINED: ''
+  PRETRAINED_LAYERS: '*'
+  IMAGE_MEAN: [0.485, 0.456, 0.406]
+  IMAGE_STD: [0.229, 0.224, 0.225]
+  SPEC:
+    DROP_RATE: 0.1
+    DROP_PATH_RATE: 0.2
+    PATCH_SIZE: [7, 3, 3, 3]
+    PATCH_STRIDE: [4, 2, 2, 2]
+    PATCH_PADDING: [3, 1, 1, 1]
+    PATCH_PRENORM: [false, true, true, true]
+    DIM_EMBED: [256, 512, 1024, 2048]
+    NUM_HEADS: [8, 16, 32, 64]
+    NUM_GROUPS: [8, 16, 32, 64]
+    DEPTHS: [1, 1, 9, 1]
+    WINDOW_SIZE: 12
+    ENABLE_CHECKPOINT: true
+
+LANG_ENCODER:
+  NAME: transformer
+  LOAD_PRETRAINED: false
+  PRETRAINED: ''
+  PRETRAINED_LAYERS: '*'
+  TOKENIZER: clip
+  CONTEXT_LENGTH: 77
+  WIDTH: 1024
+  HEADS: 16
+  LAYERS: 16
+  AUTOGRESSIVE: false
+
+UNICL_MODEL:
+  DIM_PROJECTION: 1024
+  GATHER_TENSORS: true
+  LOAD_PRETRAINED: true
+
+  # TUTORIAL STEP 2: CHOOSE MODEL PATH
+  PRETRAINED: ''
+
+  PRETRAINED_LAYERS: '*'
+
+AUG:
+  MIXUP_PROB: 0.0
+  MIXUP: 0.8
+  MIXCUT: 1.0
+  MIXCUT_MINMAX: []
+  MIXUP_SWITCH_PROB: 0.5
+  MIXUP_MODE: 'batch'
+  SCALE: [0.8, 1.0]
+  RATIO: [0.75, 1.3333333]
+  INTERPOLATION: 'bicubic'
+  TORCHVISION_AUG:
+    AUTO_AUGMENT: ta_wide
+    RE_PROB: 0.25
+    HFLIP: 0.0
+    VFLIP: 0.0
+
+LOSS:
+  LOSS: UniCL
+DATASET:
+  DATASET: 'image_text_pairs_v2'
+  TEXT_FORMAT: 'json'
+  ROOT: ''
+  TRAIN_SET: 'mimic_cxr_v2-chestxray14-chexpertv4-irma2009_v2-rsnaboneage-mura-bingmedicalfewshot'
+  DATA_FORMAT: 'tsv'
+  SAMPLER: 'default'
+  LOADER: 'default'
+  TOKEN_FILE: ''
+  #PROMPT_ENGINEERING: False
+  #SAMPLER: 'chunk'
+  #LOADER: 'azcopy'
+  #TOKEN_FILE: 'cliptrainingpairs.txt'
+  #TEST_SET: 'MarsAtrain'
+
+
+# TUTORIAL STEP 3: CHOOSE ALL BELOW EVAL PATHS (THESE ARE ALL OPTIONAL EXTRA EVALS)
+# Note how one eval is ZIP format and the other is TSV format.
+
+
+
+
+EVALDATASET_LTCXR_S100_N100_TEXT_CLASSIFIER:
+  TEXT_FORMAT: json
+  FORMAT: 'zip'
+  SPLIT: 'NIH-CXR-LT'
+  ZIP_FILE: ''
+  ZIP_MAP_FILE: ''
+  LABEL_FILE: ''
+  IMAGE_TSV: ''
+  TEXT_TSV: ''
+  CWEIGHT_FILE: ''
+  ZS_MODE: 2
+  ZS_WEIGHT: 1.0
+  KNN: 100
+#  CLASSIFICATION_SETS: ['NIH-CXR-LT']
+#  NUM_CLASSES: [20]
+
+
+
+
+# TUTORIAL STEP 4: SET THE DEFAULT ZEROSHOT EVAL (THIS IS THE MANDATORY EVAL)
+
+ZEROSHOT_EVAL_DATASET:
+  FORMAT: 'zip'
+  SPLIT: 'NIH-CXR-LT'
+  ZIP_FILE: ''
+  ZIP_MAP_FILE: ''
+  LABEL_FILE: ''
+
+
+
+EVALUATION_SPLITS: ['cls-zeroshot-eval']
+TEST:
+  BATCH_SIZE_PER_GPU: 8
+  MODEL_FILE: ''
+  CENTER_CROP: false
+TRAIN:
+  BATCH_SIZE_TOTAL: 1024
+  BATCH_SIZE_PER_GPU: 16
+
+  SHUFFLE: true
+
+WEIGHT_SMOOTHING:
+  decay: 0.999
+  use_cpu: False
+  eval_smoothed_weight: True
+
+START_LEARNING_RATE: 0.00001
+# MAX_NUM_EPOCHS: 2
+MAX_NUM_EPOCHS: 100
+OPTIMIZER: AdamW # adam
+OPTIMIZER_PARAMS:
+  weight_decay: 0.2 #0.1
+CUSTOMIZED_PARAMS_CONF:
+  NO_WEIGHT_DECAY_MODULES: ['dw', 'norm']
+  WEIGHT_DECAY_PATTERNS:
+    "\\.bias$": 0.0
+    "logit_scale": 0.0
+    "positional_embedding": 0.0
+    "token_embedding": 0.0
+
+
+
+LR_SCHEDULER: TimmScheduler
+LR_SCHEDULER_PARAMS:
+  sched: cosine
+  warmup_steps: 5
+  warmup_lr: 0.000000001
+  min_lr: 0.000000001
+
+# GRADIENT_ACCUMULATE_STEP will be updated by:
+# BATCH_SIZE_TOTAL // (BATCH_SIZE_PER_GPU * world_size)
+GRADIENT_ACCUMULATE_STEP: -1

2024.09.27/language_model/clip_tokenizer_4.16.2/special_tokens_map.json

@@ -0,0 +1,27 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "single_word": false,
+    "lstrip": false,
+    "rstrip": false,
+    "normalized": true,
+    "special": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "single_word": false,
+    "lstrip": false,
+    "rstrip": false,
+    "normalized": true,
+    "special": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "single_word": false,
+    "lstrip": false,
+    "rstrip": false,
+    "normalized": true,
+    "special": false
+  },
+  "pad_token": "<|endoftext|>"
+}

2024.09.27/language_model/clip_tokenizer_4.16.2/tokenizer_config.json

@@ -0,0 +1,38 @@
+{
+  "errors": "replace",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "single_word": false,
+    "lstrip": false,
+    "rstrip": false,
+    "normalized": true,
+    "special": false,
+    "__type": "AddedToken"
+  },
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "single_word": false,
+    "lstrip": false,
+    "rstrip": false,
+    "normalized": true,
+    "special": false,
+    "__type": "AddedToken"
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "single_word": false,
+    "lstrip": false,
+    "rstrip": false,
+    "normalized": true,
+    "special": false,
+    "__type": "AddedToken"
+  },
+  "pad_token": "<|endoftext|>",
+  "add_prefix_space": false,
+  "do_lower_case": true,
+  "name_or_path": "openai/clip-vit-base-patch32",
+  "model_max_length": 77,
+  "special_tokens_map_file": "/home/ncodella/.cache/huggingface/transformers/18a566598f286c9139f88160c99f84eec492a26bd22738fa9cb44d5b7e0a5c76.cce1206abbad28826f000510f22f354e53e66a97f7c23745a7dfe27609cc07f5",
+  "tokenizer_file": "/home/ncodella/.cache/huggingface/transformers/7811def0c53be25ba790cb67ac785669b508a8d1cf8c912b8ac046c5f08aee68.20428ea8b6821af2719b760af844a371643ff49f255c73285f6ea448e15597fe",
+  "tokenizer_class": "CLIPTokenizer"
+}

2024.09.27/vision_model/medimageinsigt-v1.0.0.pt

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

MedImageInsight/Distributed/Utils.py

@@ -0,0 +1,344 @@
+import logging
+import os
+import pickle
+import requests
+import tenacity
+import time
+import shutil
+
+import torch
+import torch.distributed as dist
+
+from PIL import Image
+from torchvision.utils import make_grid
+
+
+from fvcore.nn import FlopCountAnalysis
+from fvcore.nn import flop_count_table
+from fvcore.nn import flop_count_str
+
+logger = logging.getLogger(__name__)
+
+NORM_MODULES = [
+    torch.nn.BatchNorm1d,
+    torch.nn.BatchNorm2d,
+    torch.nn.BatchNorm3d,
+    torch.nn.SyncBatchNorm,
+    # NaiveSyncBatchNorm inherits from BatchNorm2d
+    torch.nn.GroupNorm,
+    torch.nn.InstanceNorm1d,
+    torch.nn.InstanceNorm2d,
+    torch.nn.InstanceNorm3d,
+    torch.nn.LayerNorm,
+    torch.nn.LocalResponseNorm,
+]
+
+
+def register_norm_module(cls):
+    NORM_MODULES.append(cls)
+
+    return cls
+
+
+def is_main_process():
+    rank = 0
+    if 'OMPI_COMM_WORLD_SIZE' in os.environ:
+        rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
+
+    return rank == 0
+
+
+@torch.no_grad()
+def analysis_model(model, dump_input, verbose=False):
+    model.eval()
+    flops = FlopCountAnalysis(model, dump_input)
+    total = flops.total()
+    model.train()
+    params_total = sum(p.numel() for p in model.parameters())
+    params_learned = sum(
+        p.numel() for p in model.parameters() if p.requires_grad
+    )
+    logger.info(f"flop count table:\n {flop_count_table(flops)}")
+    if verbose:
+        logger.info(f"flop count str:\n {flop_count_str(flops)}")
+    logger.info(f"  Total flops: {total / 1000 / 1000:.3f}M,")
+    logger.info(f"  Total params: {params_total / 1000 / 1000:.3f}M,")
+    logger.info(f"  Learned params: {params_learned / 1000 / 1000:.3f}M")
+
+    return total, flop_count_table(flops), flop_count_str(flops)
+
+
+def gather_tensors(tensor):
+    """
+    Performs all_gather operation on the provided tensors.
+    *** Warning ***: torch.distributed.all_gather has no gradient.
+    """
+    tensors_gather = [
+        torch.ones_like(tensor)
+        for _ in range(int(os.environ['WORLD_SIZE']))
+    ]
+
+    dist.all_gather(tensors_gather, tensor, async_op=False)
+    # need to do this to restore propagation of the gradients
+    tensors_gather[int(os.environ['RANK'])] = tensor
+    output = torch.cat(tensors_gather, dim=0)
+    return output
+
+
+def is_valid_url(url):
+    try:
+        from urllib import parse
+        return parse.urlparse(str(url)).scheme != ''
+    except Exception:
+        return False
+
+
+@tenacity.retry(stop=tenacity.stop_after_attempt(3))
+def download_file(url, filepath):
+    logger.info(f'Downloading from {url} to {filepath.absolute()}.')
+    with requests.get(url, stream=True, allow_redirects=True, timeout=60) as r:
+        if r.status_code > 200:
+            raise RuntimeError(f'Failed in downloading from {url}, status code {r.status_code}.')
+
+        with open(filepath, 'wb') as f:
+            shutil.copyfileobj(r.raw, f, length=4194304)
+
+
+class DistributionGridFactory:
+    """
+    DistributionGrid Factory for helping create, cache and share the DistributionGrid based on the usage.
+    The DistributionGrid con be shared cross modules only the when this 3 conditions:
+        1. expert parallel group size
+        2. expert parallel replica group size,
+    are the same.
+    """
+    distribution_grid_cache = {}
+
+    @classmethod
+    def get_distribution_grid(cls,
+                              expert_parallel_group_size,
+                              expert_parallel_replica_group_size,
+                              ddp_type):
+        """
+        Get the DistributionGrid by the conditions.
+        Args:
+            expert_parallel_group_size: expert parallel group size
+            expert_parallel_replica_group_size: expert parallel replica group size
+            ddp_type: distributed data parallel type. "DDP" of the recipe, only allow ddp_type is "MAINZ", "OSS" or "ShardedDDP".
+
+        Returns: new created DistributionGrid or shared DistributionGrid.
+
+        Notes: Currently get_distribution_grid only support "DDP" is "MAINZ", "OSS" or "ShardedDDP".
+        """
+        # TODO:  Support cases that "DDP" is "FSDP".
+        # For "FSDP", we use the DG of self.opt['fsdp_expert_grid'] which is initialize in DistributedTrainer directly.
+        ddp_type = ddp_type.upper()
+        assert ddp_type in ["MAINZ", "OSS", "SHARDEDDDP"], f'DistributionGrid Factory only support "DDP" is "MAINZ",' \
+                                             f' "OSS" or "ShardedDDP".' \
+                                             f' But currently "DDP" is {ddp_type}'
+
+        cached_distributed_grid = cls.distribution_grid_cache.get(
+            (expert_parallel_group_size, expert_parallel_replica_group_size), None)
+
+        if cached_distributed_grid is not None:
+            return cached_distributed_grid
+        else:
+            from ort_moe.grids import DistributionGrid
+            distributed_grid = DistributionGrid(expert_parallel_group_size=expert_parallel_group_size,
+                                                expert_parallel_replica_group_size=expert_parallel_replica_group_size)
+
+            cls.distribution_grid_cache[expert_parallel_group_size,
+                                        expert_parallel_replica_group_size] = distributed_grid
+            return distributed_grid
+
+
+def get_world_size():
+    if not dist.is_available():
+        return 1
+    if not dist.is_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not dist.is_available():
+        return 0
+    if not dist.is_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def synchronize():
+    """
+    Helper function to synchronize (barrier) among all processes when
+    using distributed training
+    """
+    if not dist.is_available():
+        return
+    if not dist.is_initialized():
+        return
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    if world_size == 1:
+        return
+
+    def _send_and_wait(r):
+        if rank == r:
+            tensor = torch.tensor(0, device="cuda")
+        else:
+            tensor = torch.tensor(1, device="cuda")
+        dist.broadcast(tensor, r)
+        while tensor.item() == 1:
+            time.sleep(1)
+
+    _send_and_wait(0)
+    # now sync on the main process
+    _send_and_wait(1)
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    # serialized to a Tensor
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to("cuda")
+
+    # obtain Tensor size of each rank
+    local_size = torch.LongTensor([tensor.numel()]).to("cuda")
+    size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.ByteTensor(size=(max_size,)).to("cuda"))
+    if local_size != max_size:
+        padding = torch.ByteTensor(size=(max_size - local_size,)).to("cuda")
+        tensor = torch.cat((tensor, padding), dim=0)
+    dist.all_gather(tensor_list, tensor)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list
+
+
+def all_gather_cpu(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors).
+    Args:
+        data: any picklable object
+        group: a torch process group. By default, will use a group which
+            contains all ranks on gloo backend.
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+
+    def _get_global_gloo_group():
+        """
+        Return a process group based on gloo backend, containing all the ranks
+        The result is cached.
+        """
+        if dist.get_backend() == "nccl":
+            return dist.new_group(backend="gloo")
+        else:
+            return dist.group.WORLD
+
+    if get_world_size() == 1:
+        return [data]
+    group = _get_global_gloo_group()  # use CPU group by default, to reduce GPU RAM usage.
+    world_size = dist.get_world_size(group)
+    if world_size == 1:
+        return [data]
+
+    output = [None for _ in range(world_size)]
+    dist.all_gather_object(output, data, group=group)
+    return output
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Args:
+        input_dict (dict): all the values will be reduced
+        average (bool): whether to do average or sum
+    Reduce the values in the dictionary from all processes so that process with rank
+    0 has the averaged results. Returns a dict with the same fields as
+    input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.no_grad():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.reduce(values, dst=0)
+        if dist.get_rank() == 0 and average:
+            # only main process gets accumulated, so only divide by
+            # world_size in this case
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict
+
+
+def broadcast_data(data):
+    if not torch.distributed.is_initialized():
+        return data
+    rank = dist.get_rank()
+    if rank == 0:
+        data_tensor = torch.tensor(data + [0], device="cuda")
+    else:
+        data_tensor = torch.tensor(data + [1], device="cuda")
+    torch.distributed.broadcast(data_tensor, 0)
+    while data_tensor.cpu().numpy()[-1] == 1:
+        time.sleep(1)
+
+    return data_tensor.cpu().numpy().tolist()[:-1]
+
+
+def reduce_sum(tensor):
+    if get_world_size() <= 1:
+        return tensor
+
+    tensor = tensor.clone()
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+    return tensor
+
+
+def save_result(result, filename):
+    output_folder = os.path.dirname(filename)
+    basename = os.path.splitext(os.path.basename(filename))[0]
+    os.makedirs(output_folder, exist_ok=True)
+
+    if isinstance(result, torch.Tensor) and result.ndim in [3,4]:
+        if result.ndim==3 and result.size(0) not in [1,3]:
+            result = make_grid(result.unsqueeze(1))
+        elif result.ndim==4:
+            result = make_grid(result)
+        else:
+            result = make_grid([result])
+
+        im = Image.fromarray(result.clamp_(0, 255).permute(1, 2, 0).to(torch.uint8).numpy())
+        im.save(os.path.join(output_folder, '{}.png'.format(basename)))
+    else:
+        torch.save(result, os.path.join(output_folder, '{}.pth'.format(basename)))

MedImageInsight/Distributed/__init__.py

@@ -0,0 +1,6 @@
+from .Utils import analysis_model
+from .Utils import is_main_process
+from .Utils import gather_tensors
+from .Utils import register_norm_module
+from .Utils import NORM_MODULES
+from .Utils import DistributionGridFactory

MedImageInsight/ImageDataLoader/__init__.py

@@ -0,0 +1,8 @@
+from .build import build_dataloader
+#from .build import build_multitask_dataloader
+from .transforms import build_transforms
+#from .imagenet.real_labels import RealLabelsImagenet
+from .constants import IMAGENET_CLASSES
+from .constants import IMAGENET_DEFAULT_TEMPLATES
+from .zipdata import ZipData
+#from .vision_dataset import VDImageTextDataset, MultiClassTorchDatasetWrapper

MedImageInsight/ImageDataLoader/blob_storage.py

@@ -0,0 +1,244 @@
+import os
+import time
+import shutil
+import logging
+import subprocess
+import os.path as op
+from typing import List
+from collections import OrderedDict
+
+import torch.distributed as distributed
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_AZCOPY_PATH = 'azcopy/azcopy'
+
+
+def disk_usage(path: str) -> float:
+    stat = shutil.disk_usage(path)
+    return stat.used / stat.total
+
+
+def is_download_successful(stdout: str) -> bool:
+    for line in stdout.split('\n'):
+        if line == "Number of Transfers Failed: 0":
+            return True
+    logger.info("Azcopy message:\n %s" % stdout)
+    return False
+
+
+def ensure_directory(path):
+    """Check existence of the given directory path. If not, create a new directory.
+
+    Args:
+        path (str): path of a given directory.
+    """
+    if path == '' or path == '.':
+        return
+    if path is not None and len(path) > 0:
+        assert not op.isfile(path), '{} is a file'.format(path)
+        if not op.exists(path) and not op.islink(path):
+            os.makedirs(path, exist_ok=True)
+        # we should always check if it succeeds.
+        assert op.isdir(op.abspath(path)), path
+
+
+class LRU(OrderedDict):
+    def __init__(self, maxsize=3):
+        self.maxsize = maxsize
+
+    def __getitem__(self, key):
+        value = super().__getitem__(key)
+        self.move_to_end(key)
+        return value
+
+    def __setitem__(self, key, value):
+        if key in self:
+            if self[key] is not None:
+                self[key].close()
+                self.move_to_end(key)
+
+        logger.debug('=> Cache {}'.format(key))
+        super().__setitem__(key, value)
+
+        if len(self) > self.maxsize:
+            oldest = next(iter(self))
+            if self[oldest] is not None:
+                self[oldest].close()
+            logger.debug('=> Purged {}'.format(oldest))
+            del self[oldest]
+
+
+class BlobStorage(OrderedDict):
+    """ Pseudo Blob Storage manager
+
+    The registered blobs are maintained in a LRU cache.
+    Limit size, evicting the least recently looked-up key when full.
+    https://docs.python.org/3/library/collections.html#collections.OrderedDict
+
+    Input argument:
+        sas_token (str): path to SAS token.
+    """
+    def __init__(self,
+                 is_train: bool,
+                 sas_token_path: str = None,
+                 azcopy_path: str = None,
+                 *args, **kwds):
+        super().__init__(*args, **kwds)
+        self.maxsize = 2 if is_train else 10    # Set maxsize to large number such val data never get purged.
+        self.is_train = is_train
+
+        if sas_token_path:
+            self.sas_token = BlobStorage.read_sas_token(sas_token_path)
+            self.base_url = self.sas_token[:self.sas_token.index("?")]
+            self.query_string = self.sas_token[self.sas_token.index("?"):]
+            self.container = BlobStorage.extract_container(self.sas_token)
+        else:
+            self.sas_token = None
+            self.base_url = None
+            self.query_string = None
+            self.container = None
+
+        logger.debug(
+            f"=> [BlobStorage] Base url: {self.base_url}"
+            f"=> [BlobStorage] Query string: {self.query_string}"
+            f"=> [BlobStorage] Container name: {self.container}"
+        )
+
+        self.azcopy_path = azcopy_path if azcopy_path else DEFAULT_AZCOPY_PATH
+        self._cached_files = LRU(3)
+
+    def __getitem__(self, key):
+        value = super().__getitem__(key)
+        self.move_to_end(key)
+        return value
+
+    def __setitem__(self, key, value):
+        if key in self:
+            self.move_to_end(key)
+        super().__setitem__(key, value)
+        # NOTE: purge the least recently used data if the disk usage is high.
+        # ITP restarts GPU clusters when disk usage reaches 80%.
+        if len(self) > self.maxsize:
+            oldest = next(iter(self))
+            del self[oldest]
+
+    @staticmethod
+    def read_sas_token(path: str) -> str:
+        with open(path, 'r') as f:
+            token = f.readline().strip()
+        return token
+
+    @staticmethod
+    def extract_container(token: str) -> str:
+        """
+        Input argument:
+            token (str): the full URI of Shared Access Signature (SAS) in the following format.
+            https://[storage_account].blob.core.windows.net/[container_name][SAS_token]
+        """
+        return os.path.basename(token.split('?')[0])
+
+    def _convert_to_blob_url(self, local_path: str):
+        return self.base_url + local_path.split("azcopy")[1] + self.query_string
+
+    def _convert_to_blob_folder_url(self, local_path: str):
+        return self.base_url + local_path.split("azcopy")[1] + "/*" + self.query_string
+
+    def fetch_blob(self, local_path: str) -> None:
+        if op.exists(local_path):
+            logger.info('=> Try to open {}'.format(local_path))
+            fp = open(local_path, 'r')
+            self._cached_files[local_path] = fp
+            logger.debug("=> %s downloaded. Skip." % local_path)
+            return
+        blob_url = self._convert_to_blob_url(local_path)
+        rank = '0' if 'RANK' not in os.environ else os.environ['RANK']
+        cmd = [self.azcopy_path, "copy", blob_url, local_path + rank]
+        curr_usage = disk_usage('/')
+        logger.info(
+            "=> Downloading %s with azcopy ... (disk usage: %.2f%%)"
+            % (local_path, curr_usage * 100)
+        )
+        proc = subprocess.run(cmd, stdout=subprocess.PIPE)
+        while not is_download_successful(proc.stdout.decode()):
+            logger.info("=> Azcopy failed to download {}. Retrying ...".format(blob_url))
+            proc = subprocess.run(cmd, stdout=subprocess.PIPE)
+        if not op.exists(local_path):
+            os.rename(local_path + rank, local_path)
+        else:
+            os.remove(local_path + rank)
+        logger.info(
+            "=> Downloaded %s with azcopy ... (disk usage: %.2f%% => %.2f%%)" %
+            (local_path, curr_usage * 100, disk_usage('/') * 100)
+        )
+
+    def fetch_blob_folder(self, local_path: str, azcopy_args: list=[]) -> None:
+        blob_url = self._convert_to_blob_folder_url(local_path)
+        cmd = [self.azcopy_path, "copy", blob_url, local_path] + azcopy_args
+        curr_usage = disk_usage('/')
+        logger.info(
+            "=> Downloading %s with azcopy args %s ... (disk usage: %.2f%%)"
+            % (local_path, ' '.join(azcopy_args), curr_usage * 100)
+        )
+        proc = subprocess.run(cmd, stdout=subprocess.PIPE)
+        while not is_download_successful(proc.stdout.decode()):
+            logger.info("=> Azcopy failed to download {} with args {}. Retrying ...".format(blob_url, ' '.join(azcopy_args)))
+            proc = subprocess.run(cmd, stdout=subprocess.PIPE)
+        logger.info(
+            "=> Downloaded %s with azcopy args %s ... (disk usage: %.2f%% => %.2f%%)" %
+            (local_path, ' '.join(azcopy_args), curr_usage * 100, disk_usage('/') * 100)
+        )
+
+    def register_local_tsv_paths(self, local_paths: List[str]) -> List[str]:
+        if self.sas_token:
+            tsv_paths_new = []
+            lineidx_paths = set()
+            linelist_paths = set()
+            for path in local_paths:
+                tsv_path_az = path.replace(self.container, 'azcopy')
+                tsv_paths_new.append(tsv_path_az)
+                logger.debug("=> Registering {}".format(tsv_path_az))
+
+                if not self.is_train:
+                    logger.info('=> Downloading {}...'.format(tsv_path_az))
+                    self.fetch_blob(tsv_path_az)
+                    logger.info('=> Downloaded {}'.format(tsv_path_az))
+
+                lineidx = op.splitext(path)[0] + '.lineidx'
+                lineidx_ = lineidx.replace(self.container, 'azcopy')
+                if self.is_train:
+                    if not op.isfile(lineidx_) and op.dirname(lineidx_) not in lineidx_paths:
+                        lineidx_paths.add(op.dirname(lineidx_))
+                else:
+                    if not op.isfile(lineidx_):
+                        ensure_directory(op.dirname(lineidx_))
+                        self.fetch_blob(lineidx_)
+
+                linelist = op.splitext(path)[0] + '.linelist'
+                linelist_ = linelist.replace(self.container, 'azcopy')
+                # .linelist does not always exist. Check existence before fetch
+                if self.is_train:
+                    if op.isfile(linelist) and not op.isfile(linelist_) and op.dirname(linelist_) not in linelist_paths:
+                        linelist_paths.add(op.dirname(linelist_))
+                else:
+                    if op.isfile(linelist) and not op.isfile(linelist_):
+                        ensure_directory(op.dirname(linelist_))
+                        self.fetch_blob(linelist_)
+
+            if self.is_train:
+                for path in lineidx_paths:
+                    self.fetch_blob_folder(path, azcopy_args=['--include-pattern', '*.lineidx'])
+
+                for path in linelist_paths:
+                    self.fetch_blob_folder(path, azcopy_args=['--include-pattern', '*.linelist'])
+
+            return tsv_paths_new
+        else:
+            return local_paths
+
+    def open(self, local_path: str):
+        if self.sas_token and 'azcopy' in local_path:
+            while not op.exists(local_path):
+                time.sleep(1)
+        fid = open(local_path, 'r')
+        return fid

MedImageInsight/ImageDataLoader/build.py

@@ -0,0 +1,260 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+import os
+import json
+import pathlib
+from os.path import basename
+
+from timm.data import create_loader
+import torch
+import torch.utils.data
+import torch.distributed as dist
+import torchvision.datasets as datasets
+from torchvision.io import read_image
+import torch.distributed as dist
+from pathlib import Path
+from yacs.config import CfgNode as CN
+
+from ..LangEncoder import build_tokenizer
+
+from .tsv import TSVImageTextDatasetV2
+from .tsv import TSVMeta
+from .transforms import build_transforms
+
+logger = logging.getLogger(__name__)
+
+
+def build_dataset(cfg, is_train):
+    if cfg['DATASET']['DATASET'] == 'image_text_pairs_v2':
+        dataset = _build_pairs_dataset_v2(cfg, is_train)
+    else:
+        raise ValueError(f'Unknown dataset: {cfg["DATASET"]["DATASET"]}')
+    return dataset
+
+
+def _get_tsv_list(cfg, is_train):
+    tmp_list = []
+    if is_train and 'TRAIN_TSV_LIST' in cfg['DATASET']:
+        tmp_list = cfg['DATASET']['TRAIN_TSV_LIST']
+    elif 'TEST_TSV_LIST' in cfg['DATASET']:
+        tmp_list = cfg['DATASET']['TEST_TSV_LIST']
+
+    tsv_list = []
+    for l in tmp_list:
+        if l.endswith('.list'):
+            with open(l, 'r') as f:
+                tsv_list.extend([i.strip() for i in f])
+        else:
+            tsv_list.append(l)
+
+    logger.info(f'tsv list: {tsv_list}')
+
+    return tsv_list
+
+
+def _get_token_file(cfg):
+    num_nodes = dist.get_world_size() // torch.cuda.device_count()
+    if isinstance(cfg['DATASET']['TOKEN_FILE'], list):
+        if num_nodes == 1:
+            logger.warning('=> Multi token files are provided, but only one node is used for training')
+            sas_token_file = cfg['DATASET']['TOKEN_FILE'][0]
+        else:
+            rank = dist.get_rank()
+            node_idx = rank // torch.cuda.device_count()
+            num_token_files = len(cfg['DATASET']['TOKEN_FILE'])
+            sas_token_file = cfg['DATASET']['TOKEN_FILE'][node_idx % num_token_files]
+    else:
+        sas_token_file = cfg['DATASET']['TOKEN_FILE']
+
+    sas_token_file = os.path.join(cfg['DATASET']['ROOT'], sas_token_file)
+
+    if (
+            cfg['DATASET']['LOADER'] == 'blobfuse'
+            or not os.path.isfile(sas_token_file)
+    ):
+        sas_token_file = None
+
+    return sas_token_file
+
+
+def _build_pairs_dataset_v2(cfg, is_train):
+    transforms = build_transforms(cfg, is_train)
+    logger.info('transforms: {}'.format(transforms))
+
+    dataset_name = cfg['DATASET']['TRAIN_SET'] \
+        if is_train else cfg['DATASET']['TEST_SET']
+
+    tokenobj = build_tokenizer(cfg['LANG_ENCODER'])
+
+    if cfg['DATASET']['DATA_FORMAT'] != 'tsv':
+        raise ValueError('Only support tsv format for pairs dataset v2')
+
+    tsv_list = _get_tsv_list(cfg, is_train)
+
+    if len(tsv_list) > 0:
+        tsv_filenames = sorted(
+            [
+                os.path.join(cfg['DATASET']['ROOT'], dataset_name, f)
+                for f in tsv_list
+            ]
+        )
+    else:
+        dataset_path = os.path.join(cfg['DATASET']['ROOT'], dataset_name)
+        tsv_files = Path(dataset_path).glob('**/*.tsv')
+
+        tsv_filenames = sorted(
+            [
+                str(path)
+                for path in tsv_files
+            ]
+        )
+
+    image_tsv_files = [
+        filename
+        for filename in tsv_filenames
+        if (
+                'image-' in basename(filename)
+                or 'image_' in basename(filename)
+                or '_image' in basename(filename)
+                or '-image' in basename(filename)
+                or 'images-' in basename(filename)
+        )
+    ]
+    text_tsv_files = [
+        filename
+        for filename in tsv_filenames
+        if (
+                'text-' in basename(filename)
+                or 'text_' in basename(filename)
+                or '_text' in basename(filename)
+                or '-text' in basename(filename)
+                or 'texts-' in basename(filename)
+        )
+    ]
+
+    logger.info(
+        "=> found %d/%d tsv file(s) to load.",
+        len(image_tsv_files), len(text_tsv_files)
+    )
+
+    num_captions = 1 \
+        if is_train else cfg['DATASET'].get('NUM_CAPTIONS', 1)
+    text_format = cfg['DATASET'].get('TEXT_FORMAT', 'json')
+
+    sas_token_file = _get_token_file(cfg)
+    logger.info("=> SAS token path: %s", sas_token_file)
+
+    metas = []
+    cfg_data = cfg['DATASET']
+    if 'CLASSIFICATION_SETS' in cfg_data and 'NUM_CLASSES' in cfg_data:
+        for source, num_classes in zip(cfg_data['CLASSIFICATION_SETS'], cfg_data['NUM_CLASSES']):
+            metas.append(
+                TSVMeta(
+                    source=source,
+                    num_classes=num_classes,
+                    task='classification'
+                )
+            )
+            logger.info('=> add meta: {}'.format(metas[-1]))
+
+    if 'coco-caption' in dataset_name:
+        logger.info('=> coco caption data is used')
+        logger.info('=> update num_captions: 5, text_format: json')
+        logger.warning('=> set sas token to None for coco evaluation')
+        sas_token_file = None
+        num_captions = 5
+        text_format = 'json'
+
+    dataset = TSVImageTextDatasetV2(
+        image_tsv_files, text_tsv_files,
+        transform=transforms,
+        tokenize=tokenobj,
+        context_length=cfg['LANG_ENCODER']['CONTEXT_LENGTH'],
+        num_captions=num_captions,
+        text_format=text_format,
+        is_train=is_train,
+        sas_token_path=sas_token_file,
+        metas=metas,
+        prompt_engineering=cfg['DATASET'].get('PROMPT_ENGINEERING', True),
+        concat_queries=cfg['DATASET'].get('CONCAT_QUERIES', False)
+    )
+
+    logger.info(
+        "=> %s set size: %d", 'train'
+        if is_train else 'val', len(dataset)
+    )
+
+    return dataset
+
+
+def build_dataloader(cfg, is_train=True, distributed=False):
+    dataset = build_dataset(cfg, is_train)
+
+    if (
+            is_train
+            and 'TIMM_AUG' in cfg['AUG']
+            and cfg['AUG']['TIMM_AUG']['USE_LOADER']
+    ):
+        logger.info('=> use timm loader for training')
+        timm_cfg = CN(init_dict=cfg['AUG']['TIMM_AUG'])
+        data_loader = create_loader(
+            dataset,
+            input_size=cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0],
+            batch_size=cfg['TRAIN']['BATCH_SIZE_PER_GPU'],
+            is_training=True,
+            use_prefetcher=True,
+            no_aug=False,
+            re_prob=timm_cfg.RE_PROB,
+            re_mode=timm_cfg.RE_MODE,
+            re_count=timm_cfg.RE_COUNT,
+            re_split=timm_cfg.RE_SPLIT,
+            scale=cfg['AUG']['SCALE'],
+            ratio=cfg['AUG']['RATIO'],
+            hflip=timm_cfg.HFLIP,
+            vflip=timm_cfg.VFLIP,
+            color_jitter=timm_cfg.COLOR_JITTER,
+            auto_augment=timm_cfg.AUTO_AUGMENT,
+            num_aug_splits=0,
+            interpolation=cfg['AUG']['INTERPOLATION'],
+            mean=cfg['IMAGE_ENCODER']['IMAGE_MEAN'],
+            std=cfg['IMAGE_ENCODER']['IMAGE_STD'],
+            num_workers=cfg['WORKERS'],
+            distributed=distributed,
+            collate_fn=None,
+            pin_memory=cfg['PIN_MEMORY'],
+            use_multi_epochs_loader=True
+        )
+    else:
+        if is_train:
+            batch_size_per_gpu = cfg['TRAIN']['BATCH_SIZE_PER_GPU']
+            shuffle = cfg['TRAIN'].get('SHUFFLE', True)
+        else:
+            batch_size_per_gpu = cfg['TEST']['BATCH_SIZE_PER_GPU']
+            shuffle = cfg['TEST'].get('SHUFFLE', False)
+
+        if distributed or cfg.get('ALWAYS_ENABLE_SAMPLER', False):
+            # sampler = build_sampler(cfg, dataset, is_train, shuffle)
+            sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle)
+            shuffle = False
+        else:
+            sampler = None
+
+        data_loader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=batch_size_per_gpu,
+            shuffle=shuffle,
+            num_workers=cfg['WORKERS'],
+            pin_memory=cfg['PIN_MEMORY'],
+            sampler=sampler,
+            drop_last=True if is_train else False,
+            prefetch_factor=cfg.get('PREFETCH_FACTOR', 2)
+        )
+
+    return data_loader
+
+
+
+

MedImageInsight/ImageDataLoader/constants.py

@@ -0,0 +1,85 @@
+IMAGENET_CLASSES = ["tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark", "electric ray", "stingray", "rooster", "hen", "ostrich", "brambling", "goldfinch", "house finch", "junco", "indigo bunting", "American robin", "bulbul", "jay", "magpie", "chickadee", "American dipper", "kite (bird of prey)", "bald eagle", "vulture", "great grey owl", "fire salamander", "smooth newt", "newt", "spotted salamander", "axolotl", "American bullfrog", "tree frog", "tailed frog", "loggerhead sea turtle", "leatherback sea turtle", "mud turtle", "terrapin", "box turtle", "banded gecko", "green iguana", "Carolina anole", "desert grassland whiptail lizard", "agama", "frilled-necked lizard", "alligator lizard", "Gila monster", "European green lizard", "chameleon", "Komodo dragon", "Nile crocodile", "American alligator", "triceratops", "worm snake", "ring-necked snake", "eastern hog-nosed snake", "smooth green snake", "kingsnake", "garter snake", "water snake", "vine snake", "night snake", "boa constrictor", "African rock python", "Indian cobra", "green mamba", "sea snake", "Saharan horned viper", "eastern diamondback rattlesnake", "sidewinder rattlesnake", "trilobite", "harvestman", "scorpion", "yellow garden spider", "barn spider", "European garden spider", "southern black widow", "tarantula", "wolf spider", "tick", "centipede", "black grouse", "ptarmigan", "ruffed grouse", "prairie grouse", "peafowl", "quail", "partridge", "african grey parrot", "macaw", "sulphur-crested cockatoo", "lorikeet", "coucal", "bee eater", "hornbill", "hummingbird", "jacamar", "toucan", "duck", "red-breasted merganser", "goose", "black swan", "tusker", "echidna", "platypus", "wallaby", "koala", "wombat", "jellyfish", "sea anemone", "brain coral", "flatworm", "nematode", "conch", "snail", "slug", "sea slug", "chiton", "chambered nautilus", "Dungeness crab", "rock crab", "fiddler crab", "red king crab", "American lobster", "spiny lobster", "crayfish", "hermit crab", "isopod", "white stork", "black stork", "spoonbill", "flamingo", "little blue heron", "great egret", "bittern bird", "crane bird", "limpkin", "common gallinule", "American coot", "bustard", "ruddy turnstone", "dunlin", "common redshank", "dowitcher", "oystercatcher", "pelican", "king penguin", "albatross", "grey whale", "killer whale", "dugong", "sea lion", "Chihuahua", "Japanese Chin", "Maltese", "Pekingese", "Shih Tzu", "King Charles Spaniel", "Papillon", "toy terrier", "Rhodesian Ridgeback", "Afghan Hound", "Basset Hound", "Beagle", "Bloodhound", "Bluetick Coonhound", "Black and Tan Coonhound", "Treeing Walker Coonhound", "English foxhound", "Redbone Coonhound", "borzoi", "Irish Wolfhound", "Italian Greyhound", "Whippet", "Ibizan Hound", "Norwegian Elkhound", "Otterhound", "Saluki", "Scottish Deerhound", "Weimaraner", "Staffordshire Bull Terrier", "American Staffordshire Terrier", "Bedlington Terrier", "Border Terrier", "Kerry Blue Terrier", "Irish Terrier", "Norfolk Terrier", "Norwich Terrier", "Yorkshire Terrier", "Wire Fox Terrier", "Lakeland Terrier", "Sealyham Terrier", "Airedale Terrier", "Cairn Terrier", "Australian Terrier", "Dandie Dinmont Terrier", "Boston Terrier", "Miniature Schnauzer", "Giant Schnauzer", "Standard Schnauzer", "Scottish Terrier", "Tibetan Terrier", "Australian Silky Terrier", "Soft-coated Wheaten Terrier", "West Highland White Terrier", "Lhasa Apso", "Flat-Coated Retriever", "Curly-coated Retriever", "Golden Retriever", "Labrador Retriever", "Chesapeake Bay Retriever", "German Shorthaired Pointer", "Vizsla", "English Setter", "Irish Setter", "Gordon Setter", "Brittany dog", "Clumber Spaniel", "English Springer Spaniel", "Welsh Springer Spaniel", "Cocker Spaniel", "Sussex Spaniel", "Irish Water Spaniel", "Kuvasz", "Schipperke", "Groenendael dog", "Malinois", "Briard", "Australian Kelpie", "Komondor", "Old English Sheepdog", "Shetland Sheepdog", "collie", "Border Collie", "Bouvier des Flandres dog", "Rottweiler", "German Shepherd Dog", "Dobermann", "Miniature Pinscher", "Greater Swiss Mountain Dog", "Bernese Mountain Dog", "Appenzeller Sennenhund", "Entlebucher Sennenhund", "Boxer", "Bullmastiff", "Tibetan Mastiff", "French Bulldog", "Great Dane", "St. Bernard", "husky", "Alaskan Malamute", "Siberian Husky", "Dalmatian", "Affenpinscher", "Basenji", "pug", "Leonberger", "Newfoundland dog", "Great Pyrenees dog", "Samoyed", "Pomeranian", "Chow Chow", "Keeshond", "brussels griffon", "Pembroke Welsh Corgi", "Cardigan Welsh Corgi", "Toy Poodle", "Miniature Poodle", "Standard Poodle", "Mexican hairless dog (xoloitzcuintli)", "grey wolf", "Alaskan tundra wolf", "red wolf or maned wolf", "coyote", "dingo", "dhole", "African wild dog", "hyena", "red fox", "kit fox", "Arctic fox", "grey fox", "tabby cat", "tiger cat", "Persian cat", "Siamese cat", "Egyptian Mau", "cougar", "lynx", "leopard", "snow leopard", "jaguar", "lion", "tiger", "cheetah", "brown bear", "American black bear", "polar bear", "sloth bear", "mongoose", "meerkat", "tiger beetle", "ladybug", "ground beetle", "longhorn beetle", "leaf beetle", "dung beetle", "rhinoceros beetle", "weevil", "fly", "bee", "ant", "grasshopper", "cricket insect", "stick insect", "cockroach", "praying mantis", "cicada", "leafhopper", "lacewing", "dragonfly", "damselfly", "red admiral butterfly", "ringlet butterfly", "monarch butterfly", "small white butterfly", "sulphur butterfly", "gossamer-winged butterfly", "starfish", "sea urchin", "sea cucumber", "cottontail rabbit", "hare", "Angora rabbit", "hamster", "porcupine", "fox squirrel", "marmot", "beaver", "guinea pig", "common sorrel horse", "zebra", "pig", "wild boar", "warthog", "hippopotamus", "ox", "water buffalo", "bison", "ram (adult male sheep)", "bighorn sheep", "Alpine ibex", "hartebeest", "impala (antelope)", "gazelle", "arabian camel", "llama", "weasel", "mink", "European polecat", "black-footed ferret", "otter", "skunk", "badger", "armadillo", "three-toed sloth", "orangutan", "gorilla", "chimpanzee", "gibbon", "siamang", "guenon", "patas monkey", "baboon", "macaque", "langur", "black-and-white colobus", "proboscis monkey", "marmoset", "white-headed capuchin", "howler monkey", "titi monkey", "Geoffroy's spider monkey", "common squirrel monkey", "ring-tailed lemur", "indri", "Asian elephant", "African bush elephant", "red panda", "giant panda", "snoek fish", "eel", "silver salmon", "rock beauty fish", "clownfish", "sturgeon", "gar fish", "lionfish", "pufferfish", "abacus", "abaya", "academic gown", "accordion", "acoustic guitar", "aircraft carrier", "airliner", "airship", "altar", "ambulance", "amphibious vehicle", "analog clock", "apiary", "apron", "trash can", "assault rifle", "backpack", "bakery", "balance beam", "balloon", "ballpoint pen", "Band-Aid", "banjo", "baluster / handrail", "barbell", "barber chair", "barbershop", "barn", "barometer", "barrel", "wheelbarrow", "baseball", "basketball", "bassinet", "bassoon", "swimming cap", "bath towel", "bathtub", "station wagon", "lighthouse", "beaker", "military hat (bearskin or shako)", "beer bottle", "beer glass", "bell tower", "baby bib", "tandem bicycle", "bikini", "ring binder", "binoculars", "birdhouse", "boathouse", "bobsleigh", "bolo tie", "poke bonnet", "bookcase", "bookstore", "bottle cap", "hunting bow", "bow tie", "brass memorial plaque", "bra", "breakwater", "breastplate", "broom", "bucket", "buckle", "bulletproof vest", "high-speed train", "butcher shop", "taxicab", "cauldron", "candle", "cannon", "canoe", "can opener", "cardigan", "car mirror", "carousel", "tool kit", "cardboard box / carton", "car wheel", "automated teller machine", "cassette", "cassette player", "castle", "catamaran", "CD player", "cello", "mobile phone", "chain", "chain-link fence", "chain mail", "chainsaw", "storage chest", "chiffonier", "bell or wind chime", "china cabinet", "Christmas stocking", "church", "movie theater", "cleaver", "cliff dwelling", "cloak", "clogs", "cocktail shaker", "coffee mug", "coffeemaker", "spiral or coil", "combination lock", "computer keyboard", "candy store", "container ship", "convertible", "corkscrew", "cornet", "cowboy boot", "cowboy hat", "cradle", "construction crane", "crash helmet", "crate", "infant bed", "Crock Pot", "croquet ball", "crutch", "cuirass", "dam", "desk", "desktop computer", "rotary dial telephone", "diaper", "digital clock", "digital watch", "dining table", "dishcloth", "dishwasher", "disc brake", "dock", "dog sled", "dome", "doormat", "drilling rig", "drum", "drumstick", "dumbbell", "Dutch oven", "electric fan", "electric guitar", "electric locomotive", "entertainment center", "envelope", "espresso machine", "face powder", "feather boa", "filing cabinet", "fireboat", "fire truck", "fire screen", "flagpole", "flute", "folding chair", "football helmet", "forklift", "fountain", "fountain pen", "four-poster bed", "freight car", "French horn", "frying pan", "fur coat", "garbage truck", "gas mask or respirator", "gas pump", "goblet", "go-kart", "golf ball", "golf cart", "gondola", "gong", "gown", "grand piano", "greenhouse", "radiator grille", "grocery store", "guillotine", "hair clip", "hair spray", "half-track", "hammer", "hamper", "hair dryer", "hand-held computer", "handkerchief", "hard disk drive", "harmonica", "harp", "combine harvester", "hatchet", "holster", "home theater", "honeycomb", "hook", "hoop skirt", "gymnastic horizontal bar", "horse-drawn vehicle", "hourglass", "iPod", "clothes iron", "carved pumpkin", "jeans", "jeep", "T-shirt", "jigsaw puzzle", "rickshaw", "joystick", "kimono", "knee pad", "knot", "lab coat", "ladle", "lampshade", "laptop computer", "lawn mower", "lens cap", "letter opener", "library", "lifeboat", "lighter", "limousine", "ocean liner", "lipstick", "slip-on shoe", "lotion", "music speaker", "loupe magnifying glass", "sawmill", "magnetic compass", "messenger bag", "mailbox", "tights", "one-piece bathing suit", "manhole cover", "maraca", "marimba", "mask", "matchstick", "maypole", "maze", "measuring cup", "medicine cabinet", "megalith", "microphone", "microwave oven", "military uniform", "milk can", "minibus", "miniskirt", "minivan", "missile", "mitten", "mixing bowl", "mobile home", "ford model t", "modem", "monastery", "monitor", "moped", "mortar and pestle", "graduation cap", "mosque", "mosquito net", "vespa", "mountain bike", "tent", "computer mouse", "mousetrap", "moving van", "muzzle", "metal nail", "neck brace", "necklace", "baby pacifier", "notebook computer", "obelisk", "oboe", "ocarina", "odometer", "oil filter", "pipe organ", "oscilloscope", "overskirt", "bullock cart", "oxygen mask", "product packet / packaging", "paddle", "paddle wheel", "padlock", "paintbrush", "pajamas", "palace", "pan flute", "paper towel", "parachute", "parallel bars", "park bench", "parking meter", "railroad car", "patio", "payphone", "pedestal", "pencil case", "pencil sharpener", "perfume", "Petri dish", "photocopier", "plectrum", "Pickelhaube", "picket fence", "pickup truck", "pier", "piggy bank", "pill bottle", "pillow", "ping-pong ball", "pinwheel", "pirate ship", "drink pitcher", "block plane", "planetarium", "plastic bag", "plate rack", "farm plow", "plunger", "Polaroid camera", "pole", "police van", "poncho", "pool table", "soda bottle", "plant pot", "potter's wheel", "power drill", "prayer rug", "printer", "prison", "projectile", "projector", "hockey puck", "punching bag", "purse", "quill", "quilt", "race car", "racket", "radiator", "radio", "radio telescope", "rain barrel", "recreational vehicle", "fishing casting reel", "reflex camera", "refrigerator", "remote control", "restaurant", "revolver", "rifle", "rocking chair", "rotisserie", "eraser", "rugby ball", "ruler measuring stick", "sneaker", "safe", "safety pin", "salt shaker", "sandal", "sarong", "saxophone", "scabbard", "weighing scale", "school bus", "schooner", "scoreboard", "CRT monitor", "screw", "screwdriver", "seat belt", "sewing machine", "shield", "shoe store", "shoji screen / room divider", "shopping basket", "shopping cart", "shovel", "shower cap", "shower curtain", "ski", "balaclava ski mask", "sleeping bag", "slide rule", "sliding door", "slot machine", "snorkel", "snowmobile", "snowplow", "soap dispenser", "soccer ball", "sock", "solar thermal collector", "sombrero", "soup bowl", "keyboard space bar", "space heater", "space shuttle", "spatula", "motorboat", "spider web", "spindle", "sports car", "spotlight", "stage", "steam locomotive", "through arch bridge", "steel drum", "stethoscope", "scarf", "stone wall", "stopwatch", "stove", "strainer", "tram", "stretcher", "couch", "stupa", "submarine", "suit", "sundial", "sunglasses", "dark glasses", "sunscreen", "suspension bridge", "mop", "sweatshirt", "swim trunks / shorts", "swing", "electrical switch", "syringe", "table lamp", "tank", "tape player", "teapot", "teddy bear", "television", "tennis ball", "thatched roof", "front curtain", "thimble", "threshing machine", "throne", "tile roof", "toaster", "tobacco shop", "toilet seat", "torch", "totem pole", "tow truck", "toy store", "tractor", "semi-trailer truck", "tray", "trench coat", "tricycle", "trimaran", "tripod", "triumphal arch", "trolleybus", "trombone", "hot tub", "turnstile", "typewriter keyboard", "umbrella", "unicycle", "upright piano", "vacuum cleaner", "vase", "vaulted or arched ceiling", "velvet fabric", "vending machine", "vestment", "viaduct", "violin", "volleyball", "waffle iron", "wall clock", "wallet", "wardrobe", "military aircraft", "sink", "washing machine", "water bottle", "water jug", "water tower", "whiskey jug", "whistle", "hair wig", "window screen", "window shade", "Windsor tie", "wine bottle", "airplane wing", "wok", "wooden spoon", "wool", "split-rail fence", "shipwreck", "sailboat", "yurt", "website", "comic book", "crossword", "traffic or street sign", "traffic light", "dust jacket", "menu", "plate", "guacamole", "consomme", "hot pot", "trifle", "ice cream", "popsicle", "baguette", "bagel", "pretzel", "cheeseburger", "hot dog", "mashed potatoes", "cabbage", "broccoli", "cauliflower", "zucchini", "spaghetti squash", "acorn squash", "butternut squash", "cucumber", "artichoke", "bell pepper", "cardoon", "mushroom", "Granny Smith apple", "strawberry", "orange", "lemon", "fig", "pineapple", "banana", "jackfruit", "cherimoya (custard apple)", "pomegranate", "hay", "carbonara", "chocolate syrup", "dough", "meatloaf", "pizza", "pot pie", "burrito", "red wine", "espresso", "tea cup", "eggnog", "mountain", "bubble", "cliff", "coral reef", "geyser", "lakeshore", "promontory", "sandbar", "beach", "valley", "volcano", "baseball player", "bridegroom", "scuba diver", "rapeseed", "daisy", "yellow lady's slipper", "corn", "acorn", "rose hip", "horse chestnut seed", "coral fungus", "agaric", "gyromitra", "stinkhorn mushroom", "earth star fungus", "hen of the woods mushroom", "bolete", "corn cob", "toilet paper"]
+
+IMAGENET_DEFAULT_TEMPLATES = [
+    '{}.',
+    'a bad photo of a {}.',
+    'a photo of many {}.',
+    'a sculpture of a {}.',
+    'a photo of the hard to see {}.',
+    'a low resolution photo of the {}.',
+    'a rendering of a {}.',
+    'graffiti of a {}.',
+    'a bad photo of the {}.',
+    'a cropped photo of the {}.',
+    'a tattoo of a {}.',
+    'the embroidered {}.',
+    'a photo of a hard to see {}.',
+    'a bright photo of a {}.',
+    'a photo of a clean {}.',
+    'a photo of a dirty {}.',
+    'a dark photo of the {}.',
+    'a drawing of a {}.',
+    'a photo of my {}.',
+    'the plastic {}.',
+    'a photo of the cool {}.',
+    'a close-up photo of a {}.',
+    'a black and white photo of the {}.',
+    'a painting of the {}.',
+    'a painting of a {}.',
+    'a pixelated photo of the {}.',
+    'a sculpture of the {}.',
+    'a bright photo of the {}.',
+    'a cropped photo of a {}.',
+    'a plastic {}.',
+    'a photo of the dirty {}.',
+    'a jpeg corrupted photo of a {}.',
+    'a blurry photo of the {}.',
+    'a photo of the {}.',
+    'a good photo of the {}.',
+    'a rendering of the {}.',
+    'a {} in a video game.',
+    'a photo of one {}.',
+    'a doodle of a {}.',
+    'a close-up photo of the {}.',
+    'a photo of a {}.',
+    'the origami {}.',
+    'the {} in a video game.',
+    'a sketch of a {}.',
+    'a doodle of the {}.',
+    'a origami {}.',
+    'a low resolution photo of a {}.',
+    'the toy {}.',
+    'a rendition of the {}.',
+    'a photo of the clean {}.',
+    'a photo of a large {}.',
+    'a rendition of a {}.',
+    'a photo of a nice {}.',
+    'a photo of a weird {}.',
+    'a blurry photo of a {}.',
+    'a cartoon {}.',
+    'art of a {}.',
+    'a sketch of the {}.',
+    'a embroidered {}.',
+    'a pixelated photo of a {}.',
+    'itap of the {}.',
+    'a jpeg corrupted photo of the {}.',
+    'a good photo of a {}.',
+    'a plushie {}.',
+    'a photo of the nice {}.',
+    'a photo of the small {}.',
+    'a photo of the weird {}.',
+    'the cartoon {}.',
+    'art of the {}.',
+    'a drawing of the {}.',
+    'a photo of the large {}.',
+    'a black and white photo of a {}.',
+    'the plushie {}.',
+    'a dark photo of a {}.',
+    'itap of a {}.',
+    'graffiti of the {}.',
+    'a toy {}.',
+    'itap of my {}.',
+    'a photo of a cool {}.',
+    'a photo of a small {}.',
+    'a tattoo of the {}.',
+]

MedImageInsight/ImageDataLoader/languages/prompt_engineering.py

@@ -0,0 +1,101 @@
+import numpy as np
+import random
+
+
+def get_prompt_templates():
+    prompt_templates = [
+        '{}.',
+        'a photo of a {}.',
+        'a bad photo of a {}.',
+        'a photo of many {}.',
+        'a sculpture of a {}.',
+        'a photo of the hard to see {}.',
+        'a low resolution photo of the {}.',
+        'a rendering of a {}.',
+        'graffiti of a {}.',
+        'a bad photo of the {}.',
+        'a cropped photo of the {}.',
+        'a tattoo of a {}.',
+        'the embroidered {}.',
+        'a photo of a hard to see {}.',
+        'a bright photo of a {}.',
+        'a photo of a clean {}.',
+        'a photo of a dirty {}.',
+        'a dark photo of the {}.',
+        'a drawing of a {}.',
+        'a photo of my {}.',
+        'the plastic {}.',
+        'a photo of the cool {}.',
+        'a close-up photo of a {}.',
+        'a black and white photo of the {}.',
+        'a painting of the {}.',
+        'a painting of a {}.',
+        'a pixelated photo of the {}.',
+        'a sculpture of the {}.',
+        'a bright photo of the {}.',
+        'a cropped photo of a {}.',
+        'a plastic {}.',
+        'a photo of the dirty {}.',
+        'a jpeg corrupted photo of a {}.',
+        'a blurry photo of the {}.',
+        'a photo of the {}.',
+        'a good photo of the {}.',
+        'a rendering of the {}.',
+        'a {} in a video game.',
+        'a photo of one {}.',
+        'a doodle of a {}.',
+        'a close-up photo of the {}.',
+        'the origami {}.',
+        'the {} in a video game.',
+        'a sketch of a {}.',
+        'a doodle of the {}.',
+        'a origami {}.',
+        'a low resolution photo of a {}.',
+        'the toy {}.',
+        'a rendition of the {}.',
+        'a photo of the clean {}.',
+        'a photo of a large {}.',
+        'a rendition of a {}.',
+        'a photo of a nice {}.',
+        'a photo of a weird {}.',
+        'a blurry photo of a {}.',
+        'a cartoon {}.',
+        'art of a {}.',
+        'a sketch of the {}.',
+        'a embroidered {}.',
+        'a pixelated photo of a {}.',
+        'itap of the {}.',
+        'a jpeg corrupted photo of the {}.',
+        'a good photo of a {}.',
+        'a plushie {}.',
+        'a photo of the nice {}.',
+        'a photo of the small {}.',
+        'a photo of the weird {}.',
+        'the cartoon {}.',
+        'art of the {}.',
+        'a drawing of the {}.',
+        'a photo of the large {}.',
+        'a black and white photo of a {}.',
+        'the plushie {}.',
+        'a dark photo of a {}.',
+        'itap of a {}.',
+        'graffiti of the {}.',
+        'a toy {}.',
+        'itap of my {}.',
+        'a photo of a cool {}.',
+        'a photo of a small {}.',
+        'a tattoo of the {}.',
+    ]
+    return prompt_templates
+
+
+def prompt_engineering(classnames):
+    prompt_templates = get_prompt_templates()
+    temp_idx = np.random.randint(len(prompt_templates))
+
+    if isinstance(classnames, list):
+        classname = random.choice(classnames)
+    else:
+        classname = classnames
+
+    return prompt_templates[temp_idx].replace('{}', classname.replace(',', '').replace('+', ' '))

MedImageInsight/ImageDataLoader/transforms/__init__.py

@@ -0,0 +1 @@
+from .build import build_transforms

MedImageInsight/ImageDataLoader/transforms/autoaugment.py

@@ -0,0 +1,447 @@
+import math
+from enum import Enum
+from typing import List, Tuple, Optional, Dict
+
+import torch
+from torch import Tensor
+
+from torchvision.transforms import functional as F
+from torchvision.transforms.functional import InterpolationMode
+
+__all__ = ["AutoAugmentPolicy", "AutoAugment", "RandAugment", "TrivialAugmentWide"]
+
+
+def _apply_op(
+    img: Tensor, op_name: str, magnitude: float, interpolation: InterpolationMode, fill: Optional[List[float]]
+):
+    if op_name == "ShearX":
+        img = F.affine(
+            img,
+            angle=0.0,
+            translate=[0, 0],
+            scale=1.0,
+            shear=[math.degrees(magnitude), 0.0],
+            interpolation=interpolation,
+            fill=fill,
+        )
+    elif op_name == "ShearY":
+        img = F.affine(
+            img,
+            angle=0.0,
+            translate=[0, 0],
+            scale=1.0,
+            shear=[0.0, math.degrees(magnitude)],
+            interpolation=interpolation,
+            fill=fill,
+        )
+    elif op_name == "TranslateX":
+        img = F.affine(
+            img,
+            angle=0.0,
+            translate=[int(magnitude), 0],
+            scale=1.0,
+            interpolation=interpolation,
+            shear=[0.0, 0.0],
+            fill=fill,
+        )
+    elif op_name == "TranslateY":
+        img = F.affine(
+            img,
+            angle=0.0,
+            translate=[0, int(magnitude)],
+            scale=1.0,
+            interpolation=interpolation,
+            shear=[0.0, 0.0],
+            fill=fill,
+        )
+    elif op_name == "Rotate":
+        img = F.rotate(img, magnitude, interpolation=interpolation, fill=fill)
+    elif op_name == "Brightness":
+        img = F.adjust_brightness(img, 1.0 + magnitude)
+    elif op_name == "Color":
+        img = F.adjust_saturation(img, 1.0 + magnitude)
+    elif op_name == "Contrast":
+        img = F.adjust_contrast(img, 1.0 + magnitude)
+    elif op_name == "Sharpness":
+        img = F.adjust_sharpness(img, 1.0 + magnitude)
+    elif op_name == "Posterize":
+        img = F.posterize(img, int(magnitude))
+    elif op_name == "Solarize":
+        img = F.solarize(img, magnitude)
+    elif op_name == "AutoContrast":
+        img = F.autocontrast(img)
+    elif op_name == "Equalize":
+        img = F.equalize(img)
+    elif op_name == "Invert":
+        img = F.invert(img)
+    elif op_name == "Identity":
+        pass
+    else:
+        raise ValueError(f"The provided operator {op_name} is not recognized.")
+    return img
+
+
+class AutoAugmentPolicy(Enum):
+    """AutoAugment policies learned on different datasets.
+    Available policies are IMAGENET, CIFAR10 and SVHN.
+    """
+
+    IMAGENET = "imagenet"
+    CIFAR10 = "cifar10"
+    SVHN = "svhn"
+
+
+# FIXME: Eliminate copy-pasted code for fill standardization and _augmentation_space() by moving stuff on a base class
+class AutoAugment(torch.nn.Module):
+    r"""AutoAugment data augmentation method based on
+    `"AutoAugment: Learning Augmentation Strategies from Data" <https://arxiv.org/pdf/1805.09501.pdf>`_.
+    If the image is torch Tensor, it should be of type torch.uint8, and it is expected
+    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If img is PIL Image, it is expected to be in mode "L" or "RGB".
+
+    Args:
+        policy (AutoAugmentPolicy): Desired policy enum defined by
+            :class:`torchvision.transforms.autoaugment.AutoAugmentPolicy`. Default is ``AutoAugmentPolicy.IMAGENET``.
+        interpolation (InterpolationMode): Desired interpolation enum defined by
+            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
+            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
+        fill (sequence or number, optional): Pixel fill value for the area outside the transformed
+            image. If given a number, the value is used for all bands respectively.
+    """
+
+    def __init__(
+        self,
+        policy: AutoAugmentPolicy = AutoAugmentPolicy.IMAGENET,
+        interpolation: InterpolationMode = InterpolationMode.NEAREST,
+        fill: Optional[List[float]] = None,
+    ) -> None:
+        super().__init__()
+        self.policy = policy
+        self.interpolation = interpolation
+        self.fill = fill
+        self.policies = self._get_policies(policy)
+
+    def _get_policies(
+        self, policy: AutoAugmentPolicy
+    ) -> List[Tuple[Tuple[str, float, Optional[int]], Tuple[str, float, Optional[int]]]]:
+        if policy == AutoAugmentPolicy.IMAGENET:
+            return [
+                (("Posterize", 0.4, 8), ("Rotate", 0.6, 9)),
+                (("Solarize", 0.6, 5), ("AutoContrast", 0.6, None)),
+                (("Equalize", 0.8, None), ("Equalize", 0.6, None)),
+                (("Posterize", 0.6, 7), ("Posterize", 0.6, 6)),
+                (("Equalize", 0.4, None), ("Solarize", 0.2, 4)),
+                (("Equalize", 0.4, None), ("Rotate", 0.8, 8)),
+                (("Solarize", 0.6, 3), ("Equalize", 0.6, None)),
+                (("Posterize", 0.8, 5), ("Equalize", 1.0, None)),
+                (("Rotate", 0.2, 3), ("Solarize", 0.6, 8)),
+                (("Equalize", 0.6, None), ("Posterize", 0.4, 6)),
+                (("Rotate", 0.8, 8), ("Color", 0.4, 0)),
+                (("Rotate", 0.4, 9), ("Equalize", 0.6, None)),
+                (("Equalize", 0.0, None), ("Equalize", 0.8, None)),
+                (("Invert", 0.6, None), ("Equalize", 1.0, None)),
+                (("Color", 0.6, 4), ("Contrast", 1.0, 8)),
+                (("Rotate", 0.8, 8), ("Color", 1.0, 2)),
+                (("Color", 0.8, 8), ("Solarize", 0.8, 7)),
+                (("Sharpness", 0.4, 7), ("Invert", 0.6, None)),
+                (("ShearX", 0.6, 5), ("Equalize", 1.0, None)),
+                (("Color", 0.4, 0), ("Equalize", 0.6, None)),
+                (("Equalize", 0.4, None), ("Solarize", 0.2, 4)),
+                (("Solarize", 0.6, 5), ("AutoContrast", 0.6, None)),
+                (("Invert", 0.6, None), ("Equalize", 1.0, None)),
+                (("Color", 0.6, 4), ("Contrast", 1.0, 8)),
+                (("Equalize", 0.8, None), ("Equalize", 0.6, None)),
+            ]
+        elif policy == AutoAugmentPolicy.CIFAR10:
+            return [
+                (("Invert", 0.1, None), ("Contrast", 0.2, 6)),
+                (("Rotate", 0.7, 2), ("TranslateX", 0.3, 9)),
+                (("Sharpness", 0.8, 1), ("Sharpness", 0.9, 3)),
+                (("ShearY", 0.5, 8), ("TranslateY", 0.7, 9)),
+                (("AutoContrast", 0.5, None), ("Equalize", 0.9, None)),
+                (("ShearY", 0.2, 7), ("Posterize", 0.3, 7)),
+                (("Color", 0.4, 3), ("Brightness", 0.6, 7)),
+                (("Sharpness", 0.3, 9), ("Brightness", 0.7, 9)),
+                (("Equalize", 0.6, None), ("Equalize", 0.5, None)),
+                (("Contrast", 0.6, 7), ("Sharpness", 0.6, 5)),
+                (("Color", 0.7, 7), ("TranslateX", 0.5, 8)),
+                (("Equalize", 0.3, None), ("AutoContrast", 0.4, None)),
+                (("TranslateY", 0.4, 3), ("Sharpness", 0.2, 6)),
+                (("Brightness", 0.9, 6), ("Color", 0.2, 8)),
+                (("Solarize", 0.5, 2), ("Invert", 0.0, None)),
+                (("Equalize", 0.2, None), ("AutoContrast", 0.6, None)),
+                (("Equalize", 0.2, None), ("Equalize", 0.6, None)),
+                (("Color", 0.9, 9), ("Equalize", 0.6, None)),
+                (("AutoContrast", 0.8, None), ("Solarize", 0.2, 8)),
+                (("Brightness", 0.1, 3), ("Color", 0.7, 0)),
+                (("Solarize", 0.4, 5), ("AutoContrast", 0.9, None)),
+                (("TranslateY", 0.9, 9), ("TranslateY", 0.7, 9)),
+                (("AutoContrast", 0.9, None), ("Solarize", 0.8, 3)),
+                (("Equalize", 0.8, None), ("Invert", 0.1, None)),
+                (("TranslateY", 0.7, 9), ("AutoContrast", 0.9, None)),
+            ]
+        elif policy == AutoAugmentPolicy.SVHN:
+            return [
+                (("ShearX", 0.9, 4), ("Invert", 0.2, None)),
+                (("ShearY", 0.9, 8), ("Invert", 0.7, None)),
+                (("Equalize", 0.6, None), ("Solarize", 0.6, 6)),
+                (("Invert", 0.9, None), ("Equalize", 0.6, None)),
+                (("Equalize", 0.6, None), ("Rotate", 0.9, 3)),
+                (("ShearX", 0.9, 4), ("AutoContrast", 0.8, None)),
+                (("ShearY", 0.9, 8), ("Invert", 0.4, None)),
+                (("ShearY", 0.9, 5), ("Solarize", 0.2, 6)),
+                (("Invert", 0.9, None), ("AutoContrast", 0.8, None)),
+                (("Equalize", 0.6, None), ("Rotate", 0.9, 3)),
+                (("ShearX", 0.9, 4), ("Solarize", 0.3, 3)),
+                (("ShearY", 0.8, 8), ("Invert", 0.7, None)),
+                (("Equalize", 0.9, None), ("TranslateY", 0.6, 6)),
+                (("Invert", 0.9, None), ("Equalize", 0.6, None)),
+                (("Contrast", 0.3, 3), ("Rotate", 0.8, 4)),
+                (("Invert", 0.8, None), ("TranslateY", 0.0, 2)),
+                (("ShearY", 0.7, 6), ("Solarize", 0.4, 8)),
+                (("Invert", 0.6, None), ("Rotate", 0.8, 4)),
+                (("ShearY", 0.3, 7), ("TranslateX", 0.9, 3)),
+                (("ShearX", 0.1, 6), ("Invert", 0.6, None)),
+                (("Solarize", 0.7, 2), ("TranslateY", 0.6, 7)),
+                (("ShearY", 0.8, 4), ("Invert", 0.8, None)),
+                (("ShearX", 0.7, 9), ("TranslateY", 0.8, 3)),
+                (("ShearY", 0.8, 5), ("AutoContrast", 0.7, None)),
+                (("ShearX", 0.7, 2), ("Invert", 0.1, None)),
+            ]
+        else:
+            raise ValueError(f"The provided policy {policy} is not recognized.")
+
+    def _augmentation_space(self, num_bins: int, image_size: List[int]) -> Dict[str, Tuple[Tensor, bool]]:
+        return {
+            # op_name: (magnitudes, signed)
+            "ShearX": (torch.linspace(0.0, 0.3, num_bins), True),
+            "ShearY": (torch.linspace(0.0, 0.3, num_bins), True),
+            "TranslateX": (torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True),
+            "TranslateY": (torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True),
+            "Rotate": (torch.linspace(0.0, 30.0, num_bins), True),
+            "Brightness": (torch.linspace(0.0, 0.9, num_bins), True),
+            "Color": (torch.linspace(0.0, 0.9, num_bins), True),
+            "Contrast": (torch.linspace(0.0, 0.9, num_bins), True),
+            "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True),
+            "Posterize": (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4)).round().int(), False),
+            "Solarize": (torch.linspace(255.0, 0.0, num_bins), False),
+            "AutoContrast": (torch.tensor(0.0), False),
+            "Equalize": (torch.tensor(0.0), False),
+            "Invert": (torch.tensor(0.0), False),
+        }
+
+    @staticmethod
+    def get_params(transform_num: int) -> Tuple[int, Tensor, Tensor]:
+        """Get parameters for autoaugment transformation
+
+        Returns:
+            params required by the autoaugment transformation
+        """
+        policy_id = int(torch.randint(transform_num, (1,)).item())
+        probs = torch.rand((2,))
+        signs = torch.randint(2, (2,))
+
+        return policy_id, probs, signs
+
+    def forward(self, img: Tensor) -> Tensor:
+        """
+            img (PIL Image or Tensor): Image to be transformed.
+
+        Returns:
+            PIL Image or Tensor: AutoAugmented image.
+        """
+        fill = self.fill
+        if isinstance(img, Tensor):
+            if isinstance(fill, (int, float)):
+                fill = [float(fill)] * F.get_image_num_channels(img)
+            elif fill is not None:
+                fill = [float(f) for f in fill]
+
+        transform_id, probs, signs = self.get_params(len(self.policies))
+
+        for i, (op_name, p, magnitude_id) in enumerate(self.policies[transform_id]):
+            if probs[i] <= p:
+                op_meta = self._augmentation_space(10, F.get_image_size(img))
+                magnitudes, signed = op_meta[op_name]
+                magnitude = float(magnitudes[magnitude_id].item()) if magnitude_id is not None else 0.0
+                if signed and signs[i] == 0:
+                    magnitude *= -1.0
+                img = _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill)
+
+        return img
+
+    def __repr__(self) -> str:
+        return self.__class__.__name__ + f"(policy={self.policy}, fill={self.fill})"
+
+
+class RandAugment(torch.nn.Module):
+    r"""RandAugment data augmentation method based on
+    `"RandAugment: Practical automated data augmentation with a reduced search space"
+    <https://arxiv.org/abs/1909.13719>`_.
+    If the image is torch Tensor, it should be of type torch.uint8, and it is expected
+    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If img is PIL Image, it is expected to be in mode "L" or "RGB".
+
+    Args:
+        num_ops (int): Number of augmentation transformations to apply sequentially.
+        magnitude (int): Magnitude for all the transformations.
+        num_magnitude_bins (int): The number of different magnitude values.
+        interpolation (InterpolationMode): Desired interpolation enum defined by
+            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
+            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
+        fill (sequence or number, optional): Pixel fill value for the area outside the transformed
+            image. If given a number, the value is used for all bands respectively.
+    """
+
+    def __init__(
+        self,
+        num_ops: int = 2,
+        magnitude: int = 9,
+        num_magnitude_bins: int = 31,
+        interpolation: InterpolationMode = InterpolationMode.NEAREST,
+        fill: Optional[List[float]] = None,
+    ) -> None:
+        super().__init__()
+        self.num_ops = num_ops
+        self.magnitude = magnitude
+        self.num_magnitude_bins = num_magnitude_bins
+        self.interpolation = interpolation
+        self.fill = fill
+
+    def _augmentation_space(self, num_bins: int, image_size: List[int]) -> Dict[str, Tuple[Tensor, bool]]:
+        return {
+            # op_name: (magnitudes, signed)
+            "Identity": (torch.tensor(0.0), False),
+            "ShearX": (torch.linspace(0.0, 0.3, num_bins), True),
+            "ShearY": (torch.linspace(0.0, 0.3, num_bins), True),
+            "TranslateX": (torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True),
+            "TranslateY": (torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True),
+            "Rotate": (torch.linspace(0.0, 30.0, num_bins), True),
+            "Brightness": (torch.linspace(0.0, 0.9, num_bins), True),
+            "Color": (torch.linspace(0.0, 0.9, num_bins), True),
+            "Contrast": (torch.linspace(0.0, 0.9, num_bins), True),
+            "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True),
+            "Posterize": (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4)).round().int(), False),
+            "Solarize": (torch.linspace(255.0, 0.0, num_bins), False),
+            "AutoContrast": (torch.tensor(0.0), False),
+            "Equalize": (torch.tensor(0.0), False),
+        }
+
+    def forward(self, img: Tensor) -> Tensor:
+        """
+            img (PIL Image or Tensor): Image to be transformed.
+
+        Returns:
+            PIL Image or Tensor: Transformed image.
+        """
+        fill = self.fill
+        if isinstance(img, Tensor):
+            if isinstance(fill, (int, float)):
+                fill = [float(fill)] * F.get_image_num_channels(img)
+            elif fill is not None:
+                fill = [float(f) for f in fill]
+
+        for _ in range(self.num_ops):
+            op_meta = self._augmentation_space(self.num_magnitude_bins, F.get_image_size(img))
+            op_index = int(torch.randint(len(op_meta), (1,)).item())
+            op_name = list(op_meta.keys())[op_index]
+            magnitudes, signed = op_meta[op_name]
+            magnitude = float(magnitudes[self.magnitude].item()) if magnitudes.ndim > 0 else 0.0
+            if signed and torch.randint(2, (1,)):
+                magnitude *= -1.0
+            img = _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill)
+
+        return img
+
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "num_ops={num_ops}"
+        s += ", magnitude={magnitude}"
+        s += ", num_magnitude_bins={num_magnitude_bins}"
+        s += ", interpolation={interpolation}"
+        s += ", fill={fill}"
+        s += ")"
+        return s.format(**self.__dict__)
+
+
+class TrivialAugmentWide(torch.nn.Module):
+    r"""Dataset-independent data-augmentation with TrivialAugment Wide, as described in
+    `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" <https://arxiv.org/abs/2103.10158>`_.
+    If the image is torch Tensor, it should be of type torch.uint8, and it is expected
+    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If img is PIL Image, it is expected to be in mode "L" or "RGB".
+
+    Args:
+        num_magnitude_bins (int): The number of different magnitude values.
+        interpolation (InterpolationMode): Desired interpolation enum defined by
+            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
+            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
+        fill (sequence or number, optional): Pixel fill value for the area outside the transformed
+            image. If given a number, the value is used for all bands respectively.
+    """
+
+    def __init__(
+        self,
+        num_magnitude_bins: int = 31,
+        interpolation: InterpolationMode = InterpolationMode.NEAREST,
+        fill: Optional[List[float]] = None,
+    ) -> None:
+        super().__init__()
+        self.num_magnitude_bins = num_magnitude_bins
+        self.interpolation = interpolation
+        self.fill = fill
+
+    def _augmentation_space(self, num_bins: int) -> Dict[str, Tuple[Tensor, bool]]:
+        return {
+            # op_name: (magnitudes, signed)
+            "Identity": (torch.tensor(0.0), False),
+            "ShearX": (torch.linspace(0.0, 0.99, num_bins), True),
+            "ShearY": (torch.linspace(0.0, 0.99, num_bins), True),
+            "TranslateX": (torch.linspace(0.0, 32.0, num_bins), True),
+            "TranslateY": (torch.linspace(0.0, 32.0, num_bins), True),
+            "Rotate": (torch.linspace(0.0, 135.0, num_bins), True),
+            "Brightness": (torch.linspace(0.0, 0.99, num_bins), True),
+            "Color": (torch.linspace(0.0, 0.99, num_bins), True),
+            "Contrast": (torch.linspace(0.0, 0.99, num_bins), True),
+            "Sharpness": (torch.linspace(0.0, 0.99, num_bins), True),
+            "Posterize": (8 - (torch.arange(num_bins) / ((num_bins - 1) / 6)).round().int(), False),
+            "Solarize": (torch.linspace(255.0, 0.0, num_bins), False),
+            "AutoContrast": (torch.tensor(0.0), False),
+            "Equalize": (torch.tensor(0.0), False),
+        }
+
+    def forward(self, img: Tensor) -> Tensor:
+        """
+            img (PIL Image or Tensor): Image to be transformed.
+
+        Returns:
+            PIL Image or Tensor: Transformed image.
+        """
+        fill = self.fill
+        if isinstance(img, Tensor):
+            if isinstance(fill, (int, float)):
+                fill = [float(fill)] * F.get_image_num_channels(img)
+            elif fill is not None:
+                fill = [float(f) for f in fill]
+
+        op_meta = self._augmentation_space(self.num_magnitude_bins)
+        op_index = int(torch.randint(len(op_meta), (1,)).item())
+        op_name = list(op_meta.keys())[op_index]
+        magnitudes, signed = op_meta[op_name]
+        magnitude = (
+            float(magnitudes[torch.randint(len(magnitudes), (1,), dtype=torch.long)].item())
+            if magnitudes.ndim > 0
+            else 0.0
+        )
+        if signed and torch.randint(2, (1,)):
+            magnitude *= -1.0
+
+        return _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill)
+
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "num_magnitude_bins={num_magnitude_bins}"
+        s += ", interpolation={interpolation}"
+        s += ", fill={fill}"
+        s += ")"
+        return s.format(**self.__dict__)

MedImageInsight/ImageDataLoader/transforms/build.py

@@ -0,0 +1,261 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import timm
+from timm.data import create_transform
+
+from yacs.config import CfgNode as CN
+from PIL import ImageFilter
+import logging
+import random
+
+import torch
+import torchvision.transforms as T
+
+
+from .autoaugment import AutoAugmentPolicy
+from .autoaugment import AutoAugment
+from .autoaugment import RandAugment
+from .autoaugment import TrivialAugmentWide
+from .threeaugment import deitIII_Solarization
+from .threeaugment import deitIII_gray_scale
+from .threeaugment import deitIII_GaussianBlur
+
+from PIL import ImageOps
+from timm.data.transforms import RandomResizedCropAndInterpolation
+
+logger = logging.getLogger(__name__)
+
+
+class GaussianBlur(object):
+    """Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709"""
+
+    def __init__(self, sigma=[.1, 2.]):
+        self.sigma = sigma
+
+    def __call__(self, x):
+        sigma = random.uniform(self.sigma[0], self.sigma[1])
+        x = x.filter(ImageFilter.GaussianBlur(radius=sigma))
+        return x
+
+
+def get_resolution(original_resolution):
+    """Takes (H,W) and returns (precrop, crop)."""
+    area = original_resolution[0] * original_resolution[1]
+    return (160, 128) if area < 96*96 else (512, 480)
+
+
+INTERPOLATION_MODES = {
+    'bilinear': T.InterpolationMode.BILINEAR,
+    'bicubic': T.InterpolationMode.BICUBIC,
+    'nearest': T.InterpolationMode.NEAREST,
+}
+
+
+def build_transforms(cfg, is_train=True):
+    # assert isinstance(cfg.DATASET.OUTPUT_SIZE, (list, tuple)), 'DATASET.OUTPUT_SIZE should be list or tuple'
+    normalize = T.Normalize(
+        mean=cfg['IMAGE_ENCODER']['IMAGE_MEAN'],
+        std=cfg['IMAGE_ENCODER']['IMAGE_STD']
+    )
+
+    transforms = None
+    if is_train:
+        if 'THREE_AUG' in cfg['AUG']:
+            img_size = cfg['IMAGE_ENCODER']['IMAGE_SIZE']
+            remove_random_resized_crop = cfg['AUG']['THREE_AUG']['SRC']
+            mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
+            primary_tfl = []
+            scale=(0.08, 1.0)
+            interpolation='bicubic'
+            if remove_random_resized_crop:
+                primary_tfl = [
+                    T.Resize(img_size, interpolation=3),
+                    T.RandomCrop(img_size, padding=4,padding_mode='reflect'),
+                    T.RandomHorizontalFlip()
+                ]
+            else:
+                primary_tfl = [
+                    RandomResizedCropAndInterpolation(
+                        img_size, scale=scale, interpolation=interpolation),
+                    T.RandomHorizontalFlip()
+                ]
+            secondary_tfl = [T.RandomChoice([gray_scale(p=1.0),
+                                             Solarization(p=1.0),
+                                             GaussianBlurDeiTv3(p=1.0)])]
+            color_jitter = cfg['AUG']['THREE_AUG']['COLOR_JITTER']
+            if color_jitter is not None and not color_jitter==0:
+                secondary_tfl.append(T.ColorJitter(color_jitter, color_jitter, color_jitter))
+            final_tfl = [
+                    T.ToTensor(),
+                    T.Normalize(
+                        mean=torch.tensor(mean),
+                        std=torch.tensor(std))
+                ]
+            return T.Compose(primary_tfl+secondary_tfl+final_tfl)
+        elif 'TIMM_AUG' in cfg['AUG'] and cfg['AUG']['TIMM_AUG']['USE_TRANSFORM']:
+            logger.info('=> use timm transform for training')
+            timm_cfg = cfg['AUG']['TIMM_AUG']
+            transforms = create_transform(
+                input_size=cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0],
+                is_training=True,
+                use_prefetcher=False,
+                no_aug=False,
+                re_prob=timm_cfg.get('RE_PROB', 0.),
+                re_mode=timm_cfg.get('RE_MODE', 'const'),
+                re_count=timm_cfg.get('RE_COUNT', 1),
+                re_num_splits= 0 if not timm_cfg.get('RE_SPLITS', False) else timm_cfg['RE_SPLITS'], # if false or 0, return 0
+                scale=cfg['AUG'].get('SCALE', None),
+                ratio=cfg['AUG'].get('RATIO', None),
+                hflip=timm_cfg.get('HFLIP', 0.5),
+                vflip=timm_cfg.get('VFLIP', 0.),
+                color_jitter=timm_cfg.get('COLOR_JITTER', 0.4),
+                auto_augment=timm_cfg.get('AUTO_AUGMENT', None),
+                interpolation=cfg['AUG']['INTERPOLATION'],
+                mean=cfg['IMAGE_ENCODER']['IMAGE_MEAN'],
+                std=cfg['IMAGE_ENCODER']['IMAGE_STD'],
+            )
+        elif 'TORCHVISION_AUG' in cfg['AUG']:
+            logger.info('=> use torchvision transform fro training')
+            crop_size = cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0]
+            interpolation = INTERPOLATION_MODES[cfg['AUG']['INTERPOLATION']]
+            trans = [
+                T.RandomResizedCrop(
+                    crop_size, scale=cfg['AUG']['SCALE'], ratio=cfg['AUG']['RATIO'],
+                    interpolation=interpolation
+                )
+            ]
+            hflip_prob = cfg['AUG']['TORCHVISION_AUG']['HFLIP']
+            auto_augment_policy = cfg['AUG']['TORCHVISION_AUG'].get('AUTO_AUGMENT', None)
+            if hflip_prob > 0:
+                trans.append(T.RandomHorizontalFlip(hflip_prob))
+            if auto_augment_policy is not None:
+                if auto_augment_policy == "ra":
+                    trans.append(RandAugment(interpolation=interpolation))
+                elif auto_augment_policy == "ta_wide":
+                    trans.append(TrivialAugmentWide(interpolation=interpolation))
+                else:
+                    aa_policy = AutoAugmentPolicy(auto_augment_policy)
+                    trans.append(AutoAugment(policy=aa_policy, interpolation=interpolation))
+            trans.extend(
+                [
+                    T.ToTensor(),
+                    normalize,
+                ]
+            )
+            random_erase_prob = cfg['AUG']['TORCHVISION_AUG']['RE_PROB']
+            random_erase_scale = cfg['AUG']['TORCHVISION_AUG'].get('RE_SCALE', 0.33)
+            if random_erase_prob > 0:
+                # NCFC (4/26/2023): Added scale parameter to random erasing for medical imaging
+                trans.append(T.RandomErasing(p=random_erase_prob, scale = (0.02, random_erase_scale)))
+
+            from torchvision.transforms import InterpolationMode
+            rotation = cfg['AUG']['TORCHVISION_AUG'].get('ROTATION', 0.0)
+            if (rotation > 0.0):
+                trans.append(T.RandomRotation(rotation, interpolation=InterpolationMode.BILINEAR))
+                logger.info(" TORCH AUG: Rotation: " + str(rotation))
+
+            transforms = T.Compose(trans)
+        elif cfg['AUG'].get('RANDOM_CENTER_CROP', False):
+            logger.info('=> use random center crop data augmenation')
+            # precrop, crop = get_resolution(cfg.TRAIN.IMAGE_SIZE)
+            crop = cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0]
+            padding = cfg['AUG'].get('RANDOM_CENTER_CROP_PADDING', 32)
+            precrop = crop + padding
+            mode = INTERPOLATION_MODES[cfg['AUG']['INTERPOLATION']]
+            transforms = T.Compose([
+                T.Resize(
+                    (precrop, precrop),
+                    interpolation=mode
+                ),
+                T.RandomCrop((crop, crop)),
+                T.RandomHorizontalFlip(),
+                T.ToTensor(),
+                normalize,
+            ])
+        elif cfg['AUG'].get('MAE_FINETUNE_AUG', False):
+            mean = cfg['IMAGE_ENCODER']['IMAGE_MEAN']
+            std = cfg['IMAGE_ENCODER']['IMAGE_STD']
+            transforms = create_transform(
+                input_size=cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0],
+                is_training=True,
+                color_jitter=cfg['AUG'].get('COLOR_JITTER', None),
+                auto_augment=cfg['AUG'].get('AUTO_AUGMENT', 'rand-m9-mstd0.5-inc1'),
+                interpolation='bicubic',
+                re_prob=cfg['AUG'].get('RE_PROB', 0.25),
+                re_mode=cfg['AUG'].get('RE_MODE', "pixel"),
+                re_count=cfg['AUG'].get('RE_COUNT', 1),
+                mean=mean,
+                std=std,
+            )
+        elif cfg['AUG'].get('MAE_PRETRAIN_AUG', False):
+            mean = cfg['IMAGE_ENCODER']['IMAGE_MEAN']
+            std = cfg['IMAGE_ENCODER']['IMAGE_STD']
+            transforms = T.Compose([
+                T.RandomResizedCrop(cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0], scale=tuple(cfg['AUG']['SCALE']), interpolation=INTERPOLATION_MODES["bicubic"]),  # 3 is bicubic
+                T.RandomHorizontalFlip(),
+                T.ToTensor(),
+                T.Normalize(mean=mean, std=std)])
+        elif cfg['AUG'].get('ThreeAugment', False): # from DeiT III
+            mean = cfg['IMAGE_ENCODER']['IMAGE_MEAN']
+            std = cfg['IMAGE_ENCODER']['IMAGE_STD']
+            img_size = cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0]
+            remove_random_resized_crop = cfg['AUG'].get('src', False)
+            mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
+            primary_tfl = []
+            scale=(0.08, 1.0)
+            interpolation='bicubic'
+            if remove_random_resized_crop:
+                primary_tfl = [
+                    T.Resize(img_size, interpolation=3), # bicubic
+                    T.RandomCrop(img_size, padding=4,padding_mode='reflect'),
+                    T.RandomHorizontalFlip()
+                ]
+            else:
+                primary_tfl = [
+                    timm.data.transforms.RandomResizedCropAndInterpolation(
+                        img_size, scale=scale, interpolation=interpolation),
+                    T.RandomHorizontalFlip()
+                ]
+
+            secondary_tfl = [T.RandomChoice([deitIII_gray_scale(p=1.0),
+                                             deitIII_Solarization(p=1.0),
+                                             deitIII_GaussianBlur(p=1.0)])]
+            color_jitter = cfg['AUG']['COLOR_JITTER']
+            secondary_tfl.append(T.ColorJitter(color_jitter, color_jitter, color_jitter))
+            final_tfl = [
+                    T.ToTensor(),
+                    T.Normalize(
+                        mean=torch.tensor(mean),
+                        std=torch.tensor(std))
+                ]
+            transforms = T.Compose(primary_tfl+secondary_tfl+final_tfl)
+        logger.info('=> training transformers: {}'.format(transforms))
+    else:
+        mode = INTERPOLATION_MODES[cfg['AUG']['INTERPOLATION']]
+        if cfg['TEST']['CENTER_CROP']:
+            transforms = T.Compose([
+                T.Resize(
+                    int(cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0] / 0.875),
+                    # the same behavior as in deit: size = int((256 / 224) * args.input_size)
+                    # 224 / 256 = 0.875
+                    interpolation=mode
+                ),
+                T.CenterCrop(cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0]),
+                T.ToTensor(),
+                normalize,
+            ])
+        else:
+            transforms = T.Compose([
+                T.Resize(
+                    (cfg['IMAGE_ENCODER']['IMAGE_SIZE'][1], cfg['IMAGE_ENCODER']['IMAGE_SIZE'][0]),
+                    interpolation=mode
+                ),
+                T.ToTensor(),
+                normalize,
+            ])
+        logger.info('=> testing transformers: {}'.format(transforms))
+
+    return transforms
+

MedImageInsight/ImageDataLoader/transforms/threeaugment.py

@@ -0,0 +1,54 @@
+import random
+from PIL import ImageFilter, ImageOps
+from torchvision import transforms
+
+
+class deitIII_GaussianBlur(object):
+    """
+    Apply Gaussian Blur to the PIL image.
+    """
+    def __init__(self, p=0.1, radius_min=0.1, radius_max=2.):
+        self.prob = p
+        self.radius_min = radius_min
+        self.radius_max = radius_max
+
+    def __call__(self, img):
+        do_it = random.random() <= self.prob
+        if not do_it:
+            return img
+
+        img = img.filter(
+            ImageFilter.GaussianBlur(
+                radius=random.uniform(self.radius_min, self.radius_max)
+            )
+        )
+        return img
+
+
+class deitIII_Solarization(object):
+    """
+    Apply Solarization to the PIL image.
+    """
+    def __init__(self, p=0.2):
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            return ImageOps.solarize(img)
+        else:
+            return img
+
+
+class deitIII_gray_scale(object):
+    """
+    Apply Solarization to the PIL image.
+    """
+    def __init__(self, p=0.2):
+        self.p = p
+        self.transf = transforms.Grayscale(3)
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            return self.transf(img)
+        else:
+            return img

MedImageInsight/ImageDataLoader/tsv.py

@@ -0,0 +1,351 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+from io import BytesIO
+import json
+import logging
+import base64
+import random
+from typing import Callable, List, Tuple, Union, NamedTuple
+from PIL import Image
+from PIL import ImageFile
+import torch.utils.data as data
+from .languages.prompt_engineering import prompt_engineering
+from .tsv_file import TSVFile, CompositeTSVFile
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+logger = logging.getLogger(__name__)
+
+
+class TSVDataset(data.Dataset):
+
+    def __init__(self,
+                 tsv_file: Union[str, List[str]],
+                 transform: Callable = None,
+                 map_file: str = None,
+                 token_file: str = None,
+                 is_train: bool = True,
+                 azcopy_path: str = None):
+        self.transform = transform
+        self._chunk_sizes = None
+        self.label2idx = self._load_map(map_file)
+        self.class_selector = list(self.label2idx.keys()) if self.label2idx else None
+
+        if isinstance(tsv_file, str):
+            if os.path.splitext(tsv_file)[1] == '.tsv':
+                self.tsv_file = TSVFile(
+                    tsv_file, class_selector=self.class_selector
+                )
+            else:
+                self.tsv_file = CompositeTSVFile(
+                    tsv_file,
+                    class_selector=self.class_selector,
+                    is_train=is_train,
+                    sas_token_path=token_file,
+                    azcopy_path=azcopy_path
+                )
+                self._chunk_sizes = self.tsv_file.get_chunk_size()
+        elif isinstance(tsv_file, list):
+            self.tsv_file = CompositeTSVFile(
+                tsv_file,
+                class_selector=self.class_selector,
+                is_train=is_train,
+                sas_token_path=token_file,
+                azcopy_path=azcopy_path
+            )
+            self._chunk_sizes = self.tsv_file.get_chunk_size()
+        else:
+            raise ValueError("Invalid input! Please check the tsv filenames")
+
+        logger.debug('=> {}\titems: {}'.format(tsv_file, len(self.tsv_file)))
+
+    def fetch_blob(self, idx):
+        image_tsv = self.tsv_file.file_list[idx]
+        self.tsv_file.blob_storage.fetch_blob(image_tsv)
+
+    def num_classes(self):
+        return len(self.class_selector)
+
+    def get_chunk_sizes(self):
+        return self._chunk_sizes
+
+    def get_class_boundaries(self):
+        # The samples of each class are organized class-by-class.
+        # _class_boundaries stores the lower- and upper-bound of each class.
+        return self.tsv_file.get_class_boundaries()
+
+    def get_filenames(self):
+        filenames = [
+            self.tsv_file.get_key(i)
+            for i in range(self.tsv_file.num_rows())
+        ]
+
+        return filenames
+
+    def _load_map(self, map_file: str):
+        if not map_file:
+            return None
+
+        label2idx = {}
+        with open(map_file) as f:
+            for line in f:
+                items = line.strip().split('\t')
+                label2idx[items[0]] = int(items[1])
+
+        return label2idx
+
+    def __getitem__(self, index: Union[int, Tuple[int, int]]):
+        items = self.tsv_file[index]
+        _, target, img = self._decode_data(items)
+
+        if self.transform:
+            img = self.transform(img)
+
+        return img, target
+
+    def _decode_data(self, items: Tuple[str, str, str]):
+        key = items[0]
+        label = self._get_label(items[1])
+        image = Image.open(BytesIO(base64.b64decode(items[2]))).convert('RGB')
+
+        return key, label, image
+
+    def _get_label(self, item: str):
+        if not self.label2idx:
+            return int(item)
+
+        js = json.loads(item)
+        return self.label2idx[js[0]['class']]
+
+    def __len__(self):
+        return len(self.tsv_file)
+
+
+class TSVMeta(NamedTuple):
+    source: str
+    num_classes: int
+    task: str
+
+
+class TSVImageTextDatasetV2(data.Dataset):
+    """
+        This class is intended for encapsulating Image/Text pair data for contrastive learning described in
+        the following paper,
+        "Learning Transferable Visual Models From Natural Language Supervision" (a.k.a CLIP)
+        V2: support image text pairs and supervised classification data
+    """
+
+    def __init__(self,
+                 image_tsv_file: Union[str, List[str]],
+                 text_tsv_file: Union[str, List[str]],
+                 transform: Callable = None,
+                 tokenize: Callable = None,
+                 context_length: int = 77,
+                 num_captions: int = 1,
+                 text_format: str = 'txt',
+                 is_train: bool = True,
+                 sas_token_path: str = None,
+                 azcopy_path: str = None,
+                 metas: List[NamedTuple] = None,
+                 prompt_engineering=True,
+                 concat_queries=False):
+        self.transform = transform
+        self.tokenize = tokenize
+        self._chunk_sizes = None
+        self.context_length = context_length
+        self.num_captions = num_captions
+        self.text_format = text_format
+        self.tsv_file_list = []
+        self.metas = metas
+        self.label_offsets = self.build_label_offsets()
+        self.prompt_engineering = prompt_engineering
+        self.concat_queries = concat_queries
+
+        if isinstance(image_tsv_file, str) and isinstance(text_tsv_file, str):
+            # single tsv file
+            if (
+                    os.path.splitext(image_tsv_file)[1].lower() == '.tsv'
+                    and os.path.splitext(text_tsv_file)[1].lower() == '.tsv'
+            ):
+                self.tsv_file_list.append((image_tsv_file, text_tsv_file))
+                self.image_tsv_file = TSVFile(
+                    image_tsv_file, if_generate_lineidx=True
+                )
+                self.text_tsv_file = TSVFile(
+                    text_tsv_file, if_generate_lineidx=True
+                )
+            else:
+                raise ValueError("Invalid input! Please check the tsv filenames.")
+        # multiple tsv files specified in a list
+        elif (
+                isinstance(image_tsv_file, list)
+                and isinstance(text_tsv_file, list)
+        ):
+            assert len(image_tsv_file) == len(text_tsv_file), \
+                "Inconsistent number of Image/Text tsv files!"
+            self.tsv_file_list = [
+                (txt, img)
+                for img, txt in zip(image_tsv_file, text_tsv_file)
+            ]
+            self.image_tsv_file = CompositeTSVFile(
+                image_tsv_file,
+                is_train=is_train,
+                sas_token_path=sas_token_path,
+                azcopy_path=azcopy_path
+            )
+            self.text_tsv_file = CompositeTSVFile(
+                text_tsv_file,
+                is_train=is_train,
+                sas_token_path=sas_token_path,
+                azcopy_path=azcopy_path
+            )
+            self._chunk_sizes = self.image_tsv_file.get_chunk_size()
+        else:
+            raise ValueError("Invalid input! Please check the tsv filenames.")
+
+        assert len(self.image_tsv_file) == len(self.text_tsv_file), \
+            "Inconsistent size of Image/Text ({}/{}) data!".format(
+                len(self.image_tsv_file), len(self.text_tsv_file)
+            )
+
+    def build_label_offsets(self):
+        if self.metas is None:
+            return None
+
+        label_offsets = {}
+        offset = 1
+        for meta in self.metas:
+            print(meta)
+            print(label_offsets)
+            label_offsets[meta.source] = offset
+            offset += meta.num_classes
+
+        return label_offsets
+
+    def fetch_blob(self, idx):
+        # image_tsv, text_tsv = self.tsv_file_list[idx]
+        image_tsv = self.image_tsv_file.file_list[idx]
+        text_tsv = self.text_tsv_file.file_list[idx]
+        self.image_tsv_file.blob_storage.fetch_blob(image_tsv)
+        self.text_tsv_file.blob_storage.fetch_blob(text_tsv)
+
+    def get_chunk_sizes(self):
+        return self._chunk_sizes
+
+    def __getitem__(self, index: Union[int, Tuple[int, int]]):
+        if index is None:
+            import torch
+            return torch.tensor([], dtype=torch.float32), \
+                torch.tensor([], dtype=torch.int64), \
+                torch.tensor([], dtype=torch.int64)
+
+        items_image = self.image_tsv_file[index]
+        items_text = self.text_tsv_file[index]
+
+        assert items_text[0] == items_image[0], \
+            'keys do not match for image and text {} vs {}'.format(
+                items_text[0], items_image[0]
+            )
+
+        _, img = self._decode_image(items_image)
+        _, txt, label = self._decode_text(items_text)
+
+        if self.transform:
+            img = self.transform(img)
+
+        tokens = self.tokenize(
+            txt, padding='max_length', truncation=True, max_length=self.context_length,
+            return_tensors='pt'
+        ) if self.tokenize else txt
+
+        tokens['input_ids'].squeeze_()
+        tokens['attention_mask'].squeeze_()
+
+        return img, tokens, label
+
+    def _decode_image(self, items: Tuple[str, str]):
+        key = items[0]
+        image = Image.open(BytesIO(base64.b64decode(items[1]))).convert('RGB')
+
+        return key, image
+
+    def _decode_text(self, items: Tuple[str, Union[str, dict]]):
+        key = items[0]
+        text = ''
+
+        if self.text_format != 'json':
+            raise ValueError('Only support json format')
+
+        # Do some reasonable handing of occasionally bad data.
+        try:
+            js = json.loads(items[1])
+        except Exception as e:
+
+            # empty dictionary
+            js = {}
+
+            # Record the data error in the log.
+            logger.info("JSON parsing error on: " + items[1])
+            logger.info(str(e))
+
+            # do not raise the exception
+            # raise e
+
+            # put some text in and continue processing data (do not kill job)
+            sstr = items[1].find("\"")
+            if (sstr < 0):
+                sstr = 0
+
+            estr = items[1][sstr:].find("\"")
+            if (estr < 0):
+                estr = len(items[1])
+
+            text = items[1][sstr:estr]
+            if (len(text) < 2):
+                text = "A picture showing some content."
+
+            label = 0
+
+        if 'captions' in js:
+            captions = js['captions']
+            if isinstance(captions, list):
+                if self.num_captions == 1:
+                    text = random.choice(captions)
+                else:
+                    text = captions
+                    if len(captions) > self.num_captions:
+                        text = captions[:self.num_captions]
+            elif isinstance(captions, str):
+                text = captions
+            else:
+                raise ValueError('captions should be str or list')
+            label = 0
+        elif 'tags' in js:
+            text = prompt_engineering(js['tags'])
+            label = 0
+        elif 'task' in js and js['task'] == 'classification':
+            if (self.prompt_engineering):
+                text = prompt_engineering(js['class_name'])
+            else:
+                text = js['class_name']
+            label = js['class_id']
+
+            if (self.label_offsets is not None):
+                if (js['source'] in self.label_offsets):
+                    label += self.label_offsets[js['source']]
+
+        if (self.concat_queries):
+            if ('queries' in js) and (len(js['queries']) > 0):
+                q = ''
+                for item in js['queries']:
+                    q = q + item + ' '
+
+                text = q + ', ' + text
+
+        return key, text, label
+
+    def __len__(self):
+        return len(self.image_tsv_file)

MedImageInsight/ImageDataLoader/tsv_file.py

@@ -0,0 +1,290 @@
+import logging
+import gc
+import os
+import os.path as op
+import json
+from typing import List
+from .blob_storage import BlobStorage, disk_usage
+
+logger = logging.getLogger(__name__)
+
+
+def generate_lineidx(filein: str, idxout: str) -> None:
+    idxout_tmp = idxout + '.tmp'
+    with open(filein, 'r') as tsvin, open(idxout_tmp, 'w') as tsvout:
+        fsize = os.fstat(tsvin.fileno()).st_size
+        fpos = 0
+        while fpos != fsize:
+            tsvout.write(str(fpos) + "\n")
+            tsvin.readline()
+            fpos = tsvin.tell()
+    os.rename(idxout_tmp, idxout)
+
+
+def read_to_character(fp, c):
+    result = []
+    while True:
+        s = fp.read(32)
+        assert s != ''
+        if c in s:
+            result.append(s[: s.index(c)])
+            break
+        else:
+            result.append(s)
+    return ''.join(result)
+
+
+class TSVFile(object):
+    def __init__(self,
+                 tsv_file: str,
+                 if_generate_lineidx: bool = True,
+                 lineidx: str = None,
+                 class_selector: List[str] = None,
+                 blob_storage: BlobStorage = None):
+        self.tsv_file = tsv_file
+        self.lineidx = op.splitext(tsv_file)[0] + '.lineidx' \
+            if not lineidx else lineidx
+        self.linelist = op.splitext(tsv_file)[0] + '.linelist'
+        self.chunks = op.splitext(tsv_file)[0] + '.chunks'
+        self._fp = None
+        self._lineidx = None
+        self._sample_indices = None
+        self._class_boundaries = None
+        self._class_selector = class_selector
+        self._blob_storage = blob_storage
+        self._len = None
+        # the process always keeps the process which opens the file.
+        # If the pid is not equal to the currrent pid, we will re-open the file.
+        self.pid = None
+        # generate lineidx if not exist
+        if not op.isfile(self.lineidx) and if_generate_lineidx:
+            generate_lineidx(self.tsv_file, self.lineidx)
+
+    def __del__(self):
+        self.gcidx()
+        if self._fp:
+            self._fp.close()
+            # physically remove the tsv file if it is retrieved by BlobStorage
+            if self._blob_storage and 'azcopy' in self.tsv_file and os.path.exists(self.tsv_file):
+                try:
+                    original_usage = disk_usage('/')
+                    os.remove(self.tsv_file)
+                    logger.info("Purged %s (disk usage: %.2f%% => %.2f%%)" %
+                                 (self.tsv_file, original_usage, disk_usage('/') * 100))
+                except:
+                    # Known issue: multiple threads attempting to delete the file will raise a FileNotFound error.
+                    # TODO: try Threadling.Lock to better handle the race condition
+                    pass
+
+    def __str__(self):
+        return "TSVFile(tsv_file='{}')".format(self.tsv_file)
+
+    def __repr__(self):
+        return str(self)
+
+    def gcidx(self):
+        logger.debug('Run gc collect')
+        self._lineidx = None
+        self._sample_indices = None
+        #self._class_boundaries = None
+        return gc.collect()
+
+    def get_class_boundaries(self):
+        return self._class_boundaries
+
+    def num_rows(self, gcf=False):
+        if (self._len is None):
+            self._ensure_lineidx_loaded()
+            retval = len(self._sample_indices)
+
+            if (gcf):
+                self.gcidx()
+
+            self._len = retval
+
+        return self._len
+
+    def seek(self, idx: int):
+        self._ensure_tsv_opened()
+        self._ensure_lineidx_loaded()
+        try:
+            pos = self._lineidx[self._sample_indices[idx]]
+        except:
+            logger.info('=> {}-{}'.format(self.tsv_file, idx))
+            raise
+        self._fp.seek(pos)
+        return [s.strip() for s in self._fp.readline().split('\t')]
+
+    def seek_first_column(self, idx: int):
+        self._ensure_tsv_opened()
+        self._ensure_lineidx_loaded()
+        pos = self._lineidx[idx]
+        self._fp.seek(pos)
+        return read_to_character(self._fp, '\t')
+
+    def get_key(self, idx: int):
+        return self.seek_first_column(idx)
+
+    def __getitem__(self, index: int):
+        return self.seek(index)
+
+    def __len__(self):
+        return self.num_rows()
+
+    def _ensure_lineidx_loaded(self):
+        if self._lineidx is None:
+            logger.debug('=> loading lineidx: {}'.format(self.lineidx))
+            with open(self.lineidx, 'r') as fp:
+                lines = fp.readlines()
+                lines = [line.strip() for line in lines]
+                self._lineidx = [int(line) for line in lines]
+
+            # read the line list if exists
+            linelist = None
+            if op.isfile(self.linelist):
+                with open(self.linelist, 'r') as fp:
+                    linelist = sorted(
+                        [
+                            int(line.strip())
+                            for line in fp.readlines()
+                        ]
+                    )
+
+            if op.isfile(self.chunks):
+                self._sample_indices = []
+                self._class_boundaries = []
+                class_boundaries = json.load(open(self.chunks, 'r'))
+                for class_name, boundary in class_boundaries.items():
+                    start = len(self._sample_indices)
+                    if class_name in self._class_selector:
+                        for idx in range(boundary[0], boundary[1] + 1):
+                            # NOTE: potentially slow when linelist is long, try to speed it up
+                            if linelist and idx not in linelist:
+                                continue
+                            self._sample_indices.append(idx)
+                    end = len(self._sample_indices)
+                    self._class_boundaries.append((start, end))
+            else:
+                if linelist:
+                    self._sample_indices = linelist
+                else:
+                    self._sample_indices = list(range(len(self._lineidx)))
+
+    def _ensure_tsv_opened(self):
+        if self._fp is None:
+            if self._blob_storage:
+                self._fp = self._blob_storage.open(self.tsv_file)
+            else:
+                self._fp = open(self.tsv_file, 'r')
+            self.pid = os.getpid()
+
+        if self.pid != os.getpid():
+            logger.debug('=> re-open {} because the process id changed'.format(self.tsv_file))
+            self._fp = open(self.tsv_file, 'r')
+            self.pid = os.getpid()
+
+
+class CompositeTSVFile:
+    def __init__(self,
+                 file_list: List[str],
+                 root: str = '.',
+                 class_selector: List[str] = None,
+                 is_train: bool = True,
+                 sas_token_path: str = None,
+                 azcopy_path: str = None):
+        self.root = root
+        self.tsvs = None
+        self.chunk_sizes = None
+        self.accum_chunk_sizes = None
+        self._class_selector = class_selector
+        self._class_boundaries = None
+        self.initialized = False
+        assert isinstance(file_list, list)
+        self.blob_storage = BlobStorage(is_train, sas_token_path, azcopy_path)
+        self.file_list = self.blob_storage.register_local_tsv_paths(file_list)
+        logger.info('=> Init CompositeTSVFile...')
+        self.initialize()
+        logger.info('=> Init CompositeTSVFile Done...')
+
+    def get_key(self, index: int):
+        idx_source, idx_row = self._calc_chunk_idx_row(index)
+        k = self.tsvs[idx_source].get_key(idx_row)
+        return '_'.join([self.file_list[idx_source], k])
+
+    def get_class_boundaries(self):
+        return self._class_boundaries
+
+    def get_chunk_size(self):
+        return self.chunk_sizes
+
+    def num_rows(self):
+        return sum(self.chunk_sizes)
+
+    def _calc_chunk_idx_row(self, index: int):
+        idx_chunk = 0
+        idx_row = index
+        while index >= self.accum_chunk_sizes[idx_chunk]:
+            idx_chunk += 1
+            idx_row = index - self.accum_chunk_sizes[idx_chunk-1]
+        return idx_chunk, idx_row
+
+    def __getitem__(self, index: int):
+        idx_source, idx_row = self._calc_chunk_idx_row(index)
+        if idx_source not in self.blob_storage:
+            self.blob_storage[idx_source] = TSVFile(
+                op.join(self.root, self.file_list[idx_source]),
+                class_selector=self._class_selector,
+                blob_storage=self.blob_storage,
+                if_generate_lineidx=True
+            )
+        return self.blob_storage[idx_source].seek(idx_row)
+
+    def __len__(self):
+        return sum(self.chunk_sizes)
+
+    def initialize(self):
+        """
+        this function has to be called in init function if cache_policy is
+        enabled. Thus, let's always call it in init funciton to make it simple.
+        """
+        if self.initialized:
+            return
+        self.tsvs = [
+            TSVFile(
+                op.join(self.root, f),
+                class_selector=self._class_selector
+            ) for f in self.file_list
+        ]
+        logger.debug("=> Calculating chunk sizes ...")
+        self.chunk_sizes = [tsv.num_rows(gcf=True) for tsv in self.tsvs]
+
+        self.accum_chunk_sizes = [0]
+        for size in self.chunk_sizes:
+            self.accum_chunk_sizes += [self.accum_chunk_sizes[-1] + size]
+        self.accum_chunk_sizes = self.accum_chunk_sizes[1:]
+
+        if (
+            self._class_selector
+            and all([tsv.get_class_boundaries() for tsv in self.tsvs])
+        ):
+            """
+            Note: When using CompositeTSVFile, make sure that the classes contained in each
+            tsv file do not overlap. Otherwise, the class boundaries won't be correct.
+            """
+            self._class_boundaries = []
+            offset = 0
+            for tsv in self.tsvs:
+                boundaries = tsv.get_class_boundaries()
+                for bound in boundaries:
+                    self._class_boundaries.append((bound[0] + offset, bound[1] + offset))
+                offset += len(tsv)
+        self.initialized = True
+
+
+def load_list_file(fname: str) -> List[str]:
+    with open(fname, 'r') as fp:
+        lines = fp.readlines()
+    result = [line.strip() for line in lines]
+    if len(result) > 0 and result[-1] == '':
+        result = result[:-1]
+    return result

MedImageInsight/ImageDataLoader/zipdata.py

@@ -0,0 +1,98 @@
+import os.path as op
+from zipfile import ZipFile, BadZipFile
+import torch.utils.data as data
+from PIL import Image
+from io import BytesIO
+import multiprocessing
+
+_VALID_IMAGE_TYPES = ['.jpg', '.jpeg', '.tiff', '.bmp', '.png']
+
+
+class ZipData(data.Dataset):
+    _IGNORE_ATTRS = {'_zip_file'}
+
+    def __init__(self, path, map_file,
+                 transform=None, target_transform=None,
+                 extensions=None):
+        self._path = path
+        if not extensions:
+            extensions = _VALID_IMAGE_TYPES
+        self._zip_file = ZipFile(path)
+        self.zip_dict = {}
+        self.samples = []
+        self.transform = transform
+        self.target_transform = target_transform
+        self.class_to_idx = {}
+        with open(map_file, 'r') as f:
+            for line in iter(f.readline, ""):
+                line = line.strip()
+                if not line:
+                    continue
+                cls_idx = [l for l in line.split('\t') if l]
+                if not cls_idx:
+                    continue
+                if (len(cls_idx) < 2):
+                    cls_idx = [l for l in line.split(' ') if l]
+                    if not cls_idx:
+                        continue
+                assert len(cls_idx) >= 2, "invalid line: {}".format(line)
+                idx = int(cls_idx[1])
+                cls = cls_idx[0]
+                del cls_idx
+                at_idx = cls.find('@')
+                assert at_idx >= 0, "invalid class: {}".format(cls)
+                cls = cls[at_idx + 1:]
+                if cls.startswith('/'):
+                    # Python ZipFile expects no root
+                    cls = cls[1:]
+                assert cls, "invalid class in line {}".format(line)
+                prev_idx = self.class_to_idx.get(cls)
+                assert prev_idx is None or prev_idx == idx, "class: {} idx: {} previously had idx: {}".format(
+                    cls, idx, prev_idx
+                )
+                self.class_to_idx[cls] = idx
+
+        for fst in self._zip_file.infolist():
+            fname = fst.filename
+            target = self.class_to_idx.get(fname)
+            if target is None:
+                continue
+            if fname.endswith('/') or fname.startswith('.') or fst.file_size == 0:
+                continue
+            ext = op.splitext(fname)[1].lower()
+            if ext in extensions:
+                self.samples.append((fname, target))
+        assert len(self), "No images found in: {} with map: {}".format(self._path, map_file)
+
+    def __repr__(self):
+        return 'ZipData({}, size={})'.format(self._path, len(self))
+
+    def __getstate__(self):
+        return {
+            key: val if key not in self._IGNORE_ATTRS else None
+            for key, val in self.__dict__.iteritems()
+        }
+
+    def __getitem__(self, index):
+        proc = multiprocessing.current_process()
+        pid = proc.pid # get pid of this process.
+        if pid not in self.zip_dict:
+            self.zip_dict[pid] = ZipFile(self._path)
+        zip_file = self.zip_dict[pid]
+
+        if index >= len(self) or index < 0:
+            raise KeyError("{} is invalid".format(index))
+        path, target = self.samples[index]
+        try:
+            sample = Image.open(BytesIO(zip_file.read(path))).convert('RGB')
+        except BadZipFile:
+            print("bad zip file")
+            return None, None
+        if self.transform is not None:
+            sample = self.transform(sample)
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+        return sample, target
+
+    def __len__(self):
+        return len(self.samples)

MedImageInsight/ImageEncoder/__init__.py

@@ -0,0 +1,8 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from .build import build_image_encoder
+
+from .coswin import *
+from .davit_v1 import *

MedImageInsight/ImageEncoder/build.py

@@ -0,0 +1,13 @@
+from .registry import image_encoders
+from .registry import is_image_encoder
+
+
+def build_image_encoder(config_encoder, verbose, **kwargs):
+    model_name = config_encoder['NAME']
+    if model_name.startswith('cls_'):
+        model_name = model_name[4:]
+
+    if not is_image_encoder(model_name):
+        raise ValueError(f'Unkown model: {model_name}')
+
+    return image_encoders(model_name)(config_encoder, verbose, **kwargs)

MedImageInsight/ImageEncoder/coswin.py

@@ -0,0 +1,779 @@
+# --------------------------------------------------------
+# CoSwin: Convolutional Swin Transformer
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ze Liu
+# Modified by Bin Xiao
+# --------------------------------------------------------
+
+import logging
+import os
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+import numpy as np
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+from .registry import register_image_encoder
+
+logger = logging.getLogger(__name__)
+
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
+
+    def flops(self, N):
+        # calculate flops for 1 window with token length of N
+        flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        return flops
+
+
+class SwinTransformerBlock(nn.Module):
+    r""" Swin Transformer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, layer_scale=False):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if self.shift_size > 0:
+            # calculate attention mask for SW-MSA
+            H, W = self.input_resolution
+            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
+            h_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            w_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            cnt = 0
+            for h in h_slices:
+                for w in w_slices:
+                    img_mask[:, h, w, :] = cnt
+                    cnt += 1
+
+            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        else:
+            attn_mask = None
+
+        self.gamma = 1.0
+        if layer_scale:
+            logger.info('=> enable layer scale')
+            self.gamma = nn.Parameter(
+                1e-4*torch.ones(dim), requires_grad=True
+            )
+
+        self.register_buffer("attn_mask", attn_mask)
+
+    def forward(self, x):
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+        else:
+            shifted_x = x
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(self.gamma*x)
+        x = x + self.drop_path(self.gamma*self.mlp(self.norm2(x)))
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # W-MSA/SW-MSA
+        nW = H * W / self.window_size / self.window_size
+        flops += nW * self.attn.flops(self.window_size * self.window_size)
+        # mlp
+        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
+        # norm2
+        flops += self.dim * H * W
+        return flops
+
+
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = x.view(B, H, W, C)
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = H * W * self.dim
+        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
+        return flops
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None,
+                 use_checkpoint=False, layer_scale=False):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(
+                dim=dim, input_resolution=input_resolution,
+                num_heads=num_heads, window_size=window_size,
+                shift_size=0 if (i % 2 == 0) else window_size // 2,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop, attn_drop=attn_drop,
+                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                norm_layer=norm_layer,
+                layer_scale=layer_scale
+            )
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            # self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
+            self.downsample = downsample(
+                input_resolution=input_resolution, patch_size=3, in_chans=dim, embed_dim=dim*2,
+                stride=2, padding=1, norm_layer=norm_layer
+            )
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+    def flops(self):
+        flops = 0
+        for blk in self.blocks:
+            flops += blk.flops()
+        if self.downsample is not None:
+            flops += self.downsample.flops()
+        return flops
+
+
+class PatchEmbed(nn.Module):
+    r""" Image to Patch Embedding
+
+    Args:
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+    def flops(self):
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class ConvEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(
+        self,
+        input_resolution=(224,224),
+        patch_size=7,
+        in_chans=3,
+        embed_dim=64,
+        stride=4,
+        padding=2,
+        norm_layer=None
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.input_resolution = input_resolution
+
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim,
+            kernel_size=patch_size,
+            stride=stride,
+            padding=padding
+        )
+        self.norm = norm_layer(embed_dim) if norm_layer else None
+
+    def forward(self, x):
+        if len(x.size()) == 3:
+            x = rearrange(
+                x, 'b (h w) c -> b c h w',
+                h=self.input_resolution[0],
+                w=self.input_resolution[1]
+            )
+
+        x = self.proj(x)
+
+        B, C, H, W = x.shape
+        x = rearrange(x, 'b c h w -> b (h w) c')
+        if self.norm:
+            x = self.norm(x)
+
+        return x
+
+
+class SwinTransformer(nn.Module):
+    r""" Swin Transformer
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+
+    Args:
+        img_size (int | tuple(int)): Input image size. Default 224
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each Swin Transformer layer.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        window_size (int): Window size. Default: 7
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
+    """
+
+    def __init__(self, img_size=224, patch_size=7, patch_padding=2, patch_stride=4, in_chans=3,
+                 num_classes=1000, embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
+                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
+                 use_checkpoint=False, layer_scale=False, **kwargs):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
+        self.mlp_ratio = mlp_ratio
+
+        # split image into non-overlapping patches
+        # self.patch_embed = PatchEmbed(
+        #     img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+        #     norm_layer=norm_layer if self.patch_norm else None)
+
+        self.patch_embed = ConvEmbed(
+            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, padding=patch_padding,
+            norm_layer=norm_layer if self.patch_norm else None
+        )
+
+        img_size = to_2tuple(img_size)
+        patches_resolution = (
+            int(np.floor(float(img_size[0]+2*patch_padding-patch_size)/patch_stride+1)),
+            int(np.floor(float(img_size[0]+2*patch_padding-patch_size)/patch_stride+1))
+        )
+        num_patches = patches_resolution[0] * patches_resolution[1]
+        # num_patches = self.patch_embed.num_patches
+        # patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        # absolute position embedding
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(
+                dim=int(embed_dim * 2 ** i_layer),
+                input_resolution=(
+                    patches_resolution[0] // (2 ** i_layer),
+                    patches_resolution[1] // (2 ** i_layer)
+                ),
+                depth=depths[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                mlp_ratio=self.mlp_ratio,
+                qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                norm_layer=norm_layer,
+                # downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                downsample=ConvEmbed if (i_layer < self.num_layers - 1) else None,
+                use_checkpoint=use_checkpoint,
+                layer_scale=layer_scale
+            )
+            self.layers.append(layer)
+
+        self.norm = norm_layer(self.num_features)
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.apply(self._init_weights)
+
+    @property
+    def dim_out(self):
+        return self.num_features
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def from_pretrained(self, pretrained='', pretrained_layers=[], verbose=True):
+        if os.path.isfile(pretrained):
+            logging.info(f'=> loading pretrained model {pretrained}')
+            pretrained_dict = torch.load(pretrained, map_location='cpu')
+
+            self.from_state_dict(pretrained_dict, pretrained_layers, verbose)
+
+    def from_state_dict(self, pretrained_dict, pretrained_layers=[], verbose=True):
+        model_dict = self.state_dict()
+        stripped_key = lambda x: x[14:] if x.startswith('image_encoder.') else x
+
+        pretrained_dict = {
+            stripped_key(k): v for k, v in pretrained_dict.items()
+            if stripped_key(k) in model_dict.keys()
+        }
+        need_init_state_dict = {}
+        for k, v in pretrained_dict.items():
+            need_init = (
+                (
+                    k.split('.')[0] in pretrained_layers
+                    or pretrained_layers[0] == '*'
+                )
+                and 'relative_position_index' not in k
+                and 'attn_mask' not in k
+            )
+
+            if need_init:
+                if verbose:
+                    logger.info(f'=> init {k} from pretrained state dict')
+
+                if 'relative_position_bias_table' in k and v.size() != model_dict[k].size():
+                    relative_position_bias_table_pretrained = v
+                    relative_position_bias_table_current = model_dict[k]
+                    L1, nH1 = relative_position_bias_table_pretrained.size()
+                    L2, nH2 = relative_position_bias_table_current.size()
+                    if nH1 != nH2:
+                        logger.info(f"Error in loading {k}, passing")
+                    else:
+                        if L1 != L2:
+                            logger.info(
+                                '=> load_pretrained: resized variant: {} to {}'
+                                .format((L1, nH1), (L2, nH2))
+                            )
+                            S1 = int(L1 ** 0.5)
+                            S2 = int(L2 ** 0.5)
+                            relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate(
+                                relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1),
+                                size=(S2, S2),
+                                mode='bicubic')
+                            v = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0)
+
+                if 'absolute_pos_embed' in k and v.size() != model_dict[k].size():
+                    absolute_pos_embed_pretrained = v
+                    absolute_pos_embed_current = model_dict[k]
+                    _, L1, C1 = absolute_pos_embed_pretrained.size()
+                    _, L2, C2 = absolute_pos_embed_current.size()
+                    if C1 != C1:
+                        logger.info(f"Error in loading {k}, passing")
+                    else:
+                        if L1 != L2:
+                            logger.info(
+                                '=> load_pretrained: resized variant: {} to {}'
+                                    .format((1, L1, C1), (1, L2, C2))
+                            )
+                            S1 = int(L1 ** 0.5)
+                            S2 = int(L2 ** 0.5)
+                            absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape(-1, S1, S1, C1)
+                            absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute(0, 3, 1, 2)
+                            absolute_pos_embed_pretrained_resized = torch.nn.functional.interpolate(
+                                absolute_pos_embed_pretrained, size=(S2, S2), mode='bicubic')
+                            v = absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1).flatten(1, 2)
+
+                need_init_state_dict[k] = v
+        self.load_state_dict(need_init_state_dict, strict=False)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        x = self.norm(x)  # B L C
+        x = self.avgpool(x.transpose(1, 2))  # B C 1
+        x = torch.flatten(x, 1)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+
+@register_image_encoder
+def image_encoder(config_encoder, verbose, **kwargs):
+    spec = config_encoder['SPEC']
+
+    coswin = SwinTransformer(
+        img_size=config_encoder['IMAGE_SIZE'],
+        patch_size=spec['PATCH_SIZE'],
+        patch_padding=spec['PATCH_PADDING'],
+        patch_stride=spec['PATCH_STRIDE'],
+        in_chans=spec['IN_CHANS'],
+        num_classes=0,
+        embed_dim=spec['EMBED_DIM'],
+        depths=spec['DEPTHS'],
+        num_heads=spec['NUM_HEADS'],
+        window_size=spec['WINDOW_SIZE'],
+        mlp_ratio=spec['MLP_RATIO'],
+        qkv_bias=spec['QKV_BIAS'],
+        qk_scale=spec.get('QK_SCALE', None),
+        drop_rate=spec['DROP_RATE'],
+        drop_path_rate=spec['DROP_PATH_RATE'],
+        ape=spec['APE'],
+        patch_norm=spec['PATCH_NORM'],
+        layer_scale=spec.get('LAYER_SCALE', False),
+        use_checkpoint=spec.get('ENABLE_CHECKPOINT', False)
+    )
+
+    if config_encoder['LOAD_PRETRAINED']:
+        coswin.from_pretrained(
+            config_encoder['PRETRAINED'],
+            config_encoder['PRETRAINED_LAYERS'],
+            verbose
+        )
+
+    return coswin

MedImageInsight/ImageEncoder/davit_v1.py

@@ -0,0 +1,727 @@
+import logging
+import os
+import copy
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from collections import OrderedDict
+
+from einops import rearrange
+from timm.models.layers import DropPath, trunc_normal_
+
+# helper methods
+from .registry import register_image_encoder
+
+import mup.init
+from mup import MuReadout, set_base_shapes
+
+logger = logging.getLogger(__name__)
+
+
+class MySequential(nn.Sequential):
+    def forward(self, *inputs):
+        for module in self._modules.values():
+            if type(inputs) == tuple:
+                inputs = module(*inputs)
+            else:
+                inputs = module(inputs)
+        return inputs
+
+
+class PreNorm(nn.Module):
+    def __init__(self, norm, fn, drop_path=None):
+        super().__init__()
+        self.norm = norm
+        self.fn = fn
+        self.drop_path = drop_path
+
+    def forward(self, x, *args, **kwargs):
+        shortcut = x
+        if self.norm != None:
+            x, size = self.fn(self.norm(x), *args, **kwargs)
+        else:
+            x, size = self.fn(x, *args, **kwargs)
+
+        if self.drop_path:
+            x = self.drop_path(x)
+
+        x = shortcut + x
+
+        return x, size
+
+
+class Mlp(nn.Module):
+    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+    """
+
+    def __init__(
+            self,
+            in_features,
+            hidden_features=None,
+            out_features=None,
+            act_layer=nn.GELU,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.net = nn.Sequential(OrderedDict([
+            ("fc1", nn.Linear(in_features, hidden_features)),
+            ("act", act_layer()),
+            ("fc2", nn.Linear(hidden_features, out_features))
+        ]))
+
+    def forward(self, x, size):
+        return self.net(x), size
+
+
+class DepthWiseConv2d(nn.Module):
+    def __init__(
+            self,
+            dim_in,
+            kernel_size,
+            padding,
+            stride,
+            bias=True,
+    ):
+        super().__init__()
+        self.dw = nn.Conv2d(
+            dim_in, dim_in,
+            kernel_size=kernel_size,
+            padding=padding,
+            groups=dim_in,
+            stride=stride,
+            bias=bias
+        )
+
+    def forward(self, x, size):
+        B, N, C = x.shape
+        H, W = size
+        assert N == H * W
+
+        x = self.dw(x.transpose(1, 2).view(B, C, H, W))
+        size = (x.size(-2), x.size(-1))
+        x = x.flatten(2).transpose(1, 2)
+        return x, size
+
+
+class ConvEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(
+            self,
+            patch_size=7,
+            in_chans=3,
+            embed_dim=64,
+            stride=4,
+            padding=2,
+            norm_layer=None,
+            pre_norm=True
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim,
+            kernel_size=patch_size,
+            stride=stride,
+            padding=padding
+        )
+
+        dim_norm = in_chans if pre_norm else embed_dim
+        self.norm = norm_layer(dim_norm) if norm_layer else None
+
+        self.pre_norm = pre_norm
+
+    def forward(self, x, size):
+        H, W = size
+        if len(x.size()) == 3:
+            if self.norm and self.pre_norm:
+                x = self.norm(x)
+            x = rearrange(
+                x, 'b (h w) c -> b c h w',
+                h=H, w=W
+            )
+
+        x = self.proj(x)
+
+        _, _, H, W = x.shape
+        x = rearrange(x, 'b c h w -> b (h w) c')
+        if self.norm and not self.pre_norm:
+            x = self.norm(x)
+
+        return x, (H, W)
+
+
+class ChannelAttention(nn.Module):
+
+    def __init__(self, dim, base_dim, groups=8, base_groups=8, qkv_bias=True, dynamic_scale=True, standparam=True):
+        super().__init__()
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+        self.dynamic_scale = dynamic_scale
+
+        self.dim = dim
+        self.groups = groups
+        self.group_dim = dim // groups
+
+        self.base_dim = base_dim
+        self.base_groups = base_groups
+        self.base_group_dim = base_dim // base_groups
+
+        self.group_wm = self.group_dim / self.base_group_dim  # Width multiplier for each group.
+        self.standparam = standparam
+
+    def forward(self, x, size):
+        B, N, C = x.shape
+        assert C == self.dim
+
+        qkv = self.qkv(x).reshape(B, N, 3, self.groups, C // self.groups).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # Shape: [B, groups, N, group_dim].
+
+        scale = N ** -0.5 if self.dynamic_scale else self.dim ** -0.5
+
+        # Change the scaling factor.
+        # Ref: examples/Transformer/model.py in muP.
+        # Note: We consider backward compatiblity and follow https://github.com/microsoft/mup/issues/18.
+        if self.standparam:
+            scale = N ** -0.5 if self.dynamic_scale else self.dim ** -0.5
+        else:
+            assert self.dynamic_scale  # Currently only support dynamic scale.
+            scale = N ** -0.5
+
+        q = q * scale
+        attention = q.transpose(-1, -2) @ k
+        attention = attention.softmax(dim=-1)
+
+        if not self.standparam:
+            # Follow https://github.com/microsoft/mup/issues/18.
+            attention = attention / self.group_wm
+
+        x = (attention @ v.transpose(-1, -2)).transpose(-1, -2)
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        return x, size
+
+
+class ChannelBlock(nn.Module):
+
+    def __init__(self, dim, base_dim, groups, base_groups, mlp_ratio=4., qkv_bias=True,
+                 drop_path_rate=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 conv_at_attn=True, conv_at_ffn=True, dynamic_scale=True, standparam=True):
+        super().__init__()
+
+        drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+
+        self.conv1 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_attn else None
+        self.channel_attn = PreNorm(
+            norm_layer(dim),
+            ChannelAttention(dim, base_dim, groups=groups, base_groups=base_groups, qkv_bias=qkv_bias,
+                             dynamic_scale=dynamic_scale, standparam=standparam),
+            drop_path
+        )
+        self.conv2 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_ffn else None
+        self.ffn = PreNorm(
+            norm_layer(dim),
+            Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer),
+            drop_path
+        )
+
+    def forward(self, x, size):
+        if self.conv1:
+            x, size = self.conv1(x, size)
+        x, size = self.channel_attn(x, size)
+
+        if self.conv2:
+            x, size = self.conv2(x, size)
+        x, size = self.ffn(x, size)
+
+        return x, size
+
+
+def window_partition(x, window_size: int):
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size: int, H: int, W: int):
+    B = windows.shape[0] // (H * W // window_size // window_size)
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+
+    def __init__(self, dim, base_dim, num_heads, base_num_heads, window_size, qkv_bias=True, standparam=True):
+
+        super().__init__()
+
+        self.window_size = window_size
+
+        self.dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.base_dim = base_dim
+        self.base_num_heads = base_num_heads
+        base_head_dim = base_dim // base_num_heads
+
+        # Change the scaling factor.
+        # Ref: examples/Transformer/model.py in muP.
+        # Note: We consider backward compatiblity and follow https://github.com/microsoft/mup/issues/17.
+        if standparam:
+            scale = float(head_dim) ** -0.5
+        else:
+            # TODO: Here we ensure backward compatibility, which may not be optimal.
+            #       We may add an argument called backward_comp. If it is set as False, we use
+            #          float(head_dim) ** -1 * math.sqrt(attn_mult)
+            #       as in the Transformer example in muP.
+            base_scale = float(base_head_dim) ** -0.5  # The same as scaling in standard parametrization.
+            head_wm = head_dim / base_head_dim  # Width multiplier for each head.
+            scale = base_scale / head_wm
+            # scale_1 = (float(base_head_dim) ** 0.5) * (float(head_dim) ** -1) # Equivalent implementation as shown in the muP paper.
+            # assert np.isclose(scale, scale_1)
+        self.scale = scale
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, size):
+
+        H, W = size
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        x = x.view(B, H, W, C)
+
+        pad_l = pad_t = 0
+        pad_r = (self.window_size - W % self.window_size) % self.window_size
+        pad_b = (self.window_size - H % self.window_size) % self.window_size
+        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
+        _, Hp, Wp, _ = x.shape
+
+        x = window_partition(x, self.window_size)
+        x = x.view(-1, self.window_size * self.window_size, C)
+
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+        attn = self.softmax(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+
+        # merge windows
+        x = x.view(
+            -1, self.window_size, self.window_size, C
+        )
+        x = window_reverse(x, self.window_size, Hp, Wp)
+
+        if pad_r > 0 or pad_b > 0:
+            x = x[:, :H, :W, :].contiguous()
+
+        x = x.view(B, H * W, C)
+
+        return x, size
+
+
+class SpatialBlock(nn.Module):
+
+    def __init__(self, dim, base_dim, num_heads, base_num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, drop_path_rate=0., act_layer=nn.GELU,
+                 norm_layer=nn.LayerNorm, conv_at_attn=True, conv_at_ffn=True, standparam=True):
+        super().__init__()
+
+        drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+
+        self.conv1 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_attn else None
+        self.window_attn = PreNorm(
+            norm_layer(dim),
+            WindowAttention(dim, base_dim, num_heads, base_num_heads, window_size, qkv_bias=qkv_bias,
+                            standparam=standparam),
+            drop_path
+        )
+        self.conv2 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_ffn else None
+        self.ffn = PreNorm(
+            norm_layer(dim),
+            Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer),
+            drop_path
+        )
+
+    def forward(self, x, size):
+        if self.conv1:
+            x, size = self.conv1(x, size)
+        x, size = self.window_attn(x, size)
+
+        if self.conv2:
+            x, size = self.conv2(x, size)
+        x, size = self.ffn(x, size)
+        return x, size
+
+
+class DaViT(nn.Module):
+    """ DaViT: Dual-Attention Transformer
+
+    Args:
+        img_size (int | tuple(int)): Input image size. Default: 224
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of spatial and channel blocks in different stages. Default: (1, 1, 3, 1)
+        patch_size (tuple(int)): Patch sizes in different stages. Default: (7, 2, 2, 2)
+        patch_stride (tuple(int)): Patch strides in different stages. Default: (4, 2, 2, 2)
+        patch_padding (tuple(int)): Patch padding sizes in different stages. Default: (3, 0, 0, 0)
+        patch_prenorm (tuple(bool)): Use pre-normalization or not in different stages. Default: (False, False, False, False)
+        embed_dims (tuple(int)): Patch embedding dimension. Default: (64, 128, 192, 256)
+        base_embed_dims (tuple(int)): Patch embedding dimension (base case for muP). Default: (64, 128, 192, 256)
+        num_heads (tuple(int)): Number of attention heads in different layers. Default: (4, 8, 12, 16)
+        base_num_heads (tuple(int)): Number of attention heads in different layers (base case for muP). Default: (4, 8, 12, 16)
+        num_groups (tuple(int)): Number of groups in channel attention in different layers. Default: (3, 6, 12, 24)
+        base_num_groups (tuple(int)): Number of groups in channel attention in different layers (base case for muP). Default: (3, 6, 12, 24)
+        window_size (int): Window size. Default: 7
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        enable_checkpoint (bool): If True, enabling checkpoint. Default: False
+        conv_at_attn (bool): If True, add convolution layer before attention. Default: True
+        conv_at_ffn (bool): If True, add convolution layer before ffn. Default: True
+        dynamic_scale (bool): If True, scale of channel attention is respect to the number of tokens. Default: True
+        standparam (bool): Use standard parametrization or mu-parametrization. Default: True (i.e., use standard paramerization)
+    """
+
+    def __init__(
+            self,
+            img_size=224,
+            in_chans=3,
+            num_classes=1000,
+            depths=(1, 1, 3, 1),
+            patch_size=(7, 2, 2, 2),
+            patch_stride=(4, 2, 2, 2),
+            patch_padding=(3, 0, 0, 0),
+            patch_prenorm=(False, False, False, False),
+            embed_dims=(64, 128, 192, 256),
+            base_embed_dims=(64, 128, 192, 256),
+            num_heads=(3, 6, 12, 24),
+            base_num_heads=(3, 6, 12, 24),
+            num_groups=(3, 6, 12, 24),
+            base_num_groups=(3, 6, 12, 24),
+            window_size=7,
+            mlp_ratio=4.,
+            qkv_bias=True,
+            drop_path_rate=0.1,
+            norm_layer=nn.LayerNorm,
+            enable_checkpoint=False,
+            conv_at_attn=True,
+            conv_at_ffn=True,
+            dynamic_scale=True,
+            standparam=True
+    ):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.num_groups = num_groups
+        self.num_stages = len(self.embed_dims)
+        self.enable_checkpoint = enable_checkpoint
+        assert self.num_stages == len(self.num_heads) == len(self.num_groups)
+
+        num_stages = len(embed_dims)
+        self.img_size = img_size
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths) * 2)]
+
+        depth_offset = 0
+        convs = []
+        blocks = []
+        for i in range(num_stages):
+            conv_embed = ConvEmbed(
+                patch_size=patch_size[i],
+                stride=patch_stride[i],
+                padding=patch_padding[i],
+                in_chans=in_chans if i == 0 else self.embed_dims[i - 1],
+                embed_dim=self.embed_dims[i],
+                norm_layer=norm_layer,
+                pre_norm=patch_prenorm[i]
+            )
+            convs.append(conv_embed)
+
+            logger.info(f'=> Depth offset in stage {i}: {depth_offset}')
+            block = MySequential(
+                *[
+                    MySequential(OrderedDict([
+                        (
+                            'spatial_block', SpatialBlock(
+                                embed_dims[i],
+                                base_embed_dims[i],
+                                num_heads[i],
+                                base_num_heads[i],
+                                window_size,
+                                drop_path_rate=dpr[depth_offset + j * 2],
+                                qkv_bias=qkv_bias,
+                                mlp_ratio=mlp_ratio,
+                                conv_at_attn=conv_at_attn,
+                                conv_at_ffn=conv_at_ffn,
+                                standparam=standparam
+                            )
+                        ),
+                        (
+                            'channel_block', ChannelBlock(
+                                embed_dims[i],
+                                base_embed_dims[i],
+                                num_groups[i],
+                                base_num_groups[i],
+                                drop_path_rate=dpr[depth_offset + j * 2 + 1],
+                                qkv_bias=qkv_bias,
+                                mlp_ratio=mlp_ratio,
+                                conv_at_attn=conv_at_attn,
+                                conv_at_ffn=conv_at_ffn,
+                                dynamic_scale=dynamic_scale,
+                                standparam=standparam
+                            )
+                        )
+                    ])) for j in range(depths[i])
+                ]
+            )
+            blocks.append(block)
+            depth_offset += depths[i] * 2
+
+        self.convs = nn.ModuleList(convs)
+        self.blocks = nn.ModuleList(blocks)
+
+        self.norms = norm_layer(self.embed_dims[-1])
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+
+        if standparam:
+            self.head = nn.Linear(self.embed_dims[-1], num_classes) if num_classes > 0 else nn.Identity()
+        else:
+            self.head = MuReadout(self.embed_dims[-1], num_classes,
+                                  readout_zero_init=True)  # Follow examples/ResNet/resnet.py in muP.
+
+        if torch.cuda.is_available():
+            self.device = torch.device(type="cuda", index=0)
+        else:
+            self.device = torch.device(type="cpu")
+
+    def custom_init_weights(self, use_original_init=True):
+        self.use_original_init = use_original_init
+        logger.info('Custom init: {}'.format('original init' if self.use_original_init else 'muP init'))
+        self.apply(self._custom_init_weights)
+
+    @property
+    def dim_out(self):
+        return self.embed_dims[-1]
+
+    def _custom_init_weights(self, m):
+        # Customized initialization for weights.
+        if self.use_original_init:
+            # Original initialization.
+            # Note: This is not SP init. We do not implement SP init here.
+            custom_trunc_normal_ = trunc_normal_
+            custom_normal_ = nn.init.normal_
+        else:
+            # muP.
+            custom_trunc_normal_ = mup.init.trunc_normal_
+            custom_normal_ = mup.init.normal_
+
+        # These initializations will overwrite the existing inializations from the modules and adjusted by set_base_shapes().
+        if isinstance(m, MuReadout):
+            pass  # Note: MuReadout is already zero initialized due to readout_zero_init=True.
+        elif isinstance(m, nn.Linear):
+            custom_trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.Conv2d):
+            custom_normal_(m.weight, std=0.02)
+            for name, _ in m.named_parameters():
+                if name in ['bias']:
+                    nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):  # Follow P24 Layernorm Weights and Biases.
+            nn.init.constant_(m.weight, 1.0)
+            nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.BatchNorm2d):  # Follow P24 Layernorm Weights and Biases.
+            nn.init.constant_(m.weight, 1.0)
+            nn.init.constant_(m.bias, 0)
+
+    def _try_remap_keys(self, pretrained_dict):
+        remap_keys = {
+            "conv_embeds": "convs",
+            "main_blocks": "blocks",
+            "0.cpe.0.proj": "spatial_block.conv1.fn.dw",
+            "0.attn": "spatial_block.window_attn.fn",
+            "0.cpe.1.proj": "spatial_block.conv2.fn.dw",
+            "0.mlp": "spatial_block.ffn.fn.net",
+            "1.cpe.0.proj": "channel_block.conv1.fn.dw",
+            "1.attn": "channel_block.channel_attn.fn",
+            "1.cpe.1.proj": "channel_block.conv2.fn.dw",
+            "1.mlp": "channel_block.ffn.fn.net",
+            "0.norm1": "spatial_block.window_attn.norm",
+            "0.norm2": "spatial_block.ffn.norm",
+            "1.norm1": "channel_block.channel_attn.norm",
+            "1.norm2": "channel_block.ffn.norm"
+        }
+
+        full_key_mappings = {}
+        for k in pretrained_dict.keys():
+            old_k = k
+            for remap_key in remap_keys.keys():
+                if remap_key in k:
+                    logger.info(f'=> Repace {remap_key} with {remap_keys[remap_key]}')
+                    k = k.replace(remap_key, remap_keys[remap_key])
+
+            full_key_mappings[old_k] = k
+
+        return full_key_mappings
+
+    def from_state_dict(self, pretrained_dict, pretrained_layers=[], verbose=True):
+        model_dict = self.state_dict()
+        stripped_key = lambda x: x[14:] if x.startswith('image_encoder.') else x
+        full_key_mappings = self._try_remap_keys(pretrained_dict)
+
+        pretrained_dict = {
+            stripped_key(full_key_mappings[k]): v.to(self.device) for k, v in pretrained_dict.items()
+            if stripped_key(full_key_mappings[k]) in model_dict.keys()
+        }
+        need_init_state_dict = {}
+        for k, v in pretrained_dict.items():
+            need_init = (
+                    k.split('.')[0] in pretrained_layers
+                    or pretrained_layers[0] == '*'
+            )
+            if need_init:
+                if verbose:
+                    logger.info(f'=> init {k} from pretrained state dict')
+
+                need_init_state_dict[k] = v.to(self.device)
+        self.load_state_dict(need_init_state_dict, strict=False)
+
+    def from_pretrained(self, pretrained='', pretrained_layers=[], verbose=True):
+        if os.path.isfile(pretrained):
+            logger.info(f'=> loading pretrained model {pretrained}')
+            pretrained_dict = torch.load(pretrained, map_location='cpu')
+
+            self.from_state_dict(pretrained_dict, pretrained_layers, verbose)
+
+    def forward_features(self, x):
+        input_size = (x.size(2), x.size(3))
+        for conv, block in zip(self.convs, self.blocks):
+            x, input_size = conv(x, input_size)
+            if self.enable_checkpoint:
+                x, input_size = checkpoint.checkpoint(block, x, input_size)
+            else:
+                x, input_size = block(x, input_size)
+
+        x = self.avgpool(x.transpose(1, 2))
+        x = torch.flatten(x, 1)
+        x = self.norms(x)
+
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+
+def create_encoder(config_encoder):
+    spec = config_encoder['SPEC']
+    standparam = spec.get('STANDPARAM', True)
+
+    if standparam:
+        # Dummy values for muP parameters.
+        base_embed_dims = spec['DIM_EMBED']
+        base_num_heads = spec['NUM_HEADS']
+        base_num_groups = spec['NUM_GROUPS']
+    else:
+        base_embed_dims = spec['BASE_DIM_EMBED']
+        base_num_heads = spec['BASE_NUM_HEADS']
+        base_num_groups = spec['BASE_NUM_GROUPS']
+
+    davit = DaViT(
+        num_classes=config_encoder['NUM_CLASSES'],
+        depths=spec['DEPTHS'],
+        embed_dims=spec['DIM_EMBED'],
+        base_embed_dims=base_embed_dims,
+        num_heads=spec['NUM_HEADS'],
+        base_num_heads=base_num_heads,
+        num_groups=spec['NUM_GROUPS'],
+        base_num_groups=base_num_groups,
+        patch_size=spec['PATCH_SIZE'],
+        patch_stride=spec['PATCH_STRIDE'],
+        patch_padding=spec['PATCH_PADDING'],
+        patch_prenorm=spec['PATCH_PRENORM'],
+        drop_path_rate=spec['DROP_PATH_RATE'],
+        img_size=config_encoder['IMAGE_SIZE'],
+        window_size=spec.get('WINDOW_SIZE', 7),
+        enable_checkpoint=spec.get('ENABLE_CHECKPOINT', False),
+        conv_at_attn=spec.get('CONV_AT_ATTN', True),
+        conv_at_ffn=spec.get('CONV_AT_FFN', True),
+        dynamic_scale=spec.get('DYNAMIC_SCALE', True),
+        standparam=standparam,
+    )
+    return davit
+
+
+def create_mup_encoder(config_encoder):
+    def gen_config(config, wm):
+        new_config = copy.deepcopy(config)
+        for name in ['DIM_EMBED', 'NUM_HEADS', 'NUM_GROUPS']:
+            base_name = 'BASE_' + name
+            new_values = [round(base_value * wm) for base_value in
+                          config['SPEC'][base_name]]  # New value = base value * width multiplier.
+            logger.info(f'config["SPEC"]["{name}"]: {new_config["SPEC"][name]} -> {new_values}')
+            new_config['SPEC'][name] = new_values
+        return new_config
+
+    logger.info('muP: Create models and set base shapes')
+    logger.info('=> Create model')
+    model = create_encoder(config_encoder)
+
+    logger.info('=> Create base model')
+    base_config = gen_config(config_encoder, wm=1.0)
+    base_model = create_encoder(base_config)
+
+    logger.info('=> Create delta model')
+    delta_config = gen_config(config_encoder, wm=2.0)
+    delta_model = create_encoder(delta_config)
+
+    logger.info('=> Set base shapes in model for training')
+    set_base_shapes(model, base=base_model, delta=delta_model)
+
+    return model
+
+
+@register_image_encoder
+def image_encoder(config_encoder, verbose, **kwargs):
+    spec = config_encoder['SPEC']
+    standparam = spec.get('STANDPARAM', True)
+
+    if standparam:
+        logger.info('Create model with standard parameterization')
+        model = create_encoder(config_encoder)
+        model.custom_init_weights(use_original_init=True)
+    else:
+        logger.info('Create model with mu parameterization')
+        model = create_mup_encoder(config_encoder)
+        model.custom_init_weights(use_original_init=False)
+
+    logger.info('Load model from pretrained checkpoint')
+    if config_encoder['LOAD_PRETRAINED']:
+        model.from_pretrained(
+            config_encoder['PRETRAINED'],
+            config_encoder['PRETRAINED_LAYERS'],
+            verbose
+        )
+
+    return model

MedImageInsight/ImageEncoder/registry.py

@@ -0,0 +1,18 @@
+_image_encoders = {}
+
+
+def register_image_encoder(fn):
+    module_name_split = fn.__module__.split('.')
+    model_name = module_name_split[-1]
+
+    _image_encoders[model_name] = fn
+
+    return fn
+
+
+def image_encoders(model_name):
+    return _image_encoders[model_name]
+
+
+def is_image_encoder(model_name):
+    return model_name in _image_encoders

MedImageInsight/LangEncoder/__init__.py

@@ -0,0 +1,13 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from .build import build_lang_encoder
+from .build import build_tokenizer
+
+from .transformer import *
+# from .hf_model import *
+# from .zcode import *
+# from .pretrain import *
+# from .tulrv6 import *
+# from .t5 import *