add submission code

README.md

@@ -1,7 +1,34 @@
-# Cross-Domain Few-Shot Segmentation
+# Adapt Before Comparision - A New Perspective on Cross-Domain Few-Shot Segmentation
 
-The work will be released as soon as review period is over
+Code for the Reproducing the Paper
 
---> Release Date: 26.02.2024
+## Preparing Data
+Because we follow PATNet and RtD, please refer to their work for prepration of the following datasets:
+- Deepglobe (PAT)
+- ISIC (PAT)
+- Chest X-Ray (Lung) (PAT)
+- FSS-1000 (PAT)
+- SUIM (RtD)
 
-Stay tuned :star2:
+You do not need to get all datasets. Just prepare the one you want to test our method with.
+
+## Python package prerequisites
+1. torch
+2. torchvision
+3. cv2
+4. numpy
+5. for others, follow the console output
+
+## Run it
+Call
+`python main.py --benchmark {} --datapath {} --nshot {}`
+
+for example
+`python main.py --benchmark deepglobe --datapath ./datasets/deepglobe/ --nshot 1`
+
+Available `benchmark` strings: deepglobe,isic,lung,fss,suim
+Easiest to prepare should be Lung or FSS.
+
+Default is quick-infer mode.
+To change this, set `config.featext.fit_every_episode = True` in the main file.
+You can change all other parameters likewise, check the available parameters in runner.makeConfig.

core/backbone.py

@@ -0,0 +1,63 @@
+from functools import reduce
+from operator import add
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import resnet
+
+class Backbone(nn.Module):
+
+    def __init__(self, typestr):
+        super(Backbone, self).__init__()
+
+        self.backbone = typestr
+
+        # feature extractor initialization
+        if typestr == 'resnet50':
+            self.feature_extractor = resnet.resnet50(weights=resnet.ResNet50_Weights.DEFAULT)
+            self.feat_channels = [256, 512, 1024, 2048]
+            self.nlayers = [3, 4, 6, 3]
+            self.feat_ids = list(range(0, 17))
+        else:
+            raise Exception('Unavailable backbone: %s' % typestr)
+        self.feature_extractor.eval()
+
+        # define model
+        self.lids = reduce(add, [[i + 1] * x for i, x in enumerate(self.nlayers)])
+        self.stack_ids = torch.tensor(self.lids).bincount()[-4:].cumsum(dim=0)
+
+        self.cross_entropy_loss = nn.CrossEntropyLoss()
+
+    def extract_feats(self, img):
+        r""" Extract input image features """
+        feats = []
+        bottleneck_ids = reduce(add, list(map(lambda x: list(range(x)), self.nlayers)))
+        # Layer 0
+        feat = self.feature_extractor.conv1.forward(img)
+        feat = self.feature_extractor.bn1.forward(feat)
+        feat = self.feature_extractor.relu.forward(feat)
+        feat = self.feature_extractor.maxpool.forward(feat)
+
+        # Layer 1-4
+        for hid, (bid, lid) in enumerate(zip(bottleneck_ids, self.lids)):
+            res = feat
+            feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].conv1.forward(feat)
+            feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].bn1.forward(feat)
+            feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].relu.forward(feat)
+            feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].conv2.forward(feat)
+            feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].bn2.forward(feat)
+            feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].relu.forward(feat)
+            feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].conv3.forward(feat)
+            feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].bn3.forward(feat)
+
+            if bid == 0:
+                res = self.feature_extractor.__getattr__('layer%d' % lid)[bid].downsample.forward(res)
+
+            feat += res
+
+            if hid + 1 in self.feat_ids:
+                feats.append(feat.clone())
+
+            feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].relu.forward(feat)
+
+        return feats

core/contrastivehead.py

