add from PIL import Image, and update get_class_colors method

README.md

@@ -85,11 +85,11 @@ For RGB-Depth semantic segmentation, the generation of HHA maps from Depth maps
 
 2. Config
 
-    Edit config file in `configs.py`, including dataset and network settings.
+    Edit config file in configs directory, including dataset and network settings.
 
 3. Run multi GPU distributed training:
     ```shell
-    $ CUDA_VISIBLE_DEVICES="GPU IDs" python -m torch.distributed.launch --nproc_per_node="GPU numbers you want to use" train.py
+    $ CUDA_VISIBLE_DEVICES="GPU IDs" python -m torch.distributed.launch --nproc_per_node="GPU numbers you want to use" train.py -f "config_file"
     ```
 
 - The tensorboard file is saved in `log_<datasetName>_<backboneSize>/tb/` directory.
@@ -98,7 +98,7 @@ For RGB-Depth semantic segmentation, the generation of HHA maps from Depth maps
 ### Evaluation
 Run the evaluation by:
 ```shell
-CUDA_VISIBLE_DEVICES="GPU IDs" python eval.py -d="Device ID" -e="epoch number or range"
+CUDA_VISIBLE_DEVICES="GPU IDs" python eval.py -f "config_file" -d="Device ID" -e="epoch number or range"
 ```
 If you want to use multi GPUs please specify multiple Device IDs (0,1,2...).
 

config.py → configs/NYUDepthv2/NYUDepthv2_mit_b2_e500_v1.py

@@ -17,7 +17,6 @@
 C.abs_dir = osp.realpath(".")
 
 # Dataset config
-"""Dataset Path"""
 C.dataset_name = 'NYUDepthv2'
 C.dataset_path = osp.join(C.root_dir, 'datasets', 'NYUDepthv2')
 C.rgb_root_folder = osp.join(C.dataset_path, 'RGB')
@@ -29,14 +28,14 @@
 # True for most dataset valid, Faslse for MFNet(?)
 C.x_root_folder = osp.join(C.dataset_path, 'HHA')
 C.x_format = '.jpg'
-C.x_is_single_channel = False # True for raw depth, thermal and aolp/dolp(not aolp/dolp tri) input
+C.x_is_single_channel = False  # True for raw depth, thermal and aolp/dolp(not aolp/dolp tri) input
 C.train_source = osp.join(C.dataset_path, "train.txt")
 C.eval_source = osp.join(C.dataset_path, "test.txt")
 C.is_test = False
 C.num_train_imgs = 795
 C.num_eval_imgs = 654
 C.num_classes = 40
