Upload 123 files

configs/multi_mo_multi_task.yaml

@@ -0,0 +1,156 @@
+train_dataset:
+  dataset:
+    name: paired-image-folders-multi-task
+    args:
+#      root_path_1: ./SAM_DATA_UNIFY/Overall_Update/split_image
+#      root_path_1: ./SAM_DATA_UNIFY2/OVERALL/split_image
+#      root_path_1: ./SAM_DATA_UNIFY2/ISAID/split_image
+#      root_path_1: [{'ISAID': './SAM_DATA_UNIFY2/ISAID/split_image', 'WHU': './SAM_DATA_UNIFY2/WHU-OPT/split_images'}]
+#      root_path_1: [{'Decoder1': "/workspace/SAM_DATA_UNIFY3/Decoder1/split_image/", 'Decoder2': "/workspace/SAM_DATA_UNIFY3/Decoder2/split_image/"}]
+      root_path_1: [{'Decoder1': "/workspace/SAM_DATA_UNIFY4/Decoder1/image/", 'Decoder2': "/workspace/SAM_DATA_UNIFY4/Decoder2/image/"}]
+#      root_path_2: ./SAM_DATA_UNIFY/Overall_Update/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY2/OVERALL/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY2/ISAID/split_gt
+#      root_path_2: [{'ISAID': './SAM_DATA_UNIFY2/ISAID/split_gt', 'WHU': './SAM_DATA_UNIFY2/WHU-OPT/split_gt'}]
+#      root_path_2: [{'Decoder1': "/workspace/SAM_DATA_UNIFY3/Decoder1/split_gt/", 'Decoder2': "/workspace/SAM_DATA_UNIFY3/Decoder2/split_gt/"}]
+      root_path_2: [{'Decoder1': "/workspace/SAM_DATA_UNIFY4/Decoder1/gt/", 'Decoder2': "/workspace/SAM_DATA_UNIFY4/Decoder2/gt/"}]
+      cache: nones
+      split_key: train
+  wrapper:
+    name: train_multi_task
+    args:
+      inp_size: 1024
+      augment: false
+#  batch_size: 2
+  batch_size: 2
+
+val_dataset:
+  dataset:
+    name: paired-image-folders-multi-task
+    args:
+#      root_path_1: ./SAM_DATA_UNIFY2/OVERALL/split_image
+#      root_path_1: [{'ISAID': './SAM_DATA_UNIFY2/ISAID/split_image', 'WHU': './SAM_DATA_UNIFY2/WHU-OPT/split_images'}]
+#      root_path_1: [{'Decoder1': "/workspace/SAM_DATA_UNIFY3/Decoder1/split_image/", 'Decoder2': "/workspace/SAM_DATA_UNIFY3/Decoder2/split_image/"}]
+      root_path_1: [{'Decoder1': "/workspace/SAM_DATA_UNIFY4/Decoder1/image/", 'Decoder2': "/workspace/SAM_DATA_UNIFY4/Decoder2/image/"}]
+#      root_path_2: ./SAM_DATA_UNIFY2/OVERALL/split_gt
+#      root_path_2: [{'ISAID': './SAM_DATA_UNIFY2/ISAID/split_gt', 'WHU': './SAM_DATA_UNIFY2/WHU-OPT/split_gt'}]
+#      root_path_2: [{'Decoder1': "/workspace/SAM_DATA_UNIFY3/Decoder1/split_gt/", 'Decoder2': "/workspace/SAM_DATA_UNIFY3/Decoder2/split_gt/"}]
+      root_path_2: [{'Decoder1': "/workspace/SAM_DATA_UNIFY4/Decoder1/gt/", 'Decoder2': "/workspace/SAM_DATA_UNIFY4/Decoder2/gt/"}]
+      cache: none
+      split_key: test
+  wrapper:
+    name: val_multi_task
+    args:
+      inp_size: 1024
+#  batch_size: 2
+  batch_size: 1
+
+test_dataset:
+  dataset:
+    name: paired-image-folders
+    args:
+
+#      root_path_1: ./SAM_DATA_UNIFY3/ISAID/split_image
+#      root_path_1: ./SAM_DATA_UNIFY3/GANFEN/split_image
+#      root_path_1: ./SAM_DATA_UNIFY3/SAR2020/split_image_ov500
+#      root_path_1: ./SAM_DATA_UNIFY3/ISAID/split_image
+#      root_path_1: ./SAM_DATA_UNIFY4/SAR2020/split_image_ov500
+#      root_path_1: ./SAM_DATA_UNIFY4/GAOFEN/split_image
+#      root_path_1: ./SAM_DATA_UNIFY4/Vaihingen/image1
+#      root_path_1: ./SAM_DATA_UNIFY4/SAR2020/split_image_ov500
+#        root_path_1: ./SAM_DATA_UNIFY4/Potsdam/image1
+#      root_path_1: ./SAM_DATA_UNIFY4/whu-opt-sar/image_sar
+      root_path_1: /workspace/AIService/FoundationModel/sam_adapter_01/TwoDecoder_data/Prompt_GUOLV_Data/prompt_all1/image
+        
+#      root_path_2: ./SAM_DATA_UNIFY3/ISAID/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY3/GANFEN/gt_decoder1
+#      root_path_2: ./SAM_DATA_UNIFY3/GANFEN/gt_decoder2
+#      root_path_2: ./SAM_DATA_UNIFY3/SAR2020/gt_decoder2
+#      root_path_2: ./SAM_DATA_UNIFY3/ISAID/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY4/SAR2020/gt_decoder2
+#      root_path_2: ./SAM_DATA_UNIFY4/GAOFEN/gt_decoder1_update
+#      root_path_2: ./SAM_DATA_UNIFY4/Vaihingen/gt2
+#      root_path_2: ./SAM_DATA_UNIFY4/Potsdam/gt1
+#      root_path_2: ./SAM_DATA_UNIFY4/SAR2020/gt_decoder2
+      root_path_2: /workspace/AIService/FoundationModel/sam_adapter_01/TwoDecoder_data/Prompt_GUOLV_Data/prompt_all1/gt
+#      root_path_2: ./SAM_DATA_UNIFY4/whu-opt-sar/gt_sar
+      cache: none
+      split_key: test
+  wrapper:
+    name: val
+    args:
+#      inp_size: 1024
+      inp_size: 1024
+  batch_size: 1
+
+#eval_type: cod
+eval_type: f1
+#sam_checkpoint: ./pretrained/sam_vit_l_0b3195.pth
+sam_checkpoint: sam_vit_h_4b8939.pth
+data_norm:
+  inp:
+    sub:
+    - 0.5
+    div:
+    - 0.5
+  gt:
+    sub:
+    - 0.5
+    div:
+    - 0.5
+  gt_rgb:
+    sub:
+    - 0.5
+    div:
+    - 0.5
+model:
+  name: sam
+  args:
+    inp_size: 1024
+#    loss: iou
+    loss: cr
+    encoder_mode:
+      name: sam
+      img_size: 1024
+      mlp_ratio: 4
+      patch_size: 16
+      qkv_bias: true
+      use_rel_pos: true
+      window_size: 14
+      out_chans: 256
+      scale_factor: 32
+      input_type: fft
+      freq_nums: 0.25
+      prompt_type: highpass
+      prompt_embed_dim: 256
+      tuning_stage: 1234
+      handcrafted_tune: true
+      embedding_tune: true
+      adaptor: adaptor
+      embed_dim: 1280
+      depth: 32
+      num_heads: 16
+      global_attn_indexes:
+      - 7
+      - 15
+      - 23
+      - 31
+optimizer:
+  name: adamw
+  args:
+#    lr: 0.0002
+#    lr: 0.00002
+    lr: 0.00008
+lr_min: 1.0e-8
+#epoch_max: 20
+epoch_max: 100
+
+multi_step_lr:
+  milestones:
+  - 1
+  gamma: 0.1
+epoch_val: 100
+epoch_save: 1
+
+#resume: 60
+#start_epoch: 60

configs/multi_mo_multi_task_sar_prompt.yaml