@@ -0,0 +1,582 @@
+import torch.nn.functional as F
+import torch
+import torch.nn as nn
+from utils import segutils
+import core.denseaffinity as dautils
+
+identity_mapping = lambda x, *args, **kwargs: x
+
+
+class ContrastiveConfig:
+    def __init__(self, config=None):
+        # Define the internal dictionary with default settings.
+        if config is None:
+            self._data = {
+                'aug': {
+                    'n_transformed_imgs': 2,
+                    'blurkernelsize': [1],  # chooses one of this kernel sizes
+                    'maxjitter': 0.0,
+                    'maxangle': 0,  # rotation
+                    # 'translate': (0,0),  # BE CAREFUL WITH TRANSLATE - if you apply it on the feature volume that has smaller spatial dims correspondences break
+                    'maxscale': 1.0,  # 1.0 = No scaling
+                    'maxshear': 20,
+                    'randomhflip': False,
+                    'apply_affine': True,
+                    'debug': False
+                },
+                'model': {
+                    'out_channels': 64,
+                    'kernel_size': 1,
+                    'prepend_relu': False,
+                    'append_normalize': False,
+                    'debug': False
+                },
+                'fitting': {
+                    'lr': 1e-2,
+                    'optimizer': torch.optim.SGD,
+                    'num_epochs': 25,
+                    'nce': {
+                        'temperature': 0.5,
+                        'debug': False
+                    },
+                    'normalize_after_fwd_pass': True,
+                    'q_nceloss': True,
+                    's_nceloss': True,
+                    'protoloss': False,
+                    'keepvarloss': True,
+                    'symmetricloss': False,
+                    'selfattentionloss': False,
+                    'o_t_contr_proto_loss': True,
+                    'debug': False
+                },
+                'featext': {
+                    'l0': 3,  # the first resnet bottleneck id to consider (0,1,2,3,4,5...15)
+                    'fit_every_episode': False
+                }
+            }
+        else:
+            self._data = config
+
+    def __getattr__(self, key):
+        # Try to get '_data' without causing a recursive call to __getattr__
+        _data = super().__getattribute__('_data') if '_data' in self.__dict__ else None
+
+        if _data is not None and key in _data:
+            if isinstance(_data[key], dict):
+                return ContrastiveConfig(_data[key])
+            return _data[key]
+
+        # If we're here, it means the key was not found in the data,
+        # so we let Python raise the appropriate AttributeError.
+        raise AttributeError(f"No setting named {key}")
+
+    def __setattr__(self, key, value):
+        # Prevent overwriting of the '_data' attribute by normal means
+        if key == '_data':
+            super().__setattr__(key, value)
+        else:
+            # Try to get '_data' without causing a recursive call to __getattr__
+            _data = super().__getattribute__('_data') if '_data' in self.__dict__ else None
+
+            if _data is not None:
+                _data[key] = value
+            else:
+                # This situation should not normally occur, handle appropriately (e.g., log an error, raise exception)
+                raise AttributeError("Unexpected")
+
+    # Optional: Representation for better debugging.
+    def __repr__(self):
+        return str(self._data)
+
+
+def dense_info_nce_loss(original_features, transformed_features, config_nce):
+    B, C, H, W = transformed_features.shape
+    o_features = original_features.expand(B, C, H, W).permute(0, 2, 3, 1).view(B, H * W, C)
+    t_features = transformed_features.permute(0, 2, 3, 1).view(B, H * W, C)
+
+    # Calculate dot product between original and transformed feature vectors for positive pairs
+    positive_logits = torch.einsum('bik,bik->bi', o_features, t_features) / config_nce.temperature
+
+    # Calculate dot product between original features and all other transformed features for negative pairs
+    all_logits = torch.einsum('bik,bjk->bij', o_features, t_features) / config_nce.temperature
+
+    if config_nce.debug: print('pos/neg:', positive_logits.mean().detach(), all_logits.mean().detach())
+
+    # Using the log-sum-exp trick
+    max_logits = torch.max(all_logits, dim=-1, keepdim=True).values
+    log_sum_exp = max_logits + torch.log(torch.sum(torch.exp(all_logits - max_logits), dim=-1, keepdim=True))
+
+    # Compute InfoNCE loss
+    loss = - (positive_logits - log_sum_exp.squeeze())
+    return loss.mean()  # [B=k*aug] or [B=k] -> scalar
+
+
+def ssim(a, b):
+    return torch.nn.CosineSimilarity()(a, b)
+
+def augwise_proto(feat_vol, mask, k, aug):
+    k, aug, c, h, w = k, aug, *feat_vol.shape[-3:]
+    feature_vectors_augwise = torch.cat(feat_vol.view(k, aug, c, h * w).unbind(0), dim=-1)
+    mask_augwise = torch.cat(segutils.downsample_mask(mask, h, w).view(k, aug, h * w).unbind(0), dim=-1)
+    assert feature_vectors_augwise.shape == (aug, c, k * h * w) and mask_augwise.shape == (
+    aug, k * h * w), "of transformed"
+
+    fg_proto, bg_proto = segutils.fg_bg_proto(feature_vectors_augwise, mask_augwise)
+    assert fg_proto.shape == bg_proto.shape == (aug, c)
+
+    return fg_proto, bg_proto
+
+
+def calc_q_pred_coarse_nodetach(qft, sft, s_mask, l0=3):
+    bsz, c, hq, wq = qft.shape
+    hs, ws = sft.shape[-2:]
+
+    sft_row = torch.cat(sft.unbind(1), -1)  # bsz,k,c,h,w -> bsz,c,h,w*k
+    smasks_downsampled = [segutils.downsample_mask(m, hs, ws) for m in s_mask.unbind(1)]
+    smask_row = torch.cat(smasks_downsampled, -1)
+
+    damat = dautils.buildDenseAffinityMat(qft, sft_row)
+    filtered = dautils.filterDenseAffinityMap(damat, smask_row)
+    q_pred_coarse = filtered.view(bsz, hq, wq)
+    return q_pred_coarse
+
+
+# input k*aug,c,h,w
+def self_attention_loss(f_base, f_transformed, mask_base, mask_transformed, k, aug):
+    c, h, w = f_base.shape[-3:]
+    pseudoquery = torch.cat(f_base.view(k, aug, c, h, w).unbind(0), -1)  # shape aug,c,h,w*k
+    pseudoquerymask = torch.cat(mask_base.view(k, aug, h, w).unbind(0), -1)  # shape aug,h,w*k
+    pseudosupport = f_transformed.view(k, aug, c, h, w).transpose(0, 1)  # shape bsz,k,c,h,w
+    pseudosupportmask = mask_transformed.view(k, aug, h, w).transpose(0, 1)  # shape bsz,k,h,w
+    # display(segutils.tensor_table(q=pseudoquery, s=pseudosupport, m=pseudosupportmask))
+    pred_map = calc_q_pred_coarse_nodetach(pseudoquery, pseudosupport, pseudosupportmask, l0=0)
+
+    loss = torch.nn.BCELoss()(pred_map.float(), pseudoquerymask.float())
+    return loss.mean()
+
+
+# features of base, transformed: [b,c,h,w]
+# if base features are aligned with transformed features, pass both same
+def ctrstive_prototype_loss(base, transformed, mask_base, mask_transformed, k, aug):
+    assert transformed.shape == base.shape, ".."
+    b, c, h, w = base.shape
+    assert b == k * aug, 'provide correct k and aug such that dim0=k*aug'
+    assert mask_base.shape == mask_transformed.shape == (b, h, w), ".."
+    fg_proto_o, bg_proto_o = augwise_proto(base, mask_base, k, aug)
+    fg_proto_t, bg_proto_t = augwise_proto(transformed, mask_transformed, k, aug)
+    # i: fg, b: bg
+    # p_b_i, p_b_j = segutils.fg_bg_proto(base.view(b,c,h*w), mask_base.view(b,h*w))
+    # p_t_i, p_t_j = segutils.fg_bg_proto(transformed.view(b,c,h*w), mask_transformed.view(b,h*w))
+    enumer = torch.exp(
+        ssim(fg_proto_o, fg_proto_t))  # 5vs5 (augvsaug), but in 5-shot: 25vs25, no, you want also augvsaug
+    denom = torch.exp(ssim(fg_proto_o, fg_proto_t)) + torch.exp(ssim(fg_proto_o, bg_proto_t))
+    assert enumer.shape == denom.shape == torch.Size([aug]), 'you want to calculate one prototype for each augmentation'
+    loss = -torch.log(enumer / denom)  # [bsz]
+    return loss.mean()
+
+
+def opposite_proto_sim_in_aug(transformed_features, mapped_s_masks, k, aug):
+    fg_proto_t, bg_proto_t = augwise_proto(transformed_features, mapped_s_masks, k, aug)
+    fg_bg_sim_t = ssim(fg_proto_t, bg_proto_t)
+    return fg_bg_sim_t.mean()
+
+
+def proto_align_val_measure(original_features, transformed_features, mapped_s_masks, k, aug):
+    fg_proto_o, _ = augwise_proto(original_features, mapped_s_masks, k, aug)
+    fg_proto_t, _ = augwise_proto(transformed_features, mapped_s_masks, k, aug)
+    fg_proto_sim = ssim(fg_proto_o, fg_proto_t)
+    return fg_proto_sim.mean()
+
+
+def atest():
+    k, aug, c, h, w = 2, 5, 8, 20, 20
+    f_base = torch.rand(k * aug, c, h, w).float()
+    f_base.requires_grad = True
+    f_transformed = torch.rand(k * aug, c, h, w).float()
+    mask_base = torch.randint(0, 2, (k * aug, h, w)).float()
+    mask_transformed = torch.randint(0, 2, (k * aug, h, w)).float()
+
+    return self_attention_loss(f_base, f_transformed, mask_base, mask_transformed, k, aug)
+
+def keep_var_loss(original_features, transformed_features):
+    meandiff = original_features.mean((-2, -1)) - transformed_features.mean((-2, -1))
+    vardiff = original_features.var((-2, -1)) - transformed_features.var((-2, -1))
+    keepvarloss = torch.abs(meandiff).mean() + torch.abs(
+        vardiff).mean()  # [k*aug,c] -> [scalar] or  [aug,c] -> [scalar]
+    return keepvarloss
+
+class ContrastiveFeatureTransformer(nn.Module):
+    def __init__(self, in_channels, config_model):
+        super(ContrastiveFeatureTransformer, self).__init__()
+
+        out_channels, kernel_size = config_model.out_channels, config_model.kernel_size
+        # Add a convolutional layer and a batch normalization layer for learning
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=(kernel_size - 1) // 2)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.linear = nn.Conv2d(out_channels, out_channels, 1)
+
+        self.prepend_relu = config_model.prepend_relu
+        self.append_normalize = config_model.append_normalize
+        self.debug = config_model.debug
+
+    def forward(self, x):
+        if self.prepend_relu:
+            x = nn.ReLU()(x)
+        x = self.conv(x)
+        x = self.bn(x)
+        x = nn.ReLU()(x)
+        x = self.linear(x)
+        if self.append_normalize:
+            x = F.normalize(x, p=2, dim=1)
+        return x
+
+    # fits the model for one semantic class, therefore does not work with batches
+    # mapped_qfeat_vol, aug_qfeat_vols: [aug,c,h,w]
+    # mapped_sfeat_vol, aug_sfeat_vols: [k*aug,c,h,w]
+    # augmented_smasks: [k*aug,h,w]
+    def fit(self, mapped_qfeat_vol, aug_qfeat_vols, mapped_sfeat_vol, aug_sfeat_vols, augmented_smasks, config_fit):
+        f_norm = F.normalize if config_fit.normalize_after_fwd_pass else identity_mapping
+        optimizer = config_fit.optimizer(self.parameters(), lr=config_fit.lr)
+        for epoch in range(config_fit.num_epochs):
+            # Pass original and transformed image batches through the model
+
+            # Q
+            original_features = f_norm(self(mapped_qfeat_vol), p=2, dim=1)  # fwd pass non-augmented
+            transformed_features = f_norm(self(aug_qfeat_vols), p=2, dim=1)  # fwd pass augmented
+
+            qloss = dense_info_nce_loss(original_features, transformed_features,
+                                        config_fit.nce) if config_fit.q_nceloss else 0
+            if config_fit.keepvarloss:  # 1. idea: Let query and support have the same feature distribution (mean/var per channel)
+                qloss += keep_var_loss(original_features, transformed_features)
+            # S
+            original_features = f_norm(self(mapped_sfeat_vol), p=2, dim=1)  # fwd pass non-augmented
+            transformed_features = f_norm(self(aug_sfeat_vols), p=2, dim=1)  # fwd pass augmented
+
+            sloss = dense_info_nce_loss(original_features, transformed_features,
+                                        config_fit.nce) if config_fit.s_nceloss else 0
+            if config_fit.keepvarloss:
+                sloss += keep_var_loss(original_features, transformed_features)
+
+            # 2. class-aware loss: opposite classes should get opposite features
+            # for prototype calculation, we want only one prototype per class
+            # so we average over features of entire k
+            # but calculate prototype for each augmentation individually [k*aug,c,h,w]->[aug,c,k*h*w]->[aug,c]
+            kaug, c, h, w = transformed_features.shape
+            aug = aug_qfeat_vols.shape[0]
+            k = kaug // aug
+            if config_fit.protoloss:
+                assert not config_fit.o_t_contr_proto_loss, 'only one of the proto losses should be used'
+                opposite_proto_sim = opposite_proto_sim_in_aug(transformed_features, augmented_smasks, k, aug)
+                if config_fit.debug and (epoch == config_fit.num_epochs - 1 or epoch == 0): print(
+                    'proto-sim intER-class transf<->transf', opposite_proto_sim.item())
+                proto_loss = opposite_proto_sim
+            elif config_fit.selfattentionloss:
+                proto_loss = self_attention_loss(original_features, transformed_features, augmented_smasks,
+                                                 augmented_smasks, k, aug)
+                if config_fit.debug and (epoch == config_fit.num_epochs - 1 or epoch == 0): print(
+                    'self-att non-transf<->transformed bce', proto_loss.item())
+            elif config_fit.o_t_contr_proto_loss:
+                o_t_contr_proto_loss = ctrstive_prototype_loss(original_features, transformed_features,
+                                                               augmented_smasks, augmented_smasks, k, aug)
+                if config_fit.debug and (epoch == config_fit.num_epochs - 1 or epoch == 0): print(
+                    'proto-contr non-transf<->transformed', o_t_contr_proto_loss.item())
+                proto_loss = o_t_contr_proto_loss
+            else:
+                proto_loss = 0
+
+            if config_fit.debug and (epoch == config_fit.num_epochs - 1 or epoch == 0):
+                proto_align_val = proto_align_val_measure(original_features, transformed_features, augmented_smasks, k,
+                                                          aug)
+                print('proto-sim intRA-class non-transf<->transformed (for validation)', proto_align_val.item())
+
+            # 3. do not let only one image fit well - regularization
+            q_s_loss_diff = torch.abs(qloss - sloss) if config_fit.symmetricloss else 0
+
+            # Aggregate loss
+            loss = qloss + sloss + q_s_loss_diff + proto_loss
+            assert loss.isfinite().all(), f"invalid contrastive loss:{loss}"
+
+            # Backpropagation and optimization
+            if config_fit.debug and (epoch == config_fit.num_epochs - 1 or epoch == 0):
+                def gradient_magnitude(loss_term):
+                    optimizer.zero_grad()
+                    loss_term.backward(retain_graph=True)
+                    magn = torch.abs(self.conv.weight.grad.mean()) + torch.abs(self.linear.weight.grad.mean())
+                    return magn
+
+                q_loss_grad_magnitude = gradient_magnitude(qloss)
+                s_loss_grad_magnitude = gradient_magnitude(sloss)
+                proto_loss_grad_magnitude = gradient_magnitude(proto_loss)
+                q_s_loss_diff_grad_magnitude = gradient_magnitude(q_s_loss_diff)
+                display(segutils.tensor_table(q_loss_grad_magnitude=q_loss_grad_magnitude,
+                                              s_loss_grad_magnitude=s_loss_grad_magnitude,
+                                              proto_loss_grad_magnitude=proto_loss_grad_magnitude,
+                                              q_s_loss_diff_grad_magnitude=q_s_loss_diff_grad_magnitude))
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            if config_fit.debug and epoch % 10 == 0: print('loss', loss.detach())
+
+
+import numpy as np
+import torch.nn.functional as F
+from torchvision.transforms.functional import affine
+from torchvision.transforms import GaussianBlur, ColorJitter
+
+
+class AffineProxy:
+    def __init__(self, angle, translate, scale, shear):
+        self.affine_params = {
+            'angle': angle,
+            'translate': translate,
+            'scale': scale,
+            'shear': shear
+        }
+
+    def apply(self, img):
+        return affine(img, angle=self.affine_params['angle'], translate=self.affine_params['translate'],
+                      scale=self.affine_params['scale'], shear=self.affine_params['shear'])
+
+
+# def affine_proxy(angle, translate, scale, shear):
+#     def inner(img):
+#         return affine(img, angle=angle, translate=translate, scale=scale, shear=shear)
+
+#     return inner
+
+class Augmen:
+    def __init__(self, config_aug):
+        self.config = config_aug
+        self.blurs, self.jitters, self.affines = self.setup_augmentations()
+
+    def copy_construct(self, blurs, jitters, affines, config_aug):
+        self.config = config_aug
+        self.blurs, self.jitters, self.affines = blurs, jitters, affines
+
+    def setup_augmentations(self):
+        blurkernelsize = self.config.blurkernelsize
+        maxjitter = self.config.maxjitter
+
+        maxangle = self.config.maxangle
+        translate = (0, 0)
+        maxscale = self.config.maxscale
+        maxshear = self.config.maxshear
+
+        blurs = []
+        jitters = []
+        affine_trans = []
+        for i in range(self.config.n_transformed_imgs):
+            # Randomize kernel size for GaussianBlur
+            kernel_size = np.random.choice(torch.tensor(blurkernelsize), (1,)).item()
+            blur = GaussianBlur(kernel_size)
+            blurs.append(blur)
+
+            # Randomize values for ColorJitter
+            brightness_val = torch.rand(1).item() * maxjitter  # up to <maxjitter> change
+            contrast_val = torch.rand(1).item() * maxjitter
+            saturation_val = torch.rand(1).item() * maxjitter
+            jitter = ColorJitter(brightness=brightness_val, contrast=contrast_val, saturation=saturation_val)
+            jitters.append(jitter)
+
+            # Random values for each iteration
+            angle = torch.randint(-maxangle, maxangle + 1, (1,)).item()
+            shear = [torch.randint(-maxshear, maxshear + 1, (1,)).item() for _ in range(2)]
+            scale = torch.rand(1).item() * (1 - maxscale) + maxscale
+            affine_trans.append(AffineProxy(angle=angle, translate=translate, scale=scale, shear=shear))
+
+        return (blurs, jitters, affine_trans)  # tuple of lists
+
+    def augment(self, original_image, orignal_mask):
+        transformed_imgs = []
+        transformed_masks = []
+        for blur, jitter, affine_trans in zip(self.blurs, self.jitters, self.affines):
+            # Apply non-geometric transformations
+            t_img = blur(original_image)
+            t_img = jitter(t_img)
+            t_mask = orignal_mask.clone()
+
+            if self.config.apply_affine:
+                t_img = affine_trans.apply(t_img)
+                t_mask = affine_trans.apply(t_mask)
+
+            transformed_imgs.append(t_img)
+            transformed_masks.append(t_mask)
+        return torch.stack(transformed_imgs, dim=1), torch.stack(transformed_masks, dim=1)
+
+    # [bsz,ch,h,w] -> [bsz,aug,ch,h,w], where aug is the number of augmentated images
+    def applyAffines(self, feat_vol):
+        return torch.stack([trans.apply(feat_vol) for trans in self.affines], dim=1)
+
+
+class CTrBuilder:
+    # call init 1st, pass all config parameters (initatiate a ContrastiveConfig class in your code)
+    def __init__(self, config, augmentator=None):
+        if augmentator is None:
+            augmentator = Augmen(config.aug)
+        self.augmentator = augmentator
+
+        self.augimgs = self.AugImgStack(augmentator)
+
+        self.hasfit = False
+        self.config = config
+
+    class AugImgStack():
+        def __init__(self, augmentator):
+            self.augmentator = augmentator
+            self.q, self.s, self.s_mask = None, None, None
+
+        def init(self, s_img):
+            # c is color channels here, not feature channels
+            bsz, k, aug, c, h, w = *s_img.shape[:2], self.augmentator.config.n_transformed_imgs, *s_img.shape[-3:]
+            self.q = torch.empty(bsz, aug, c, h, w).to(s_img.device)
+            self.s = torch.empty(bsz, k, aug, c, h, w).to(s_img.device)
+            self.s_mask = torch.empty(bsz, k, aug, h, w).to(s_img.device)
+
+        def show(self):
+            bsz_, k_, aug_ = self.s.shape[:3]
+            for b in range(bsz_):
+                display('aug x queries', segutils.pilImageRow(*[segutils.norm(img) for img in self.q[b]]))
+                for k in range(k_):
+                    print('k=', k, ' aug x (s, smask):')
+                    display(segutils.pilImageRow(*[segutils.norm(img) for img in self.s[b, k]]))
+                    display(segutils.pilImageRow(*self.s_mask[b, k]))
+
+    def showAugmented(self):
+        self.augimgs.show()
+
+    # 2nd call makeAugmented
+    def makeAugmented(self, q_img, s_img, s_mask):
+        # 2. Augmentation
+        # 2.1 Apply transformations to images
+        self.augimgs.init(s_img)
+        self.augimgs.q, _ = self.augmentator.augment(q_img, s_mask)
+
+        for k in range(s_img.shape[1]):
+            s_aug_imgs, s_aug_masks = self.augmentator.augment(s_img[:, k], s_mask[:, k])
+            self.augimgs.s[:, k] = s_aug_imgs
+            self.augimgs.s_mask[:, k] = s_aug_masks
+        if self.config.aug.debug: self.augimgs.show()
+
+    # 3rd call build_and_fit
+    def build_and_fit(self, q_feat, s_feat, q_feataug, s_feataug, s_maskaug=None):
+        if s_maskaug is None: s_maskaug = self.augimgs.s_mask
+        self.ctrs = self.buildContrastiveTransformers(q_feat, s_feat, q_feataug, s_feataug, s_maskaug)
+        self.hasfit = True
+
+    def buildContrastiveTransformers(self, qfeat_alllayers, sfeat_alllayers, query_feats_aug, support_feats_aug,
+                                     supp_aug_mask, s_mask=None):
+        contrastive_transformers = []
+        l0 = self.config.featext.l0
+        # [bsz,k,aug,h,w] -> [k*aug,h,w]
+        s_aug_mask = supp_aug_mask.view(-1, *supp_aug_mask.shape[-2:])
+        # iterate over feature layers
+        for (qfeat, sfeat, qfeataug, sfeataug) in zip(qfeat_alllayers[l0:], sfeat_alllayers[l0:], query_feats_aug[l0:],
+                                                      support_feats_aug[l0:]):
+            bsz, k, aug, ch, h, w = sfeataug.shape
+            # we fit it for exactly one class, so use no batches
+            assert bsz == 1, "bsz should be 1"
+            assert supp_aug_mask.shape[1] == sfeat.shape[
+                1] == k, f'augmented support shot-dimension mismatch:{s_aug_mask.shape[1]=},{sfeat.shape[1]=},(bsz,k,aug,ch,h,w)={bsz, k, aug, ch, h, w}'
+            assert supp_aug_mask.shape[2] == qfeataug.shape[1] == aug, 'augmented shot-dimension mismatch'
+            # [bsz,c,h,w] -> [1,c,h,w]
+            qfeat = qfeat.view(-1, *qfeat.shape[-3:])
+            # [bsz,k,c,h,w] -> [k,c,h,w]
+            sfeat = sfeat.view(-1, *sfeat.shape[-3:])
+            # [bsz,aug,c,h,w] -> [aug,c,h,w]
+            qfeataug = qfeataug.view(-1, *qfeataug.shape[-3:])
+            # [bsz,k,aug,c,h,w] -> [k*aug,c,h,w]
+            sfeataug = sfeataug.view(-1, *qfeataug.shape[-3:])
+
+            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+            contrastive_head = ContrastiveFeatureTransformer(in_channels=ch, config_model=self.config.model).to(device)
+
+            # 3. Feature volumes from untransformed image need to be geometrically mapped to allow for dense matching
+            mapped_qfeat = self.augmentator.applyAffines(qfeat)
+            assert mapped_qfeat.shape[1] == aug, "should be 1,aug,c,h,w"
+            mapped_qfeat = mapped_qfeat.view(-1, *qfeat.shape[-3:])  # ->[aug,c,h,w]
+            mapped_sfeat = self.augmentator.applyAffines(sfeat)
+            assert mapped_sfeat.shape[1] == aug and mapped_sfeat.shape[0] == k, "should be k,aug,c,h,w"
+            mapped_sfeat = mapped_sfeat.view(-1, *sfeat.shape[-3:])  # ->[k*aug,c,h,w]
+
+            contrastive_head.fit(mapped_qfeat, qfeataug, mapped_sfeat, sfeataug,
+                                 segutils.downsample_mask(s_aug_mask, h, w), self.config.fitting)
+
+            contrastive_transformers.append(contrastive_head)
+            # show how support image and its augmentations would produce a affinity map
+            if s_mask != None:
+                display(segutils.to_pil(segutils.norm(dautils.filterDenseAffinityMap(
+                    dautils.buildDenseAffinityMat(contrastive_head(sfeat), contrastive_head(sfeataug[:1])),
+                    segutils.downsample_mask(s_mask, h, w)).view(1, h, w))))
+                display(segutils.to_pil(segutils.norm(dautils.filterDenseAffinityMap(
+                    dautils.buildDenseAffinityMat(contrastive_head(qfeat), contrastive_head(sfeat)),
+                    segutils.downsample_mask(s_mask, h, w)).view(1, h, w))))
+        return contrastive_transformers
+
+    # You have fitted the contrastive transformers, now apply the transform and then pass to the downstream DCAMA
+    # you just need to append the empty layers you exluded ([:3]), they're also skipped in dcama
+    # Obtain the result of the contrastive head, which will be the new query and support feat representation
+    def getTaskAdaptedFeats(self, layerwise_feats):
+        if (self.ctrs == None): print("error: call buildContrastiveTransformers() first")
+        task_adapted_feats = []
+
+        for idx in range(len(layerwise_feats)):
+            if idx < self.config.featext.l0:
+                task_adapted_feats.append(None)
+            else:
+                input_shape = layerwise_feats[idx].shape
+                idxth_feat = layerwise_feats[idx].view(-1, *input_shape[-3:])
+                forward_pass_res = self.ctrs[idx - self.config.featext.l0](idxth_feat)
+                target_shape = *input_shape[:-3], *forward_pass_res.shape[
+                                                   -3:]  # borrow channel dim from result, but bsz,k dims from input
+                task_adapted_feats.append(forward_pass_res.view(target_shape))
+
+        return task_adapted_feats
+
+
+class FeatureMaker:
+    def __init__(self, feat_extraction_method, class_ids, config=ContrastiveConfig()):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.featextractor = feat_extraction_method
+        self.c_trs = {ctr: CTrBuilder(config) for ctr in class_ids}
+        self.norm_bb_feats = False
+
+    def extract_bb_feats(self, img):
+        with torch.no_grad():
+            return self.featextractor(img)
+
+    def create_and_fit(self, c_tr, q_img, s_img, s_mask, q_feat, s_feat):
+        print('doing contrastive')
+        c_tr.makeAugmented(q_img, s_img, s_mask)
+
+        bsz, k, c, h, w = s_img.shape
+        aug = c_tr.augmentator.config.n_transformed_imgs
+        # [bsz,aug,c,h,w]->[bsz*aug,c,h,w] squeeze for forward pass
+        q_feataug = self.extract_bb_feats(c_tr.augimgs.q.view(-1, c, h, w))  # returns layer-list
+        # then restore
+        q_feataug = [l.view(bsz, aug, *l.shape[1:]) for l in q_feataug]
+        # [bsz,k,aug,c,h,w]->[bsz*k*aug,c,h,w]->[bsz,k,aug,c,h,w]
+        s_feataug = self.extract_bb_feats(c_tr.augimgs.s.view(-1, c, h, w))
+        s_feataug = [l.view(bsz, k, aug, *l.shape[1:]) for l in s_feataug]
+
+        c_tr.build_and_fit(q_feat, s_feat, q_feataug, s_feataug)
+
+    def taskAdapt(self, q_img, s_img, s_mask, class_id):
+        ch_norm = lambda t: t / torch.linalg.norm(t, dim=1)
+        q_feat = self.extract_bb_feats(q_img)
+        bsz, k, c, h, w = s_img.shape
+        s_feat = self.extract_bb_feats(s_img.view(-1, c, h, w))
+        if self.norm_bb_feats:
+            q_feat = [ch_norm(l) for l in q_feat]
+            s_feat = [ch_norm(l) for l in q_feat]
+        s_feat = [l.view(bsz, k, *l.shape[1:]) for l in s_feat]
+
+        c_tr = self.c_trs[class_id]  # select the relevant ctr for this class
+
+        if c_tr.hasfit is False or c_tr.config.featext.fit_every_episode:  # create and fit a contrastive transformer if not existing yet
+            self.create_and_fit(c_tr, q_img, s_img, s_mask, q_feat, s_feat)
+
+        q_feat_t, s_feat_t = c_tr.getTaskAdaptedFeats(q_feat), c_tr.getTaskAdaptedFeats(
+            s_feat)  # tocheck: do they require_grad here?
+        return q_feat_t, s_feat_t