-C.class_names =  ['wall','floor','cabinet','bed','chair','sofa','table','door','window','bookshelf','picture','counter','blinds',
+C.class_names = ['wall','floor','cabinet','bed','chair','sofa','table','door','window','bookshelf','picture','counter','blinds',
     'desk','shelves','curtain','dresser','pillow','mirror','floor mat','clothes','ceiling','books','refridgerator',
     'television','paper','towel','shower curtain','box','whiteboard','person','night stand','toilet',
     'sink','lamp','bathtub','bag','otherstructure','otherfurniture','otherprop']
@@ -49,7 +48,7 @@
 C.norm_std = np.array([0.229, 0.224, 0.225])
 
 """ Settings for network, this would be different for each kind of model"""
-C.backbone = 'mit_b2' # Remember change the path below.
+C.backbone = 'mit_b2'  # Remember change the path below.
 C.pretrained_model = C.root_dir + '/pretrained/segformer/mit_b2.pth'
 C.decoder = 'MLPDecoder'
 C.decoder_embed_dim = 512
@@ -62,7 +61,7 @@
 C.weight_decay = 0.01
 C.batch_size = 8
 C.nepochs = 500
-C.niters_per_epoch = C.num_train_imgs // C.batch_size  + 1
+C.niters_per_epoch = C.num_train_imgs // C.batch_size + 1
 C.num_workers = 16
 C.train_scale_array = [0.5, 0.75, 1, 1.25, 1.5, 1.75]
 C.warm_up_epoch = 10
@@ -74,9 +73,9 @@
 """Eval Config"""
 C.eval_iter = 25
 C.eval_stride_rate = 2 / 3
-C.eval_scale_array = [1] # [0.75, 1, 1.25] # 
-C.eval_flip = False # True # 
-C.eval_crop_size = [480, 640] # [height weight]
+C.eval_scale_array = [1]  # [0.75, 1, 1.25] #
+C.eval_flip = False  # True #
+C.eval_crop_size = [480, 640]  # [height weight]
 
 """Store Config"""
 C.checkpoint_start_epoch = 250
@@ -88,7 +87,8 @@ def add_path(path):
         sys.path.insert(0, path)
 add_path(osp.join(C.root_dir))
 
-C.log_dir = osp.abspath('log_' + C.dataset_name + '_' + C.backbone)
+config_name=os.path.basename(__file__).split(".")[0]
+C.log_dir = osp.abspath('logs/log_' + config_name)
 C.tb_dir = osp.abspath(osp.join(C.log_dir, "tb"))
 C.log_dir_link = C.log_dir
 C.checkpoint_dir = osp.abspath(osp.join(C.log_dir, "checkpoint"))
@@ -106,5 +106,5 @@ def add_path(path):
         '-tb', '--tensorboard', default=False, action='store_true')
     args = parser.parse_args()
 
-    if args.tensorboard:
-        open_tensorboard()
+    # if args.tensorboard:
+    #     open_tensorboard()

dataloader/RGBXDataset.py

@@ -51,7 +51,7 @@ def __getitem__(self, index):
             x = self._open_image(x_path, cv2.IMREAD_GRAYSCALE)
             x = cv2.merge([x, x, x])
         else:
-            x =  self._open_image(x_path, cv2.COLOR_BGR2RGB)
+            x = self._open_image(x_path, cv2.COLOR_BGR2RGB)
 
         if self.preprocess is not None:
             rgb, gt, x = self.preprocess(rgb, gt, x)

dataloader/dataloader.py

@@ -3,9 +3,10 @@
 import numpy as np
 from torch.utils import data
 import random
-from config import config
+# from config import config
 from utils.transforms import generate_random_crop_pos, random_crop_pad_to_shape, normalize
 
+
 def random_mirror(rgb, gt, modal_x):
     if random.random() >= 0.5:
         rgb = cv2.flip(rgb, 1)
@@ -14,6 +15,7 @@ def random_mirror(rgb, gt, modal_x):
 
     return rgb, gt, modal_x
 
+
 def random_scale(rgb, gt, modal_x, scales):
     scale = random.choice(scales)
     sh = int(rgb.shape[0] * scale)
@@ -24,20 +26,22 @@ def random_scale(rgb, gt, modal_x, scales):
 
     return rgb, gt, modal_x, scale
 
+
 class TrainPre(object):
-    def __init__(self, norm_mean, norm_std):
-        self.norm_mean = norm_mean
-        self.norm_std = norm_std
+    def __init__(self, config):
+        self.config = config
+        self.norm_mean = config.norm_mean
+        self.norm_std = config.norm_std
 
     def __call__(self, rgb, gt, modal_x):
         rgb, gt, modal_x = random_mirror(rgb, gt, modal_x)
-        if config.train_scale_array is not None:
-            rgb, gt, modal_x, scale = random_scale(rgb, gt, modal_x, config.train_scale_array)
+        if self.config.train_scale_array is not None:
+            rgb, gt, modal_x, scale = random_scale(rgb, gt, modal_x, self.config.train_scale_array)
 
         rgb = normalize(rgb, self.norm_mean, self.norm_std)
         modal_x = normalize(modal_x, self.norm_mean, self.norm_std)
 