@@ -0,0 +1,174 @@
+train_dataset:
+  dataset:
+    name: paired-image-folders
+    args:
+#      root_path_1: ./ISAID/train/trainprompt/sub_images
+#      root_path_1: ./ISAID/train/trainprompt/images
+      root_path_1: ./SAR_prompt/image
+#      root_path_1: ./SAM_DATA_UNIFY2/OVERALL/split_image
+#      root_path_1: ./SAM_DATA_UNIFY2/ISAID/split_image
+#      root_path_1: [{'ISAID': './SAM_DATA_UNIFY2/ISAID/split_image', 'WHU': './SAM_DATA_UNIFY2/WHU-OPT/split_images'}]
+#      root_path_1: [{'Decoder1': "/workspace/SAM_DATA_UNIFY3/Decoder1/split_image/", 'Decoder2': "/workspace/SAM_DATA_UNIFY3/Decoder2/split_image/"}]
+#     root_path_1: [{'Decoder1': "/workspace/SAM_DATA_UNIFY4/Potsdam/image1/", 'Decoder2': "/workspace/SAM_DATA_UNIFY4/Decoder2/image/"}]
+#      root_path_2: ./ISAID/train/trainprompt/sub_gt
+      root_path_2: ./SAR_prompt/gt
+#      root_path_2: ./SAM_DATA_UNIFY2/OVERALL/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY2/ISAID/split_gt
+#      root_path_2: [{'ISAID': './SAM_DATA_UNIFY2/ISAID/split_gt', 'WHU': './SAM_DATA_UNIFY2/WHU-OPT/split_gt'}]
+#      root_path_2: [{'Decoder1': "/workspace/SAM_DATA_UNIFY3/Decoder1/split_gt/", 'Decoder2': "/workspace/SAM_DATA_UNIFY3/Decoder2/split_gt/"}]
+#      root_path_2: [{'Decoder1': "/workspace/SAM_DATA_UNIFY4/Potsdam/gt1/", 'Decoder2': "/workspace/SAM_DATA_UNIFY4/Decoder2/gt/"}]
+      cache: none
+      split_key: train
+  wrapper:
+    name: train
+    args:
+      inp_size: 1024
+      augment: false
+#  batch_size: 2
+  batch_size: 1
+
+val_dataset:
+  dataset:
+    name: paired-image-folders
+    args:
+#      root_path_1: ./ISAID/train/trainprompt/images
+      root_path_1: ./SAR_prompt/image
+#      root_path_1: [{'ISAID': './SAM_DATA_UNIFY2/ISAID/split_image', 'WHU': './SAM_DATA_UNIFY2/WHU-OPT/split_images'}]
+#      root_path_1: [{'Decoder1': "/workspace/SAM_DATA_UNIFY3/Decoder1/split_image/", 'Decoder2': "/workspace/SAM_DATA_UNIFY3/Decoder2/split_image/"}]
+#      root_path_1: [{'Decoder1': "/workspace/SAM_DATA_UNIFY4/Potsdam/image1/", 'Decoder2': "/workspace/SAM_DATA_UNIFY4/Decoder2/image/"}]
+#      root_path_2: ./ISAID/train/trainprompt/gt
+      root_path_2: ./SAR_prompt/gt
+#      root_path_2: [{'ISAID': './SAM_DATA_UNIFY2/ISAID/split_gt', 'WHU': './SAM_DATA_UNIFY2/WHU-OPT/split_gt'}]
+#      root_path_2: [{'Decoder1': "/workspace/SAM_DATA_UNIFY3/Decoder1/split_gt/", 'Decoder2': "/workspace/SAM_DATA_UNIFY3/Decoder2/split_gt/"}]
+#      root_path_2: [{'Decoder1': "/workspace/SAM_DATA_UNIFY4/Potsdam/gt1/", 'Decoder2': "/workspace/SAM_DATA_UNIFY4/Decoder2/gt/"}]
+      cache: none
+      split_key: test
+  wrapper:
+    name: val
+    args:
+      inp_size: 1024
+#  batch_size: 2
+  batch_size: 1
+
+test_dataset:
+  dataset:
+    name: paired-image-folders
+    args:
+#      root_path_1: ./ISAID/train/trainprompt/images
+#      root_path_1: ./ISAID/train/trainprompt/sub_images
+      root_path_1:  ./save/SAR_prompt/image
+#      root_path_1: ./SAM_DATA_UNIFY/Vaihingen/split_image
+#      root_path_1: ./SAM_DATA_UNIFY/SAR2020/split_image_ov500
+#      root_path_1: ./SAM_DATA_UNIFY/POLARIS_SAR/split_image
+#      root_path_1: ./SAM_DATA_UNIFY/Overall_Update/split_image
+#      root_path_1: ./SAM_DATA_UNIFY2/ISAID/split_image
+#      root_path_1: ./SAM_DATA_UNIFY2/whu-sar-test/split_image
+#      root_path_1: ./SAM_DATA_UNIFY2/WHU-SAR/split_image
+#      root_path_1: ./SAM_DATA_UNIFY2/WHU_ALL/split_image
+#      root_path_1: ./SAM_DATA_UNIFY3/WHU_SAR/split_image
+#      root_path_1: ./SAM_DATA_UNIFY3/WHU_OPT/split_image
+#      root_path_1: ./SAM_DATA_UNIFY3/ISAID/split_image
+#      root_path_1: ./SAM_DATA_UNIFY3/GANFEN/split_image
+#      root_path_1: ./SAM_DATA_UNIFY4/SAR2020/split_image_ov500
+
+#      root_path_2: ./ISAID/train/trainprompt/gt
+#      root_path_2: ./ISAID/train/trainprompt/sub_gt
+      root_path_2:  ./save/SAR_prompt/gt
+#      root_path_2: ./SAM_DATA_UNIFY/Vaihingen/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY2/ISAID/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY/POLARIS_SAR/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY/Overall_Update/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY2/ISAID/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY2/whu-sar-test/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY2/WHU-SAR/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY2/WHU_ALL/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY3/WHU_SAR/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY3/WHU_OPT/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY3/ISAID/split_gt
+#      root_path_2: ./SAM_DATA_UNIFY3/GANFEN/gt_decoder1
+#      root_path_2: ./SAM_DATA_UNIFY3/GANFEN/gt_decoder2
+#      root_path_2: ./SAM_DATA_UNIFY4/SAR2020/gt_decoder2
+      cache: none
+      split_key: test
+  wrapper:
+    name: val
+    args:
+#      inp_size: 1024
+      inp_size: 1024
+  batch_size: 1
+
+#eval_type: cod
+eval_type: f1
+#sam_checkpoint: ./pretrained/sam_vit_l_0b3195.pth
+#sam_checkpoint: sam_vit_h_4b8939.pth
+sam_checkpoint: ./save/_multi_mo_multi_task_0626/model_epoch_last.pth
+#sam_checkpoint: ./save/_multi_mo_multi_task_0626/model_epoch_last.pth
+data_norm:
+  inp:
+    sub:
+    - 0.5
+    div:
+    - 0.5
+  gt:
+    sub:
+    - 0.5
+    div:
+    - 0.5
+  gt_rgb:
+    sub:
+    - 0.5
+    div:
+    - 0.5
+model:
+  name: sam
+  args:
+    inp_size: 1024
+#    loss: iou
+    loss: cr
+    encoder_mode:
+      name: sam
+      img_size: 1024
+      mlp_ratio: 4
+      patch_size: 16
+      qkv_bias: true
+      use_rel_pos: true
+      window_size: 14
+      out_chans: 256
+      scale_factor: 32
+      input_type: fft
+      freq_nums: 0.25
+      prompt_type: highpass
+      prompt_embed_dim: 256
+      tuning_stage: 1234
+      handcrafted_tune: true
+      embedding_tune: true
+      adaptor: adaptor
+      embed_dim: 1280
+      depth: 32
+      num_heads: 16
+      global_attn_indexes:
+      - 7
+      - 15
+      - 23
+      - 31
+optimizer:
+  name: adamw
+  args:
+#    lr: 0.0002
+#    lr: 0.00002
+#    lr: 0.00004
+#    lr: 0.00008
+    lr: 0.0002
+lr_min: 1.0e-8
+#epoch_max: 20
+epoch_max: 200
+
+multi_step_lr:
+  milestones:
+  - 1
+  gamma: 0.1
+epoch_val: 200
+epoch_save: 1
+
+#resume: 60
+#start_epoch: 60

datasets/__init__.py

@@ -0,0 +1,3 @@
+from .datasets import register, make
+from . import image_folder
+from . import wrappers

datasets/data_loader_multi_tasks.py

@@ -0,0 +1,26 @@
+
+def build_loader_simmim(config):
+    ############ single model #####################
+    # transform = SimMIMTransform(config)
+    # dataset = ImageFolder(config.DATA.DATA_PATH, transform)
+    # sampler = DistributedSampler(dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank(), shuffle=True)
+    # dataloader = DataLoader(dataset, config.DATA.BATCH_SIZE, sampler=sampler, num_workers=config.DATA.NUM_WORKERS, pin_memory=True, drop_last=True, collate_fn=collate_fn)
+
+    ############## multi model ####################
+    datasets = []
+    ### 数据增强 ######
+    model_paths = config.DATA.TYPE_PATH[0]
+    for i in model_paths.keys():
+        a = config.DATA.SCALE[0][i].split(',')
+        scale_model = (float(a[0].split('(')[1]) ,float(a[1].split(')')[0]))
+        transform = SimMIMTransform(config, config.DATA.NORM[0][i], scale_model)
+        dataset = CachedImageFolder(model_paths[i], transform = transform, model = i)
+        datasets.append(dataset)
+    multi_task_train_dataset = MultiTaskDataset(datasets)
+    print(len(datasets))
+    multi_task_batch_sampler = DistrubutedMultiTaskBatchSampler(datasets, batch_size=config.DATA.BATCH_SIZE, num_replicas=dist.get_world_size(), rank=dist.get_rank(), mix_opt=0, extra_task_ratio=0, drop_last=True ,shuffle =True)
+    dataloader = DataLoader(multi_task_train_dataset, batch_sampler=multi_task_batch_sampler, num_workers=config.DATA.NUM_WORKERS, pin_memory=True, collate_fn=collate_fn)
+    # dataloader = DataLoader(multi_task_train_dataset, batch_sampler=multi_task_batch_sampler, pin_memory=True, collate_fn=collate_fn)
+    print(len(dataloader))
+
+    return dataloader

datasets/data_simmim_pt.py