core/denseaffinity.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn.functional as F
+import math
+from utils import segutils
+
+
+def buildHyperCol(feat_pyram):
+    # concatenate along channel dim
+    # upsample spatial size to largest feat vol space available
+    target_size = feat_pyram[0].shape[-2:]
+    upsampled = []
+    for layer in feat_pyram:
+        # if idx < self.stack_ids[0]: continue
+        upsampled.append(F.interpolate(layer, size=target_size, mode='bilinear', align_corners=False))
+    return torch.cat(upsampled, dim=1)
+
+
+# accepts both:
+# s_feat_vol: [bsz,k,c,h,w]->[bsz,c,h,w*k]
+# s_mask: [bsz,k,h,w]->[bsz,h,w*k]
+def paste_supports_together(supports):
+    return torch.cat(supports.unbind(dim=1), dim=-1)
+
+
+# Attention regular:
+# 1. Dot product
+# 2. Divide by square root of key length (#nchannels)
+# 3. Softmax
+# 4. Multiply with V (mask)
+
+def buildDenseAffinityMat(qfeat_volume, sfeat_volume, softmax_arg2=True):  # bsz,C,H,W
+    qfeat_volume, sfeat_volume = qfeat_volume.permute(0, 2, 3, 1), sfeat_volume.permute(0, 2, 3, 1)
+    bsz, H, Wq, C = qfeat_volume.shape
+    Ws = sfeat_volume.shape[2]
+    # [px,C][C,px]=[px,px]
+    dense_affinity_mat = torch.matmul(qfeat_volume.view(bsz, H * Wq, C),
+                                      sfeat_volume.view(bsz, H * Ws, C).transpose(1, 2))
+    if softmax_arg2 is False: return dense_affinity_mat
+    dense_affinity_mat_softmax = (dense_affinity_mat / math.sqrt(C)).softmax(
+        dim=-1)  # each query pixel's affinities sum up to 1 over support pxls
+    return dense_affinity_mat_softmax
+
+
+# filter with support mask following DAM
+def filterDenseAffinityMap(dense_affinity_mat, downsampled_smask):
+    # for each query pixel, aggregate all correlations where the support mask ==1
+    # [px,px][px,1]=[px,1]
+    bsz, HWq, HWs = dense_affinity_mat.shape
+    # let mean(V)=1 -> sum(V)=len(V) -> d_mask / mean(d_mask)
+    # downsampled_smask_norm = downsampled_smask / downsampled_smask.mean()
+    q_coarse = torch.matmul(dense_affinity_mat, downsampled_smask.view(bsz, HWs, 1))
+    return q_coarse.view(bsz, HWq)
+
+
+def upsample(volume, h, w):
+    return F.interpolate(volume, size=(h, w), mode='bilinear', align_corners=False)
+
+class DAMatComparison:
+
+    def algo_mean(self, q_pred_coarses_t, s_mask=None):
+        return q_pred_coarses_t.mean(1)
+
+    def calc_q_pred_coarses(self, q_feat_t, s_feat_t, s_mask, l0=3):
+        q_pred_coarses = []
+        h0, w0 = q_feat_t[l0].shape[-2:]
+        for (qft, sft) in zip(q_feat_t[l0:], s_feat_t[l0:]):
+            qft, sft = qft.detach(), sft.detach()
+            bsz, c, hq, wq = qft.shape
+            hs, ws = sft.shape[-2:]
+
+            sft_row = torch.cat(sft.unbind(1), -1)  # bsz,k,c,h,w -> bsz,c,h,w*k
+            smasks_downsampled = [segutils.downsample_mask(m, hs, ws) for m in s_mask.unbind(1)]
+            smask_row = torch.cat(smasks_downsampled, -1)
+
+            damat = buildDenseAffinityMat(qft, sft_row)
+            filtered = filterDenseAffinityMap(damat, smask_row)
+            q_pred_coarse = upsample(filtered.view(bsz, 1, hq, wq), h0, w0).squeeze(1)
+            q_pred_coarses.append(q_pred_coarse)
+        return torch.stack(q_pred_coarses, dim=1)
+
+    def forward(self, q_feat_t, s_feat_t, s_mask, upsample=True, debug=False):
+        q_pred_coarses_t = self.calc_q_pred_coarses(q_feat_t, s_feat_t, s_mask)
+
+        if debug: display(segutils.pilImageRow(*q_pred_coarses_t.unbind(1), q_pred_coarses_t.mean(1)))
+
+        # select the algorithm
+        postprocessing_algorithm = self.algo_mean
+        # do the postprocessing
+        logit_mask = postprocessing_algorithm(q_pred_coarses_t, s_mask)
+        if upsample:  # if query and support have different shape, then you must do upsampling yourself afterwards
+            logit_mask = segutils.downsample_mask(logit_mask, *s_mask.shape[-2:])
+
+        return logit_mask

core/runner.py

@@ -0,0 +1,165 @@
+from data.dataset import FSSDataset
+from core.backbone import Backbone
+from eval.logger import Logger, AverageMeter
+from eval.evaluation import Evaluator
+from utils import commonutils as utils
+import utils.segutils as segutils
+import core.contrastivehead as ctrutils
+import core.denseaffinity as dautils
+import torch
+
+class args:
+    backbone = 'resnet50'
+    logpath = '/kaggle/working/logs'
+    nworker = 0
+    bsz = 1
+    benchmark='' #e.g. deepglobe,isic,etc.
+    datapath='' #path to the selected dataset
+    fold = 0
+    nshot = 1
+
+class SingleSampleEval:
+    def __init__(self, batch, feat_maker, debug=False):
+        self.damat_comp = dautils.DAMatComparison()
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.batch = batch
+        self.feat_maker = feat_maker
+        self.debug = debug
+        self.thresh_method = 'pred_mean'
+
+    def taskAdapt(self, detach=True):
+        b = self.batch
+        if self.device.type == 'cuda': b = utils.to_cuda(b)
+        self.q_img, self.s_img, self.s_mask, self.class_id = b['query_img'], b['support_imgs'], b['support_masks'], b[
+            'class_id'].item()
+        self.task_adapted = self.feat_maker.taskAdapt(self.q_img, self.s_img, self.s_mask, self.class_id)
+
+    def compare_feats(self):
+        if self.task_adapted is None:
+            print("error, do task adaption first")
+            return None
+        self.logit_mask = self.damat_comp.forward(self.task_adapted[0], self.task_adapted[1], self.s_mask)
+        return self.logit_mask
+
+    def threshold(self, method=None):
+        if self.logit_mask is None:
+            print("error, calculate logit mask first (do forward pass)")
+        if method is None:
+            method = self.thresh_method
+        self.thresh = calcthresh(self.logit_mask, self.s_mask, method)
+        self.pred_mask = (self.logit_mask > self.thresh).float()
+        return self.thresh, self.pred_mask
+
+    def apply_crf(self):
+        return apply_crf(self.q_img, self.logit_mask, thresh_fn(self.thresh_method))
+
+    # this method calls above components sequentially
+    def forward(self):
+        self.taskAdapt()
+
+        self.logit_mask = self.compare_feats()
+
+        self.thresh, self.pred_mask = self.threshold()
+
+        return self.logit_mask, self.pred_mask
+
+    def calc_metrics(self):
+        # assert torch.logical_or(self.logit_mask<0, self.logit_mask>1).sum()==0, display(tensor_table(logit_mask=self.logit_mask))
+        self.area_inter, self.area_union = Evaluator.classify_prediction(self.pred_mask, self.batch)
+        self.fgratio_pred = self.pred_mask.float().mean()
+        self.fgratio_gt = self.batch['query_mask'].float().mean()
+        return self.area_inter[1] / self.area_union[1]  # fg-iou
+
+    def plots(self):
+        display(pilImageRow(norm(self.logit_mask[0]), (self.logit_mask[0] > self.thresh).float(), self.pred_mask,
+                            self.batch['query_mask'][:1], norm(self.q_img[0]), norm(self.s_img[0, 0])))
+        display(segutils.tensor_table(probs=self.logit_mask))
+
+        print('s_mask.mean, pred_mask.mean, thresh:', self.s_mask.mean().item(), self.logit_mask.mean().item(),
+              self.thresh.item())
+
+class AverageMeterWrapper:
+    def __init__(self, dataloader, device='cpu', initlogger=True):
+        if initlogger: Logger.initialize(args, training=False)
+        self.average_meter = AverageMeter(dataloader.dataset, device)
+        self.device=device
+        self.dataloader = dataloader
+        self.write_batch_idx = 50
+    def update(self, sseval):
+        self.average_meter.update(sseval.area_inter, sseval.area_union, torch.tensor(sseval.class_id).to(self.device), loss=None)
+    def update_manual(self, area_inter, area_union, class_id):
+        if isinstance(class_id, int): class_id = torch.tensor(class_id).to(self.device)
+        self.average_meter.update(area_inter, area_union, class_id, loss=None)
+    def write(self, i):
+        self.average_meter.write_process(i, len(self.dataloader), 0, self.write_batch_idx)
+
+def makeDataloader():
+
+    FSSDataset.initialize(img_size=400, datapath=args.datapath)
+    dataloader = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'test', args.nshot)
+    return dataloader
+
+
+def makeConfig():
+    config = ctrutils.ContrastiveConfig()
+    config.fitting.protoloss = False
+    config.fitting.o_t_contr_proto_loss = True
+    config.fitting.selfattentionloss = False
+    config.fitting.keepvarloss = True
+    config.fitting.symmetricloss = False
+    config.fitting.q_nceloss = True
+    config.fitting.s_nceloss = True
+    config.fitting.num_epochs = 25
+    config.fitting.lr = 1e-2
+    config.fitting.debug = False
+    config.model.out_channels = 64
+    config.featext.fit_every_episode = False
+    config.aug.blurkernelsize = [1]
+    config.aug.n_transformed_imgs = 2
+    config.aug.maxjitter = 0.0
+    config.aug.maxangle = 0
+    config.aug.maxscale = 1
+    config.aug.maxshear = 20
+    config.aug.apply_affine = True
+    config.aug.debug = False
+    return config
+
+
+def makeFeatureMaker(dataset, config, device='cpu', randseed=2, feat_extr_method=None):
+    utils.fix_randseed(randseed)
+    if feat_extr_method is None:
+        feat_extr_method = Backbone(args.backbone).to(device).extract_feats
+    feat_maker = ctrutils.FeatureMaker(feat_extr_method, dataset.class_ids, config)
+    utils.fix_randseed(randseed)
+    feat_maker.norm_bb_feats = False
+    return feat_maker
+def apply_crf(rgb_img, fg_pred, thresh_fn,iterations=5): #5 on deployment, 1 on support-aug test for speedup
+    crf = segutils.CRF(gaussian_stdxy=(1,1), gaussian_compat=2,
+                 bilateral_stdxy=(35,35), bilateral_compat=1, stdrgb=(13,13,13))
+    q = crf.iterrefine(iterations, rgb_img, fg_pred, thresh_fn)
+    return q.argmax(1)
+
+def calcthresh(fused_pred, s_masks, method='otsus'):
+    if method=='iterotsus':
+        thresh = segutils.iterative_otsus(fused_pred,s_masks,maxiters=5)[0]
+        return thresh
+    elif method=='1iterotsus':
+        thresh = segutils.iterative_otsus(fused_pred,s_masks,maxiters=1)[0]
+        return thresh
+    elif method=='otsus':
+        thresh = segutils.otsus(fused_pred)[0]
+        return thresh
+    # elif method=='via_triclass':
+    #     thresh = segutils.otsus(fused_pred, mode='via_triclass')[0]
+    elif method=='pred_mean':
+        otsu_thresh = segutils.otsus(fused_pred)[0]
+        thresh = torch.max(otsu_thresh, fused_pred.mean())
+    # elif method=='3kmeans':
+    #     k3 = segutils.KMeans(fused_pred.float().view(1,-1), k=3)
+    #     thresh = k3.compute_thresholds()[0][-1]
+    return thresh
+
+def thresh_fn(method):
+    def inner(fused_pred, s_masks=None):
+        return calcthresh(fused_pred, s_masks, method)
+    return inner

data/coco.py