-        crop_size = (config.image_height, config.image_width)
+        crop_size = (self.config.image_height, self.config.image_width)
         crop_pos = generate_random_crop_pos(rgb.shape[:2], crop_size)
 
         p_rgb, _ = random_crop_pad_to_shape(rgb, crop_pos, crop_size, 0)
@@ -49,11 +53,13 @@ def __call__(self, rgb, gt, modal_x):
 
         return p_rgb, p_gt, p_modal_x
 
+
 class ValPre(object):
     def __call__(self, rgb, gt, modal_x):
         return rgb, gt, modal_x
 
-def get_train_loader(engine, dataset):
+
+def get_train_loader(engine, dataset, config=None):
     data_setting = {'rgb_root': config.rgb_root_folder,
                     'rgb_format': config.rgb_format,
                     'gt_root': config.gt_root_folder,
@@ -66,7 +72,7 @@ def get_train_loader(engine, dataset):
                     'train_source': config.train_source,
                     'eval_source': config.eval_source,
                     'class_names': config.class_names}
-    train_preprocess = TrainPre(config.norm_mean, config.norm_std)
+    train_preprocess = TrainPre(config)
 
     train_dataset = dataset(data_setting, "train", train_preprocess, config.batch_size * config.niters_per_epoch)
 

engine/engine.py

@@ -61,6 +61,8 @@ def __init__(self, custom_parser=None):
 
     def inject_default_parser(self):
         p = self.parser
+        p.add_argument("-f", "--config_file", default=None, type=str,
+                            help="plz input your experiment description file", )
         p.add_argument('-d', '--devices', default='',
                        help='set data parallel training')
         p.add_argument('-c', '--continue', type=extant_file,

engine/evaluator.py

@@ -50,29 +50,34 @@ def run(self, model_path, model_indice, log_file, log_file_link):
             models = [model_indice, ]
         elif "-" in model_indice:
             start_epoch = int(model_indice.split("-")[0])
-            end_epoch = model_indice.split("-")[1]
-
-            models = os.listdir(model_path)
-            models.remove("epoch-last.pth")
-            sorted_models = [None] * len(models)
-            model_idx = [0] * len(models)
-
-            for idx, m in enumerate(models):
-                num = m.split(".")[0].split("-")[1]
-                model_idx[idx] = num
-                sorted_models[idx] = m
-            model_idx = np.array([int(i) for i in model_idx])
-
-            down_bound = model_idx >= start_epoch
-            up_bound = [True] * len(sorted_models)
-            if end_epoch:
-                end_epoch = int(end_epoch)
-                assert start_epoch < end_epoch
-                up_bound = model_idx <= end_epoch
-            bound = up_bound * down_bound
-            model_slice = np.array(sorted_models)[bound]
-            models = [os.path.join(model_path, model) for model in
-                      model_slice]
+            end_epoch = int(model_indice.split("-")[1])
+
+            models = []
+            for i in range(start_epoch, end_epoch + 1):
+                models.append(os.path.join(model_path, 'epoch-%s.pth' % i))
+
+            # models = os.listdir(model_path)
+            # models.remove("epoch-last.pth")
+            # sorted_models = [None] * len(models)
+            # model_idx = [0] * len(models)
+            #
+            # for idx, m in enumerate(models):
+            #     num = m.split(".")[0].split("-")[1]
+            #     model_idx[idx] = num
+            #     sorted_models[idx] = m
+            # model_idx = np.array([int(i) for i in model_idx])
+            # model_idx.sort()
+            #
+            # down_bound = model_idx >= start_epoch
+            # up_bound = [True] * len(sorted_models)
+            # if end_epoch:
+            #     end_epoch = int(end_epoch)
+            #     assert start_epoch < end_epoch
+            #     up_bound = model_idx <= end_epoch
+            # bound = up_bound * down_bound
+            # model_slice = np.array(sorted_models)[bound]
+            # models = [os.path.join(model_path, model) for model in
+            #           model_slice]
         else:
             if os.path.exists(model_path):
                 models = [os.path.join(model_path, 'epoch-%s.pth' % model_indice), ]
@@ -306,7 +311,10 @@ def process_image(self, img, crop_size=None):
     def sliding_eval_rgbX(self, img, modal_x, crop_size, stride_rate, device=None):
         crop_size = to_2tuple(crop_size)
         ori_rows, ori_cols, _ = img.shape
-        processed_pred = np.zeros((ori_rows, ori_cols, self.class_num))
+        if self.class_num < 2:
+            processed_pred = np.zeros((ori_rows, ori_cols))
+        else:
+            processed_pred = np.zeros((ori_rows, ori_cols, self.class_num))
 
         for s in self.multi_scales:
             img_scale = cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_LINEAR)