@@ -0,0 +1,271 @@
+# --------------------------------------------------------
+# SimMIM
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Zhenda Xie
+# --------------------------------------------------------
+
+import math
+import random
+import numpy as np
+
+import torch
+import torch.distributed as dist
+import torchvision.transforms as T
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.utils.data._utils.collate import default_collate
+from torchvision.datasets import ImageFolder
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from torch.utils.data import Dataset, BatchSampler
+from torchvision.io import read_image
+from .cached_image_folder import CachedImageFolder
+
+class MultiTaskDataset(Dataset):
+    """
+    useage example:
+    train_datasets = [SemData_Single(), SemData_Single()]
+    multi_task_train_dataset = MultiTaskDataset(train_datasets)
+    multi_task_batch_sampler = MultiTaskBatchSampler(train_datasets, batch_size=4, mix_opt=0, extra_task_ratio=0, drop_last=True)
+    multi_task_train_data = DataLoader(multi_task_train_dataset, batch_sampler=multi_task_batch_sampler)
+    for i, (task_id, input, target) in enumerate(multi_task_train_data):
+        pre = model(input)
+    """
+    def __init__(self, datasets):
+        self._datasets = datasets
+        task_id_2_data_set_dic = {}
+        for i, dataset in  enumerate(datasets):
+            task_id = i
+            assert task_id not in task_id_2_data_set_dic, "Duplicate task_id %s" % task_id
+            task_id_2_data_set_dic[task_id] = dataset
+
+        self._task_id_2_data_set_dic = task_id_2_data_set_dic
+
+    def __len__(self):
+        return sum(len(dataset) for dataset in self._datasets)
+
+    def __getitem__(self, idx):
+        task_id, sample_id = idx
+        return self._task_id_2_data_set_dic[task_id][sample_id]
+
+class DistrubutedMultiTaskBatchSampler(BatchSampler):
+    """
+    datasets: class the class of the Dataset
+    batch_size: int
+    mix_opt: int mix_opt ==0 shuffle all_task; mix_opt ==1 shuffle extra_task
+    extra_task_ratio(float, optional): the rate between task one and extra task
+    drop_last (bool, optional): set to ``True`` to drop the last incomplete batch,
+            if the dataset size is not divisible by the batch size. If ``False`` and
+            the size of dataset is not divisible by the batch size, then the last batch
+            will be smaller. (default: ``True``)
+    """
+    def __init__(self, datasets, batch_size, num_replicas, rank, mix_opt=0, extra_task_ratio=0, drop_last=True,shuffle = True):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        if rank >= num_replicas or rank < 0:
+            raise ValueError(
+                "Invalid rank {}, rank should be in the interval"
+                " [0, {}]".format(rank, num_replicas - 1))
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        assert mix_opt in [0, 1], 'mix_opt must equal 0 or 1'
+        assert extra_task_ratio >= 0, 'extra_task_ratio must greater than 0'
+        self._datasets = datasets
+        self._batch_size = batch_size
+        self._mix_opt = mix_opt
+        self._extra_task_ratio = extra_task_ratio
+        self._drop_last = drop_last
+        train_data_list = []
+        self.shuffle = shuffle
+        for dataset in datasets:
+            print(len(dataset))
+            train_data_list.append(self._get_index_batches(len(dataset), batch_size, self._drop_last))
+
+        ######### 一个列表里存n个dataset的数据，数据也以列表形式存在，一个dataset的列表里面把数据划分成了不同的batch的index
+        self._train_data_list = train_data_list
+        self.total_len = sum(len(train_data) for train_data in self._train_data_list)
+        
+        ######### DDP ######################
+        if self._drop_last and self.total_len % self.num_replicas != 0:  # type: ignore[arg-type]
+            self.num_samples = math.ceil(
+                (self.total_len - self.num_replicas) / self.num_replicas  # type: ignore[arg-type]
+            )
+        else:
+            self.num_samples = math.ceil(self.total_len / self.num_replicas)  # type: ignore[arg-type]
+        
+        self.total_size = self.num_samples * self.num_replicas
+        self.epoch = 0
+        self.seed = 0
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+    @staticmethod
+    def _get_index_batches(dataset_len, batch_size, drop_last):
+        # index_batches = [list(range(i, min(i+batch_size, dataset_len))) for i in range(0, dataset_len, batch_size)]
+        index = list(range(dataset_len))
+        if drop_last and dataset_len % batch_size:
+            del index[-(dataset_len % batch_size):]
+        index_batches = [index[i:i+batch_size] for i in range(0, len(index), batch_size)]
+        return index_batches
+
+    def __len__(self):
+        # return sum(len(train_data) for train_data in self._train_data_list)
+        return self.num_samples
+
+    def __iter__(self):
+        all_iters = [iter(item) for item in self._train_data_list]
+        all_indices = self._gen_task_indices(self._train_data_list, self._mix_opt, self._extra_task_ratio)
+        
+        ######### DDP ######################
+        random.shuffle(all_indices)
+        all_indices = all_indices[self.rank:self.total_size:self.num_replicas]
+        assert len(all_indices) == self.num_samples
+        
+        for local_task_idx in all_indices:
+            # task_id = self._datasets[local_task_idx].get_task_id() 
+            batch = next(all_iters[local_task_idx])
+            # batch = batch[self.rank:len(batch):self.num_replicas]
+            # print(local_task_idx)
+            yield [(local_task_idx, sample_id) for sample_id in batch]
+            # yield iter(batch)
+
+    @staticmethod
+    def _gen_task_indices(train_data_list, mix_opt, extra_task_ratio):
+
+        ########## accoding to the number of models ###########
+        all_indices = []
+        for i in range(len(train_data_list)):
+            all_indices += [i] * len(train_data_list[i])
+        # print(all_indices)
+        return all_indices
+    # def set_epoch(self, epoch)
+    #     r"""
+    #     Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
+    #     use a different random ordering for each epoch. Otherwise, the next iteration of this
+    #     sampler will yield the same ordering.
+
+    #     Args:
+    #         epoch (int): Epoch number.
+    #     """
+    #     self.epoch = epoch
+    
+
+class MaskGenerator:
+    def __init__(self, input_size=192, mask_patch_size=32, model_patch_size=4, mask_ratio=0.6):
+        self.input_size = input_size
+        self.mask_patch_size = mask_patch_size
+        self.model_patch_size = model_patch_size
+        self.mask_ratio = mask_ratio
+        
+        assert self.input_size % self.mask_patch_size == 0
+        assert self.mask_patch_size % self.model_patch_size == 0
+        
+        self.rand_size = self.input_size // self.mask_patch_size
+        self.scale = self.mask_patch_size // self.model_patch_size
+        
+        self.token_count = self.rand_size ** 2
+        self.mask_count = int(np.ceil(self.token_count * self.mask_ratio))
+        
+    def __call__(self):
+        mask_idx = np.random.permutation(self.token_count)[:self.mask_count]
+        mask = np.zeros(self.token_count, dtype=int)
+        mask[mask_idx] = 1
+        
+        mask = mask.reshape((self.rand_size, self.rand_size))
+        mask = mask.repeat(self.scale, axis=0).repeat(self.scale, axis=1)
+        
+        return mask
+
+
+class ZeroOneNormalize(object):
+    def __call__(self, img):
+        return img.float().div(255)
+
+class SimMIMTransform:
+    def __init__(self, config, NORM, SCALE):
+        self.transform_img = T.Compose([
+            # T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+            # T.RandomResizedCrop(config.DATA.IMG_SIZE, scale=(0.67, 1.), ratio=(3. / 4., 4. / 3.)),
+            # T.RandomHorizontalFlip(),
+            # T.ToTensor(),
+            # T.Normalize(mean=torch.tensor(IMAGENET_DEFAULT_MEAN),std=torch.tensor(IMAGENET_DEFAULT_STD)),
+
+            T.RandomResizedCrop(config.DATA.IMG_SIZE, scale=SCALE, ratio=(3. / 4., 4. / 3.)),
+            T.RandomHorizontalFlip(),
+            ZeroOneNormalize(),
+            T.Normalize(mean=torch.tensor(NORM[0]),std=torch.tensor(NORM[1])),
+        ])
+ 
+        if config.MODEL.TYPE in ['swin', 'swinv2']:
+            model_patch_size=config.MODEL.SWIN.PATCH_SIZE
+        else:
+            raise NotImplementedError
+        
+        self.mask_generator = MaskGenerator(
+            input_size=config.DATA.IMG_SIZE,
+            mask_patch_size=config.DATA.MASK_PATCH_SIZE,
+            model_patch_size=model_patch_size,
+            mask_ratio=config.DATA.MASK_RATIO,
+        )
+    
+    def __call__(self, img):
+        img = self.transform_img(img)
+        mask = self.mask_generator()
+        
+        return img, mask
+
+def collate_fn(batch):
+    # print(len(batch))
+    # print('*'*10)
+    # print(batch[0][0])
+    # print('#'*10)
+    # print(batch[0][1])
+    # batch = list(filter(lambda x: x[0][0] is not None, batch))
+    # if len(batch) == 0: return torch.Tensor()
+
+    if not isinstance(batch[0][0], tuple):
+        return default_collate(batch)
+    else:
+        batch_num = len(batch)
+        ret = []
+        for item_idx in range(len(batch[0][0])):
+            if batch[0][0][item_idx] is None:
+                ret.append(None)
+            else:
+                ret.append(default_collate([batch[i][0][item_idx] for i in range(batch_num)]))
+        ret.append(default_collate([batch[i][1] for i in range(batch_num)]))
+        return ret
+
+
+def build_loader_simmim(config):
+    ############ single model #####################
+    # transform = SimMIMTransform(config)
+    # dataset = ImageFolder(config.DATA.DATA_PATH, transform)
+    # sampler = DistributedSampler(dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank(), shuffle=True)
+    # dataloader = DataLoader(dataset, config.DATA.BATCH_SIZE, sampler=sampler, num_workers=config.DATA.NUM_WORKERS, pin_memory=True, drop_last=True, collate_fn=collate_fn)
+
+    ############## multi model ####################
+    datasets = []
+    ### 数据增强 ######
+    model_paths = config.DATA.TYPE_PATH[0]
+    for i in model_paths.keys():
+        a = config.DATA.SCALE[0][i].split(',')
+        scale_model = (float(a[0].split('(')[1]),float(a[1].split(')')[0]))
+        transform = SimMIMTransform(config, config.DATA.NORM[0][i], scale_model)
+        dataset = CachedImageFolder(model_paths[i], transform = transform, model = i)
+        datasets.append(dataset)
+    multi_task_train_dataset = MultiTaskDataset(datasets)
+    print(len(datasets))
+    multi_task_batch_sampler = DistrubutedMultiTaskBatchSampler(datasets, batch_size=config.DATA.BATCH_SIZE, num_replicas=dist.get_world_size(), rank=dist.get_rank(), mix_opt=0, extra_task_ratio=0, drop_last=True,shuffle =True)
+    dataloader = DataLoader(multi_task_train_dataset, batch_sampler=multi_task_batch_sampler, num_workers=config.DATA.NUM_WORKERS, pin_memory=True, collate_fn=collate_fn)
+    # dataloader = DataLoader(multi_task_train_dataset, batch_sampler=multi_task_batch_sampler, pin_memory=True, collate_fn=collate_fn)
+    print(len(dataloader))
+    
+    return dataloader

datasets/datasets.py

@@ -0,0 +1,21 @@
+import copy
+
+
+datasets = {}
+
+
+def register(name):
+    def decorator(cls):
+        datasets[name] = cls
+        return cls
+    return decorator
+
+
+def make(dataset_spec, args=None):
+    if args is not None:
+        dataset_args = copy.deepcopy(dataset_spec['args'])
+        dataset_args.update(args)
+    else:
+        dataset_args = dataset_spec['args']
+    dataset = datasets[dataset_spec['name']](**dataset_args)
+    return dataset

datasets/image_folder.py