@@ -0,0 +1,111 @@
+r""" COCO-20i few-shot semantic segmentation dataset """
+import os
+import pickle
+
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+import torch
+import PIL.Image as Image
+import numpy as np
+
+
+class DatasetCOCO(Dataset):
+    def __init__(self, datapath, fold, transform, split, shot, use_original_imgsize=False):
+        self.split = 'val' if split in ['val', 'test'] else 'trn'
+        self.fold = fold
+        self.nfolds = 4
+        self.nclass = 80
+        self.benchmark = 'coco'
+        self.shot = shot
+        self.split_coco = split if split == 'val2014' else 'train2014'
+        self.base_path = os.path.join(datapath, 'COCO2014')
+        self.transform = transform
+        self.use_original_imgsize = use_original_imgsize
+
+        self.class_ids = self.build_class_ids()
+        self.img_metadata_classwise = self.build_img_metadata_classwise()
+        self.img_metadata = self.build_img_metadata()
+
+    def __len__(self):
+        return len(self.img_metadata) if self.split == 'trn' else 1000
+
+    def __getitem__(self, idx):
+        # ignores idx during training & testing and perform uniform sampling over object classes to form an episode
+        # (due to the large size of the COCO dataset)
+        query_img, query_mask, support_imgs, support_masks, query_name, support_names, class_sample, org_qry_imsize = self.load_frame()
+
+        query_img = self.transform(query_img)
+        query_mask = query_mask.float()
+        if not self.use_original_imgsize:
+            query_mask = F.interpolate(query_mask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze()
+
+        support_imgs = torch.stack([self.transform(support_img) for support_img in support_imgs])
+        for midx, smask in enumerate(support_masks):
+            support_masks[midx] = F.interpolate(smask.unsqueeze(0).unsqueeze(0).float(), support_imgs.size()[-2:], mode='nearest').squeeze()
+        support_masks = torch.stack(support_masks)
+
+        batch = {'query_img': query_img,
+                 'query_mask': query_mask,
+                 'query_name': query_name,
+
+                 'org_query_imsize': org_qry_imsize,
+
+                 'support_imgs': support_imgs,
+                 'support_masks': support_masks,
+                 'support_names': support_names,
+                 'class_id': torch.tensor(class_sample)}
+
+        return batch
+
+    def build_class_ids(self):
+        nclass_trn = self.nclass // self.nfolds
+        class_ids_val = [self.fold + self.nfolds * v for v in range(nclass_trn)]
+        class_ids_trn = [x for x in range(self.nclass) if x not in class_ids_val]
+        class_ids = class_ids_trn if self.split == 'trn' else class_ids_val
+
+        return class_ids
+
+    def build_img_metadata_classwise(self):
+        with open('./data/splits/coco/%s/fold%d.pkl' % (self.split, self.fold), 'rb') as f:
+            img_metadata_classwise = pickle.load(f)
+        return img_metadata_classwise
+
+    def build_img_metadata(self):
+        img_metadata = []
+        for k in self.img_metadata_classwise.keys():
+            img_metadata += self.img_metadata_classwise[k]
+        return sorted(list(set(img_metadata)))
+
+    def read_mask(self, name):
+        mask_path = os.path.join(self.base_path, 'annotations', name)
+        mask = torch.tensor(np.array(Image.open(mask_path[:mask_path.index('.jpg')] + '.png')))
+        return mask
+
+    def load_frame(self):
+        class_sample = np.random.choice(self.class_ids, 1, replace=False)[0]
+        query_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+        query_img = Image.open(os.path.join(self.base_path, query_name)).convert('RGB')
+        query_mask = self.read_mask(query_name)
+
+        org_qry_imsize = query_img.size
+
+        query_mask[query_mask != class_sample + 1] = 0
+        query_mask[query_mask == class_sample + 1] = 1
+
+        support_names = []
+        while True:  # keep sampling support set if query == support
+            support_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+            if query_name != support_name: support_names.append(support_name)
+            if len(support_names) == self.shot: break
+
+        support_imgs = []
+        support_masks = []
+        for support_name in support_names:
+            support_imgs.append(Image.open(os.path.join(self.base_path, support_name)).convert('RGB'))
+            support_mask = self.read_mask(support_name)
+            support_mask[support_mask != class_sample + 1] = 0
+            support_mask[support_mask == class_sample + 1] = 1
+            support_masks.append(support_mask)
+
+        return query_img, query_mask, support_imgs, support_masks, query_name, support_names, class_sample, org_qry_imsize
+

data/dataset.py

@@ -0,0 +1,52 @@
+r""" Dataloader builder for few-shot semantic segmentation dataset  """
+from torch.utils.data.distributed import DistributedSampler as Sampler
+from torch.utils.data import DataLoader
+from torchvision import transforms
+
+from data.pascal import DatasetPASCAL
+from data.coco import DatasetCOCO
+from data.fss import DatasetFSS
+from data.deepglobe import DatasetDeepglobe
+from data.isic import DatasetISIC
+from data.lung import DatasetLung
+from data.fss import DatasetFSS
+from data.suim import DatasetSUIM
+
+
+class FSSDataset:
+
+    @classmethod
+    def initialize(cls, img_size, datapath):
+
+        cls.datasets = {
+            'pascal': DatasetPASCAL,
+            'coco': DatasetCOCO,
+            'fss': DatasetFSS,
+            'deepglobe': DatasetDeepglobe,
+            'isic': DatasetISIC,
+            'lung': DatasetLung,
+            'suim': DatasetSUIM
+        }
+
+        cls.img_mean = [0.485, 0.456, 0.406]
+        cls.img_std = [0.229, 0.224, 0.225]
+        cls.datapath = datapath
+
+        cls.transform = transforms.Compose([transforms.Resize(size=(img_size, img_size)),
+                                            transforms.ToTensor(),
+                                            transforms.Normalize(cls.img_mean, cls.img_std)])
+
+    @classmethod
+    def build_dataloader(cls, benchmark, bsz, nworker, fold, split, shot=1):
+        nworker = nworker if split == 'trn' else 0
+
+        dataset = cls.datasets[benchmark](cls.datapath, fold=fold,
+                                          transform=cls.transform,
+                                          split=split, shot=shot)
+        # Force randomness during training for diverse episode combinations
+        # Freeze randomness during testing for reproducibility
+        #train_sampler = Sampler(dataset) if split == 'trn' else None
+        dataloader = DataLoader(dataset, batch_size=bsz, shuffle=split=='trn', num_workers=nworker,
+                                pin_memory=True)
+
+        return dataloader

data/deepglobe.py

@@ -0,0 +1,119 @@
+r""" FSS-1000 few-shot semantic segmentation dataset """
+import os
+import glob
+
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+import torch
+import PIL.Image as Image
+import numpy as np
+
+
+class DatasetDeepglobe(Dataset):
+    def __init__(self, datapath, fold, transform, split, shot, num_val=600):
+        self.split = split
+        self.benchmark = 'deepglobe'
+        self.shot = shot
+        self.num_val = num_val
+
+        self.base_path = os.path.join(datapath)
+        self.to_annpath = lambda p: p.replace('jpg', 'png').replace('origin', 'groundtruth')
+
+        self.categories = ['1','2','3','4','5','6']
+
+        self.class_ids = range(0, 6)
+        self.img_metadata_classwise, self.num_images = self.build_img_metadata_classwise()
+
+        self.transform = transform
+
+    def __len__(self):
+        # if it is the target domain, then also test on entire dataset
+        return self.num_images if self.split !='val' else self.num_val
+
+    def __getitem__(self, idx):
+        query_name, support_names, class_sample = self.sample_episode(idx)
+        query_img, query_mask, support_imgs, support_masks = self.load_frame(query_name, support_names)
+
+        query_img = self.transform(query_img)
+        query_mask = F.interpolate(query_mask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze()
+
+        support_imgs = torch.stack([self.transform(support_img) for support_img in support_imgs])
+
+        support_masks_tmp = []
+        for smask in support_masks:
+            smask = F.interpolate(smask.unsqueeze(0).unsqueeze(0).float(), support_imgs.size()[-2:], mode='nearest').squeeze()
+            support_masks_tmp.append(smask)
+        support_masks = torch.stack(support_masks_tmp)
+
+        batch = {'query_img': query_img,
+                 'query_mask': query_mask,
+                 'support_set': (support_imgs, support_masks),
+                 'support_classes': torch.tensor([class_sample]), # adapt to Nway
+
+                 'query_name': query_name, # REMOVE
+                 'support_imgs': support_imgs, # REMOVE
+                 'support_masks': support_masks, # REMOVE
+                 'support_names': support_names, # REMOVE
+                 'class_id': torch.tensor(class_sample)} # REMOVE
+
+        return batch
+
+
+    def load_frame(self, query_name, support_names):
+        query_img = Image.open(query_name).convert('RGB')
+        support_imgs = [Image.open(name).convert('RGB') for name in support_names]
+
+        query_id = query_name.split('/')[-1].split('.')[0]
+        ann_path = os.path.join(self.base_path, query_name.split('/')[-4], 'test', 'groundtruth')
+        query_name = os.path.join(ann_path, query_id) + '.png'
+        support_ids = [name.split('/')[-1].split('.')[0] for name in support_names]
+        support_names = [os.path.join(ann_path, sid) + '.png' for name, sid in zip(support_names, support_ids)]
+
+        query_mask = self.read_mask(query_name)
+        support_masks = [self.read_mask(name) for name in support_names]
+
+        return query_img, query_mask, support_imgs, support_masks
+
+    def read_mask(self, img_name):
+        mask = torch.tensor(np.array(Image.open(img_name).convert('L')))
+        mask[mask < 128] = 0
+        mask[mask >= 128] = 1
+        return mask
+
+    def sample_episode(self, idx):
+        class_id = idx % len(self.class_ids)
+        class_sample = self.categories[class_id]
+
+        query_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+        support_names = []
+        while True:  # keep sampling support set if query == support
+            support_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+            if query_name != support_name: support_names.append(support_name)
+            if len(support_names) == self.shot: break
+
+        return query_name, support_names, class_id
+
+
+    # def build_img_metadata(self):
+    #     img_metadata = []
+    #     for cat in self.categories:
+    #         os.path.join(self.base_path, cat)
+    #         img_paths = sorted([path for path in glob.glob('%s/*' % os.path.join(self.base_path, cat, 'test', 'origin'))])
+    #         for img_path in img_paths:
+    #             if os.path.basename(img_path).split('.')[1] == 'jpg':
+    #                 img_metadata.append(img_path)
+    #     return img_metadata
+
+    def build_img_metadata_classwise(self):
+        num_images=0
+        img_metadata_classwise = {}
+        for cat in self.categories:
+            img_metadata_classwise[cat] = []
+
+        for cat in self.categories:
+            img_paths = sorted([path for path in glob.glob('%s/*' % os.path.join(self.base_path, cat, 'test', 'origin'))])
+            for img_path in img_paths:
+                if os.path.basename(img_path).split('.')[1] == 'jpg':
+                    img_metadata_classwise[cat] += [img_path]
+                    num_images += 1
+        return img_metadata_classwise, num_images

data/fss.py

@@ -0,0 +1,114 @@
+r""" FSS-1000 few-shot semantic segmentation dataset """
+import os
+import glob
+
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+import torch
+import PIL.Image as Image
+import numpy as np
+
+
+class DatasetFSS(Dataset):
+    def __init__(self, datapath, fold, transform, split, shot, use_original_imgsize=False):
+        self.split = split
+        self.benchmark = 'fss'
+        self.shot = shot
+
+        self.base_path = os.path.join(datapath, 'FSS-1000')
+
+        # Given predefined test split, load randomly generated training/val splits:
+        # (reference regarding trn/val/test splits: https://github.com/HKUSTCV/FSS-1000/issues/7))
+        with open('./data/splits/fss/%s.txt' % split, 'r') as f:
+            self.categories = f.read().split('\n')[:-1]
+        self.categories = sorted(self.categories)
+
+        self.class_ids = self.build_class_ids()
+        self.img_metadata = self.build_img_metadata()
+
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.img_metadata)
+
+    def __getitem__(self, idx):
+        query_name, support_names, class_sample = self.sample_episode(idx)
+        query_img, query_mask, support_imgs, support_masks = self.load_frame(query_name, support_names)
+
+        query_img = self.transform(query_img)
+        query_mask = F.interpolate(query_mask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze()
+
+        support_imgs = torch.stack([self.transform(support_img) for support_img in support_imgs])
+
+        support_masks_tmp = []
+        for smask in support_masks:
+            smask = F.interpolate(smask.unsqueeze(0).unsqueeze(0).float(), support_imgs.size()[-2:], mode='nearest').squeeze()
+            support_masks_tmp.append(smask)
+        support_masks = torch.stack(support_masks_tmp)
+
+        batch = {'query_img': query_img,
+                 'query_mask': query_mask,
+                 'query_name': query_name,
+
+                 'support_imgs': support_imgs,
+                 'support_masks': support_masks,
+                 'support_names': support_names,
+
+                 'class_id': torch.tensor(class_sample)}
+
+        return batch
+
+    def load_frame(self, query_name, support_names):
+        query_img = Image.open(query_name).convert('RGB')
+        support_imgs = [Image.open(name).convert('RGB') for name in support_names]
+
+        query_id = query_name.split('/')[-1].split('.')[0]
+        query_name = os.path.join(os.path.dirname(query_name), query_id) + '.png'
+        support_ids = [name.split('/')[-1].split('.')[0] for name in support_names]
+        support_names = [os.path.join(os.path.dirname(name), sid) + '.png' for name, sid in zip(support_names, support_ids)]
+
+        query_mask = self.read_mask(query_name)
+        support_masks = [self.read_mask(name) for name in support_names]
+
+        return query_img, query_mask, support_imgs, support_masks
+
+    def read_mask(self, img_name):
+        mask = torch.tensor(np.array(Image.open(img_name).convert('L')))
+        mask[mask < 128] = 0
+        mask[mask >= 128] = 1
+        return mask
+
+    def sample_episode(self, idx):
+        query_name = self.img_metadata[idx]
+        class_sample = self.categories.index(query_name.split('/')[-2])
+        if self.split == 'val':
+            class_sample += 520
+        elif self.split == 'test':
+            class_sample += 760
+
+        support_names = []
+        while True:  # keep sampling support set if query == support
+            support_name = np.random.choice(range(1, 11), 1, replace=False)[0]
+            support_name = os.path.join(os.path.dirname(query_name), str(support_name)) + '.jpg'
+            if query_name != support_name: support_names.append(support_name)
+            if len(support_names) == self.shot: break
+
+        return query_name, support_names, class_sample
+
+    def build_class_ids(self):
+        if self.split == 'trn':
+            class_ids = range(0, 520)
+        elif self.split == 'val':
+            class_ids = range(520, 760)
+        elif self.split == 'test':
+            class_ids = range(760, 1000)
+        return class_ids
+
+    def build_img_metadata(self):
+        img_metadata = []
+        for cat in self.categories:
+            img_paths = sorted([path for path in glob.glob('%s/*' % os.path.join(self.base_path, cat))])
+            for img_path in img_paths:
+                if os.path.basename(img_path).split('.')[1] == 'jpg':
+                    img_metadata.append(img_path)
+        return img_metadata

data/isic.py

@@ -0,0 +1,113 @@
+r""" ISIC few-shot semantic segmentation dataset """
+import os
+import glob
+
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+import torch
+import PIL.Image as Image
+import numpy as np
+
+
+class DatasetISIC(Dataset):
+    def __init__(self, datapath, fold, transform, split, shot, num_val=600):
+        self.split = split
+        self.benchmark = 'isic'
+        self.shot = shot
+        self.num_val = num_val
+
+        self.base_path = os.path.join(datapath)
+        self.categories = ['1', '2', '3']
+
+        self.class_ids = range(0, 3)
+        self.img_metadata_classwise,self.num_images = self.build_img_metadata_classwise()
+
+        self.transform = transform
+
+    def __len__(self):
+        return self.num_images if self.split != 'val' else self.num_val
+
+    def __getitem__(self, idx):
+        query_name, support_names, class_sample = self.sample_episode(idx)
+        query_img, query_mask, support_imgs, support_masks = self.load_frame(query_name, support_names)
+
+        query_img = self.transform(query_img)
+        query_mask = F.interpolate(query_mask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze()
+
+        support_imgs = torch.stack([self.transform(support_img) for support_img in support_imgs])
+
+        support_masks_tmp = []
+        for smask in support_masks:
+            smask = F.interpolate(smask.unsqueeze(0).unsqueeze(0).float(), support_imgs.size()[-2:], mode='nearest').squeeze()
+            support_masks_tmp.append(smask)
+        support_masks = torch.stack(support_masks_tmp)
+
+        batch = {'query_img': query_img,
+                 'query_mask': query_mask,
+                 'query_name': query_name,
+
+                 'support_imgs': support_imgs,
+                 'support_masks': support_masks,
+                 'support_names': support_names,
+
+                 'class_id': torch.tensor(class_sample)}
+
+        return batch
+
+    def load_frame(self, query_name, support_names):
+        query_img = Image.open(query_name).convert('RGB')
+        support_imgs = [Image.open(name).convert('RGB') for name in support_names]
+
+        query_id = query_name.split('/')[-1].split('.')[0]
+        ann_path = os.path.join(self.base_path, 'ISIC2018_Task1_Training_GroundTruth')
+        query_name = os.path.join(ann_path, query_id) + '_segmentation.png'
+        support_ids = [name.split('/')[-1].split('.')[0] for name in support_names]
+        support_names = [os.path.join(ann_path, sid) + '_segmentation.png' for name, sid in zip(support_names, support_ids)]
+
+        query_mask = self.read_mask(query_name)
+        support_masks = [self.read_mask(name) for name in support_names]
+
+        return query_img, query_mask, support_imgs, support_masks
+
+    def read_mask(self, img_name):
+        mask = torch.tensor(np.array(Image.open(img_name).convert('L')))
+        mask[mask < 128] = 0
+        mask[mask >= 128] = 1
+        return mask
+
+    def sample_episode(self, idx):
+        class_id = idx % len(self.class_ids)
+        class_sample = self.categories[class_id]
+
+        query_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+        support_names = []
+        while True:  # keep sampling support set if query == support
+            support_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+            if query_name != support_name: support_names.append(support_name)
+            if len(support_names) == self.shot: break
+
+        return query_name, support_names, class_id
+
+    def build_img_metadata(self):
+        img_metadata = []
+        for cat in self.categories:
+            os.path.join(self.base_path, cat)
+            img_paths = sorted([path for path in glob.glob('%s/*' % os.path.join(self.base_path, 'ISIC2018_Task1-2_Training_Input', cat))])
+            for img_path in img_paths:
+                if os.path.basename(img_path).split('.')[1] == 'jpg':
+                    img_metadata.append(img_path)
+        return img_metadata
+
+    def build_img_metadata_classwise(self):
+        num_images=0
+        img_metadata_classwise = {}
+        for cat in self.categories:
+            img_metadata_classwise[cat] = []
+
+        for cat in self.categories:
+            img_paths = sorted([path for path in glob.glob('%s/*' % os.path.join(self.base_path, 'ISIC2018_Task1-2_Training_Input', cat))])
+            for img_path in img_paths:
+                if os.path.basename(img_path).split('.')[1] == 'jpg':
+                    img_metadata_classwise[cat] += [img_path]
+                    num_images += 1
+        return img_metadata_classwise, num_images

data/lung.py

@@ -0,0 +1,116 @@
+r""" Chest X-ray few-shot semantic segmentation dataset """
+import os
+import glob
+
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+import torch
+import PIL.Image as Image
+import numpy as np
+
+
+class DatasetLung(Dataset):
+    def __init__(self, datapath, fold, transform, split, shot=1, num_val=600):
+        self.benchmark = 'lung'
+        self.shot = shot
+        self.split = split
+        self.num_val = num_val
+
+        self.base_path = os.path.join(datapath)
+        self.img_path = os.path.join(self.base_path, 'CXR_png')
+        self.ann_path = os.path.join(self.base_path, 'masks')
+
+        self.categories = ['1']
+
+        self.class_ids = range(0, 1)
+        self.img_metadata_classwise, self.num_images = self.build_img_metadata_classwise()
+
+        self.transform = transform
+
+    def __len__(self):
+        return self.num_images if self.split != 'val' else self.num_val
+
+    def __getitem__(self, idx):
+        query_name, support_names, class_sample = self.sample_episode(idx)
+        query_img, query_mask, support_imgs, support_masks = self.load_frame(query_name, support_names)
+
+        query_img = self.transform(query_img)
+        query_mask = F.interpolate(query_mask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze()
+
+        support_imgs = torch.stack([self.transform(support_img) for support_img in support_imgs])
+
+        support_masks_tmp = []
+        for smask in support_masks:
+            smask = F.interpolate(smask.unsqueeze(0).unsqueeze(0).float(), support_imgs.size()[-2:], mode='nearest').squeeze()
+            support_masks_tmp.append(smask)
+        support_masks = torch.stack(support_masks_tmp)
+
+        batch = {'query_img': query_img,
+                 'query_mask': query_mask,
+                 'query_name': query_name,
+                 'support_imgs': support_imgs,
+                 'support_masks': support_masks,
+                 'class_id': torch.tensor(class_sample),
+                 'support_names': support_names,
+
+                 'support_set': [support_imgs, support_masks],
+                 'support_classes': torch.tensor([class_sample])
+                 }
+
+        return batch
+
+    def load_frame(self, query_name, support_names):
+        query_mask = self.read_mask(query_name)
+        support_masks = [self.read_mask(name) for name in support_names]
+
+        query_id = query_name[:-9] + '.png'
+        query_img = Image.open(os.path.join(self.img_path, os.path.basename(query_id))).convert('RGB')
+
+        support_ids = [os.path.basename(name)[:-9] + '.png' for name in support_names]
+        support_names = [os.path.join(self.img_path, sid) for sid in support_ids]
+        support_imgs = [Image.open(name).convert('RGB') for name in support_names]
+
+        return query_img, query_mask, support_imgs, support_masks
+
+    def read_mask(self, img_name):
+        mask = torch.tensor(np.array(Image.open(img_name).convert('L')))
+        mask[mask < 128] = 0
+        mask[mask >= 128] = 1
+        return mask
+
+    def sample_episode(self, idx):
+        class_id = idx % len(self.class_ids)
+        class_sample = self.categories[class_id]
+
+        query_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+        support_names = []
+        while True:  # keep sampling support set if query == support
+            support_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+            if query_name != support_name: support_names.append(support_name)
+            if len(support_names) == self.shot: break
+
+        return query_name, support_names, class_id
+
+    def build_img_metadata(self):
+        img_metadata = []
+        for cat in self.categories:
+            os.path.join(self.base_path, cat)
+            img_paths = sorted([path for path in glob.glob('%s/*' % os.path.join(self.img_path, cat))])
+            for img_path in img_paths:
+                if os.path.basename(img_path).split('.')[1] == 'png':
+                    img_metadata.append(img_path)
+        return img_metadata
+
+    def build_img_metadata_classwise(self):
+        num_images=0
+        img_metadata_classwise = {}
+        for cat in self.categories:
+            img_metadata_classwise[cat] = []
+
+        for cat in self.categories:
+            img_paths = sorted([path for path in glob.glob('%s/*' % self.ann_path)])
+            for img_path in img_paths:
+                if os.path.basename(img_path).split('.')[1] == 'png':
+                    img_metadata_classwise[cat] += [img_path]
+                    num_images+=1
+        return img_metadata_classwise, num_images

data/pascal.py

@@ -0,0 +1,148 @@
+r""" PASCAL-5i few-shot semantic segmentation dataset """
+import os
+
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+import torch
+import PIL.Image as Image
+import numpy as np
+
+
+class DatasetPASCAL(Dataset):
+    def __init__(self, datapath, fold, transform, split, shot, use_original_imgsize=False):
+        self.split = 'val' if split in ['val', 'test'] else 'trn'
+        self.fold = fold
+        self.nfolds = 4
+        self.nclass = 20
+        self.benchmark = 'pascal'
+        self.shot = shot
+        self.use_original_imgsize = use_original_imgsize
+
+        self.img_path = os.path.join(datapath, 'VOC2012/JPEGImages/')
+        self.ann_path = os.path.join(datapath, 'VOC2012/SegmentationClassAug/')
+        self.transform = transform
+
+        self.class_ids = self.build_class_ids()
+        self.img_metadata = self.build_img_metadata()
+        self.img_metadata_classwise = self.build_img_metadata_classwise()
+
+    def __len__(self):
+        return len(self.img_metadata) if self.split == 'trn' else 1000
+
+    def __getitem__(self, idx):
+        idx %= len(self.img_metadata)  # for testing, as n_images < 1000
+        query_name, support_names, class_sample = self.sample_episode(idx)
+        query_img, query_cmask, support_imgs, support_cmasks, org_qry_imsize = self.load_frame(query_name, support_names)
+
+        query_img = self.transform(query_img)
+        if not self.use_original_imgsize:
+            query_cmask = F.interpolate(query_cmask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze()
+        query_mask, query_ignore_idx = self.extract_ignore_idx(query_cmask.float(), class_sample)
+
+        support_imgs = torch.stack([self.transform(support_img) for support_img in support_imgs])
+
+        support_masks = []
+        support_ignore_idxs = []
+        for scmask in support_cmasks:
+            scmask = F.interpolate(scmask.unsqueeze(0).unsqueeze(0).float(), support_imgs.size()[-2:], mode='nearest').squeeze()
+            support_mask, support_ignore_idx = self.extract_ignore_idx(scmask, class_sample)
+            support_masks.append(support_mask)
+            support_ignore_idxs.append(support_ignore_idx)
+        support_masks = torch.stack(support_masks)
+        support_ignore_idxs = torch.stack(support_ignore_idxs)
+
+        batch = {'query_img': query_img,
+                 'query_mask': query_mask,
+                 'query_name': query_name,
+                 'query_ignore_idx': query_ignore_idx,
+
+                 'org_query_imsize': org_qry_imsize,
+
+                 'support_imgs': support_imgs,
+                 'support_masks': support_masks,
+                 'support_names': support_names,
+                 'support_ignore_idxs': support_ignore_idxs,
+
+                 'class_id': torch.tensor(class_sample)}
+
+        return batch
+
+    def extract_ignore_idx(self, mask, class_id):
+        boundary = (mask / 255).floor()
+        mask[mask != class_id + 1] = 0
+        mask[mask == class_id + 1] = 1
+
+        return mask, boundary
+
+    def load_frame(self, query_name, support_names):
+        query_img = self.read_img(query_name)
+        query_mask = self.read_mask(query_name)
+        support_imgs = [self.read_img(name) for name in support_names]
+        support_masks = [self.read_mask(name) for name in support_names]
+
+        org_qry_imsize = query_img.size
+
+        return query_img, query_mask, support_imgs, support_masks, org_qry_imsize
+
+    def read_mask(self, img_name):
+        r"""Return segmentation mask in PIL Image"""
+        mask = torch.tensor(np.array(Image.open(os.path.join(self.ann_path, img_name) + '.png')))
+        return mask
+
+    def read_img(self, img_name):
+        r"""Return RGB image in PIL Image"""
+        return Image.open(os.path.join(self.img_path, img_name) + '.jpg')
+
+    def sample_episode(self, idx):
+        query_name, class_sample = self.img_metadata[idx]
+
+        support_names = []
+        while True:  # keep sampling support set if query == support
+            support_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+            if query_name != support_name: support_names.append(support_name)
+            if len(support_names) == self.shot: break
+
+        return query_name, support_names, class_sample
+
+    def build_class_ids(self):
+        nclass_trn = self.nclass // self.nfolds
+        class_ids_val = [self.fold * nclass_trn + i for i in range(nclass_trn)]
+        class_ids_trn = [x for x in range(self.nclass) if x not in class_ids_val]
+
+        if self.split == 'trn':
+            return class_ids_trn
+        else:
+            return class_ids_val
+
+    def build_img_metadata(self):
+
+        def read_metadata(split, fold_id):
+            fold_n_metadata = os.path.join('data/splits/pascal/%s/fold%d.txt' % (split, fold_id))
+            with open(fold_n_metadata, 'r') as f:
+                fold_n_metadata = f.read().split('\n')[:-1]
+            fold_n_metadata = [[data.split('__')[0], int(data.split('__')[1]) - 1] for data in fold_n_metadata]
+            return fold_n_metadata
+
+        img_metadata = []
+        if self.split == 'trn':  # For training, read image-metadata of "the other" folds
+            for fold_id in range(self.nfolds):
+                if fold_id == self.fold:  # Skip validation fold
+                    continue
+                img_metadata += read_metadata(self.split, fold_id)
+        elif self.split == 'val':  # For validation, read image-metadata of "current" fold
+            img_metadata = read_metadata(self.split, self.fold)
+        else:
+            raise Exception('Undefined split %s: ' % self.split)
+
+        print('Total (%s) images are : %d' % (self.split, len(img_metadata)))
+
+        return img_metadata
+
+    def build_img_metadata_classwise(self):
+        img_metadata_classwise = {}
+        for class_id in range(self.nclass):
+            img_metadata_classwise[class_id] = []
+
+        for img_name, img_class in self.img_metadata:
+            img_metadata_classwise[img_class] += [img_name]
+        return img_metadata_classwise

data/splits/fss/test.txt

@@ -0,0 +1,240 @@
+bus
+hotel_slipper
+burj_al
+reflex_camera
+abe's_flyingfish
+oiltank_car
+doormat
+fish_eagle
+barber_shaver
+motorbike
+feather_clothes
+wandering_albatross
+rice_cooker
+delta_wing
+fish
+nintendo_switch
+bustard
+diver
+minicooper
+cathedrale_paris
+big_ben
+combination_lock
+villa_savoye
+american_alligator
+gym_ball
+andean_condor
+leggings
+pyramid_cube
+jet_aircraft
+meatloaf
+reel
+swan
+osprey
+crt_screen
+microscope
+rubber_eraser
+arrow
+monkey
+mitten
+spiderman
+parthenon
+bat
+chess_king
+sulphur_butterfly
+quail_egg
+oriole
+iron_man
+wooden_boat
+anise
+steering_wheel
+groenendael
+dwarf_beans
+pteropus
+chalk_brush
+bloodhound
+moon
+english_foxhound
+boxing_gloves
+peregine_falcon
+pyraminx
+cicada
+screw
+shower_curtain
+tredmill
+bulb
+bell_pepper
+lemur_catta
+doughnut
+twin_tower
+astronaut
+nintendo_3ds
+fennel_bulb
+indri
+captain_america_shield
+kunai
+broom
+iphone
+earphone1
+flying_squirrel
+onion
+vinyl
+sydney_opera_house
+oyster
+harmonica
+egg
+breast_pump
+guitar
+potato_chips
+tunnel
+cuckoo
+rubick_cube
+plastic_bag
+phonograph
+net_surface_shoes
+goldfinch
+ipad
+mite_predator
+coffee_mug
+golden_plover
+f1_racing
+lapwing
+nintendo_gba
+pizza
+rally_car
+drilling_platform
+cd
+fly
+magpie_bird
+leaf_fan
+little_blue_heron
+carriage
+moist_proof_pad
+flying_snakes
+dart_target
+warehouse_tray
+nintendo_wiiu
+chiffon_cake
+bath_ball
+manatee
+cloud
+marimba
+eagle
+ruler
+soymilk_machine
+sled
+seagull
+glider_flyingfish
+doublebus
+transport_helicopter
+window_screen
+truss_bridge
+wasp
+snowman
+poached_egg
+strawberry
+spinach
+earphone2
+downy_pitch
+taj_mahal
+rocking_chair
+cablestayed_bridge
+sealion
+banana_boat
+pheasant
+stone_lion
+electronic_stove
+fox
+iguana
+rugby_ball
+hang_glider
+water_buffalo
+lotus
+paper_plane
+missile
+flamingo
+american_chamelon
+kart
+chinese_knot
+cabbage_butterfly
+key
+church
+tiltrotor
+helicopter
+french_fries
+water_heater
+snow_leopard
+goblet
+fan
+snowplow
+leafhopper
+pspgo
+black_bear
+quail
+condor
+chandelier
+hair_razor
+white_wolf
+toaster
+pidan
+pyramid
+chicken_leg
+letter_opener
+apple_icon
+porcupine
+chicken
+stingray
+warplane
+windmill
+bamboo_slip
+wig
+flying_geckos
+stonechat
+haddock
+australian_terrier
+hover_board
+siamang
+canton_tower
+santa_sledge
+arch_bridge
+curlew
+sushi
+beet_root
+accordion
+leaf_egg
+stealth_aircraft
+stork
+bucket
+hawk
+chess_queen
+ocarina
+knife
+whippet
+cantilever_bridge
+may_bug
+wagtail
+leather_shoes
+wheelchair
+shumai
+speedboat
+vacuum_cup
+chess_knight
+pumpkin_pie
+wooden_spoon
+bamboo_dragonfly
+ganeva_chair
+soap
+clearwing_flyingfish
+pencil_sharpener1
+cricket
+photocopier
+nintendo_sp
+samarra_mosque
+clam
+charge_battery
+flying_frog
+ferrari911
+polo_shirt
+echidna
+coin
+tower_pisa

data/splits/fss/trn.txt

@@ -0,0 +1,520 @@
+fountain
+taxi
+assult_rifle
+radio
+comb
+box_turtle
+igloo
+head_cabbage
+cottontail
+coho
+ashtray
+joystick
+sleeping_bag
+jackfruit
+trailer_truck
+shower_cap
+ibex
+kinguin
+squirrel
+ac_wall
+sidewinder
+remote_control
+marshmallow
+bolotie
+polar_bear
+rock_beauty
+tokyo_tower
+wafer
+red_bayberry
+electronic_toothbrush
+hartebeest
+cassette
+oil_filter
+bomb
+walnut
+toilet_tissue
+memory_stick
+wild_boar
+cableways
+chihuahua
+envelope
+bison
+poker
+pubg_lvl3helmet
+indian_cobra
+staffordshire
+park_bench
+wombat
+black_grouse
+submarine
+washer
+agama
+coyote
+feeder
+sarong
+buckingham_palace
+frog
+steam_locomotive
+acorn
+german_pointer
+obelisk
+polecat
+black_swan
+butterfly
+mountain_tent
+gorilla
+sloth_bear
+aubergine
+stinkhorn
+stole
+owl
+mooli
+pool_table
+collar
+lhasa_apso
+ambulance
+spade
+pufferfish
+paint_brush
+lark
+golf_ball
+hock
+fork
+drake
+bee_house
+mooncake
+wok
+cocacola
+water_bike
+ladder
+psp
+bassoon
+bear
+border_terrier
+petri_dish
+pill_bottle
+aircraft_carrier
+panther
+canoe
+baseball_player
+turtle
+espresso
+throne
+cornet
+coucal
+eletrical_switch
+bra
+snail
+backpack
+jacamar
+scroll_brush
+gliding_lizard
+raft
+pinwheel
+grasshopper
+green_mamba
+eft_newt
+computer_mouse
+vine_snake
+recreational_vehicle
+llama
+meerkat
+chainsaw
+ferret
+garbage_can
+kangaroo
+litchi
+carbonara
+housefinch
+modem
+tebby_cat
+thatch
+face_powder
+tomb
+apple
+ladybug
+killer_whale
+rocket
+airship
+surfboard
+lesser_panda
+jordan_logo
+banana
+nail_scissor
+swab
+perfume
+punching_bag
+victor_icon
+waffle_iron
+trimaran
+garlic
+flute
+langur
+starfish
+parallel_bars
+dandie_dinmont
+cosmetic_brush
+screwdriver
+brick_card
+balance_weight
+hornet
+carton
+toothpaste
+bracelet
+egg_tart
+pencil_sharpener2
+swimming_glasses
+howler_monkey
+camel
+dragonfly
+lionfish
+convertible
+mule
+usb
+conch
+papaya
+garbage_truck
+dingo
+radiator
+solar_dish
+streetcar
+trilobite
+bouzouki
+ringlet_butterfly
+space_shuttle
+waffle
+american_staffordshire
+violin
+flowerpot
+forklift
+manx
+sundial
+snowmobile
+chickadee_bird
+ruffed_grouse
+brick_tea
+paddle
+stove
+carousel
+spatula
+beaker
+gas_pump
+lawn_mower
+speaker
+tank
+tresher
+kappa_logo
+hare
+tennis_racket
+shopping_cart
+thimble
+tractor
+anemone_fish
+trolleybus
+steak
+capuchin
+red_breasted_merganser
+golden_retriever
+light_tube
+flatworm
+melon_seed
+digital_watch
+jacko_lantern
+brown_bear
+cairn
+mushroom
+chalk
+skull
+stapler
+potato
+telescope
+proboscis
+microphone
+torii
+baseball_bat
+dhole
+excavator
+fig
+snake
+bradypod
+pepitas
+prairie_chicken
+scorpion
+shotgun
+bottle_cap
+file_cabinet
+grey_whale
+one-armed_bandit
+banded_gecko
+flying_disc
+croissant
+toothbrush
+miniskirt
+pokermon_ball
+gazelle
+grey_fox
+esport_chair
+necklace
+ptarmigan
+watermelon
+besom
+pomelo
+radio_telescope
+studio_couch
+black_stork
+vestment
+koala
+brambling
+muscle_car
+window_shade
+space_heater
+sunglasses
+motor_scooter
+ladyfinger
+pencil_box
+titi_monkey
+chicken_wings
+mount_fuji
+giant_panda
+dart
+fire_engine
+running_shoe
+dumbbell
+donkey
+loafer
+hard_disk
+globe
+lifeboat
+medical_kit
+brain_coral
+paper_towel
+dugong
+seatbelt
+skunk
+military_vest
+cocktail_shaker
+zucchini
+quad_drone
+ocicat
+shih-tzu
+teapot
+tile_roof
+cheese_burger
+handshower
+red_wolf
+stop_sign
+mouse
+battery
+adidas_logo2
+earplug
+hummingbird
+brush_pen
+pistachio
+hamster
+air_strip
+indian_elephant
+otter
+cucumber
+scabbard
+hawthorn
+bullet_train
+leopard
+whale
+cream
+chinese_date
+jellyfish
+lobster
+skua
+single_log
+chicory
+bagel
+beacon
+pingpong_racket
+spoon
+yurt
+wallaby
+egret
+christmas_stocking
+mcdonald_uncle
+wrench
+spark_plug
+triceratops
+wall_clock
+jinrikisha
+pickup
+rhinoceros
+swimming_trunk
+band-aid
+spotted_salamander
+leeks
+marmot
+warthog
+cello
+stool
+chest
+toilet_plunger
+wardrobe
+cannon
+adidas_logo1
+drumstick
+lady_slipper
+puma_logo
+great_wall
+white_shark
+witch_hat
+vending_machine
+wreck
+chopsticks
+garfish
+african_elephant
+children_slide
+hornbill
+zebra
+boa_constrictor
+armour
+pineapple
+angora
+brick
+car_wheel
+wallet
+boston_bull
+hyena
+lynx
+crash_helmet
+terrapin_turtle
+persian_cat
+shift_gear
+cactus_ball
+fur_coat
+plate
+pen
+okra
+mario
+airedale
+cowboy_hat
+celery
+macaque
+candle
+goose
+raccoon
+brasscica
+almond
+maotai_bottle
+soccer_ball
+sports_car
+tobacco_pipe
+water_polo
+eggnog
+hook
+ostrich
+patas
+table_lamp
+teddy
+mongoose
+spoonbill
+redheart
+crane
+dinosaur
+kitchen_knife
+seal
+baboon
+golfcart
+roller_coaster
+avocado
+birdhouse
+yorkshire_terrier
+saluki
+basketball
+buckler
+harvester
+afghan_hound
+beam_bridge
+guinea_pig
+lorikeet
+shakuhachi
+motarboard
+statue_liberty
+police_car
+sulphur_crested
+gourd
+sombrero
+mailbox
+adhensive_tape
+night_snake
+bushtit
+mouthpiece
+beaver
+bathtub
+printer
+cumquat
+orange
+cleaver
+quill_pen
+panpipe
+diamond
+gypsy_moth
+cauliflower
+lampshade
+cougar
+traffic_light
+briefcase
+ballpoint
+african_grey
+kremlin
+barometer
+peacock
+paper_crane
+sunscreen
+tofu
+bedlington_terrier
+snowball
+carrot
+tiger
+mink
+cristo_redentor
+ladle
+keyboard
+maraca
+monitor
+water_snake
+can_opener
+mud_turtle
+bald_eagle
+carp
+cn_tower
+egyptian_cat
+hen_of_the_woods
+measuring_cup
+roller_skate
+kite
+sandwich_cookies
+sandwich
+persimmon
+chess_bishop
+coffin
+ruddy_turnstone
+prayer_rug
+rain_barrel
+neck_brace
+nematode
+rosehip
+dutch_oven
+goldfish
+blossom_card
+dough
+trench_coat
+sponge
+stupa
+wash_basin
+electric_fan
+spring_scroll
+potted_plant
+sparrow
+car_mirror
+gecko
+diaper
+leatherback_turtle
+strainer
+guacamole
+microwave

data/splits/fss/val.txt

@@ -0,0 +1,240 @@
+handcuff
+mortar
+matchstick
+wine_bottle
+dowitcher
+triumphal_arch
+gyromitra
+hatchet
+airliner
+broccoli
+olive
+pubg_lvl3backpack
+calculator
+toucan
+shovel
+sewing_machine
+icecream
+woodpecker
+pig
+relay_stick
+mcdonald_sign
+cpu
+peanut
+pumpkin
+sturgeon
+hammer
+hami_melon
+squirrel_monkey
+shuriken
+power_drill
+pingpong_ball
+crocodile
+carambola
+monarch_butterfly
+drum
+water_tower
+panda
+toilet_brush
+pay_phone
+yonex_icon
+cricketball
+revolver
+chimpanzee
+crab
+corn
+baseball
+rabbit
+croquet_ball
+artichoke
+abacus
+harp
+bell
+gas_tank
+scissors
+vase
+upright_piano
+typewriter
+bittern
+impala
+tray
+fire_hydrant
+beer_bottle
+sock
+soup_bowl
+spider
+cherry
+macaw
+toilet_seat
+fire_balloon
+french_ball
+fox_squirrel
+volleyball
+cornmeal
+folding_chair
+pubg_airdrop
+beagle
+skateboard
+narcissus
+whiptail
+cup
+arabian_camel
+badger
+stopwatch
+ab_wheel
+ox
+lettuce
+monocycle
+redshank
+vulture
+whistle
+smoothing_iron
+mashed_potato
+conveyor
+yoga_pad
+tow_truck
+siamese_cat
+cigar
+white_stork
+sniper_rifle
+stretcher
+tulip
+handkerchief
+basset
+iceberg
+gibbon
+lacewing
+thrush
+cheetah
+bighorn_sheep
+espresso_maker
+pretzel
+english_setter
+sandbar
+cheese
+daisy
+arctic_fox
+briard
+colubus
+balance_beam
+coffeepot
+soap_dispenser
+yawl
+consomme
+parking_meter
+cactus
+turnstile
+taro
+fire_screen
+digital_clock
+rose
+pomegranate
+bee_eater
+schooner
+ski_mask
+jay_bird
+plaice
+red_fox
+syringe
+camomile
+pickelhaube
+blenheim_spaniel
+pear
+parachute
+common_newt
+bowtie
+cigarette
+oscilloscope
+laptop
+african_crocodile
+apron
+coconut
+sandal
+kwanyin
+lion
+eel
+balloon
+crepe
+armadillo
+kazoo
+lemon
+spider_monkey
+tape_player
+ipod
+bee
+sea_cucumber
+suitcase
+television
+pillow
+banjo
+rock_snake
+partridge
+platypus
+lycaenid_butterfly
+pinecone
+conversion_plug
+wolf
+frying_pan
+timber_wolf
+bluetick
+crayon
+giant_schnauzer
+orang
+scarerow
+kobe_logo
+loguat
+saxophone
+ceiling_fan
+cardoon
+equestrian_helmet
+louvre_pyramid
+hotdog
+ironing_board
+razor
+nagoya_castle
+loggerhead_turtle
+lipstick
+cradle
+strongbox
+raven
+kit_fox
+albatross
+flat-coated_retriever
+beer_glass
+ice_lolly
+sungnyemun
+totem_pole
+vacuum
+bolete
+mango
+ginger
+weasel
+cabbage
+refrigerator
+school_bus
+hippo
+tiger_cat
+saltshaker
+piano_keyboard
+windsor_tie
+sea_urchin
+microsd
+barbell
+swim_ring
+bulbul_bird
+water_ouzel
+ac_ground
+sweatshirt
+umbrella
+hair_drier
+hammerhead_shark
+tomato
+projector
+cushion
+dishwasher
+three-toed_sloth
+tiger_shark
+har_gow
+baby
+thor's_hammer
+nike_logo

data/suim.py

@@ -0,0 +1,119 @@
+r""" FSS-1000 few-shot semantic segmentation dataset """
+import os
+import glob
+
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+import torch
+import PIL.Image as Image
+import numpy as np
+
+
+class DatasetSUIM(Dataset):
+    def __init__(self, datapath, fold, transform, split, shot, num_val=600):
+        self.split = split
+        self.benchmark = 'suim'
+        self.shot = shot
+        self.num_val = num_val
+
+        self.base_path = os.path.join(datapath)
+        self.img_path = os.path.join(self.base_path, 'images')
+        self.ann_path = os.path.join(self.base_path, 'masks')
+
+        self.categories = ['FV','HD','PF','RI','RO','SR','WR']
+
+        self.class_ids = range(len(self.categories))
+        self.img_metadata_classwise, self.num_images = self.build_img_metadata_classwise()
+
+        self.transform = transform
+
+    def __len__(self):
+        # if it is the target domain, then also test on entire dataset
+        return self.num_images if self.split !='val' else self.num_val
+
+    def __getitem__(self, idx):
+        query_name, support_names, class_sample = self.sample_episode(idx)
+        query_img, query_mask, support_imgs, support_masks = self.load_frame(query_name, support_names)
+
+        query_img = self.transform(query_img)
+        query_mask = F.interpolate(query_mask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze()
+
+        support_imgs = torch.stack([self.transform(support_img) for support_img in support_imgs])
+
+        support_masks_tmp = []
+        for smask in support_masks:
+            smask = F.interpolate(smask.unsqueeze(0).unsqueeze(0).float(), support_imgs.size()[-2:], mode='nearest').squeeze()
+            support_masks_tmp.append(smask)
+        support_masks = torch.stack(support_masks_tmp)
+
+        batch = {'query_img': query_img,
+                 'query_mask': query_mask,
+                 'support_set': (support_imgs, support_masks),
+                 'support_classes': torch.tensor([class_sample]), # adapt to Nway
+
+                 'query_name': query_name, # REMOVE
+                 'support_imgs': support_imgs, # REMOVE
+                 'support_masks': support_masks, # REMOVE
+                 'support_names': support_names, # REMOVE
+                 'class_id': torch.tensor(class_sample)} # REMOVE
+
+        return batch
+
+
+    def load_frame(self, query_mask_path, support_mask_paths):
+        def maskpath_to_imgpath(maskpath):
+            filename, imgext = maskpath.split('/')[-1].split('.')[0], '.jpg'
+            return os.path.join(self.img_path, filename) + imgext
+
+        query_img = Image.open(maskpath_to_imgpath(query_mask_path)).convert('RGB')
+
+        support_imgs = [Image.open(maskpath_to_imgpath(s_mask_path)).convert('RGB') for s_mask_path in support_mask_paths]
+
+        query_mask = self.read_mask(query_mask_path)
+        support_masks = [self.read_mask(s_mask_path) for s_mask_path in support_mask_paths]
+
+        return query_img, query_mask, support_imgs, support_masks
+
+    def read_mask(self, img_name):
+        mask = torch.tensor(np.array(Image.open(img_name).convert('L')))
+        mask[mask < 128] = 0
+        mask[mask >= 128] = 1
+        return mask
+
+    def sample_episode(self, idx):
+        class_id = idx % len(self.class_ids)
+        class_sample = self.categories[class_id]
+
+        query_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+        support_names = []
+        while True:  # keep sampling support set if query == support
+            support_name = np.random.choice(self.img_metadata_classwise[class_sample], 1, replace=False)[0]
+            if query_name != support_name: support_names.append(support_name)
+            if len(support_names) == self.shot: break
+
+        return query_name, support_names, class_id
+
+
+    # def build_img_metadata(self):
+    #     img_metadata = []
+    #     for cat in self.categories:
+    #         os.path.join(self.base_path, cat)
+    #         img_paths = sorted([path for path in glob.glob('%s/*' % os.path.join(self.base_path, cat, 'test', 'origin'))])
+    #         for img_path in img_paths:
+    #             if os.path.basename(img_path).split('.')[1] == 'jpg':
+    #                 img_metadata.append(img_path)
+    #     return img_metadata
+
+    def build_img_metadata_classwise(self):
+        num_images=0
+        img_metadata_classwise = {}
+        for cat in self.categories:
+            img_metadata_classwise[cat] = []
+
+        for cat in self.categories:
+            mask_paths = sorted([path for path in glob.glob('%s/*' % os.path.join(self.base_path, 'masks', cat))])
+            for mask_path in mask_paths:
+                if self.read_mask(mask_path).count_nonzero() > 0: #no empty masks
+                    img_metadata_classwise[cat] += [mask_path]
+                    num_images += 1
+        return img_metadata_classwise, num_images

eval/evaluation.py

@@ -0,0 +1,39 @@
+r""" Evaluate mask prediction """
+import torch
+
+
+class Evaluator:
+    r""" Computes intersection and union between prediction and ground-truth """
+    @classmethod
+    def initialize(cls):
+        cls.ignore_index = 255
+
+    @classmethod
+    def classify_prediction(cls, pred_mask, batch):
+        gt_mask = batch.get('query_mask')
+
+        # Apply ignore_index in PASCAL-5i masks (following evaluation scheme in PFE-Net (TPAMI 2020))
+        query_ignore_idx = batch.get('query_ignore_idx')
+        if query_ignore_idx is not None:
+            assert torch.logical_and(query_ignore_idx, gt_mask).sum() == 0
+            query_ignore_idx *= cls.ignore_index
+            gt_mask = gt_mask + query_ignore_idx
+            pred_mask[gt_mask == cls.ignore_index] = cls.ignore_index
+
+        # compute intersection and union of each episode in a batch
+        area_inter, area_pred, area_gt = [],  [], []
+        for _pred_mask, _gt_mask in zip(pred_mask, gt_mask):
+            _inter = _pred_mask[_pred_mask == _gt_mask]
+            if _inter.size(0) == 0:  # as torch.histc returns error if it gets empty tensor (pytorch 1.5.1)
+                _area_inter = torch.tensor([0, 0], device=_pred_mask.device)
+            else:
+                _area_inter = torch.histc(_inter, bins=2, min=0, max=1)
+            area_inter.append(_area_inter)
+            area_pred.append(torch.histc(_pred_mask, bins=2, min=0, max=1))
+            area_gt.append(torch.histc(_gt_mask, bins=2, min=0, max=1))
+        area_inter = torch.stack(area_inter).t()
+        area_pred = torch.stack(area_pred).t()
+        area_gt = torch.stack(area_gt).t()
+        area_union = area_pred + area_gt - area_inter
+
+        return area_inter, area_union

eval/logger.py

@@ -0,0 +1,149 @@
+r""" Logging during training/testing """
+import datetime
+import logging
+import os
+
+from tensorboardX import SummaryWriter
+import torch
+
+class AverageMeter:
+    r""" Stores loss, evaluation results """
+    def __init__(self, dataset, device='cuda'):
+        self.benchmark = dataset.benchmark
+        if self.benchmark == 'pascal':
+            self.class_ids_interest = dataset.class_ids
+            self.class_ids_interest = torch.tensor(self.class_ids_interest).to(device)
+            self.nclass = 20
+        elif self.benchmark == 'fss':
+            self.class_ids_interest = dataset.class_ids
+            self.class_ids_interest = torch.tensor(self.class_ids_interest).to(device)
+            self.nclass = 1000
+        elif self.benchmark == 'deepglobe':
+            self.class_ids_interest = dataset.class_ids
+            self.class_ids_interest = torch.tensor(self.class_ids_interest).to(device)
+            self.nclass = 6
+        elif self.benchmark == 'isic':
+            self.class_ids_interest = dataset.class_ids
+            self.class_ids_interest = torch.tensor(self.class_ids_interest).to(device)
+            self.nclass = 3
+        elif self.benchmark == 'lung':
+            self.class_ids_interest = dataset.class_ids
+            self.class_ids_interest = torch.tensor(self.class_ids_interest).to(device)
+            self.nclass = 1
+        elif self.benchmark == 'suim':
+            self.class_ids_interest = dataset.class_ids
+            self.class_ids_interest = torch.tensor(self.class_ids_interest).to(device)
+            self.nclass = 7
+        else:
+            raise Exception('Unknown dataset: %s' % dataset)
+
+        self.intersection_buf = torch.zeros([2, self.nclass]).float().to(device)
+        self.union_buf = torch.zeros([2, self.nclass]).float().to(device)
+        self.ones = torch.ones_like(self.union_buf)
+        self.loss_buf = []
+
+    def update(self, inter_b, union_b, class_id, loss):
+        self.intersection_buf.index_add_(1, class_id, inter_b.float())
+        self.union_buf.index_add_(1, class_id, union_b.float())
+        if loss is None:
+            loss = torch.tensor(0.0)
+        self.loss_buf.append(loss)
+
+    def compute_iou(self):
+        iou = self.intersection_buf.float() / \
+              torch.max(torch.stack([self.union_buf, self.ones]), dim=0)[0]
+        iou = iou.index_select(1, self.class_ids_interest)
+        miou = iou[1].mean() * 100
+
+        fb_iou = (self.intersection_buf.index_select(1, self.class_ids_interest).sum(dim=1) /
+                  self.union_buf.index_select(1, self.class_ids_interest).sum(dim=1)).mean() * 100
+
+        return miou, fb_iou
+
+    def write_result(self, split, epoch):
+        iou,fb_iou = self.compute_iou()
+
+        loss_buf = torch.stack(self.loss_buf)
+        msg = '\n*** %s ' % split
+        msg += '[@Epoch %02d] ' % epoch
+        msg += 'Avg L: %6.5f  ' % loss_buf.mean()
+        msg += 'mIoU: %5.2f   ' % iou
+        msg += 'FB-IoU: %5.2f   ' % fb_iou
+
+        msg += '***\n'
+        Logger.info(msg)
+
+    def write_process(self, batch_idx, datalen, epoch, write_batch_idx=20):
+        if batch_idx % write_batch_idx == 0:
+            msg = '[Epoch: %02d] ' % epoch if epoch != -1 else ''
+            msg += '[Batch: %04d/%04d] ' % (batch_idx+1, datalen)
+            iou,fb_iou = self.compute_iou()
+            if epoch != -1:
+                loss_buf = torch.stack(self.loss_buf)
+                msg += 'L: %6.5f  ' % loss_buf[-1]
+                msg += 'Avg L: %6.5f  ' % loss_buf.mean()
+            msg += 'mIoU: %5.2f  |  ' % iou
+            msg += 'FB-IoU: %5.2f' % fb_iou
+            Logger.info(msg)
+
+
+class Logger:
+    r""" Writes evaluation results of training/testing """
+    @classmethod
+    def initialize(cls, args, training):
+        logtime = datetime.datetime.now().__format__('_%m%d_%H%M%S')
+        logpath = args.logpath if training else args.logpath + '_TEST_' + logtime # changed lopath created for test
+        if logpath == '': logpath = logtime
+
+        cls.logpath = os.path.join('logs', logpath + '.log')
+        cls.benchmark = args.benchmark
+        print("logdir: ",cls.logpath)
+        os.makedirs(cls.logpath)
+
+        logging.basicConfig(filemode='w',
+                            filename=os.path.join(cls.logpath, 'log.txt'),
+                            level=logging.INFO,
+                            format='%(message)s',
+                            datefmt='%m-%d %H:%M:%S')
+
+        # Console log config
+        console = logging.StreamHandler()
+        console.setLevel(logging.INFO)
+        formatter = logging.Formatter('%(message)s')
+        console.setFormatter(formatter)
+        logging.getLogger('').addHandler(console)
+
+        # Tensorboard writer
+        cls.tbd_writer = SummaryWriter(os.path.join(cls.logpath, 'tbd/runs'))
+
+        # Log arguments
+        logging.info('\n:=========== Adapt Before Comparison - A New Perspective on Cross-Domain Few-Shot Segmentation ===========')
+        for arg_key in args.__dict__:
+            logging.info('| %20s: %-24s' % (arg_key, str(args.__dict__[arg_key])))
+        logging.info(':================================================\n')
+
+    @classmethod
+    def info(cls, msg):
+        r""" Writes log message to log.txt """
+        logging.info(msg)
+
+    @classmethod
+    def save_model_miou(cls, model, epoch, val_miou):
+        torch.save(model.state_dict(), os.path.join(cls.logpath, 'best_model.pt'))
+        cls.info('Model saved @%d w/ val. mIoU: %5.2f.\n' % (epoch, val_miou))
+
+    @classmethod
+    def log_params(cls, model):
+        backbone_param = 0
+        learner_param = 0
+        for k in model.state_dict().keys():
+            n_param = model.state_dict()[k].view(-1).size(0)
+            if k.split('.')[0] in 'backbone':
+                if k.split('.')[1] in ['classifier', 'fc']:  # as fc layers are not used in HSNet
+                    continue
+                backbone_param += n_param
+            else:
+                learner_param += n_param
+        Logger.info('Backbone # param.: %d' % backbone_param)
+        Logger.info('Learnable # param.: %d' % learner_param)
+        Logger.info('Total # param.: %d' % (backbone_param + learner_param))

main.py

@@ -0,0 +1,37 @@
+from core import runner
+import torch
+import argparse
+
+def parse_opts():
+    r"""arguments"""
+    parser = argparse.ArgumentParser(description='Adapt Before Comparison - A New Perspective on Cross-Domain Few-Shot Segmentation')
+
+    # common
+    parser.add_argument('--benchmark', type=str, default='lung', choices=['fss', 'deepglobe', 'lung', 'isic', 'fss', 'lung'])
+    parser.add_argument('--datapath', type=str)
+    parser.add_argument('--nshot', type=int, default=1)
+
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    args = parse_opts()
+    print(args)
+    runner.args.benchmark = args.benchmark
+    runner.args.datapath = args.datapath
+    runner.args.nshot = args.nshot
+
+    dataloader = runner.makeDataloader()
+    config = runner.makeConfig()
+    feat_maker = runner.makeFeatureMaker(dataloader.dataset, config, device=device)
+    average_meter = runner.AverageMeterWrapper(dataloader, device)
+
+    for idx, batch in enumerate(dataloader):
+        sseval = runner.SingleSampleEval(batch, feat_maker)
+        sseval.forward()
+        sseval.calc_metrics()
+        average_meter.update(sseval)
+        average_meter.write(idx)
+    print('Result m|FB:', average_meter.average_meter.compute_iou())

utils/commonutils.py

@@ -0,0 +1,32 @@
+r""" Helper functions """
+import random
+
+import torch
+import numpy as np
+
+
+def fix_randseed(seed):
+    r""" Set random seeds for reproducibility """
+    if seed is None:
+        seed = int(random.random() * 1e5)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+
+def mean(x):
+    return sum(x) / len(x) if len(x) > 0 else 0.0
+
+
+def to_cuda(batch):
+    for key, value in batch.items():
+        if isinstance(value, torch.Tensor):
+            batch[key] = value.cuda()
+    return batch
+
+
+def to_cpu(tensor):
+    return tensor.detach().clone().cpu()

utils/segutils.py

@@ -0,0 +1,584 @@
+import os
+import cv2
+import numpy as np
+import torch
+from torchvision import transforms
+from PIL import Image, ImageDraw
+import torch.nn.functional as F
+
+norm = lambda t: (t - t.min()) / (t.max() - t.min())
+denorm = lambda t, min_, max_: t * (max_ - min_) + min_
+
+percentilerange = lambda t, perc: t.min() + perc * (t.max() - t.min())
+midrange = lambda t: percentilerange(t, .5)
+
+downsample_mask = lambda mask, H, W: F.interpolate(mask.unsqueeze(1), size=(H, W), mode='bilinear',
+                                                   align_corners=False).squeeze(1)
+
+
+# downsampled_mask: [bsz,vecs], vecs can be H*W for example
+# s_feat_volume: [bsz,c,vecs]
+# returns [bsz,c], [bsz,c,vecs]
+def fg_bg_proto(sfeat_volume, downsampled_smask):
+    B, C, vecs = sfeat_volume.shape
+    reshaped_mask = downsampled_smask.expand(B, vecs).unsqueeze(1)  # ->[B,1,vecs]
+
+    masked_fg = reshaped_mask * sfeat_volume
+    fg_proto = torch.sum(masked_fg, dim=-1) / (torch.sum(reshaped_mask, dim=-1) + 1e-8)
+
+    masked_bg = (1 - reshaped_mask) * sfeat_volume
+    bg_proto = torch.sum(masked_bg, dim=-1) / (torch.sum(1 - reshaped_mask, dim=-1) + 1e-8)
+    assert fg_proto.shape == (B, C), ":o"
+    return fg_proto, bg_proto
+
+
+# intersection = lambda pred, target: (pred * target).float().sum()
+# union = lambda pred, target: (pred + target).clamp(0, 1).float().sum()
+#
+#
+# def iou(pred, target):  # binary only, input bsz,h,w
+#     i, u = intersection(pred, target), union(pred, target)
+#     iou = (i + 1e-8) / (u + 1e-8)
+#     return iou.item()
+#
+#
+# class SimpleAvgMeter:
+#     def __init__(self, n_classes, device=torch.device('cuda')):
+#         self.n_lasses = n_classes
+#         self.intersection_buf = torch.zeros(n_classes).to(device)
+#         self.union_buf = torch.zeros(n_classes).to(device)
+#
+#     def update(self, pred, target, class_id):
+#         self.intersection_buf[class_id] += intersection(pred, target)
+#         self.union_buf[class_id] += union(pred, target)
+#
+#     def IoU(self, class_id):
+#         return self.intersection_buf[class_id] / self.union_buf[class_id] * 100
+#
+#     def cls_mIoU(self, class_ids):
+#         return (self.intersection_buf[class_ids] / self.union_buf[class_ids]).mean() * 100
+#
+#     def compute_mIoU(self):
+#         noentry = self.union_buf == 0
+#         if noentry.sum() > 0: print("SimpleAvgMeter warning: ", noentry.sum(), "elements of", self.nclasses,
+#                                     "have no empty.")
+#         return self.cls_mIoU(~noentry)
+
+# class KMeans():
+#     # expects input to be in shape [bsz, -1]
+#     def __init__(self, data, k=2, num_iterations=10):
+#         self.k = k
+#         self.device = data.device
+#         self.centroids = self._init_centroids(data)
+#
+#         for _ in range(num_iterations):
+#             labels = self._assign_clusters(data)
+#             self._update_centroids(data, labels)
+#
+#         self.labels = self._assign_clusters(data)  # Final cluster assignment
+#
+#     def _init_centroids(self, data):
+#         # Randomly initialize centroids
+#         centroids = []
+#         min_values = data.min(dim=1, keepdim=True).values
+#         range_values = (data.max(dim=1, keepdim=True).values - min_values)
+#
+#         for _ in range(self.k):
+#             random_values = torch.rand((data.shape[0], 1)).to(self.device)
+#             centroids.append(min_values + random_values * range_values)
+#
+#         return torch.cat(centroids, dim=1)
+#
+#     def _assign_clusters(self, data):
+#         # Calculate distances between data points and centroids
+#         distances = torch.abs(data.unsqueeze(2) - self.centroids)  # Expand data tensor to calculate distances
+#         # Determine the closest centroid for each data point
+#         labels = torch.argmin(distances, dim=2)
+#         # Sort labels so that the largest mean data point has the highest label
+#         cluster_means = [data[labels == k].mean() for k in range(self.k)]
+#         sorted_labels = {k: rank for rank, k in enumerate(sorted(range(self.k), key=lambda k: cluster_means[k]))}
+#         labels = torch.tensor([sorted_labels[label.item()] for label in labels.flatten()]).reshape_as(labels).to(
+#             self.device)
+#
+#         return labels
+#
+#     def _update_centroids(self, data, labels):
+#         # Calculate new centroids as the mean of the data points closest to each centroid
+#         mask = torch.nn.functional.one_hot(labels, num_classes=self.k).to(torch.float32)
+#         summed_data = torch.bmm(mask.transpose(1, 2), data.unsqueeze(2))  # Sum data points per centroid
+#         self.centroids = summed_data.squeeze() / mask.sum(dim=1, keepdim=True)  # Normalize to get the mean
+#
+#     def compute_thresholds(self):
+#         # Flatten the centroids along the middle dimension
+#         flat_centroids = self.centroids.view(self.centroids.size(0), -1)
+#
+#         # Sort the flattened centroids
+#         sorted_centroids, _ = torch.sort(flat_centroids, dim=1)
+#
+#         # Compute the midpoints between consecutive centroids
+#         thresholds = (sorted_centroids[:, :-1] + sorted_centroids[:, 1:]) / 2.0
+#
+#         return thresholds
+#
+#     def inference(self, data):
+#         # Assign data points to the nearest centroid
+#         return self._assign_clusters(data)
+
+# def iterative_triclass_thresholding(image, max_iterations=100, tolerance=25):
+#     # Ensure image is grayscale
+#     if len(image.shape) == 3:
+#         image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+#
+#     # Initialize iteration parameters
+#     TBD_region = image.copy()
+#     iteration = 0
+#     prev_threshold = 0
+#
+#     while iteration < max_iterations:
+#         iteration += 1
+#
+#         # Step 1: Apply Otsu's thresholding on the TBD region
+#         current_threshold, _ = cv2.threshold(TBD_region, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+#
+#         # Check stopping criteria
+#         if abs(current_threshold - prev_threshold) < tolerance:
+#             break
+#         prev_threshold = current_threshold
+#
+#         # Step 2: Calculate means for upper and lower regions
+#         upper_region = TBD_region[TBD_region > current_threshold]
+#         lower_region = TBD_region[TBD_region <= current_threshold]
+#
+#         if len(upper_region) == 0 or len(lower_region) == 0:
+#             break  # No further division possible
+#
+#         mean_upper = np.mean(upper_region)
+#         mean_lower = np.mean(lower_region)
+#
+#         # Step 3: Update temporary foreground, background, and TBD regions
+#         TBD_region[(TBD_region > mean_upper)] = 255  # Temporary foreground F
+#         TBD_region[(TBD_region < mean_lower)] = 0  # Temporary background B
+#
+#         # Extracting the new TBD region (between mean_lower and mean_upper)
+#         mask = (TBD_region > mean_lower) & (TBD_region < mean_upper)
+#         TBD_region = TBD_region[mask]  # Apply mask to extract region
+#
+#     # Final classification after convergence or max iterations
+#     final_foreground = (image > current_threshold).astype(np.uint8) * 255
+#     final_background = (image <= current_threshold).astype(np.uint8) * 255
+#
+#     return current_threshold, final_foreground
+
+def otsus(batched_tensor_image, drop_least=0.05, mode='ordinary'):
+    bsz = batched_tensor_image.size(0)
+    binary_tensors = []
+    thresholds = []
+
+    for i in range(bsz):
+        # Convert the tensor to numpy array
+        numpy_image = batched_tensor_image[i].cpu().numpy()
+
+        # Rescale to [0, 255] and convert to uint8 type for OpenCV compatibility
+        npmin, npmax = numpy_image.min(), numpy_image.max()
+        numpy_image = (norm(numpy_image) * 255).astype(np.uint8)
+
+        # Drop values that are in the lowest percentiles
+        truncated_vals = numpy_image[numpy_image >= int(255 * drop_least)]
+
+        # Apply Otsu's thresholding
+        if mode == 'via_triclass':
+            thresh_value, _ = iterative_triclass_thresholding(truncated_vals)
+        else:
+            thresh_value, _ = cv2.threshold(truncated_vals, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+
+        # Apply the computed threshold on the original image
+        binary_image = (numpy_image > thresh_value).astype(np.uint8) * 255
+
+        # Convert the result back to a tensor and append to the list
+        binary_tensors.append(torch.from_numpy(binary_image).float() / 255)
+
+        thresholds.append(torch.tensor(denorm(thresh_value / 255, npmin, npmax)) \
+                          .to(batched_tensor_image.device, dtype=batched_tensor_image.dtype))
+
+    # Convert list of tensors back to a single batched tensor
+    binary_tensor_batch = torch.stack(binary_tensors, dim=0)
+    thresh_batch = torch.stack(thresholds, dim=0)
+    return thresh_batch, binary_tensor_batch
+
+
+def iterative_otsus(probab_mask, s_mask, maxiters=5, mode='ordinary',
+                    debug=False):  # verify that it works correctly when batch_size >1
+    it = 1
+    otsuthresh = 0
+    assert probab_mask.min() >= 0 and probab_mask.max() <= 1, 'you should pass probabilites'
+    while True:
+        clipped = torch.where(probab_mask < otsuthresh, 0, probab_mask)
+        otsuthresh, newmask = otsus(clipped.detach(), drop_least=.02, mode=mode)
+        if otsuthresh >= s_mask.mean():
+            return otsuthresh.to(probab_mask.device), newmask.to(probab_mask.device)
+        if it >= maxiters:
+            if debug:
+                print('reached maxiter:', it, 'with thresh', otsuthresh.item(), \
+                      'removed', int(((clipped == 0).sum() / clipped.numel()).item() * 10000) / 100, \
+                      '% at lower and and new min,max is', clipped[clipped > 0].min().item(), clipped.max().item())
+                display(pilImageRow(norm(probab_mask[0]), s_mask[0], maxwidth=300))
+            return s_mask.mean(), (probab_mask > s_mask.mean()).float()  # otsuthresh
+        it += 1
+
+
+# def upgrade_scipy():
+#     os.system('!pip install - -upgrade scipy')
+#
+#
+# def slicRGB(q_img, n_segments=50, compactness=10., sigma=1, mask=None, debug=False):
+#     import skimage.segmentation as skseg
+#
+#     rgb_labels = skseg.slic(q_img, n_segments=n_segments, compactness=compactness, sigma=sigma, mask=mask,
+#                             enforce_connectivity=True)
+#
+#     if debug:
+#         plt.imshow(skseg.mark_boundaries(q_img, rgb_labels))
+#         plt.show()
+#
+#     return rgb_labels
+#
+#
+#
+# def slicRGBP(q_img, fg_pred, n_segments=30, compactness=0.1, sigma=1, mask=None, debug=False):
+#     import skimage.segmentation as skseg
+#
+#     def concat_rgb_pred(rgbimg, pred):
+#         h, w = rgbimg.shape[:2]
+#         return np.concatenate((rgbimg, pred.reshape(h, w, 1)), axis=-1)
+#
+#     rgbp_img = concat_rgb_pred(q_img, fg_pred)
+#     rgbp_labels = skseg.slic(rgbp_img, n_segments=n_segments, compactness=compactness, mask=mask, sigma=sigma,
+#                              enforce_connectivity=True)
+#
+#     if debug:
+#         rgb_labels = skseg.slic(q_img, n_segments=n_segments, compactness=10., sigma=sigma, mask=mask,
+#                                 enforce_connectivity=True)
+#         pred_labels = skseg.slic(fg_pred, n_segments=n_segments, compactness=compactness, sigma=sigma, mask=mask,
+#                                  channel_axis=None, enforce_connectivity=True)
+#
+#         rows, cols = 1, 3
+#         fig, ax = plt.subplots(rows, cols, figsize=(10, 10), sharex=True, sharey=True)
+#         ax[0].imshow(skseg.mark_boundaries(q_img, rgbp_labels))
+#         ax[1].imshow(skseg.mark_boundaries(q_img, rgb_labels))
+#         ax[2].imshow(skseg.mark_boundaries(q_img, pred_labels))
+#         plt.show()
+#
+#     return rgbp_labels
+#
+#
+# def calc_cluster_means(label_id_map, fg_prob):
+#     fg_pred_clustered = np.zeros_like(fg_prob)
+#     label_ids = np.unique(label_id_map)
+#     for lab_id in label_ids:
+#         cluster = fg_prob[label_id_map == lab_id]
+#         fg_pred_clustered[label_id_map == lab_id] = cluster.mean()
+#     return fg_pred_clustered
+
+
+def install_pydensecrf():
+    os.system('pip install git+https://github.com/lucasb-eyer/pydensecrf.git')
+
+
+class CRF:
+    def __init__(self, gaussian_stdxy=(3, 3), gaussian_compat=3,
+                 bilateral_stdxy=(80, 80), bilateral_compat=10, stdrgb=(13, 13, 13)):
+        self.gaussian_stdxy = gaussian_stdxy
+        self.gaussian_compat = gaussian_compat
+        self.bilateral_stdxy = bilateral_stdxy
+        self.bilateral_compat = bilateral_compat
+        self.stdrgb = stdrgb
+        self.iters = 5
+        self.debug = False
+
+    def refine(self, image_tensor, fg_probs, soft_thresh=None, T=1):
+
+        """
+        Refine segmentation using DenseCRF.
+
+        Args:
+            - image_tensor (tensor): Original image, shape [1, 3, H, W].
+            - fg_probs (tensor): Fg probabilities from the network, shape [1, H, W]
+            - soft_thresh: The preferred threshold for fg_probs for segmenting into binary prediction mask
+            - T: a temperature for softmax/sigmoid
+
+        Returns:
+            - Refined segmentation mask, shape [1, H, W].
+        """
+        try:
+            import pydensecrf.densecrf as dcrf
+            from pydensecrf.utils import unary_from_softmax, create_pairwise_bilateral
+        except ImportError as e:
+            print("pydensecrf not found. Installing...")
+            install_pydensecrf()  # Ensure this function installs pydensecrf and handles any potential errors during installation.
+
+        # After installation, retry importing. This is placed inside the except block to avoid repeating the import statements.
+        try:
+            import pydensecrf.densecrf as dcrf
+            from pydensecrf.utils import unary_from_softmax, create_pairwise_bilateral
+        except ImportError as e:
+            print("Failed to import after installation. Please check the installation of pydensecrf.")
+            raise  # This will raise the last exception that was handled by the except block
+
+        # We find the segmentation threshold that splits fg-bg
+        if soft_thresh is None:
+            soft_thresh, _ = otsus(fg_probs)
+        image_tensor, fg_probs, soft_thresh = image_tensor.cpu(), fg_probs.cpu(), soft_thresh.cpu()
+        # Then we presume at this threshold the probability should be 0.5
+        # probability 0 should stay 0, 1 should stay 1
+        # sigmoid=lambda x: 1/(1 + np.exp(-x))
+        fg_probs = torch.sigmoid(T * (fg_probs - soft_thresh))
+        probs = torch.stack([1 - fg_probs, fg_probs], dim=1)  # crf expects both classes as input
+        if self.debug:
+            print('softthresh', soft_thresh)
+            print('fg_probs min max', fg_probs.min(), fg_probs.max())
+        # C: Number of classes
+        bsz, C, H, W = probs.shape
+
+        refined_masks = []
+        image_numpy = np.ascontiguousarray( \
+            (255 * image_tensor.permute(0, 2, 3, 1)).numpy().astype(np.uint8))
+        probs_numpy = probs.numpy()
+        for (image, prob) in zip(image_numpy, probs_numpy):
+            # Unary potentials
+            unary = np.ascontiguousarray(unary_from_softmax(prob))
+            d = dcrf.DenseCRF2D(W, H, C)
+            d.setUnaryEnergy(unary)
+
+            # Add pairwise potentials
+            d.addPairwiseGaussian(sxy=self.gaussian_stdxy, compat=self.gaussian_compat)
+            d.addPairwiseBilateral(sxy=self.bilateral_stdxy, srgb=self.stdrgb,
+                                   rgbim=image, compat=self.bilateral_compat)
+
+            # Perform inference
+            Q = d.inference(self.iters)
+            if self.debug:
+                print('Q:', np.array(Q).shape, np.array(Q)[0].mean(), np.array(Q).mean())
+            result = np.reshape(Q, (2, H, W))  # np.argmax(Q, axis=0).reshape((H, W))
+            refined_masks.append(result)
+
+        return torch.from_numpy(np.stack(refined_masks, axis=0))
+
+    #     def iterrefine(self, iters, image_tensor, fg_probs, soft_thresh=None, T=1):
+    #         q1 = fg_probs
+    #         for iter in range(iters):
+    #             print(q1.shape)
+    #             q1 = self.refine(image_tensor, q1, soft_thresh=None, T=1)[:,1]
+    #         return q1
+
+    def iterrefine(self, iters, q_img, fg_probs, thresh_fn, debug=False):
+        pred = fg_probs.unsqueeze(1).expand(1, 2, *fg_probs.shape[-2:])
+        for it in range(iters):
+            thresh = thresh_fn(pred[:, 1])[0]
+
+            if debug and i % 10 == 0:
+                print('thresh', thresh)
+                display(to_pil(pred[0, 1]))
+
+            pred = self.refine(q_img, pred[:, 1], soft_thresh=thresh)
+        return pred
+
+
+#
+# class Subplot:
+#     def __init__(self):
+#         self.vertical_lines = []
+#         self.histograms = []
+#         self.gaussian_curves = []
+#         self.colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
+#         self.title = ''
+#
+#     class Element:
+#         def __init__(self, x=None, y=None, label=''):
+#             if x is not None:
+#                 self.x = Subplot.to_np(x)
+#             if y is not None:
+#                 self.y = Subplot.to_np(y)
+#
+#             self.label = label
+#
+#     @staticmethod
+#     def to_np(t):
+#         return t.detach().cpu().numpy()
+#
+#     def add_vertical(self, x, label=''):
+#         self.vertical_lines.append(Subplot.Element(x=x, label=label))
+#         return self
+#
+#     def add_histogram(self, samples, label=''):
+#         self.histograms.append(Subplot.Element(x=samples, label=label))
+#         return self
+#
+#     def add_gaussian(self, gaussian):
+#         samples, mu, var = gaussian.samples, gaussian.mean, gaussian.covs
+#         # Generate a range of x values
+#         x_values = np.linspace(samples.min(), samples.max(), 100)
+#         x_values = np.linspace(samples.min(), samples.max(), 100)
+#
+#         # Compute Gaussian values for these x values
+#         gaussian1_values = gaussian.gaussian_pdf(x_values, mu[0].item(), var[0].item())
+#         gaussian2_values = gaussian.gaussian_pdf(x_values, mu[1].item(), var[1].item())
+#         self.gaussian_curves.append(Subplot.Element(x_values, gaussian1_values))
+#         self.gaussian_curves.append(Subplot.Element(x_values, gaussian2_values))
+#         return self
+#
+#
+# class PredHistos2():
+#     def __init__(self, n_cols=1):
+#         self.fig, self.axes = plt.subplots(nrows=1, ncols=n_cols, figsize=(10, 4))
+#         self.n_cols = n_cols
+#         if n_cols == 1:
+#             self.builder = Subplot()
+#         self.subplots = [Subplot() for x in range(n_cols)]
+#         self.alpha = 0.5
+#         self.bins = 200
+#
+#     def reload(self, n_cols=1):
+#         self.fig, self.axes = plt.subplots(nrows=1, ncols=n_cols, figsize=(10, 4))
+#
+#     def aggr(self, ax, sub):
+#         for hist, col in zip(sub.histograms, sub.colors):
+#             ax.hist(hist.x, self.bins, density=True, color=col, alpha=self.alpha, label=hist.label)
+#         for vline, col in zip(sub.vertical_lines, sub.colors):
+#             ax.axvline(x=vline.x, color=col, label=vline.label, linestyle='--')
+#         for gaussian, col in zip(sub.gaussian_curves, sub.colors):
+#             ax.plot(gaussian.x, gaussian.y, gaussian.label, col)
+#         ax.legend()
+#
+#     def plot(self, name=''):
+#
+#         if self.n_cols == 1:
+#             self.aggr(plt, self.builder)
+#         else:
+#             for ax, sub in zip(self.axes, self.subplots):
+#                 self.aggr(ax, sub)
+#                 ax.set_title(sub.title)
+#
+#         plt.legend()
+#         plt.title(name)
+#         plt.show()
+#
+#
+# from sklearn.mixture import GaussianMixture
+# import scipy.optimize as opt
+# from scipy.optimize import fsolve
+#
+#
+# class GMM:
+#     def __init__(self, q_pred_coarse, name='gaussian', n_components=2):
+#         samples = q_pred_coarse.detach().cpu().numpy()
+#         self.samples = samples.reshape(-1, 1)
+#
+#         # Fit a mixture of 2 Gaussians using EM
+#         gmm = GaussianMixture(n_components)
+#         gmm.fit(samples)
+#         self.means = gmm.means_.flatten()
+#         self.covs = gmm.covariances_.flatten()
+#         self.weights = gmm.weights_
+#         self.label = name
+#
+#     def intersect(self):
+#         # Use fsolve to find the intersection
+#         gaussian_intersect, = fsolve(difference, self.means.mean(), args=(
+#         self.means[0].item(), self.covs[0].item(), self.means[1].item(), self.means[1].item()))
+#         return gaussian_intersect
+#
+#
+# class PredHistoSNS():
+#     def __init__(self, n_cols=1):
+#         import seaborn as sns
+#         sns.set_theme(style="whitegrid")  # Set the Seaborn theme. You can change the style as needed.
+#         self.fig, self.axes = plt.subplots(nrows=1, ncols=n_cols, figsize=(10, 4))
+#         self.n_cols = n_cols
+#         if n_cols == 1:
+#             self.axes = [self.axes]  # Wrap the single axis in a list to simplify the loop logic later.
+#             self.builder = Subplot()  # This is assuming Subplot is a properly defined class.
+#         self.subplots = [Subplot() for _ in range(n_cols)]  # Use underscore for unused loop variable.
+#         self.alpha = 0.5
+#         self.bins = 200
+#
+#     def reload(self, n_cols=1):
+#         self.fig, self.axes = plt.subplots(nrows=1, ncols=n_cols, figsize=(10, 4))
+#
+#     def aggr(self, ax, sub):
+#         import seaborn as sns
+#         for hist, col in zip(sub.histograms, sub.colors):
+#             sns.histplot(hist.x, bins=self.bins, kde=False, color=col, ax=ax, alpha=self.alpha, label=hist.label)
+#         for vline, col in zip(sub.vertical_lines, sub.colors):
+#             ax.axvline(x=vline.x, color=col, label=vline.label, linestyle='--')
+#         for gaussian, col in zip(sub.gaussian_curves, sub.colors):
+#             sns.lineplot(x=gaussian.x, y=gaussian.y, label=gaussian.label, color=col, ax=ax)
+#         ax.legend()
+#
+#     def plot(self, name=''):
+#
+#         if self.n_cols == 1:
+#             self.aggr(self.axes[0], self.builder)
+#         else:
+#             for ax, sub in zip(self.axes, self.subplots):
+#                 self.aggr(ax, sub)
+#                 ax.set_title(sub.title)
+#
+#         plt.show()
+#
+#
+# def overlay_mask(image, mask, color=[255, 0, 0], alpha=0.2):
+#     """
+#     Apply an overlay of a binary mask onto an image using a specified color.
+#
+#     :param image: A PyTorch tensor of the image (C x H x W) with pixel values in [0, 1].
+#     :param mask: A PyTorch tensor of the mask (H x W) with binary values (0 or 1).
+#     :param color: A list of 3 elements representing the RGB values of the overlay color.
+#     :param alpha: A float representing the transparency of the overlay (0 to 1).
+#     :return: An overlayed image tensor.
+#     """
+#     # Ensure the mask is binary
+#     mask = (mask > 0).float()
+#
+#     # Create an RGB version of the mask
+#     mask_rgb = torch.tensor(color).view(3, 1, 1) / 255.0  # Normalize the color vector
+#     mask_rgb = mask_rgb * mask
+#
+#     # Overlay the mask onto the image
+#     overlayed_image = (1 - alpha) * image + alpha * mask_rgb
+#
+#     # Ensure the resulting tensor values are between 0 and 1
+#     overlayed_image = torch.clamp(overlayed_image, 0, 1)
+#
+#     return overlayed_image
+#
+#
+# import pandas as pd
+
+# to_pil = lambda t: transforms.ToPILImage()(t) if t.shape[-1] > 4 else transforms.ToPILImage()(t.permute(2, 0, 1))
+#
+#
+# def pilImageRow(*imgs, maxwidth=800, bordercolor=0x000000):
+#     imgs = [to_pil(im.float()) for im in imgs]
+#     dst = Image.new('RGB', (sum(im.width for im in imgs), imgs[0].height))
+#     for i, im in enumerate(imgs):
+#         loc = [x0, y0, x1, y1] = [i * im.width, 0, (i + 1) * im.width, im.height]
+#         dst.paste(im, (x0, y0))
+#         ImageDraw.Draw(dst).rectangle(loc, width=2, outline=bordercolor)
+#     factorToBig = dst.width / maxwidth
+#     dst = dst.resize((int(dst.width / factorToBig), int(dst.height / factorToBig)))
+#     return dst
+#
+#
+# def tensor_table(**kwargs):
+#     tensor_overview = {}
+#     for name, tensor in kwargs.items():
+#         if callable(tensor):
+#             print(name, [tensor(t) for _, t in kwargs.items() if isinstance(t, torch.Tensor)])
+#         else:
+#             tensor_overview[name] = {
+#                 'min': tensor.min().item(),
+#                 'max': tensor.max().item(),
+#                 'shape': tensor.shape,
+#             }
+#     return pd.DataFrame.from_dict(tensor_overview, orient='index')
+