@@ -319,7 +327,8 @@ def sliding_eval_rgbX(self, img, modal_x, crop_size, stride_rate, device=None):
             processed_pred += self.scale_process_rgbX(img_scale, modal_x_scale, (ori_rows, ori_cols),
                                                         crop_size, stride_rate, device)
 
-        pred = processed_pred.argmax(2)
+        if self.class_num > 1:
+            pred = processed_pred.argmax(2)
 
         return pred
 

eval.py

@@ -6,18 +6,21 @@
 import torch
 import torch.nn as nn
 
-from config import config
+# from config import config
+from utils.config_utils import get_config_by_file
 from utils.pyt_utils import ensure_dir, link_file, load_model, parse_devices
-from utils.visualize import print_iou, show_img
+from utils.visualize import print_iou, show_img, get_class_colors
 from engine.evaluator import Evaluator
 from engine.logger import get_logger
 from utils.metric import hist_info, compute_score
 from dataloader.RGBXDataset import RGBXDataset
 from models.builder import EncoderDecoder as segmodel
 from dataloader.dataloader import ValPre
+from PIL import Image
 
 logger = get_logger()
 
+
 class SegEvaluator(Evaluator):
     def func_per_iteration(self, data, device):
         img = data['data']
@@ -36,7 +39,7 @@ def func_per_iteration(self, data, device):
 
             # save colored result
             result_img = Image.fromarray(pred.astype(np.uint8), mode='P')
-            class_colors = get_class_colors()
+            class_colors = get_class_colors(config.num_classes)
             palette_list = list(np.array(class_colors).flat)
             if len(palette_list) < 768:
                 palette_list += [0] * (768 - len(palette_list))
@@ -75,8 +78,11 @@ def compute_metric(self, results):
                                 dataset.class_names, show_no_back=False)
         return result_line
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--config_file", default=None, type=str,
+        help="plz input your experiment description file",)
     parser.add_argument('-e', '--epochs', default='last', type=str)
     parser.add_argument('-d', '--devices', default='0', type=str)
     parser.add_argument('-v', '--verbose', default=False, action='store_true')
@@ -86,6 +92,7 @@ def compute_metric(self, results):
 
     args = parser.parse_args()
     all_dev = parse_devices(args.devices)
+    config = get_config_by_file(args.config_file)
 
     network = segmodel(cfg=config, criterion=None, norm_layer=nn.BatchNorm2d)
     data_setting = {'rgb_root': config.rgb_root_folder,

models/builder.py

@@ -47,6 +47,10 @@ def __init__(self, cfg=None, criterion=nn.CrossEntropyLoss(reduction='mean', ign
             self.channels = [32, 64, 160, 256]
             from .encoders.dual_segformer import mit_b0 as backbone
             self.backbone = backbone(norm_fuse=norm_layer)
+        elif cfg.backbone == 'mit_b2_s':
+            logger.info('Using backbone: Segformer-B2')
+            from .encoders.segformer import mit_b2 as backbone
+            self.backbone = backbone(norm_fuse=norm_layer)
         else:
             logger.info('Using backbone: Segformer-B2')
             from .encoders.dual_segformer import mit_b2 as backbone