@@ -0,0 +1,370 @@
+import os
+import json
+from PIL import Image
+
+import pickle
+import imageio
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+import random
+from datasets import register
+
+import math
+import torch.distributed as dist
+from torch.utils.data import BatchSampler
+
+from torch.utils.data._utils.collate import default_collate
+
+@register('image-folder')
+class ImageFolder(Dataset):
+    def __init__(self, path,  split_file=None, split_key=None, first_k=None, size=None,
+                 repeat=1, cache='none', mask=False):
+        self.repeat = repeat
+        self.cache = cache
+        self.path = path
+        self.Train = False
+        self.split_key = split_key
+
+        self.size = size
+        self.mask = mask
+        if self.mask:
+            self.img_transform = transforms.Compose([
+                transforms.Resize((self.size, self.size), interpolation=Image.NEAREST),
+                transforms.ToTensor(),
+            ])
+        else:
+            self.img_transform = transforms.Compose([
+                transforms.Resize((self.size, self.size)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+            ])
+
+        if split_file is None:
+            filenames = sorted(os.listdir(path))
+        else:
+            with open(split_file, 'r') as f:
+                filenames = json.load(f)[split_key]
+        if first_k is not None:
+            filenames = filenames[:first_k]
+
+        self.files = []
+
+        for filename in filenames:
+            file = os.path.join(path, filename)
+            self.append_file(file)
+
+    def append_file(self, file):
+        if self.cache == 'none':
+            self.files.append(file)
+        elif self.cache == 'in_memory':
+            self.files.append(self.img_process(file))
+
+    def __len__(self):
+        return len(self.files) * self.repeat
+
+    def __getitem__(self, idx):
+        x = self.files[idx % len(self.files)]
+
+        if self.cache == 'none':
+            return self.img_process(x)
+        elif self.cache == 'in_memory':
+            return x
+
+    def img_process(self, file):
+        if self.mask:
+            # return Image.open(file).convert('L')
+            return file
+        else:
+            return Image.open(file).convert('RGB')
+
+@register('paired-image-folders')
+class PairedImageFolders(Dataset):
+
+    def __init__(self, root_path_1, root_path_2, **kwargs):
+        self.dataset_1 = ImageFolder(root_path_1, **kwargs)
+        self.dataset_2 = ImageFolder(root_path_2, **kwargs, mask=True)
+
+    def __len__(self):
+        return len(self.dataset_1)
+
+    def __getitem__(self, idx):
+        return self.dataset_1[idx], self.dataset_2[idx]
+
+class ImageFolder_multi_task(Dataset):
+    def __init__(self, path,  split_file=None, split_key=None, first_k=None, size=None,
+                 repeat=1, cache='none', mask=False):
+        self.repeat = repeat
+        self.cache = cache
+        self.path = path
+        self.Train = False
+        self.split_key = split_key
+
+        self.size = size
+        self.mask = mask
+        if self.mask:
+            self.img_transform = transforms.Compose([
+                transforms.Resize((self.size, self.size), interpolation=Image.NEAREST),
+                transforms.ToTensor(),
+            ])
+        else:
+            self.img_transform = transforms.Compose([
+                transforms.Resize((self.size, self.size)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+            ])
+
+        if split_file is None:
+            filenames = sorted(os.listdir(path))
+        else:
+            with open(split_file, 'r') as f:
+                filenames = json.load(f)[split_key]
+        if first_k is not None:
+            filenames = filenames[:first_k]
+
+        self.files = []
+
+        for filename in filenames:
+            file = os.path.join(path, filename)
+            self.append_file(file)
+
+    def append_file(self, file):
+        if self.cache == 'none':
+            self.files.append(file)
+        elif self.cache == 'in_memory':
+            self.files.append(self.img_process(file))
+
+    def __len__(self):
+        return len(self.files) * self.repeat
+
+    def __getitem__(self, idx):
+        x = self.files[idx % len(self.files)]
+
+        if self.cache == 'none':
+            return self.img_process(x)
+        elif self.cache == 'in_memory':
+            return x
+
+    def img_process(self, file):
+        if self.mask:
+            # return Image.open(file).convert('L')
+            return file
+        else:
+            return Image.open(file).convert('RGB')
+
+@register('paired-image-folders-multi-task')
+class PairedImageFolders_multi_task(Dataset):
+
+    def __init__(self, root_path_1, root_path_2, model=None, **kwargs):
+
+        self.dataset_1 = ImageFolder_multi_task(root_path_1, **kwargs)
+        self.dataset_2 = ImageFolder_multi_task(root_path_2, **kwargs, mask=True)
+
+    def __len__(self):
+        return len(self.dataset_1)
+
+    def __getitem__(self, idx):
+        return self.dataset_1[idx], self.dataset_2[idx]
+
+
+
+
+# class MultiTaskDataset(Dataset):
+#     """
+#     useage example:
+#     train_datasets = [SemData_Single(), SemData_Single()]
+#     multi_task_train_dataset = MultiTaskDataset(train_datasets)
+#     multi_task_batch_sampler = MultiTaskBatchSampler(train_datasets, batch_size=4, mix_opt=0, extra_task_ratio=0, drop_last=True)
+#     multi_task_train_data = DataLoader(multi_task_train_dataset, batch_sampler=multi_task_batch_sampler)
+#     for i, (task_id, input, target) in enumerate(multi_task_train_data):
+#         pre = model(input)
+#     """
+#     def __init__(self, datasets_image, datasets_gt):
+#         self._datasets = datasets_image
+#         task_id_2_image_set_dic = {}
+#         for i, dataset in  enumerate(datasets_image):
+#             task_id = i
+#             assert task_id not in task_id_2_image_set_dic, "Duplicate task_id %s" % task_id
+#             task_id_2_image_set_dic[task_id] = dataset
+#         self.datasets_1 = task_id_2_image_set_dic
+#
+#         task_id_2_gt_set_dic = {}
+#         for i, dataset in enumerate(datasets_gt):
+#             task_id = i
+#             assert task_id not in task_id_2_gt_set_dic, "Duplicate task_id %s" % task_id
+#             task_id_2_gt_set_dic[task_id] = dataset
+#         self.dataset_2 = task_id_2_gt_set_dic
+#
+#
+#     def __len__(self):
+#         return sum(len(dataset) for dataset in self._datasets)
+#
+#     def __getitem__(self, idx):
+#         task_id, sample_id = idx
+#         # return self._task_id_2_data_set_dic[task_id][sample_id]
+#         return self.dataset_1[task_id][sample_id], self.dataset_2[task_id][sample_id]
+
+class MultiTaskDataset(Dataset):
+    """
+    useage example:
+    train_datasets = [SemData_Single(), SemData_Single()]
+    multi_task_train_dataset = MultiTaskDataset(train_datasets)
+    multi_task_batch_sampler = MultiTaskBatchSampler(train_datasets, batch_size=4, mix_opt=0, extra_task_ratio=0, drop_last=True)
+    multi_task_train_data = DataLoader(multi_task_train_dataset, batch_sampler=multi_task_batch_sampler)
+    for i, (task_id, input, target) in enumerate(multi_task_train_data):
+        pre = model(input)
+    """
+    def __init__(self, datasets):
+        self._datasets = datasets
+        task_id_2_data_set_dic = {}
+        for i, dataset in enumerate(datasets):
+            task_id = i
+            assert task_id not in task_id_2_data_set_dic, "Duplicate task_id %s" % task_id
+            task_id_2_data_set_dic[task_id] = dataset
+
+        self._task_id_2_data_set_dic = task_id_2_data_set_dic
+
+    def __len__(self):
+        return sum(len(dataset) for dataset in self._datasets)
+
+    def __getitem__(self, idx):
+        task_id, sample_id = idx
+        # print('----', idx, task_id, sample_id)
+        return self._task_id_2_data_set_dic[task_id][sample_id]
+
+def collate_fn(batch):
+    # print(len(batch))
+    # print('*'*10)
+    # print(batch[0][0])
+    # print('#'*10)
+    # print(batch[0][1])
+    # batch = list(filter(lambda x: x[0][0] is not None, batch))
+    # if len(batch) == 0: return torch.Tensor()
+    print('******------',batch)
+    if not isinstance(batch[0][0], tuple):
+        return default_collate(batch)
+    else:
+        batch_num = len(batch)
+        ret = []
+        for item_idx in range(len(batch[0][0])):
+            if batch[0][0][item_idx] is None:
+                ret.append(None)
+            else:
+                ret.append(default_collate([batch[i][0][item_idx] for i in range(batch_num)]))
+        ret.append(default_collate([batch[i][1] for i in range(batch_num)]))
+        return ret
+
+class DistrubutedMultiTaskBatchSampler(BatchSampler):
+    """
+    datasets: class the class of the Dataset
+    batch_size: int
+    mix_opt: int mix_opt ==0 shuffle all_task; mix_opt ==1 shuffle extra_task
+    extra_task_ratio(float, optional): the rate between task one and extra task
+    drop_last (bool, optional): set to ``True`` to drop the last incomplete batch,
+            if the dataset size is not divisible by the batch size. If ``False`` and
+            the size of dataset is not divisible by the batch size, then the last batch
+            will be smaller. (default: ``True``)
+    """
+
+    def __init__(self, datasets, batch_size, num_replicas, rank, mix_opt=0, extra_task_ratio=0, drop_last=True,
+                 shuffle=True):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        if rank >= num_replicas or rank < 0:
+            raise ValueError(
+                "Invalid rank {}, rank should be in the interval"
+                " [0, {}]".format(rank, num_replicas - 1))
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        assert mix_opt in [0, 1], 'mix_opt must equal 0 or 1'
+        assert extra_task_ratio >= 0, 'extra_task_ratio must greater than 0'
+        # self._datasets = datasets
+        self._batch_size = batch_size
+        self._mix_opt = mix_opt
+        self._extra_task_ratio = extra_task_ratio
+        self._drop_last = drop_last
+        train_data_list = []
+        self.shuffle = shuffle
+        for dataset in datasets:
+            print(len(dataset))
+            train_data_list.append(self._get_index_batches(len(dataset), batch_size, self._drop_last))
+
+        ######### 一个列表里存n个dataset的数据，数据也以列表形式存在，一个dataset的列表里面把数据划分成了不同的batch的index
+        self._train_data_list = train_data_list
+        self.total_len = sum(len(train_data) for train_data in self._train_data_list)
+
+        ######### DDP ######################
+        if self._drop_last and self.total_len % self.num_replicas != 0:  # type: ignore[arg-type]
+            self.num_samples = math.ceil(
+                (self.total_len - self.num_replicas) / self.num_replicas  # type: ignore[arg-type]
+            )
+        else:
+            self.num_samples = math.ceil(self.total_len / self.num_replicas)  # type: ignore[arg-type]
+
+        self.total_size = self.num_samples * self.num_replicas
+        self.epoch = 0
+        self.seed = 0
+
+    def set_epoch(self, epoch):
+        # print('&&&&****')
+        self.epoch = epoch
+
+    @staticmethod
+    def _get_index_batches(dataset_len, batch_size, drop_last):
+        # index_batches = [list(range(i, min(i+batch_size, dataset_len))) for i in range(0, dataset_len, batch_size)]
+        index = list(range(dataset_len))
+        if drop_last and dataset_len % batch_size:
+            del index[-(dataset_len % batch_size):]
+        index_batches = [index[i:i + batch_size] for i in range(0, len(index), batch_size)]
+        return index_batches
+
+    def __len__(self):
+        # return sum(len(train_data) for train_data in self._train_data_list)
+        return self.num_samples
+
+    def __iter__(self):
+        all_iters = [iter(item) for item in self._train_data_list]
+        all_indices = self._gen_task_indices(self._train_data_list, self._mix_opt, self._extra_task_ratio)
+
+        ######### DDP ######################
+        random.shuffle(all_indices)
+        all_indices = all_indices[self.rank:self.total_size:self.num_replicas]
+        assert len(all_indices) == self.num_samples
+
+        for local_task_idx in all_indices:
+            # task_id = self._datasets[local_task_idx].get_task_id()
+            batch = next(all_iters[local_task_idx])
+            # batch = batch[self.rank:len(batch):self.num_replicas]
+            # print(local_task_idx)
+            yield [(local_task_idx, sample_id) for sample_id in batch]
+            # yield iter(batch)
+
+    @staticmethod
+    def _gen_task_indices(train_data_list, mix_opt, extra_task_ratio):
+
+        ########## accoding to the number of models ###########
+        all_indices = []
+        for i in range(len(train_data_list)):
+            all_indices += [i] * len(train_data_list[i])
+        # print(all_indices)
+        return all_indices
+    # def set_epoch(self, epoch)
+    #     r"""
+    #     Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
+    #     use a different random ordering for each epoch. Otherwise, the next iteration of this
+    #     sampler will yield the same ordering.
+
+    #     Args:
+    #         epoch (int): Epoch number.
+    #     """
+    #     self.epoch = epoch

datasets/wrappers.py

@@ -0,0 +1,231 @@
+
+import functools
+import random
+import math
+from PIL import Image
+import cv2
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+import torchvision
+
+from datasets import register
+import cv2
+from math import pi
+from torchvision.transforms import InterpolationMode
+
+import torch.nn.functional as F
+def to_mask(mask):
+    return transforms.ToTensor()(
+        transforms.Grayscale(num_output_channels=1)(
+            transforms.ToPILImage()(mask)))
+
+
+def resize_fn(img, size):
+    return transforms.ToTensor()(
+        transforms.Resize(size)(
+            transforms.ToPILImage()(img)))
+
+
+@register('val')
+class ValDataset(Dataset):
+    def __init__(self, dataset, inp_size=None, augment=False):
+        self.dataset = dataset
+        self.inp_size = inp_size
+        self.augment = augment
+
+        self.img_transform = transforms.Compose([
+                # transforms.Resize((inp_size, inp_size)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+            ])
+        self.mask_transform = transforms.Compose([
+                transforms.Resize((inp_size, inp_size), interpolation=Image.NEAREST),
+                transforms.ToTensor(),
+            ])
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        img, mask = self.dataset[idx]
+        mask_name = mask
+        a = self.img_transform(img)
+        # b = self.mask_transform(mask)
+
+        # print(idx, mask.filename)
+        # b = cv2.imread(mask.filename,cv2.IMREAD_UNCHANGED)
+        b = cv2.imread(mask,cv2.IMREAD_UNCHANGED)
+        return {
+            'inp': self.img_transform(img),
+            'gt': torch.tensor(b),
+            'name': mask_name,
+            'filp': False
+            # 'idx': idx
+        }
+
+
+@register('train')
+class TrainDataset(Dataset):
+    def __init__(self, dataset, size_min=None, size_max=None, inp_size=None,
+                 augment=False, gt_resize=None):
+        self.dataset = dataset
+        self.size_min = size_min
+        if size_max is None:
+            size_max = size_min
+        self.size_max = size_max
+        self.augment = augment
+        self.gt_resize = gt_resize
+
+        self.inp_size = inp_size
+        self.img_transform = transforms.Compose([
+                transforms.Resize((self.inp_size, self.inp_size)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+            ])
+        self.inverse_transform = transforms.Compose([
+                transforms.Normalize(mean=[0., 0., 0.],
+                                     std=[1/0.229, 1/0.224, 1/0.225]),
+                transforms.Normalize(mean=[-0.485, -0.456, -0.406],
+                                     std=[1, 1, 1])
+            ])
+        self.mask_transform = transforms.Compose([
+                transforms.Resize((self.inp_size, self.inp_size)),
+                transforms.ToTensor(),
+            ])
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        # print('lodd****',idx,self.dataset[idx])
+        img, mask = self.dataset[idx]
+        mask_name = mask
+        # print('befor mask', mask)
+        #new add
+        # print(idx, mask.filename, img.size)
+
+        # mask = cv2.imread(mask.filename, cv2.IMREAD_UNCHANGED)
+        mask = cv2.imread(mask, cv2.IMREAD_UNCHANGED)
+        # print('befor mask', mask)
+        # random filp
+        if random.random() < 0.5:
+            img = img.transpose(Image.FLIP_LEFT_RIGHT)
+            # mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
+            mask = cv2.flip(mask, 1)
+
+        img = transforms.Resize((self.inp_size, self.inp_size))(img)
+        # mask = transforms.Resize((self.inp_size, self.inp_size), interpolation=InterpolationMode.NEAREST)(mask)
+        mask = torch.from_numpy(mask)
+        # print('behind mask', mask)
+        return {
+            'inp': self.img_transform(img),
+            # 'gt': self.mask_transform(mask)
+            'gt': mask,
+            'name': mask_name,
+            # 'idx': idx
+        }
+
+@register('train_multi_task')
+class TrainDataset(Dataset):
+    def __init__(self, dataset, size_min=None, size_max=None, inp_size=None,
+                 augment=False, gt_resize=None):
+        self.dataset = dataset
+        self.size_min = size_min
+        if size_max is None:
+            size_max = size_min
+        self.size_max = size_max
+        self.augment = augment
+        self.gt_resize = gt_resize
+
+        self.inp_size = inp_size
+        self.img_transform = transforms.Compose([
+                transforms.Resize((self.inp_size, self.inp_size)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+            ])
+        self.inverse_transform = transforms.Compose([
+                transforms.Normalize(mean=[0., 0., 0.],
+                                     std=[1/0.229, 1/0.224, 1/0.225]),
+                transforms.Normalize(mean=[-0.485, -0.456, -0.406],
+                                     std=[1, 1, 1])
+            ])
+        self.mask_transform = transforms.Compose([
+                transforms.Resize((self.inp_size, self.inp_size)),
+                transforms.ToTensor(),
+            ])
+
+    def __len__(self):
+        return len(self.dataset)
+        # return sum(len(dataset) for dataset in self.datasets)
+
+    def __getitem__(self, idx):
+        # print('lodd****',idx,self.dataset[idx])
+        # print('+++++',idx)
+        img, mask = self.dataset[idx]
+        # print('befor mask', mask)
+        #new add
+        # print('****',idx, mask)
+        mask_name = mask
+        mask = cv2.imread(mask, cv2.IMREAD_UNCHANGED)
+
+        # print('****',mask)
+        # print('befor mask', mask)
+        # random filp
+        if random.random() < 0.5:
+            img = img.transpose(Image.FLIP_LEFT_RIGHT)
+            # mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
+            mask = cv2.flip(mask, 1)
+
+        img = transforms.Resize((self.inp_size, self.inp_size))(img)
+        # mask = transforms.Resize((self.inp_size, self.inp_size), interpolation=InterpolationMode.NEAREST)(mask)
+        mask = torch.from_numpy(mask)
+        # print('behind mask', mask)
+        return {
+            'inp': self.img_transform(img),
+            # 'gt': self.mask_transform(mask)
+            'gt': mask,
+            'name': mask_name
+        }
+
+
+@register('val_multi_task')
+class ValDataset(Dataset):
+    def __init__(self, dataset, inp_size=None, augment=False):
+        self.dataset = dataset
+        self.inp_size = inp_size
+        self.augment = augment
+
+        self.img_transform = transforms.Compose([
+                transforms.Resize((inp_size, inp_size)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+            ])
+        self.mask_transform = transforms.Compose([
+                transforms.Resize((inp_size, inp_size), interpolation=Image.NEAREST),
+                transforms.ToTensor(),
+            ])
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        img, mask = self.dataset[idx]
+        a = self.img_transform(img)
+        # b = self.mask_transform(mask)
+        mask_name = mask
+        # print(idx, mask.filename)
+        # b = cv2.imread(mask.filename,cv2.IMREAD_UNCHANGED)
+        b = cv2.imread(mask, cv2.IMREAD_UNCHANGED)
+        return {
+            'inp': self.img_transform(img),
+            'gt': torch.tensor(b),
+            'name': mask_name
+            # 'idx': idx
+        }

models/__init__.py

@@ -0,0 +1,4 @@
+from .models import register, make
+from . import sam
+from . import sam_single
+

models/block.py

@@ -0,0 +1,128 @@
+from __future__ import print_function
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+
+
+class MergeAndConv(nn.Module):
+
+    def __init__(self, ic, oc, inner=32):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(ic, inner, kernel_size=3, stride=1, padding=1)
+        self.bn = nn.BatchNorm2d(inner)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(inner, oc, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        x = self.conv2(self.bn(self.relu(self.conv1(x))))
+        x = torch.sigmoid(x)
+        return x
+
+
+class SideClassifer(nn.Module):
+    def __init__(self, ic, n_class=1, M=2, kernel_size=1):
+        super().__init__()
+
+        sides = []
+        for i in range(M):
+            sides.append(nn.Conv2d(ic, n_class, kernel_size=kernel_size))
+
+        self.sides = nn.ModuleList(sides)
+
+    def forward(self, x):
+        return [fn(x) for fn in self.sides]
+
+
+class UpsampleSKConv(nn.Module):
+    """docstring for UpsampleSKConvPlus"""
+
+    def __init__(self, ic, oc, reduce=4):
+        super(UpsampleSKConv, self).__init__()
+
+        self.relu = nn.ReLU(inplace=True)
+        self.prev = nn.Conv2d(ic, ic // reduce, kernel_size=3, stride=1, padding=1)
+        self.bn = nn.BatchNorm2d(ic // reduce)
+
+        self.next = nn.Conv2d(ic // reduce, oc, kernel_size=1, stride=1)
+        self.bn2 = nn.BatchNorm2d(oc)
+
+        self.sk = SKSPP(ic // reduce, ic // reduce, M=4)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2)
+
+        x = self.bn(self.relu(self.prev(x)))
+
+        x = self.sk(x)
+
+        x = self.bn2(self.relu(self.next(x)))
+
+        return x
+
+
+class SKSPP(nn.Module):
+    def __init__(self, features, WH, M=2, G=1, r=16, stride=1, L=32):
+        """ Constructor
+        Args:
+            features: input channel dimensionality.
+            WH: input spatial dimensionality, used for GAP kernel size.
+            M: the number of branchs.
+            G: num of convolution groups.
+            r: the radio for compute d, the length of z.
+            stride: stride, default 1.
+            L: the minimum dim of the vector z in paper, default 32.
+        """
+        super(SKSPP, self).__init__()
+        d = max(int(features / r), L)
+        self.M = M  # original
+        self.features = features
+        self.convs = nn.ModuleList([])
+
+        # 1,3,5,7 padding:[0,1,2,3]
+        for i in range(1, M):
+            self.convs.append(nn.Sequential(
+                nn.Conv2d(features, features, kernel_size=1 + i * 2, dilation=1 + i * 2, stride=stride,
+                          padding=((1 + i * 2) * (i * 2) + 1) // 2, groups=G),
+                nn.BatchNorm2d(features),
+                nn.ReLU(inplace=False)
+            ))
+        # self.gap = nn.AvgPool2d(int(WH/stride))
+        self.fc = nn.Linear(features, d)
+        self.fcs = nn.ModuleList([])
+        for i in range(M):
+            self.fcs.append(
+                nn.Linear(d, features)
+            )
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+
+        feas = torch.unsqueeze(x, dim=1)
+
+        # F->conv1x1->conv3x3->conv5x5->conv7x7
+
+        for i, conv in enumerate(self.convs):
+            x = conv(x)
+            # if i == 0:
+            #     feas = fea
+            # else:
+            feas = torch.cat([feas, torch.unsqueeze(x, dim=1)], dim=1)
+
+        fea_U = torch.sum(feas, dim=1)
+        fea_s = fea_U.mean(-1).mean(-1)
+        fea_z = self.fc(fea_s)
+
+        for i, fc in enumerate(self.fcs):
+            vector = fc(fea_z).unsqueeze_(dim=1)
+            if i == 0:
+                attention_vectors = vector
+            else:
+                attention_vectors = torch.cat([attention_vectors, vector], dim=1)
+
+        attention_vectors = self.softmax(attention_vectors)
+        attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1)
+        fea_v = (feas * attention_vectors).sum(dim=1)
+        return fea_v

models/bn_helper.py

@@ -0,0 +1,16 @@
+import torch
+import functools
+
+if torch.__version__.startswith('0'):
+    from .sync_bn.inplace_abn.bn import InPlaceABNSync
+    BatchNorm2d = functools.partial(InPlaceABNSync, activation='none')
+    BatchNorm2d_class = InPlaceABNSync
+    relu_inplace = False
+else:
+    BatchNorm2d_class = BatchNorm2d = torch.nn.SyncBatchNorm
+    relu_inplace = True
+
+import torch
+BatchNorm2d = torch.nn.BatchNorm2d
+BatchNorm2d_class = BatchNorm2d
+relu_inplace = False

models/iou_loss.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+###################################################################
+# ########################## iou loss #############################
+###################################################################
+class IOU(torch.nn.Module):
+    def __init__(self):
+        super(IOU, self).__init__()
+
+    def _iou(self, pred, target):
+        pred = torch.sigmoid(pred)
+        inter = (pred * target).sum(dim=(2, 3))
+        union = (pred + target).sum(dim=(2, 3)) - inter
+        iou = 1 - (inter / union)
+
+        return iou.mean()
+
+    def forward(self, pred, target):
+        return self._iou(pred, target)

models/mmseg/__init__.py

@@ -0,0 +1,33 @@
+import mmcv
+
+from .version import __version__, version_info
+
+# MMCV_MIN = '1.1.4'
+# MMCV_MAX = '1.3.0'
+
+MMCV_MIN = '1.1.4'
+MMCV_MAX = '1.7.0'
+
+
+def digit_version(version_str):
+    digit_version = []
+    for x in version_str.split('.'):
+        if x.isdigit():
+            digit_version.append(int(x))
+        elif x.find('rc') != -1:
+            patch_version = x.split('rc')
+            digit_version.append(int(patch_version[0]) - 1)
+            digit_version.append(int(patch_version[1]))
+    return digit_version
+
+
+mmcv_min_version = digit_version(MMCV_MIN)
+mmcv_max_version = digit_version(MMCV_MAX)
+mmcv_version = digit_version(mmcv.__version__)
+
+
+assert (mmcv_min_version <= mmcv_version <= mmcv_max_version), \
+    f'MMCV=={mmcv.__version__} is used but incompatible. ' \
+    f'Please install mmcv>={mmcv_min_version}, <={mmcv_max_version}.'
+
+__all__ = ['__version__', 'version_info']

models/mmseg/apis/__init__.py

@@ -0,0 +1,9 @@
+from .inference import inference_segmentor, init_segmentor, show_result_pyplot
+from .test import multi_gpu_test, single_gpu_test
+from .train import get_root_logger, set_random_seed, train_segmentor
+
+__all__ = [
+    'get_root_logger', 'set_random_seed', 'train_segmentor', 'init_segmentor',
+    'inference_segmentor', 'multi_gpu_test', 'single_gpu_test',
+    'show_result_pyplot'
+]

models/mmseg/apis/inference.py

@@ -0,0 +1,118 @@
+import matplotlib.pyplot as plt
+import mmcv
+import torch
+from mmcv.parallel import collate, scatter
+from mmcv.runner import load_checkpoint
+
+from mmseg.datasets.pipelines import Compose
+from mmseg.models import build_segmentor
+
+
+def init_segmentor(config, checkpoint=None, device='cuda:0'):
+    """Initialize a segmentor from config file.
+
+    Args:
+        config (str or :obj:`mmcv.Config`): Config file path or the config
+            object.
+        checkpoint (str, optional): Checkpoint path. If left as None, the model
+            will not load any weights.
+        device (str, optional) CPU/CUDA device option. Default 'cuda:0'.
+            Use 'cpu' for loading model on CPU.
+    Returns:
+        nn.Module: The constructed segmentor.
+    """
+    if isinstance(config, str):
+        config = mmcv.Config.fromfile(config)
+    elif not isinstance(config, mmcv.Config):
+        raise TypeError('config must be a filename or Config object, '
+                        'but got {}'.format(type(config)))
+    config.model.pretrained = None
+    config.model.train_cfg = None
+    model = build_segmentor(config.model, test_cfg=config.get('test_cfg'))
+    if checkpoint is not None:
+        checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
+        model.CLASSES = checkpoint['meta']['CLASSES']
+        model.PALETTE = checkpoint['meta']['PALETTE']
+    model.cfg = config  # save the config in the model for convenience
+    model.to(device)
+    model.eval()
+    return model
+
+
+class LoadImage:
+    """A simple pipeline to load image."""
+
+    def __call__(self, results):
+        """Call function to load images into results.
+
+        Args:
+            results (dict): A result dict contains the file name
+                of the image to be read.
+
+        Returns:
+            dict: ``results`` will be returned containing loaded image.
+        """
+
+        if isinstance(results['img'], str):
+            results['filename'] = results['img']
+            results['ori_filename'] = results['img']
+        else:
+            results['filename'] = None
+            results['ori_filename'] = None
+        img = mmcv.imread(results['img'])
+        results['img'] = img
+        results['img_shape'] = img.shape
+        results['ori_shape'] = img.shape
+        return results
+
+
+def inference_segmentor(model, img):
+    """Inference image(s) with the segmentor.
+
+    Args:
+        model (nn.Module): The loaded segmentor.
+        imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
+            images.
+
+    Returns:
+        (list[Tensor]): The segmentation result.
+    """
+    cfg = model.cfg
+    device = next(model.parameters()).device  # model device
+    # build the data pipeline
+    test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
+    test_pipeline = Compose(test_pipeline)
+    # prepare data
+    data = dict(img=img)
+    data = test_pipeline(data)
+    data = collate([data], samples_per_gpu=1)
+    if next(model.parameters()).is_cuda:
+        # scatter to specified GPU
+        data = scatter(data, [device])[0]
+    else:
+        data['img_metas'] = [i.data[0] for i in data['img_metas']]
+
+    # forward the model
+    with torch.no_grad():
+        result = model(return_loss=False, rescale=True, **data)
+    return result
+
+
+def show_result_pyplot(model, img, result, palette=None, fig_size=(15, 10)):
+    """Visualize the segmentation results on the image.
+
+    Args:
+        model (nn.Module): The loaded segmentor.
+        img (str or np.ndarray): Image filename or loaded image.
+        result (list): The segmentation result.
+        palette (list[list[int]]] | None): The palette of segmentation
+            map. If None is given, random palette will be generated.
+            Default: None
+        fig_size (tuple): Figure size of the pyplot figure.
+    """
+    if hasattr(model, 'module'):
+        model = model.module
+    img = model.show_result(img, result, palette=palette, show=False)
+    plt.figure(figsize=fig_size)
+    plt.imshow(mmcv.bgr2rgb(img))
+    plt.show()

models/mmseg/apis/test.py

@@ -0,0 +1,235 @@
+import os.path as osp
+import pickle
+import shutil
+import tempfile
+
+import mmcv
+import numpy as np
+import torch
+import torch.distributed as dist
+from mmcv.image import tensor2imgs
+from mmcv.runner import get_dist_info
+from IPython import embed
+from mmseg.ops import resize
+
+def np2tmp(array, temp_file_name=None):
+    """Save ndarray to local numpy file.
+
+    Args:
+        array (ndarray): Ndarray to save.
+        temp_file_name (str): Numpy file name. If 'temp_file_name=None', this
+            function will generate a file name with tempfile.NamedTemporaryFile
+            to save ndarray. Default: None.
+
+    Returns:
+        str: The numpy file name.
+    """
+
+    if temp_file_name is None:
+        temp_file_name = tempfile.NamedTemporaryFile(
+            suffix='.npy', delete=False).name
+    np.save(temp_file_name, array)
+    return temp_file_name
+
+
+def single_gpu_test(model,
+                    data_loader,
+                    show=False,
+                    out_dir=None,
+                    efficient_test=False):
+    """Test with single GPU.
+
+    Args:
+        model (nn.Module): Model to be tested.
+        data_loader (utils.data.Dataloader): Pytorch data loader.
+        show (bool): Whether show results during infernece. Default: False.
+        out_dir (str, optional): If specified, the results will be dumped into
+            the directory to save output results.
+        efficient_test (bool): Whether save the results as local numpy files to
+            save CPU memory during evaluation. Default: False.
+
+    Returns:
+        list: The prediction results.
+    """
+
+    model.eval()
+    results = []
+    dataset = data_loader.dataset
+    prog_bar = mmcv.ProgressBar(len(dataset))
+    for i, data in enumerate(data_loader):
+        with torch.no_grad():
+            result = model(return_loss=False, **data)
+
+        if show or out_dir:
+            img_tensor = data['img'][0]
+            img_metas = data['img_metas'][0].data[0]
+            imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg'])
+            assert len(imgs) == len(img_metas)
+
+            for img, img_meta in zip(imgs, img_metas):
+                h, w, _ = img_meta['img_shape']
+                img_show = img[:h, :w, :]
+
+                ori_h, ori_w = img_meta['ori_shape'][:-1]
+                img_show = mmcv.imresize(img_show, (ori_w, ori_h))
+
+                if out_dir:
+                    out_file = osp.join(out_dir, img_meta['ori_filename'])
+                else:
+                    out_file = None
+
+                model.module.show_result(
+                    img_show,
+                    result,
+                    palette=dataset.PALETTE,
+                    show=show,
+                    out_file=out_file)
+
+        if isinstance(result, list):
+            if efficient_test:
+                result = [np2tmp(_) for _ in result]
+            results.extend(result)
+        else:
+            if efficient_test:
+                result = np2tmp(result)
+            results.append(result)
+
+        batch_size = data['img'][0].size(0)
+        for _ in range(batch_size):
+            prog_bar.update()
+    return results
+
+
+def multi_gpu_test(model,
+                   data_loader,
+                   tmpdir=None,
+                   gpu_collect=False,
+                   efficient_test=False):
+    """Test model with multiple gpus.
+
+    This method tests model with multiple gpus and collects the results
+    under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'
+    it encodes results to gpu tensors and use gpu communication for results
+    collection. On cpu mode it saves the results on different gpus to 'tmpdir'
+    and collects them by the rank 0 worker.
+
+    Args:
+        model (nn.Module): Model to be tested.
+        data_loader (utils.data.Dataloader): Pytorch data loader.
+        tmpdir (str): Path of directory to save the temporary results from
+            different gpus under cpu mode.
+        gpu_collect (bool): Option to use either gpu or cpu to collect results.
+        efficient_test (bool): Whether save the results as local numpy files to
+            save CPU memory during evaluation. Default: False.
+
+    Returns:
+        list: The prediction results.
+    """
+
+    model.eval()
+    results = []
+    dataset = data_loader.dataset
+    rank, world_size = get_dist_info()
+    if rank == 0:
+        prog_bar = mmcv.ProgressBar(len(dataset))
+    for i, data in enumerate(data_loader):
+        with torch.no_grad():
+            result = model(return_loss=False, rescale=True, **data)
+
+        if isinstance(result, list):
+            if efficient_test:
+                result = [np2tmp(_) for _ in result]
+            results.extend(result)
+        else:
+            if efficient_test:
+                result = np2tmp(result)
+            results.append(result)
+
+        if rank == 0:
+            batch_size = data['img'][0].size(0)
+            for _ in range(batch_size * world_size):
+                prog_bar.update()
+
+    # collect results from all ranks
+    if gpu_collect:
+        results = collect_results_gpu(results, len(dataset))
+    else:
+        results = collect_results_cpu(results, len(dataset), tmpdir)
+    return results
+
+
+def collect_results_cpu(result_part, size, tmpdir=None):
+    """Collect results with CPU."""
+    rank, world_size = get_dist_info()
+    # create a tmp dir if it is not specified
+    if tmpdir is None:
+        MAX_LEN = 512
+        # 32 is whitespace
+        dir_tensor = torch.full((MAX_LEN, ),
+                                32,
+                                dtype=torch.uint8,
+                                device='cuda')
+        if rank == 0:
+            tmpdir = tempfile.mkdtemp()
+            tmpdir = torch.tensor(
+                bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda')
+            dir_tensor[:len(tmpdir)] = tmpdir
+        dist.broadcast(dir_tensor, 0)
+        tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip()
+    else:
+        mmcv.mkdir_or_exist(tmpdir)
+    # dump the part result to the dir
+    mmcv.dump(result_part, osp.join(tmpdir, 'part_{}.pkl'.format(rank)))
+    dist.barrier()
+    # collect all parts
+    if rank != 0:
+        return None
+    else:
+        # load results of all parts from tmp dir
+        part_list = []
+        for i in range(world_size):
+            part_file = osp.join(tmpdir, 'part_{}.pkl'.format(i))
+            part_list.append(mmcv.load(part_file))
+        # sort the results
+        ordered_results = []
+        for res in zip(*part_list):
+            ordered_results.extend(list(res))
+        # the dataloader may pad some samples
+        ordered_results = ordered_results[:size]
+        # remove tmp dir
+        shutil.rmtree(tmpdir)
+        return ordered_results
+
+
+def collect_results_gpu(result_part, size):
+    """Collect results with GPU."""
+    rank, world_size = get_dist_info()
+    # dump result part to tensor with pickle
+    part_tensor = torch.tensor(
+        bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')
+    # gather all result part tensor shape
+    shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
+    shape_list = [shape_tensor.clone() for _ in range(world_size)]
+    dist.all_gather(shape_list, shape_tensor)
+    # padding result part tensor to max length
+    shape_max = torch.tensor(shape_list).max()
+    part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')
+    part_send[:shape_tensor[0]] = part_tensor
+    part_recv_list = [
+        part_tensor.new_zeros(shape_max) for _ in range(world_size)
+    ]
+    # gather all result part
+    dist.all_gather(part_recv_list, part_send)
+
+    if rank == 0:
+        part_list = []
+        for recv, shape in zip(part_recv_list, shape_list):
+            part_list.append(
+                pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()))
+        # sort the results
+        ordered_results = []
+        for res in zip(*part_list):
+            ordered_results.extend(list(res))
+        # the dataloader may pad some samples
+        ordered_results = ordered_results[:size]
+        return ordered_results

models/mmseg/apis/train.py

@@ -0,0 +1,115 @@
+import random
+import warnings
+
+import numpy as np
+import torch
+from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
+from mmcv.runner import build_optimizer, build_runner
+
+from mmseg.core import DistEvalHook, EvalHook
+from mmseg.datasets import build_dataloader, build_dataset
+from mmseg.utils import get_root_logger
+
+
+def set_random_seed(seed, deterministic=False):
+    """Set random seed.
+    Args:
+        seed (int): Seed to be used.
+        deterministic (bool): Whether to set the deterministic option for
+            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`
+            to True and `torch.backends.cudnn.benchmark` to False.
+            Default: False.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+def train_segmentor(model,
+                    dataset,
+                    cfg,
+                    distributed=False,
+                    validate=False,
+                    timestamp=None,
+                    meta=None):
+    """Launch segmentor training."""
+    logger = get_root_logger(cfg.log_level)
+
+    # prepare data loaders
+    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
+    data_loaders = [
+        build_dataloader(
+            ds,
+            cfg.data.samples_per_gpu,
+            cfg.data.workers_per_gpu,
+            # cfg.gpus will be ignored if distributed
+            len(cfg.gpu_ids),
+            dist=distributed,
+            seed=cfg.seed,
+            drop_last=True) for ds in dataset
+    ]
+
+    # put model on gpus
+    if distributed:
+        find_unused_parameters = cfg.get('find_unused_parameters', False)
+        # Sets the `find_unused_parameters` parameter in
+        # torch.nn.parallel.DistributedDataParallel
+        model = MMDistributedDataParallel(
+            model.cuda(),
+            device_ids=[torch.cuda.current_device()],
+            broadcast_buffers=False,
+            find_unused_parameters=find_unused_parameters)
+    else:
+        model = MMDataParallel(
+            model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)
+
+    # build runner
+    optimizer = build_optimizer(model, cfg.optimizer)
+
+    if cfg.get('runner') is None:
+        cfg.runner = {'type': 'IterBasedRunner', 'max_iters': cfg.total_iters}
+        warnings.warn(
+            'config is now expected to have a `runner` section, '
+            'please set `runner` in your config.', UserWarning)
+
+    runner = build_runner(
+        cfg.runner,
+        default_args=dict(
+            model=model,
+            batch_processor=None,
+            optimizer=optimizer,
+            work_dir=cfg.work_dir,
+            logger=logger,
+            meta=meta))
+
+    # register hooks
+    runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config,
+                                   cfg.checkpoint_config, cfg.log_config,
+                                   cfg.get('momentum_config', None))
+
+    # an ugly walkaround to make the .log and .log.json filenames the same
+    runner.timestamp = timestamp
+
+    # register eval hooks
+    if validate:
+        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
+        val_dataloader = build_dataloader(
+            val_dataset,
+            samples_per_gpu=1,
+            workers_per_gpu=cfg.data.workers_per_gpu,
+            dist=distributed,
+            shuffle=False)
+        eval_cfg = cfg.get('evaluation', {})
+        eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
+        eval_hook = DistEvalHook if distributed else EvalHook
+        runner.register_hook(eval_hook(val_dataloader, **eval_cfg))
+
+    if cfg.resume_from:
+        runner.resume(cfg.resume_from)
+    elif cfg.load_from:
+        runner.load_checkpoint(cfg.load_from)
+    runner.run(data_loaders, cfg.workflow)

models/mmseg/core/__init__.py

@@ -0,0 +1,3 @@
+from .evaluation import *  # noqa: F401, F403
+from .seg import *  # noqa: F401, F403
+from .utils import *  # noqa: F401, F403

models/mmseg/core/evaluation/__init__.py

@@ -0,0 +1,8 @@
+from .class_names import get_classes, get_palette
+from .eval_hooks import DistEvalHook, EvalHook
+from .metrics import eval_metrics, mean_dice, mean_iou
+
+__all__ = [
+    'EvalHook', 'DistEvalHook', 'mean_dice', 'mean_iou', 'eval_metrics',
+    'get_classes', 'get_palette'
+]

models/mmseg/core/evaluation/class_names.py

@@ -0,0 +1,152 @@
+import mmcv
+
+
+def cityscapes_classes():
+    """Cityscapes class names for external use."""
+    return [
+        'road', 'sidewalk', 'building', 'wall', 'fence', 'pole',
+        'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky',
+        'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle',
+        'bicycle'
+    ]
+
+
+def ade_classes():
+    """ADE20K class names for external use."""
+    return [
+        'wall', 'building', 'sky', 'floor', 'tree', 'ceiling', 'road', 'bed ',
+        'windowpane', 'grass', 'cabinet', 'sidewalk', 'person', 'earth',
+        'door', 'table', 'mountain', 'plant', 'curtain', 'chair', 'car',
+        'water', 'painting', 'sofa', 'shelf', 'house', 'sea', 'mirror', 'rug',
+        'field', 'armchair', 'seat', 'fence', 'desk', 'rock', 'wardrobe',
+        'lamp', 'bathtub', 'railing', 'cushion', 'base', 'box', 'column',
+        'signboard', 'chest of drawers', 'counter', 'sand', 'sink',
+        'skyscraper', 'fireplace', 'refrigerator', 'grandstand', 'path',
+        'stairs', 'runway', 'case', 'pool table', 'pillow', 'screen door',
+        'stairway', 'river', 'bridge', 'bookcase', 'blind', 'coffee table',
+        'toilet', 'flower', 'book', 'hill', 'bench', 'countertop', 'stove',
+        'palm', 'kitchen island', 'computer', 'swivel chair', 'boat', 'bar',
+        'arcade machine', 'hovel', 'bus', 'towel', 'light', 'truck', 'tower',
+        'chandelier', 'awning', 'streetlight', 'booth', 'television receiver',
+        'airplane', 'dirt track', 'apparel', 'pole', 'land', 'bannister',
+        'escalator', 'ottoman', 'bottle', 'buffet', 'poster', 'stage', 'van',
+        'ship', 'fountain', 'conveyer belt', 'canopy', 'washer', 'plaything',
+        'swimming pool', 'stool', 'barrel', 'basket', 'waterfall', 'tent',
+        'bag', 'minibike', 'cradle', 'oven', 'ball', 'food', 'step', 'tank',
+        'trade name', 'microwave', 'pot', 'animal', 'bicycle', 'lake',
+        'dishwasher', 'screen', 'blanket', 'sculpture', 'hood', 'sconce',
+        'vase', 'traffic light', 'tray', 'ashcan', 'fan', 'pier', 'crt screen',
+        'plate', 'monitor', 'bulletin board', 'shower', 'radiator', 'glass',
+        'clock', 'flag'
+    ]
+
+
+def voc_classes():
+    """Pascal VOC class names for external use."""
+    return [
+        'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
+        'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse',
+        'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train',
+        'tvmonitor'
+    ]
+
+
+def cityscapes_palette():
+    """Cityscapes palette for external use."""
+    return [[128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156],
+            [190, 153, 153], [153, 153, 153], [250, 170, 30], [220, 220, 0],
+            [107, 142, 35], [152, 251, 152], [70, 130, 180], [220, 20, 60],
+            [255, 0, 0], [0, 0, 142], [0, 0, 70], [0, 60, 100], [0, 80, 100],
+            [0, 0, 230], [119, 11, 32]]
+
+
+def ade_palette():
+    """ADE20K palette for external use."""
+    return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
+            [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
+            [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
+            [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
+            [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
+            [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
+            [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
+            [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
+            [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
+            [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
+            [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
+            [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
+            [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
+            [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
+            [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
+            [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
+            [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
+            [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
+            [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
+            [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
+            [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
+            [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
+            [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
+            [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
+            [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
+            [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
+            [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
+            [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
+            [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
+            [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
+            [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
+            [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
+            [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
+            [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
+            [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
+            [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
+            [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
+            [102, 255, 0], [92, 0, 255]]
+
+
+def voc_palette():
+    """Pascal VOC palette for external use."""
+    return [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
+            [128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0],
+            [192, 0, 0], [64, 128, 0], [192, 128, 0], [64, 0, 128],
+            [192, 0, 128], [64, 128, 128], [192, 128, 128], [0, 64, 0],
+            [128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128]]
+
+
+dataset_aliases = {
+    'cityscapes': ['cityscapes'],
+    'ade': ['ade', 'ade20k'],
+    'voc': ['voc', 'pascal_voc', 'voc12', 'voc12aug']
+}
+
+
+def get_classes(dataset):
+    """Get class names of a dataset."""
+    alias2name = {}
+    for name, aliases in dataset_aliases.items():
+        for alias in aliases:
+            alias2name[alias] = name
+
+    if mmcv.is_str(dataset):
+        if dataset in alias2name:
+            labels = eval(alias2name[dataset] + '_classes()')
+        else:
+            raise ValueError(f'Unrecognized dataset: {dataset}')
+    else:
+        raise TypeError(f'dataset must a str, but got {type(dataset)}')
+    return labels
+
+
+def get_palette(dataset):
+    """Get class palette (RGB) of a dataset."""
+    alias2name = {}
+    for name, aliases in dataset_aliases.items():
+        for alias in aliases:
+            alias2name[alias] = name
+
+    if mmcv.is_str(dataset):
+        if dataset in alias2name:
+            labels = eval(alias2name[dataset] + '_palette()')
+        else:
+            raise ValueError(f'Unrecognized dataset: {dataset}')
+    else:
+        raise TypeError(f'dataset must a str, but got {type(dataset)}')
+    return labels

models/mmseg/core/evaluation/eval_hooks.py

@@ -0,0 +1,107 @@
+import os.path as osp
+
+from mmcv.runner import Hook
+from torch.utils.data import DataLoader
+
+
+class EvalHook(Hook):
+    """Evaluation hook.
+
+    Attributes:
+        dataloader (DataLoader): A PyTorch dataloader.
+        interval (int): Evaluation interval (by epochs). Default: 1.
+    """
+
+    def __init__(self, dataloader, interval=1, by_epoch=False, **eval_kwargs):
+        if not isinstance(dataloader, DataLoader):
+            raise TypeError('dataloader must be a pytorch DataLoader, but got '
+                            f'{type(dataloader)}')
+        self.dataloader = dataloader
+        self.interval = interval
+        self.by_epoch = by_epoch
+        self.eval_kwargs = eval_kwargs
+
+    def after_train_iter(self, runner):
+        """After train epoch hook."""
+        if self.by_epoch or not self.every_n_iters(runner, self.interval):
+            return
+        from mmseg.apis import single_gpu_test
+        runner.log_buffer.clear()
+        results = single_gpu_test(runner.model, self.dataloader, show=False)
+        self.evaluate(runner, results)
+
+    def after_train_epoch(self, runner):
+        """After train epoch hook."""
+        if not self.by_epoch or not self.every_n_epochs(runner, self.interval):
+            return
+        from mmseg.apis import single_gpu_test
+        runner.log_buffer.clear()
+        results = single_gpu_test(runner.model, self.dataloader, show=False)
+        self.evaluate(runner, results)
+
+    def evaluate(self, runner, results):
+        """Call evaluate function of dataset."""
+        eval_res = self.dataloader.dataset.evaluate(
+            results, logger=runner.logger, **self.eval_kwargs)
+        for name, val in eval_res.items():
+            runner.log_buffer.output[name] = val
+        runner.log_buffer.ready = True
+
+
+class DistEvalHook(EvalHook):
+    """Distributed evaluation hook.
+
+    Attributes:
+        dataloader (DataLoader): A PyTorch dataloader.
+        interval (int): Evaluation interval (by epochs). Default: 1.
+        tmpdir (str | None): Temporary directory to save the results of all
+            processes. Default: None.
+        gpu_collect (bool): Whether to use gpu or cpu to collect results.
+            Default: False.
+    """
+
+    def __init__(self,
+                 dataloader,
+                 interval=1,
+                 gpu_collect=False,
+                 by_epoch=False,
+                 **eval_kwargs):
+        if not isinstance(dataloader, DataLoader):
+            raise TypeError(
+                'dataloader must be a pytorch DataLoader, but got {}'.format(
+                    type(dataloader)))
+        self.dataloader = dataloader
+        self.interval = interval
+        self.gpu_collect = gpu_collect
+        self.by_epoch = by_epoch
+        self.eval_kwargs = eval_kwargs
+
+    def after_train_iter(self, runner):
+        """After train epoch hook."""
+        if self.by_epoch or not self.every_n_iters(runner, self.interval):
+            return
+        from mmseg.apis import multi_gpu_test
+        runner.log_buffer.clear()
+        results = multi_gpu_test(
+            runner.model,
+            self.dataloader,
+            tmpdir=osp.join(runner.work_dir, '.eval_hook'),
+            gpu_collect=self.gpu_collect)
+        if runner.rank == 0:
+            print('\n')
+            self.evaluate(runner, results)
+
+    def after_train_epoch(self, runner):
+        """After train epoch hook."""
+        if not self.by_epoch or not self.every_n_epochs(runner, self.interval):
+            return
+        from mmseg.apis import multi_gpu_test
+        runner.log_buffer.clear()
+        results = multi_gpu_test(
+            runner.model,
+            self.dataloader,
+            tmpdir=osp.join(runner.work_dir, '.eval_hook'),
+            gpu_collect=self.gpu_collect)
+        if runner.rank == 0:
+            print('\n')
+            self.evaluate(runner, results)

models/mmseg/core/evaluation/metrics.py

@@ -0,0 +1,229 @@
+import mmcv
+import numpy as np
+
+
+def intersect_and_union(pred_label,
+                        label,
+                        num_classes,
+                        ignore_index,
+                        label_map=dict(),
+                        reduce_zero_label=False):
+    """Calculate intersection and Union.
+
+    Args:
+        pred_label (ndarray): Prediction segmentation map.
+        label (ndarray): Ground truth segmentation map.
+        num_classes (int): Number of categories.
+        ignore_index (int): Index that will be ignored in evaluation.
+        label_map (dict): Mapping old labels to new labels. The parameter will
+            work only when label is str. Default: dict().
+        reduce_zero_label (bool): Wether ignore zero label. The parameter will
+            work only when label is str. Default: False.
+
+     Returns:
+         ndarray: The intersection of prediction and ground truth histogram
+             on all classes.
+         ndarray: The union of prediction and ground truth histogram on all
+             classes.
+         ndarray: The prediction histogram on all classes.
+         ndarray: The ground truth histogram on all classes.
+    """
+
+    if isinstance(pred_label, str):
+        pred_label = np.load(pred_label)
+
+    if isinstance(label, str):
+        label = mmcv.imread(label, flag='unchanged', backend='pillow')
+    # modify if custom classes
+    if label_map is not None:
+        for old_id, new_id in label_map.items():
+            label[label == old_id] = new_id
+    if reduce_zero_label:
+        # avoid using underflow conversion
+        label[label == 0] = 255
+        label = label - 1
+        label[label == 254] = 255
+
+    mask = (label != ignore_index)
+    pred_label = pred_label[mask]
+    label = label[mask]
+
+    intersect = pred_label[pred_label == label]
+    area_intersect, _ = np.histogram(
+        intersect, bins=np.arange(num_classes + 1))
+    area_pred_label, _ = np.histogram(
+        pred_label, bins=np.arange(num_classes + 1))
+    area_label, _ = np.histogram(label, bins=np.arange(num_classes + 1))
+    area_union = area_pred_label + area_label - area_intersect
+
+    return area_intersect, area_union, area_pred_label, area_label
+
+
+def total_intersect_and_union(results,
+                              gt_seg_maps,
+                              num_classes,
+                              ignore_index,
+                              label_map=dict(),
+                              reduce_zero_label=False):
+    """Calculate Total Intersection and Union.
+
+    Args:
+        results (list[ndarray]): List of prediction segmentation maps.
+        gt_seg_maps (list[ndarray]): list of ground truth segmentation maps.
+        num_classes (int): Number of categories.
+        ignore_index (int): Index that will be ignored in evaluation.
+        label_map (dict): Mapping old labels to new labels. Default: dict().
+        reduce_zero_label (bool): Wether ignore zero label. Default: False.
+
+     Returns:
+         ndarray: The intersection of prediction and ground truth histogram
+             on all classes.
+         ndarray: The union of prediction and ground truth histogram on all
+             classes.
+         ndarray: The prediction histogram on all classes.
+         ndarray: The ground truth histogram on all classes.
+    """
+
+    num_imgs = len(results)
+    assert len(gt_seg_maps) == num_imgs
+    total_area_intersect = np.zeros((num_classes, ), dtype=np.float)
+    total_area_union = np.zeros((num_classes, ), dtype=np.float)
+    total_area_pred_label = np.zeros((num_classes, ), dtype=np.float)
+    total_area_label = np.zeros((num_classes, ), dtype=np.float)
+    for i in range(num_imgs):
+        area_intersect, area_union, area_pred_label, area_label = \
+            intersect_and_union(results[i], gt_seg_maps[i], num_classes,
+                                ignore_index, label_map, reduce_zero_label)
+        total_area_intersect += area_intersect
+        total_area_union += area_union
+        total_area_pred_label += area_pred_label
+        total_area_label += area_label
+    return total_area_intersect, total_area_union, \
+        total_area_pred_label, total_area_label
+
+
+def mean_iou(results,
+             gt_seg_maps,
+             num_classes,
+             ignore_index,
+             nan_to_num=None,
+             label_map=dict(),
+             reduce_zero_label=False):
+    """Calculate Mean Intersection and Union (mIoU)
+
+    Args:
+        results (list[ndarray]): List of prediction segmentation maps.
+        gt_seg_maps (list[ndarray]): list of ground truth segmentation maps.
+        num_classes (int): Number of categories.
+        ignore_index (int): Index that will be ignored in evaluation.
+        nan_to_num (int, optional): If specified, NaN values will be replaced
+            by the numbers defined by the user. Default: None.
+        label_map (dict): Mapping old labels to new labels. Default: dict().
+        reduce_zero_label (bool): Wether ignore zero label. Default: False.
+
+     Returns:
+         float: Overall accuracy on all images.
+         ndarray: Per category accuracy, shape (num_classes, ).
+         ndarray: Per category IoU, shape (num_classes, ).
+    """
+
+    all_acc, acc, iou = eval_metrics(
+        results=results,
+        gt_seg_maps=gt_seg_maps,
+        num_classes=num_classes,
+        ignore_index=ignore_index,
+        metrics=['mIoU'],
+        nan_to_num=nan_to_num,
+        label_map=label_map,
+        reduce_zero_label=reduce_zero_label)
+    return all_acc, acc, iou
+
+
+def mean_dice(results,
+              gt_seg_maps,
+              num_classes,
+              ignore_index,
+              nan_to_num=None,
+              label_map=dict(),
+              reduce_zero_label=False):
+    """Calculate Mean Dice (mDice)
+
+    Args:
+        results (list[ndarray]): List of prediction segmentation maps.
+        gt_seg_maps (list[ndarray]): list of ground truth segmentation maps.
+        num_classes (int): Number of categories.
+        ignore_index (int): Index that will be ignored in evaluation.
+        nan_to_num (int, optional): If specified, NaN values will be replaced
+            by the numbers defined by the user. Default: None.
+        label_map (dict): Mapping old labels to new labels. Default: dict().
+        reduce_zero_label (bool): Wether ignore zero label. Default: False.
+
+     Returns:
+         float: Overall accuracy on all images.
+         ndarray: Per category accuracy, shape (num_classes, ).
+         ndarray: Per category dice, shape (num_classes, ).
+    """
+
+    all_acc, acc, dice = eval_metrics(
+        results=results,
+        gt_seg_maps=gt_seg_maps,
+        num_classes=num_classes,
+        ignore_index=ignore_index,
+        metrics=['mDice'],
+        nan_to_num=nan_to_num,
+        label_map=label_map,
+        reduce_zero_label=reduce_zero_label)
+    return all_acc, acc, dice
+
+
+def eval_metrics(results,
+                 gt_seg_maps,
+                 num_classes,
+                 ignore_index,
+                 metrics=['mIoU'],
+                 nan_to_num=None,
+                 label_map=dict(),
+                 reduce_zero_label=False):
+    """Calculate evaluation metrics
+    Args:
+        results (list[ndarray]): List of prediction segmentation maps.
+        gt_seg_maps (list[ndarray]): list of ground truth segmentation maps.
+        num_classes (int): Number of categories.
+        ignore_index (int): Index that will be ignored in evaluation.
+        metrics (list[str] | str): Metrics to be evaluated, 'mIoU' and 'mDice'.
+        nan_to_num (int, optional): If specified, NaN values will be replaced
+            by the numbers defined by the user. Default: None.
+        label_map (dict): Mapping old labels to new labels. Default: dict().
+        reduce_zero_label (bool): Wether ignore zero label. Default: False.
+     Returns:
+         float: Overall accuracy on all images.
+         ndarray: Per category accuracy, shape (num_classes, ).
+         ndarray: Per category evalution metrics, shape (num_classes, ).
+    """
+
+    if isinstance(metrics, str):
+        metrics = [metrics]
+    allowed_metrics = ['mIoU', 'mDice']
+    if not set(metrics).issubset(set(allowed_metrics)):
+        raise KeyError('metrics {} is not supported'.format(metrics))
+    total_area_intersect, total_area_union, total_area_pred_label, \
+        total_area_label = total_intersect_and_union(results, gt_seg_maps,
+                                                     num_classes, ignore_index,
+                                                     label_map,
+                                                     reduce_zero_label)
+    all_acc = total_area_intersect.sum() / total_area_label.sum()
+    acc = total_area_intersect / total_area_label
+    ret_metrics = [all_acc, acc]
+    for metric in metrics:
+        if metric == 'mIoU':
+            iou = total_area_intersect / total_area_union
+            ret_metrics.append(iou)
+        elif metric == 'mDice':
+            dice = 2 * total_area_intersect / (
+                total_area_pred_label + total_area_label)
+            ret_metrics.append(dice)
+    if nan_to_num is not None:
+        ret_metrics = [
+            np.nan_to_num(metric, nan=nan_to_num) for metric in ret_metrics
+        ]
+    return ret_metrics

models/mmseg/core/seg/__init__.py

@@ -0,0 +1,4 @@
+from .builder import build_pixel_sampler
+from .sampler import BasePixelSampler, OHEMPixelSampler
+
+__all__ = ['build_pixel_sampler', 'BasePixelSampler', 'OHEMPixelSampler']

models/mmseg/core/seg/builder.py

@@ -0,0 +1,8 @@
+from mmcv.utils import Registry, build_from_cfg
+
+PIXEL_SAMPLERS = Registry('pixel sampler')
+
+
+def build_pixel_sampler(cfg, **default_args):
+    """Build pixel sampler for segmentation map."""
+    return build_from_cfg(cfg, PIXEL_SAMPLERS, default_args)

models/mmseg/core/seg/sampler/__init__.py

@@ -0,0 +1,4 @@
+from .base_pixel_sampler import BasePixelSampler
+from .ohem_pixel_sampler import OHEMPixelSampler
+
+__all__ = ['BasePixelSampler', 'OHEMPixelSampler']

models/mmseg/core/seg/sampler/base_pixel_sampler.py

@@ -0,0 +1,13 @@
+from abc import ABCMeta, abstractmethod
+
+
+class BasePixelSampler(metaclass=ABCMeta):
+    """Base class of pixel sampler."""
+
+    def __init__(self, **kwargs):
+        pass
+
+    @abstractmethod
+    def sample(self, seg_logit, seg_label):
+        """Placeholder for sample function."""
+        pass

models/mmseg/core/seg/sampler/ohem_pixel_sampler.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn.functional as F
+
+from ..builder import PIXEL_SAMPLERS
+from .base_pixel_sampler import BasePixelSampler
+
+
+@PIXEL_SAMPLERS.register_module()
+class OHEMPixelSampler(BasePixelSampler):
+    """Online Hard Example Mining Sampler for segmentation.
+
+    Args:
+        context (nn.Module): The context of sampler, subclass of
+            :obj:`BaseDecodeHead`.
+        thresh (float, optional): The threshold for hard example selection.
+            Below which, are prediction with low confidence. If not
+            specified, the hard examples will be pixels of top ``min_kept``
+            loss. Default: None.
+        min_kept (int, optional): The minimum number of predictions to keep.
+            Default: 100000.
+    """
+
+    def __init__(self, context, thresh=None, min_kept=100000):
+        super(OHEMPixelSampler, self).__init__()
+        self.context = context
+        assert min_kept > 1
+        self.thresh = thresh
+        self.min_kept = min_kept
+
+    def sample(self, seg_logit, seg_label):
+        """Sample pixels that have high loss or with low prediction confidence.
+
+        Args:
+            seg_logit (torch.Tensor): segmentation logits, shape (N, C, H, W)
+            seg_label (torch.Tensor): segmentation label, shape (N, 1, H, W)
+
+        Returns:
+            torch.Tensor: segmentation weight, shape (N, H, W)
+        """
+        with torch.no_grad():
+            assert seg_logit.shape[2:] == seg_label.shape[2:]
+            assert seg_label.shape[1] == 1
+            seg_label = seg_label.squeeze(1).long()
+            batch_kept = self.min_kept * seg_label.size(0)
+            valid_mask = seg_label != self.context.ignore_index
+            seg_weight = seg_logit.new_zeros(size=seg_label.size())
+            valid_seg_weight = seg_weight[valid_mask]
+            if self.thresh is not None:
+                seg_prob = F.softmax(seg_logit, dim=1)
+
+                tmp_seg_label = seg_label.clone().unsqueeze(1)
+                tmp_seg_label[tmp_seg_label == self.context.ignore_index] = 0
+                seg_prob = seg_prob.gather(1, tmp_seg_label).squeeze(1)
+                sort_prob, sort_indices = seg_prob[valid_mask].sort()
+
+                if sort_prob.numel() > 0:
+                    min_threshold = sort_prob[min(batch_kept,
+                                                  sort_prob.numel() - 1)]
+                else:
+                    min_threshold = 0.0
+                threshold = max(min_threshold, self.thresh)
+                valid_seg_weight[seg_prob[valid_mask] < threshold] = 1.
+            else:
+                losses = self.context.loss_decode(
+                    seg_logit,
+                    seg_label,
+                    weight=None,
+                    ignore_index=self.context.ignore_index,
+                    reduction_override='none')
+                # faster than topk according to https://github.com/pytorch/pytorch/issues/22812  # noqa
+                _, sort_indices = losses[valid_mask].sort(descending=True)
+                valid_seg_weight[sort_indices[:batch_kept]] = 1.
+
+            seg_weight[valid_mask] = valid_seg_weight
+
+            return seg_weight

models/mmseg/core/utils/__init__.py

@@ -0,0 +1,3 @@
+from .misc import add_prefix
+
+__all__ = ['add_prefix']

models/mmseg/core/utils/misc.py

@@ -0,0 +1,17 @@
+def add_prefix(inputs, prefix):
+    """Add prefix for dict.
+
+    Args:
+        inputs (dict): The input dict with str keys.
+        prefix (str): The prefix to add.
+
+    Returns:
+
+        dict: The dict with keys updated with ``prefix``.
+    """
+
+    outputs = dict()
+    for name, value in inputs.items():
+        outputs[f'{prefix}.{name}'] = value
+
+    return outputs