├── images ├── .gitignore └── Main Figure.png ├── Codes ├── src │ ├── core │ │ ├── schedulers │ │ │ ├── __init__.py │ │ │ ├── iteration_polyLR.py │ │ │ ├── Poly_LR.py │ │ │ ├── torch_poly_lr_decay.py │ │ │ ├── OneCycle_LR.py │ │ │ └── cyclic_cos_annealing_lr.py │ │ ├── __init__.py │ │ ├── data │ │ │ ├── __init__.py │ │ │ └── samplers │ │ │ │ ├── __init__.py │ │ │ │ ├── common.py │ │ │ │ ├── batch_sampler.py │ │ │ │ └── multiscale_batch_samplers.py │ │ ├── optimizers │ │ │ ├── __init__.py │ │ │ └── adamw.py │ │ ├── loss │ │ │ ├── __init__.py │ │ │ └── segmentation_loss.py │ │ ├── callbacks │ │ │ ├── __init__.py │ │ │ └── progressbar.py │ │ ├── metrics │ │ │ ├── __init__.py │ │ │ ├── AverageMeter.py │ │ │ └── pl_segmentation_score.py │ │ ├── models │ │ │ ├── encoder │ │ │ │ ├── __init__.py │ │ │ │ ├── splat.py │ │ │ │ └── resnext.py │ │ │ ├── segmentation_decoder │ │ │ │ ├── __init__.py │ │ │ │ └── ftnet.py │ │ │ ├── __init__.py │ │ │ └── segbase.py │ │ └── utils │ │ │ ├── __init__.py │ │ │ ├── json_extension.py │ │ │ ├── tabulate_metrics.py │ │ │ ├── mean_std.py │ │ │ ├── img_utils.py │ │ │ ├── utils.py │ │ │ ├── filesystem.py │ │ │ ├── collect_env.py │ │ │ ├── logger.py │ │ │ ├── visualize.py │ │ │ └── optimizer_scheduler_helper.py │ ├── lightning_scripts │ │ ├── .gitignore │ │ ├── run_performace_10_times.sh │ │ ├── trainers │ │ │ ├── __init__.py │ │ │ └── thermal_edge_trainer.py │ │ ├── __init__.py │ │ └── performance_test.py │ ├── datasets │ │ ├── utils │ │ │ ├── splits.mat │ │ │ ├── Cityscape_folderMap.py │ │ │ ├── MFNDataset_folderMap.py │ │ │ ├── scutseg_foldermap.py │ │ │ └── SODA_folderMap.py │ │ ├── __init__.py │ │ ├── dataloaders │ │ │ ├── __init__.py │ │ │ ├── segbase.py │ │ │ ├── mfn.py │ │ │ ├── scutseg.py │ │ │ └── soda.py │ │ ├── encode_originals.py │ │ ├── edge_generation │ │ │ ├── main_cityscape_visible.m │ │ │ ├── main.m │ │ │ └── seg2edge.m │ │ └── calculate_mean_std.py │ └── bash │ │ ├── Paper_runs.sh │ │ ├── Train_and_test.sh │ │ └── Train_and_test_all.sh ├── Trained_Models │ └── .gitignore └── pretrained_models │ └── .gitignore ├── Dataset └── .gitignore ├── requirements.txt ├── .gitignore ├── LICENSE └── README.md /images/.gitignore: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /Codes/src/core/schedulers/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /Dataset/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | */ 3 | !.gitignore 4 | -------------------------------------------------------------------------------- /Codes/src/core/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | -------------------------------------------------------------------------------- /Codes/src/lightning_scripts/.gitignore: -------------------------------------------------------------------------------- 1 | lightning_logs/ 2 | -------------------------------------------------------------------------------- /Codes/Trained_Models/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | */ 3 | !.gitignore 4 | -------------------------------------------------------------------------------- /Codes/pretrained_models/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | */ 3 | !.gitignore 4 | -------------------------------------------------------------------------------- /Codes/src/core/data/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | from .samplers import * 4 | -------------------------------------------------------------------------------- /Codes/src/core/optimizers/__init__.py: -------------------------------------------------------------------------------- 1 | from .adabound import * 2 | from .adamw import * 3 | -------------------------------------------------------------------------------- /Codes/src/core/loss/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | from .segmentation_loss import * 4 | -------------------------------------------------------------------------------- /images/Main Figure.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyaskamathkm/FTNet/HEAD/images/Main Figure.png -------------------------------------------------------------------------------- /Codes/src/core/callbacks/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | 4 | from .progressbar import * 5 | -------------------------------------------------------------------------------- /Codes/src/core/metrics/__init__.py: -------------------------------------------------------------------------------- 1 | from .AverageMeter import * 2 | from .pl_segmentation_score import * 3 | -------------------------------------------------------------------------------- /Codes/src/datasets/utils/splits.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shreyaskamathkm/FTNet/HEAD/Codes/src/datasets/utils/splits.mat -------------------------------------------------------------------------------- /Codes/src/core/models/encoder/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | from .resnext import * 4 | from .resnetv1b import * 5 | -------------------------------------------------------------------------------- /Codes/src/core/data/samplers/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | from .batch_sampler import * 3 | from .common import * 4 | from .multiscale_batch_samplers import * 5 | -------------------------------------------------------------------------------- /Codes/src/core/models/segmentation_decoder/__init__.py: -------------------------------------------------------------------------------- 1 | """Model Zoo""" 2 | # https://github.com/qubvel/segmentation_models.pytorch 3 | from core.models.segmentation_decoder.ftnet import * 4 | -------------------------------------------------------------------------------- /Codes/src/lightning_scripts/run_performace_10_times.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | for (( c=1; c<=10; c++ )) 3 | do 4 | echo "Welcome $c times" 5 | python performance_test.py --model deeplabv3 --backbone resnet50 6 | done 7 | -------------------------------------------------------------------------------- /Codes/src/lightning_scripts/trainers/__init__.py: -------------------------------------------------------------------------------- 1 | from .thermal_edge_trainer import SegmentationLightningModel as thermal_edge_trainer 2 | from .thermal_edge_trainer_train_only import SegmentationTrainLightningModel as thermal_edge_trainer_train_only 3 | -------------------------------------------------------------------------------- /Codes/src/core/utils/__init__.py: -------------------------------------------------------------------------------- 1 | """Utility functions.""" 2 | 3 | from .collect_env import * 4 | from .download import * 5 | from .filesystem import * 6 | from .utils import * 7 | from .logger import * 8 | from .img_utils import * 9 | from .optimizer_scheduler_helper import * 10 | from .json_extension import * 11 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | download 2 | fvcore 3 | lib 4 | matplotlib 5 | numpy 6 | opencv_python 7 | pandas 8 | Pillow 9 | rfconv 10 | scikit_image 11 | scikit_learn 12 | seaborn==0.11.2 13 | skimage 14 | tabulate 15 | termcolor 16 | tqdm 17 | wandb 18 | torch==1.6.0 19 | torchmetrics==0.5.1 20 | torchvision==0.10.0 21 | ptflops==0.6.6 22 | pycocotools==2.0.2 23 | pytorch_lightning==1.3.4 24 | -------------------------------------------------------------------------------- /Codes/src/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .dataloaders import * 2 | from .encode_originals import * 3 | 4 | 5 | datasets = { 6 | 'cityscapes_thermal_combine': CityscapesCombineThermalDataset, 7 | 'cityscapes_thermal_split': CityscapesThermalsSplitDataset, 8 | 'soda': SODADataset, 9 | 'mfn': MFNDataset, 10 | 'scutseg': SCUTSEGDataset} 11 | 12 | 13 | def get_segmentation_dataset(name, **kwargs): 14 | """Segmentation Datasets""" 15 | return datasets[name.lower()](**kwargs) 16 | -------------------------------------------------------------------------------- /Codes/src/core/models/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | try: 3 | import segmentation_decoder as seg_dec 4 | except: 5 | import core.models.segmentation_decoder as seg_dec 6 | from typing import Any 7 | 8 | __all__ = ['get_segmentation_model'] 9 | 10 | 11 | def get_segmentation_model(model_name: str, **kwargs: Any): 12 | 13 | model = seg_dec.__dict__['get_' + model_name.lower()] 14 | return model(**kwargs) 15 | 16 | 17 | if __name__ == '__main__': 18 | model = get_segmentation_model('ftnet') 19 | -------------------------------------------------------------------------------- /Codes/src/datasets/dataloaders/__init__.py: -------------------------------------------------------------------------------- 1 | try: 2 | from .cityscapes_thermal import CityscapesCombineThermalDataset 3 | from .soda import SODADataset 4 | from .mfn import MFNDataset 5 | from .cityscapes_thermal_split import CityscapesThermalsSplitDataset 6 | from .scutseg import SCUTSEGDataset 7 | 8 | 9 | except: 10 | from datasets.dataloaders.cityscapes_thermal import CityscapesCombineThermalDataset 11 | from datasets.dataloaders.soda import SODADataset 12 | from datasets.dataloaders.mfn import MFNDataset 13 | from datasets.dataloaders.cityscapes_thermal_split import CityscapesThermalsSplitDataset 14 | from datasets.dataloaders.scutseg import SCUTSEGDataset 15 | -------------------------------------------------------------------------------- /Codes/src/bash/Paper_runs.sh: -------------------------------------------------------------------------------- 1 | sbatch --job-name='FTN50' --gres=gpu:v100:2 --export=BACKBONE='resnext50_32x4d' RUN_me.sh 2 | sbatch --job-name='FTN101' --gres=gpu:v100:2 --export=BACKBONE='resnext101_32x8d' RUN_me.sh 3 | 4 | sbatch --job-name='RX50_20' --gres=gpu:v100:2 --export=BACKBONE='resnext50_32x4d',ALPHA='20' RUN_me_AlphaLoss.sh 5 | sbatch --job-name='RX50_10' --gres=gpu:v100:2 --export=BACKBONE='resnext50_32x4d',ALPHA='10' RUN_me_AlphaLoss.sh 6 | sbatch --job-name='RX50_30' --gres=gpu:v100:2 --export=BACKBONE='resnext50_32x4d',ALPHA='30' RUN_me_AlphaLoss.sh 7 | sbatch --job-name='RX50_15' --gres=gpu:v100:2 --export=BACKBONE='resnext50_32x4d',ALPHA='15' RUN_me_AlphaLoss.sh 8 | sbatch --job-name='RX50_5' --gres=gpu:v100:2 --export=BACKBONE='resnext50_32x4d',ALPHA='5' RUN_me_AlphaLoss.sh 9 | -------------------------------------------------------------------------------- /Codes/src/core/schedulers/iteration_polyLR.py: -------------------------------------------------------------------------------- 1 | """Popular Learning Rate Schedulers""" 2 | from __future__ import division 3 | 4 | import math 5 | from bisect import bisect_right 6 | 7 | import torch 8 | 9 | __all__ = ['IterationPolyLR'] 10 | 11 | 12 | class IterationPolyLR(torch.optim.lr_scheduler._LRScheduler): 13 | def __init__(self, optimizer, target_lr=0, max_iters=0, power=0.9, last_epoch=-1): 14 | self.target_lr = target_lr 15 | self.max_iters = max_iters 16 | self.power = power 17 | super(IterationPolyLR, self).__init__(optimizer, last_epoch) 18 | 19 | def get_lr(self): 20 | N = self.max_iters 21 | T = self.last_epoch 22 | factor = pow(1 - T / N, self.power) 23 | # https://blog.csdn.net/mieleizhi0522/article/details/83113824 24 | return [self.target_lr + (base_lr - self.target_lr) * factor for base_lr in self.base_lrs] 25 | -------------------------------------------------------------------------------- /Codes/src/core/schedulers/Poly_LR.py: -------------------------------------------------------------------------------- 1 | from torch.optim.lr_scheduler import _LRScheduler 2 | 3 | __all__ = ['Poly'] 4 | 5 | 6 | class Poly(_LRScheduler): 7 | def __init__(self, optimizer, num_epochs, steps_per_epoch=0, warmup_epochs=0, power=0.9, last_epoch=-1): 8 | self.steps_per_epoch = steps_per_epoch 9 | self.cur_iter = 0 10 | self.power = power 11 | self.N = num_epochs * steps_per_epoch 12 | self.warmup_iters = warmup_epochs * steps_per_epoch 13 | super(Poly, self).__init__(optimizer, last_epoch) 14 | 15 | def get_lr(self): 16 | T = self.last_epoch * self.steps_per_epoch + self.cur_iter 17 | factor = pow((1 - 1.0 * T / self.N), self.power) 18 | if self.warmup_iters > 0 and T < self.warmup_iters: 19 | factor = 1.0 * T / self.warmup_iters 20 | 21 | self.cur_iter %= self.steps_per_epoch 22 | self.cur_iter += 1 23 | return [base_lr * factor for base_lr in self.base_lrs] 24 | -------------------------------------------------------------------------------- /Codes/src/datasets/encode_originals.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | from core.utils.visualize import get_color_pallete 4 | import cv2 5 | import matplotlib.pyplot as plt 6 | import numpy as np 7 | 8 | mask_path = '/mnt/ACF29FC2F29F8F68/Work/Deep_Learning/Thermal_Segmentation/Dataset/InfraredSemanticLabel-20210430T150555Z-001/SODA/mask/test/' 9 | mask_color_path = '/mnt/ACF29FC2F29F8F68/Work/Deep_Learning/Thermal_Segmentation/Dataset/InfraredSemanticLabel-20210430T150555Z-001/SODA/mask_color/test/' 10 | 11 | 12 | if not os.path.exists(mask_color_path): 13 | os.makedirs(mask_color_path) 14 | 15 | 16 | def is_image_file(filename): 17 | return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg", ".bmp", ".tif", "tiff"]) 18 | 19 | 20 | masks = {os.path.splitext(os.path.basename(x))[0]: x for x in glob.glob( 21 | os.path.join(mask_path, '*'), recursive=True) if is_image_file(x)} 22 | 23 | 24 | for name, path in masks.items(): 25 | img = cv2.imread(path)[:, :, 0] 26 | temp = get_color_pallete(img, dataset='soda') 27 | plt.imsave(mask_color_path + 28 | '/original_mask_{}.png'.format(name), np.array(temp)) 29 | -------------------------------------------------------------------------------- /Codes/src/lightning_scripts/__init__.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import os 3 | import sys 4 | 5 | # Add script directory to sys.path. 6 | # This is complicated due to the fact that __file__ is not always defined. 7 | 8 | 9 | def GetScriptDirectory(): 10 | if hasattr(GetScriptDirectory, "dir"): 11 | return GetScriptDirectory.dir 12 | module_path = "" 13 | try: 14 | # The easy way. Just use __file__. 15 | # Unfortunately, __file__ is not available when cx_Freeze is used or in IDLE. 16 | module_path = __file__ 17 | except NameError: 18 | if len(sys.argv) > 0 and len(sys.argv[0]) > 0 and os.path.isabs(sys.argv[0]): 19 | module_path = sys.argv[0] 20 | else: 21 | module_path = os.path.abspath(inspect.getfile(GetScriptDirectory)) 22 | if not os.path.exists(module_path): 23 | # If cx_Freeze is used the value of the module_path variable at this point is in the following format. 24 | # {PathToExeFile}\{NameOfPythonSourceFile}. This makes it necessary to strip off the file name to get the correct 25 | # path. 26 | module_path = os.path.dirname(module_path) 27 | GetScriptDirectory.dir = os.path.dirname(module_path) 28 | return GetScriptDirectory.dir 29 | 30 | 31 | abspath = os.path.abspath(os.path.join(GetScriptDirectory(), "..")) 32 | sys.path.append(abspath) 33 | print(GetScriptDirectory()) 34 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.pyc 4 | *.py[cod] 5 | *$py.class 6 | 7 | # C extensions 8 | *.so 9 | 10 | # Distribution / packaging 11 | .Python 12 | env/ 13 | build/ 14 | develop-eggs/ 15 | dist/ 16 | downloads/ 17 | eggs/ 18 | .eggs/ 19 | lib/ 20 | lib64/ 21 | parts/ 22 | sdist/ 23 | var/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | 28 | 29 | # Folders 30 | Codes/Results/ 31 | Codes/src/.spyproject 32 | Codes/Results_/ 33 | Test_Results/ 34 | Codes/src/pretrained_models/ 35 | 36 | 37 | # PyInstaller 38 | # Usually these files are written by a python script from a template 39 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 40 | *.manifest 41 | *.spec 42 | 43 | # Error and output files 44 | *.dat 45 | 46 | # Installer logs 47 | pip-log.txt 48 | pip-delete-this-directory.txt 49 | 50 | # Unit test / coverage reports 51 | htmlcov/ 52 | .tox/ 53 | .coverage 54 | .coverage.* 55 | .cache 56 | nosetests.xml 57 | coverage.xml 58 | *,cover 59 | 60 | # Translations 61 | *.mo 62 | *.pot 63 | 64 | # Django stuff: 65 | *.log 66 | 67 | # Sphinx documentation 68 | docs/_build/ 69 | 70 | # PyBuilder 71 | target/ 72 | 73 | # PyTorch 74 | *.pt 75 | *.pdf 76 | *.png 77 | *.swp 78 | .vscode 79 | 80 | 81 | .spyproject/ 82 | 83 | Codes/src/core/models/crisscross_decoder_dilated.py 84 | 85 | Codes/src/debug/lightning_logs/ 86 | Codes/src/debug/log_dict.py 87 | Codes/src/debug/uuid.py 88 | 89 | 90 | *.lprof 91 | *.pstats 92 | 93 | *.insyncdl 94 | 95 | 96 | -------------------------------------------------------------------------------- /Codes/src/datasets/edge_generation/main_cityscape_visible.m: -------------------------------------------------------------------------------- 1 | clc;clear all;close all; 2 | % , 'soda' 3 | datasets = {'cityscapes'}; 4 | path = 'E:\Work\Deep_Learning\Thermal_Segmentation\Dataset/' 5 | labels_path = {'Cityscapes_thermal\CITYSCAPE_5000/', 'InfraredSemanticLabel-20210430T150555Z-001\SODA/','Cityscapes\gtFine/' }; 6 | 7 | % numWorker = 6; % Number of matlab workers for parallel computing 8 | % delete(gcp('nocreate')); 9 | % parpool('local', numWorker); 10 | for idxSet = 1:length(datasets) 11 | if strcmp(datasets{idxSet}, 'cityscapes_thermal') 12 | setList = {'train'}; 13 | data_path = labels_path{1} 14 | else 15 | setList = {'train', 'val', 'test'}; 16 | data_path = labels_path{3} 17 | end 18 | 19 | for set_idx = 1:length(setList) 20 | save_path = [path data_path 'edges/' setList{set_idx} ]; 21 | if(exist(save_path, 'file')==0) 22 | mkdir(save_path); 23 | end 24 | fileList = dir([path data_path 'mask/' setList{set_idx} ]); 25 | fileList=fileList(~ismember({fileList.name},{'.','..'})); 26 | for idxFile = 1 :length(fileList) 27 | mask = imread([fileList(idxFile).folder '/' fileList(idxFile).name]); 28 | edgeMapBin = seg2edge(mask, 1, []', 'regular'); % Avoid generating edges on "rectification border" (labelId==2) and "out of roi" (labelId==3) 29 | imwrite(edgeMapBin, [save_path '/' fileList(idxFile).name]) 30 | end 31 | end 32 | end 33 | 34 | 35 | 36 | 37 | -------------------------------------------------------------------------------- /Codes/src/datasets/edge_generation/main.m: -------------------------------------------------------------------------------- 1 | clc;clear all;close all; 2 | % , 'soda' 3 | datasets = {'cityscapes','soda','scutseg','mfn'}; 4 | path = 'D:\processed_dataset/'; 5 | labels_path = {'CITYSCAPE_5000/', 'SODA/', 'SCUTSEG/','MFN/', }; 6 | radius = [2,1,1,1]; 7 | % numWorker = 6; % Number of matlab workers for parallel computing 8 | % delete(gcp('nocreate')); 9 | % parpool('local', numWorker); 10 | for idxSet = 1:length(datasets) 11 | if strcmp(datasets{idxSet}, 'cityscapes') 12 | setList = {'train'}; 13 | data_path = labels_path{1}; 14 | r = radius(1); 15 | elseif strcmp(datasets{idxSet}, 'soda') 16 | setList = {'train', 'val', 'test'}; 17 | data_path = labels_path{2}; 18 | r = radius(2); 19 | elseif strcmp(datasets{idxSet}, 'scutseg') 20 | setList = {'train', 'test'}; 21 | data_path = labels_path{3}; 22 | r = radius(3); 23 | elseif strcmp(datasets{idxSet}, 'mfn') 24 | setList = {'train', 'val', 'test'}; 25 | data_path = labels_path{4}; 26 | r = radius(4); 27 | end 28 | 29 | for set_idx = 1:length(setList) 30 | save_path = [path data_path 'edges/' setList{set_idx} ]; 31 | if(exist(save_path, 'file')==0) 32 | mkdir(save_path); 33 | end 34 | fileList = dir([path data_path 'mask/' setList{set_idx} ]); 35 | fileList=fileList(~ismember({fileList.name},{'.','..'})); 36 | for idxFile = 1 :length(fileList) 37 | mask = imread([fileList(idxFile).folder '/' fileList(idxFile).name]); 38 | %SODA=1; Cityscapes = 2; 39 | edgeMapBin = seg2edge(mask, r, []', 'regular'); % Avoid generating edges on "rectification border" (labelId==2) and "out of roi" (labelId==3) 40 | imwrite(edgeMapBin, [save_path '/' fileList(idxFile).name]) 41 | end 42 | end 43 | end 44 | 45 | 46 | 47 | 48 | -------------------------------------------------------------------------------- /Codes/src/core/schedulers/torch_poly_lr_decay.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | ''' 3 | Code adapted from 4 | https://github.com/cmpark0126/pytorch-polynomial-lr-decay 5 | 6 | ''' 7 | 8 | from torch.optim.lr_scheduler import _LRScheduler 9 | 10 | __all__ = ['PolynomialLRDecay'] 11 | 12 | 13 | class PolynomialLRDecay(_LRScheduler): 14 | """Polynomial learning rate decay until step reach to max_decay_step 15 | 16 | Args: 17 | optimizer (Optimizer): Wrapped optimizer. 18 | max_decay_steps: after this step, we stop decreasing learning rate 19 | end_learning_rate: scheduler stoping learning rate decay, value of learning rate must be this value 20 | power: The power of the polynomial. 21 | """ 22 | 23 | def __init__(self, optimizer, max_decay_steps, end_learning_rate=0.0001, power=1.0): 24 | if max_decay_steps <= 1.: 25 | raise ValueError('max_decay_steps should be greater than 1.') 26 | self.max_decay_steps = max_decay_steps 27 | self.end_learning_rate = end_learning_rate 28 | self.power = power 29 | self.last_step = 0 30 | super().__init__(optimizer) 31 | 32 | def get_lr(self): 33 | if self.last_step > self.max_decay_steps: 34 | return [self.end_learning_rate for _ in self.base_lrs] 35 | 36 | return [(base_lr - self.end_learning_rate) * 37 | ((1 - self.last_step / self.max_decay_steps) ** (self.power)) + 38 | self.end_learning_rate for base_lr in self.base_lrs] 39 | 40 | def step(self, step=None): 41 | if step is None: 42 | step = self.last_step + 1 43 | self.last_step = step if step != 0 else 1 44 | if self.last_step <= self.max_decay_steps: 45 | decay_lrs = [(base_lr - self.end_learning_rate) * 46 | ((1 - self.last_step / self.max_decay_steps) ** (self.power)) + 47 | self.end_learning_rate for base_lr in self.base_lrs] 48 | for param_group, lr in zip(self.optimizer.param_groups, decay_lrs): 49 | param_group['lr'] = lr 50 | -------------------------------------------------------------------------------- /Codes/src/core/utils/json_extension.py: -------------------------------------------------------------------------------- 1 | import types 2 | import argparse 3 | import json 4 | from .utils import to_python_float 5 | # ============================================================================= 6 | # https://discuss.pytorch.org/t/typeerror-tensor-is-not-json-serializable/36065/3 7 | # ============================================================================= 8 | 9 | 10 | def _to_json_dict_with_strings(dictionary): 11 | """ 12 | Convert dict to dict with leafs only being strings. So it recursively makes keys to strings 13 | if they are not dictionaries. 14 | 15 | Use case: 16 | - saving dictionary of tensors (convert the tensors to strins!) 17 | - saving arguments from script (e.g. argparse) for it to be pretty 18 | 19 | e.g. 20 | 21 | """ 22 | if type(dictionary) != dict: 23 | return str(to_python_float(dictionary)) 24 | d = {k: _to_json_dict_with_strings(v) for k, v in dictionary.items()} 25 | return d 26 | 27 | 28 | def to_json(dic): 29 | if type(dic) is dict: 30 | dic = dict(dic) 31 | else: 32 | dic = dic.__dict__ 33 | return _to_json_dict_with_strings(dic) 34 | 35 | 36 | def save_to_json_pretty(dic, path, mode='w', indent=4, sort_keys=True): 37 | 38 | with open(path, mode) as f: 39 | json.dump(to_json(dic), f, indent=indent, sort_keys=sort_keys) 40 | 41 | 42 | def my_pprint(dic): 43 | """ 44 | 45 | @param dic: 46 | @return: 47 | 48 | Note: this is not the same as pprint. 49 | """ 50 | # make all keys strings recursively with their naitve str function 51 | dic = to_json(dic) 52 | # pretty print 53 | pretty_dic = json.dumps(dic, indent=4, sort_keys=True) 54 | print(pretty_dic) 55 | # print(json.dumps(dic, indent=4, sort_keys=True)) 56 | # return pretty_dic 57 | 58 | 59 | if __name__ == '__main__': 60 | import torch 61 | # import json # results in non serializabe errors for torch.Tensors 62 | from pprint import pprint 63 | 64 | dic = {'x': torch.randn(1, 3), 'rec': {'y': torch.randn(1, 3)}} 65 | 66 | my_pprint(dic) 67 | pprint(dic) 68 | -------------------------------------------------------------------------------- /Codes/src/core/utils/tabulate_metrics.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # To distribute train, valid, test for SODA 3 | # ============================================================================= 4 | 5 | import os 6 | import pandas as pd 7 | from sklearn.model_selection import train_test_split 8 | from glob import glob 9 | import shutil 10 | import argparse 11 | 12 | 13 | def is_txt_file(filename): 14 | return any(filename.endswith(extension) for extension in [".txt"]) 15 | 16 | 17 | def str2bool(v): 18 | if v.lower() in ('yes', 'true', 't', 'y', '1'): 19 | return True 20 | elif v.lower() in ('no', 'false', 'f', 'n', '0'): 21 | return False 22 | else: 23 | raise argparse.ArgumentTypeError('Boolean value expected.') 24 | 25 | 26 | # %% 27 | # ============================================================================= 28 | # Copying files to different folders 29 | # ============================================================================= 30 | def copying(tiles, path_label, basepath, fileset_path): 31 | tiles.set_index(path_label, inplace=True) 32 | for img_path in tiles.index: 33 | print('Path = {}'.format(img_path)) 34 | dst_path = os.path.join(basepath, fileset_path) 35 | 36 | shutil.copy(img_path, dst_path) 37 | 38 | # %% 39 | # ============================================================================= 40 | # Creating folders 41 | # ============================================================================= 42 | 43 | 44 | def tabulate_metrics(directory): 45 | subfolders = [f.path for f in os.scandir(directory) if f.is_dir()] 46 | for x in glob(os.path.join(subfolders, 'soda/Segmented images/test/', 'Average.txt')): 47 | with open(x) as f: 48 | lines = f.readlines() 49 | 50 | 51 | if __name__ == '__main__': 52 | import argparse 53 | parser = argparse.ArgumentParser() 54 | parser.add_argument('directory', 55 | type=str, 56 | help='the path to list') 57 | args = parser.parse_args() 58 | tabulate_metrics(args.directory) 59 | -------------------------------------------------------------------------------- /Codes/src/datasets/calculate_mean_std.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import glob 3 | import os 4 | import time 5 | from multiprocessing.dummy import Pool as ThreadPool 6 | import json 7 | import cv2 8 | import numpy as np 9 | import tqdm 10 | 11 | from skimage.exposure import equalize_adapthist 12 | 13 | from PIL import Image 14 | from PIL import ImageFilter 15 | from torchvision import transforms 16 | 17 | 18 | def im2double(im): 19 | info = np.iinfo(im.dtype) # Get the data type of the input image 20 | return (im.astype(np.float) / info.max) # Divide all values by the largest possible value in the datatype 21 | 22 | 23 | parser = argparse.ArgumentParser(description="Finding Mean and Std Deviation") 24 | parser.add_argument("--test_folder", 25 | type=str, 26 | default="./Dataset/SCUT/SCUTSEG/image/train/", 27 | help="Dataset for checking the model") 28 | args = parser.parse_args() 29 | args.test_folder = args.test_folder.replace("\\", "/") 30 | 31 | 32 | def is_image_file(filename): 33 | return any( 34 | filename.endswith(extension) 35 | for extension in [".png", ".jpg", ".jpeg", ".bmp", ".tif", "tiff"] 36 | ) 37 | 38 | 39 | image_list = [x 40 | for x in glob.glob(os.path.join(args.test_folder, "*"), recursive=True) 41 | if is_image_file(x) 42 | ] 43 | 44 | 45 | def caluate_MSTD(paths): 46 | img = cv2.imread(paths, 1) 47 | img = im2double(img) 48 | mean = np.mean(img, axis=(0, 1)) 49 | std = np.std(img, axis=(0, 1)) 50 | return mean, std 51 | 52 | 53 | start = time.time() 54 | with ThreadPool(6) as p: 55 | patch_grids = list( 56 | tqdm.tqdm(p.imap(caluate_MSTD, image_list), total=len(image_list)) 57 | ) 58 | 59 | all_img_means = [x[0] for x in patch_grids] 60 | all_img_std = [x[1] for x in patch_grids] 61 | 62 | mean_per_channel = np.mean(np.array(all_img_means), axis=0) 63 | std_per_channel = np.mean(np.array(all_img_std), axis=0) 64 | 65 | 66 | print("Total Time = {}".format(time.time() - start)) 67 | 68 | print(mean_per_channel) 69 | print(std_per_channel) 70 | 71 | np.savez("./mfn", mean=mean_per_channel, std=std_per_channel) 72 | -------------------------------------------------------------------------------- /Codes/src/core/utils/mean_std.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # https://news.ycombinator.com/item?id=16609033 3 | # ============================================================================= 4 | from __future__ import absolute_import, division, print_function 5 | 6 | import os 7 | 8 | cur_path = os.path.abspath(os.path.dirname("__file__")) 9 | root_path = os.path.split(cur_path)[0] 10 | import sys 11 | 12 | sys.path.append(root_path) 13 | import random 14 | 15 | import lib.data.dataloader.qtransforms as qtrans 16 | import numpy 17 | import torch 18 | import torch.optim 19 | import torchvision.transforms as transforms 20 | from lib.data.dataloader import get_classification_dataset 21 | from lib.data.dataloader.prefetch_generator.dataloader_v2 import \ 22 | DataLoaderX as DataLoader 23 | from lib.utils.utils import plot_tensors 24 | from options import parse_args 25 | 26 | args = parse_args() 27 | 28 | # ============================================================================= 29 | # Setting up 30 | # ============================================================================= 31 | # Data loading code 32 | traindir = os.path.join(args.dataset_path, 'train') 33 | valdir = os.path.join(args.dataset_path, 'valid') 34 | 35 | # ============================================================================= 36 | # 0.2126 * R + 0.7152 * G + 0.0722 * B 37 | # ============================================================================= 38 | 39 | train_dataset = get_classification_dataset(name=args.dataset, 40 | root=traindir, 41 | train=True, 42 | download=True, 43 | transform=qtrans.ToQTensor()) 44 | train_loader = DataLoader(train_dataset, 45 | batch_size=64, 46 | shuffle=False, 47 | num_workers=args.workers, 48 | pin_memory=True) 49 | 50 | mean, std = [], [] 51 | for i, (input, target) in enumerate(train_loader): 52 | mean.append(input.mean(dim=[0, 2, 3])) 53 | mean.append(input.std(dim=[0, 2, 3])) 54 | -------------------------------------------------------------------------------- /Codes/src/core/data/samplers/common.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Iterator 2 | 3 | import torch 4 | import torch.utils.data as data 5 | from core.data.samplers.batch_sampler import IterationBasedBatchSampler 6 | from core.data.samplers.multiscale_batch_samplers import ( 7 | IterationBasedMultiscaleBatchSampler, MultiscaleBatchSampler) 8 | 9 | __all__ = ['make_data_sampler', 'make_batch_data_sampler', 'make_multiscale_batch_data_sampler'] 10 | 11 | 12 | def make_data_sampler(dataset: Callable, shuffle: bool, distributed: bool) -> Iterator: 13 | if distributed: 14 | return torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle) 15 | if shuffle: 16 | sampler = data.sampler.RandomSampler(dataset) 17 | else: 18 | sampler = data.sampler.SequentialSampler(dataset) 19 | return sampler 20 | 21 | 22 | def make_batch_data_sampler(sampler: Iterator, 23 | batch_size: int, 24 | num_iters: int = None, 25 | start_iter: int = 0) -> Iterator: 26 | batch_sampler = data.sampler.BatchSampler(sampler, batch_size, drop_last=True) 27 | if num_iters is not None: 28 | batch_sampler = IterationBasedBatchSampler(batch_sampler, num_iters, start_iter) 29 | return batch_sampler 30 | 31 | 32 | def make_multiscale_batch_data_sampler(sampler: Iterator, 33 | batch_size: int = 2, 34 | multiscale_step: int = 1, 35 | scales: int = 1, 36 | num_iters: int = None, 37 | start_iter: int = 0) -> Iterator: 38 | 39 | batch_sampler = MultiscaleBatchSampler(sampler=sampler, 40 | batch_size=batch_size, 41 | drop_last=True, 42 | multiscale_step=multiscale_step, 43 | scales=scales) 44 | if num_iters is not None: 45 | batch_sampler = IterationBasedMultiscaleBatchSampler(batch_sampler, num_iters, start_iter) 46 | return batch_sampler 47 | 48 | 49 | if __name__ == '__main__': 50 | pass 51 | -------------------------------------------------------------------------------- /Codes/src/datasets/edge_generation/seg2edge.m: -------------------------------------------------------------------------------- 1 | % This function takes an input segment and produces binary bdrys. 2 | % Multi-channel input segments are supported by the function. 3 | %%% 4 | %Code Adapted from 5 | %https://github.com/Lavender105/DFF/blob/152397cec4a3dac2aa86e92a65cc27e6c8016ab9/lib/matlab/modules/data/seg2edge.m 6 | %%% 7 | function [idxEdge] = seg2edge(seg, radius, labelIgnore, edge_type) 8 | % Get dimensions 9 | [height, width, chn] = size(seg); 10 | if(~isempty(labelIgnore)) 11 | if(chn~=size(labelIgnore, 2)) 12 | error('Channel dimension not matching ignored label dimension!') 13 | end 14 | end 15 | 16 | % Set the considered neighborhood 17 | radius_search = max(ceil(radius), 1); 18 | [X, Y] = meshgrid(1:width, 1:height); 19 | [x, y] = meshgrid(-radius_search:radius_search, -radius_search:radius_search); 20 | 21 | % Columnize everything 22 | X = X(:); Y = Y(:); 23 | x = x(:); y = y(:); 24 | if(chn == 1) 25 | seg = seg(:); 26 | else 27 | seg = reshape(seg, [height*width chn]); 28 | end 29 | 30 | % Build circular neighborhood 31 | idxNeigh = sqrt(x.^2 + y.^2) <= radius; 32 | x = x(idxNeigh); y = y(idxNeigh); 33 | numPxlImg = length(X); 34 | numPxlNeigh = length(x); 35 | 36 | % Compute Gaussian weight 37 | idxEdge = false(numPxlImg, 1); 38 | for i = 1:numPxlNeigh 39 | XNeigh = X+x(i); 40 | YNeigh = Y+y(i); 41 | idxValid = find( XNeigh >= 1 & XNeigh <= width & YNeigh >=1 & YNeigh <= height ); 42 | 43 | XCenter = X(idxValid); 44 | YCenter = Y(idxValid); 45 | XNeigh = XNeigh(idxValid); 46 | YNeigh = YNeigh(idxValid); 47 | LCenter = seg(sub2ind([height width], YCenter, XCenter), :); 48 | LNeigh = seg(sub2ind([height width], YNeigh, XNeigh), :); 49 | 50 | if(strcmp(edge_type, 'regular')) 51 | idxDiff = find(any(LCenter~=LNeigh, 2)); 52 | elseif(strcmp(edge_type, 'inner')) 53 | idxDiff = find(any(LCenter~=LNeigh, 2) & any(LCenter~=0, 2) & all(LNeigh==0, 2) ); 54 | elseif(strcmp(edge_type, 'outer')) 55 | idxDiff = find(any(LCenter~=LNeigh, 2) & all(LCenter==0, 2) & any(LNeigh~=0, 2) ); 56 | else 57 | error('Wrong edge type input!'); 58 | end 59 | 60 | LCenterEdge = LCenter(idxDiff, :); 61 | LNeighEdge = LNeigh(idxDiff, :); 62 | idxIgnore2 = false(length(idxDiff), 1); 63 | for j = 1:size(labelIgnore, 1) 64 | idxIgnore2 = idxIgnore2 | ( all(bsxfun(@eq, LCenterEdge, labelIgnore(j, :)), 2) | all(bsxfun(@eq, LNeighEdge, labelIgnore(j, :)), 2) ); 65 | end 66 | 67 | idxDiffGT = idxDiff(~idxIgnore2); 68 | idxEdge(idxValid(idxDiffGT)) = true; 69 | end 70 | idxEdge = reshape(idxEdge, [height, width]); 71 | -------------------------------------------------------------------------------- /Codes/src/core/utils/img_utils.py: -------------------------------------------------------------------------------- 1 | from typing import List, Tuple 2 | 3 | import cv2 4 | import numpy as np 5 | import torch 6 | 7 | 8 | def resize_by_scaled_short_side( 9 | image: torch.tensor, 10 | base_size: int, 11 | scale: float, 12 | ) -> torch.tensor: 13 | """ Equivalent to ResizeShort(), but functional, instead of OOP paradigm, and w/ scale param. 14 | 15 | Args: 16 | image: Numpy array of shape () 17 | scale: scaling factor for image 18 | 19 | Returns: 20 | image_scaled: 21 | """ 22 | n, c, h, w = image.shape 23 | short_size = round(scale * base_size) 24 | new_h = short_size 25 | new_w = short_size 26 | # Preserve the aspect ratio 27 | if h > w: 28 | new_h = round(short_size / float(w) * h) 29 | else: 30 | new_w = round(short_size / float(h) * w) 31 | image_scaled = torch.nn.functional.interpolate(image, (new_w, new_h), mode='bilinear') 32 | return image_scaled 33 | 34 | 35 | def pad_to_crop_sz( 36 | image: torch.tensor, 37 | crop_h: int, 38 | crop_w: int, 39 | mean: Tuple[float, float, float] 40 | ) -> Tuple[torch.tensor, int, int]: 41 | """ 42 | Network input should be at least crop size, so we pad using mean values if 43 | provided image is too small. No rescaling is performed here. 44 | We use cv2.copyMakeBorder to copy the source image into the middle of a 45 | destination image. The areas to the left, to the right, above and below the 46 | copied source image will be filled with extrapolated pixels, in this case the 47 | provided mean pixel intensity. 48 | 49 | Args: 50 | image: 51 | crop_h: integer representing crop height 52 | crop_w: integer representing crop width 53 | 54 | Returns: 55 | image: Numpy array of shape (crop_h x crop_w) representing a 56 | square image, with short side of square is at least crop size. 57 | pad_h_half: half the number of pixels used as padding along height dim 58 | pad_w_half" half the number of pixels used as padding along width dim 59 | """ 60 | n, ch, orig_h, orig_w = image.shape 61 | 62 | pad_h = max(crop_h - orig_h, 0) 63 | pad_w = max(crop_w - orig_w, 0) 64 | pad_h_half = int(pad_h / 2) 65 | pad_w_half = int(pad_w / 2) 66 | if pad_h > 0 or pad_w > 0: 67 | temp = [] 68 | for c in range(ch): 69 | temp.append(torch.nn.functional.pad(input=image[:, c, :, :], 70 | pad=(pad_w_half, pad_w - pad_w_half, pad_h_half, pad_h - pad_h_half), 71 | mode='constant', 72 | value=mean[c])) # (padding_left, padding_right, \text{padding\_top}, \text{padding\_bottom})padding_top, padding_bottom) 73 | image = torch.stack(temp, dim=1) 74 | 75 | return image, pad_h_half, pad_w_half 76 | -------------------------------------------------------------------------------- /Codes/src/lightning_scripts/performance_test.py: -------------------------------------------------------------------------------- 1 | # https://deci.ai/the-correct-way-to-measure-inference-time-of-deep-neural-networks/ 2 | import os 3 | import math 4 | import torch 5 | import cv2 6 | import glob 7 | from lightning_scripts.options import parse_args 8 | import time 9 | import numpy as np 10 | from core.models import get_segmentation_model 11 | import torch.nn as nn 12 | from core.utils import collect_env_info, get_rank, setup_logger 13 | BatchNorm2d = nn.BatchNorm2d 14 | 15 | 16 | class ImageTester(): 17 | def __init__(self, hparams): 18 | self.hparams = hparams 19 | self.device = torch.device('cuda') 20 | self.logger = setup_logger(name="pytorch_lightning", 21 | save_dir=None, 22 | distributed_rank=get_rank(), 23 | color=True, 24 | abbrev_name=None, 25 | ) 26 | 27 | self.model = get_segmentation_model(model_name=self.hparams.model, 28 | dataset=self.hparams.dataset, 29 | backbone=self.hparams.backbone, 30 | norm_layer=BatchNorm2d, 31 | no_of_filters=self.hparams.no_of_filters, 32 | pretrained_base=self.hparams.pretrained_base, 33 | edge_extracts=self.hparams.edge_extracts, 34 | num_blocks=self.hparams.num_blocks) 35 | self.model.to(self.device) 36 | self.dummy_input = torch.randn(1, 3, 512, 512, dtype=torch.float).to(self.device) 37 | 38 | def test(self): 39 | torch.set_grad_enabled(False) 40 | 41 | self.logger.info('\nEvaluation on Images:') 42 | self.model.eval() 43 | 44 | # INIT LOGGERS 45 | starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) 46 | repetitions = 300 47 | timings = np.zeros((repetitions, 1)) 48 | # GPU-WARM-UP 49 | for _ in range(10): 50 | _ = self.model(self.dummy_input) 51 | # MEASURE PERFORMANCE 52 | with torch.no_grad(): 53 | for rep in range(repetitions): 54 | starter.record() 55 | _ = self.model(self.dummy_input) 56 | ender.record() 57 | # WAIT FOR GPU SYNC 58 | torch.cuda.synchronize() 59 | curr_time = starter.elapsed_time(ender) 60 | timings[rep] = curr_time 61 | mean_syn = np.sum(timings) / repetitions 62 | std_syn = np.std(timings) 63 | print(mean_syn) 64 | 65 | def prepare(self, *args): 66 | def _prepare(tensor): 67 | if self.hparams.precision == 'half': 68 | tensor = tensor.half() 69 | return tensor.to(self.device) 70 | return [_prepare(a) for a in args] 71 | 72 | 73 | if __name__ == '__main__': 74 | 75 | hparams = parse_args() 76 | test_model = ImageTester(hparams) 77 | test_model.test() 78 | -------------------------------------------------------------------------------- /Codes/src/core/loss/segmentation_loss.py: -------------------------------------------------------------------------------- 1 | ''' 2 | MixSoftmaxCrossEntropyLoss adapted from 3 | https://github.com/Tramac/awesome-semantic-segmentation-pytorch/blob/master/core/utils/loss.py 4 | ''' 5 | from typing import Any, Callable 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | __all__ = ['MixSoftmaxCrossEntropyLoss', 11 | 'EdgeNetLoss', 12 | 'get_segmentation_loss'] 13 | 14 | 15 | class MixSoftmaxCrossEntropyLoss(nn.CrossEntropyLoss): 16 | def __init__(self, aux: bool = False, 17 | aux_weight: float = 0.2, 18 | ignore_index: int = -1, 19 | **kwargs: Any) -> None: 20 | super(MixSoftmaxCrossEntropyLoss, self).__init__(ignore_index=ignore_index) 21 | self.aux = aux 22 | self.aux_weight = aux_weight 23 | 24 | def _aux_forward(self, 25 | *inputs: torch.Tensor, 26 | **kwargs: Any): 27 | *preds, target = tuple(inputs) 28 | 29 | loss = super(MixSoftmaxCrossEntropyLoss, self).forward(preds[0], target) 30 | for i in range(1, len(preds)): 31 | aux_loss = super(MixSoftmaxCrossEntropyLoss, self).forward(preds[i], target) 32 | loss += self.aux_weight * aux_loss 33 | return loss 34 | 35 | def forward(self, *inputs: torch.Tensor, **kwargs: Any) -> torch.Tensor: 36 | if isinstance(inputs[0], tuple): 37 | preds, target = tuple(inputs) 38 | inputs = tuple(list(preds) + [target]) 39 | 40 | if self.aux: 41 | return self._aux_forward(*inputs) 42 | else: 43 | return super(MixSoftmaxCrossEntropyLoss, self).forward(*inputs) 44 | 45 | 46 | class EdgeNetLoss(nn.Module): 47 | # https://github.com/gasparian/PicsArtHack-binary-segmentation/blob/ecab001f334949d5082a79b8fbd1dc2fdb8b093e/utils.py#L217 48 | def __init__(self, bce_weight: float = None, loss_weight: int = 1, **kwargs: Any) -> None: 49 | super(EdgeNetLoss, self).__init__() 50 | self.cross_entropy = MixSoftmaxCrossEntropyLoss(**kwargs) 51 | self.loss_weight = loss_weight 52 | 53 | def auto_weight_bce(self, y_hat_log, y): 54 | if y.ndim == 3: 55 | y = y.unsqueeze(1) 56 | with torch.no_grad(): 57 | beta = y.mean(dim=[2, 3], keepdims=True) 58 | logit_1 = F.logsigmoid(y_hat_log) 59 | logit_0 = F.logsigmoid(-y_hat_log) 60 | loss = -(1 - beta) * logit_1 * y - beta * logit_0 * (1 - y) 61 | return loss.mean() 62 | 63 | def forward(self, *inputs: torch.Tensor, **kwargs: Any) -> torch.Tensor: 64 | if isinstance(inputs[0], tuple): 65 | pred_class_map, pred_edge_map = inputs[0] 66 | target_maps, target_edges = inputs[1] 67 | 68 | loss1 = self.loss_weight * (self.auto_weight_bce(pred_edge_map, target_edges.float())) 69 | loss2 = self.cross_entropy(pred_class_map, target_maps) 70 | print("loss_weight: {} loss1: {} loss2: {} ".format(self.loss_weight, loss1, loss2)) 71 | return loss1 + loss2 72 | 73 | 74 | def get_segmentation_loss(model, **kwargs): 75 | model = model.lower() 76 | if 'ftnet' in model: 77 | return EdgeNetLoss(**kwargs) 78 | else: 79 | return MixSoftmaxCrossEntropyLoss(**kwargs) 80 | -------------------------------------------------------------------------------- /Codes/src/core/data/samplers/batch_sampler.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | ''' 3 | Code Adapted from: 4 | https://github.com/CaoWGG/multi-scale-training 5 | 6 | ''' 7 | import numpy as np 8 | from torch.utils.data import RandomSampler, Sampler, SequentialSampler 9 | 10 | __all__ = ['BatchSampler', 'IterationBasedBatchSampler'] 11 | 12 | 13 | class BatchSampler(object): 14 | def __init__(self, sampler, batch_size, drop_last, multiscale_step=None, img_sizes=None): 15 | if not isinstance(sampler, Sampler): 16 | raise ValueError("sampler should be an instance of " 17 | "torch.utils.data.Sampler, but got sampler={}" 18 | .format(sampler)) 19 | if not isinstance(drop_last, bool): 20 | raise ValueError("drop_last should be a boolean value, but got " 21 | "drop_last={}".format(drop_last)) 22 | self.sampler = sampler 23 | self.batch_size = batch_size 24 | self.drop_last = drop_last 25 | if multiscale_step is not None and multiscale_step < 1: 26 | raise ValueError("multiscale_step should be > 0, but got " 27 | "multiscale_step={}".format(multiscale_step)) 28 | if multiscale_step is not None and img_sizes is None: 29 | raise ValueError("img_sizes must a list, but got img_sizes={} ".format(img_sizes)) 30 | 31 | self.multiscale_step = multiscale_step 32 | self.img_sizes = img_sizes 33 | 34 | def __iter__(self): 35 | num_batch = 0 36 | batch = [] 37 | size = np.random.choice(self.img_sizes) 38 | for idx in self.sampler: 39 | batch.append([idx, size]) 40 | if len(batch) == self.batch_size: 41 | yield batch 42 | num_batch += 1 43 | batch = [] 44 | if self.multiscale_step and num_batch % self.multiscale_step == 0: 45 | size = np.random.choice(self.img_sizes) 46 | 47 | if len(batch) > 0 and not self.drop_last: 48 | yield batch 49 | 50 | def __len__(self): 51 | if self.drop_last: 52 | return len(self.sampler) // self.batch_size 53 | else: 54 | return (len(self.sampler) + self.batch_size - 1) // self.batch_size 55 | 56 | 57 | class IterationBasedBatchSampler(BatchSampler): 58 | """ 59 | Wraps a BatchSampler, resampling from it until 60 | a specified number of iterations have been sampled 61 | """ 62 | 63 | def __init__(self, batch_sampler, num_iterations, start_iter=0): 64 | self.batch_sampler = batch_sampler 65 | self.num_iterations = num_iterations 66 | self.start_iter = start_iter 67 | 68 | def __iter__(self): 69 | iteration = self.start_iter 70 | while iteration <= self.num_iterations: 71 | # if the underlying sampler has a set_epoch method, like 72 | # DistributedSampler, used for making each process see 73 | # a different split of the dataset, then set it 74 | if hasattr(self.batch_sampler.sampler, "set_epoch"): 75 | self.batch_sampler.sampler.set_epoch(iteration) 76 | for batch in self.batch_sampler: 77 | iteration += 1 78 | if iteration > self.num_iterations: 79 | break 80 | yield batch 81 | 82 | def __len__(self): 83 | return self.num_iterations 84 | -------------------------------------------------------------------------------- /Codes/src/bash/Train_and_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | MODEL=ftnet 3 | BACKBONE=resnext50_32x4d 4 | GPUS=2 5 | FILTERS=128 6 | EDGES='3' 7 | NBLOCKS=2 8 | NODES=1 9 | ALPHA=20 10 | 11 | DATASET=cityscapes_thermal_combine 12 | RUNDIR=./../../../Training_Paper/Lightning/"$MODEL"_"$BACKBONE"_"$FILTERS"_"$EDGES"_"$NBLOCKS"/"$DATASET"/ 13 | mkdir -p "$RUNDIR" 14 | 15 | ## Pretraining 16 | python ./../lightning_scripts/main.py \ 17 | --mode 'train' \ 18 | --no-of-filters "$FILTERS" \ 19 | --edge-extracts "$EDGES" \ 20 | --loss-weight "$ALPHA" \ 21 | --num-blocks "$NBLOCKS" \ 22 | --train_only True \ 23 | --model "$MODEL" \ 24 | --pretrained-base True \ 25 | --backbone "$BACKBONE" \ 26 | --dataset "$DATASET" \ 27 | --dataset-path './../../../../Thermal_Segmentation/Dataset/' \ 28 | --base-size '520' \ 29 | --crop-size '480' \ 30 | --train-batch-size 16 \ 31 | --val-batch-size 16 \ 32 | --epochs 100 \ 33 | --optimizer 'SGD' \ 34 | --lr 0.01 \ 35 | --scheduler-type 'poly_warmstartup' \ 36 | --warmup-iters 0 \ 37 | --warmup-factor 0.3333 \ 38 | --warmup-method 'linear' \ 39 | --save-images-as-subplots False \ 40 | --save-images False \ 41 | --debug False \ 42 | --workers 16 \ 43 | --momentum 0.9 \ 44 | --weight-decay 0.0001 \ 45 | --beta1 0.9 \ 46 | --beta2 0.999 \ 47 | --epsilon 1e-8 \ 48 | --seed 0 \ 49 | --gpus $GPUS \ 50 | --num-nodes $NODES \ 51 | --distributed-backend 'ddp' \ 52 | --wandb-name-ext '' \ 53 | --save-dir "$RUNDIR" >"$RUNDIR"/log_train_cityscape.txt 54 | 55 | 56 | ## FineTuning 57 | DATASET2='soda' 58 | 59 | 60 | RUNDIR2=./../../../Training_Paper/Lightning/"$MODEL"_"$BACKBONE"_"$FILTERS"_"$EDGES"_"$NBLOCKS"/"$DATASET2"/ 61 | mkdir -p "$RUNDIR2" 62 | 63 | python ./../lightning_scripts/main.py \ 64 | --mode 'train' \ 65 | --train_only False \ 66 | --no-of-filters "$FILTERS" \ 67 | --num-blocks "$NBLOCKS" \ 68 | --edge-extracts "$EDGES" \ 69 | --loss-weight "$ALPHA" \ 70 | --model "$MODEL" \ 71 | --pretrain-checkpoint "$RUNDIR"/ckpt/last.ckpt \ 72 | --pretrained-base False \ 73 | --backbone "$BACKBONE" \ 74 | --dataset "$DATASET2" \ 75 | --dataset-path './../../../../Thermal_Segmentation/Dataset/' \ 76 | --base-size '520' \ 77 | --crop-size '480' \ 78 | --train-batch-size 16 \ 79 | --val-batch-size 16 \ 80 | --test-batch-size 1 \ 81 | --epochs 100 \ 82 | --optimizer 'SGD' \ 83 | --lr 0.001 \ 84 | --scheduler-type 'poly_warmstartup' \ 85 | --warmup-iters 0 \ 86 | --warmup-factor 0.3333 \ 87 | --warmup-method 'linear' \ 88 | --save-images-as-subplots False \ 89 | --save-images False \ 90 | --debug False \ 91 | --workers 16 \ 92 | --momentum 0.9 \ 93 | --weight-decay 0.0001 \ 94 | --beta1 0.9 \ 95 | --beta2 0.999 \ 96 | --epsilon 1e-8 \ 97 | --seed 0 \ 98 | --gpus $GPUS \ 99 | --num-nodes $NODES \ 100 | --distributed-backend 'ddp' \ 101 | --wandb-name-ext '' \ 102 | --save-dir "$RUNDIR2" 103 | 104 | 105 | ## Testing 106 | python ./../lightning_scripts/main.py \ 107 | --mode 'test' \ 108 | --model "$MODEL" \ 109 | --edge-extracts "$EDGES" \ 110 | --loss-weight "$ALPHA" \ 111 | --num-blocks "$NBLOCKS" \ 112 | --backbone "$BACKBONE" \ 113 | --no-of-filters "$FILTERS" \ 114 | --test-monitor 'val_mIOU' \ 115 | --test-monitor-path "$RUNDIR2"/ckpt/ \ 116 | --pretrained-base False \ 117 | --dataset "$DATASET2" \ 118 | --dataset-path './../../../../Thermal_Segmentation/Dataset/' \ 119 | --test-batch-size 1 \ 120 | --save-images True \ 121 | --save-images-as-subplots False \ 122 | --debug False \ 123 | --workers 16 \ 124 | --seed 0 \ 125 | --gpus 1 \ 126 | --num-nodes 1 \ 127 | --distributed-backend 'dp' \ 128 | --wandb-name-ext '' \ 129 | --save-dir "$RUNDIR2"/Best_MIOU/ 130 | -------------------------------------------------------------------------------- /Codes/src/core/models/segbase.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | ''' 3 | Adapted from 4 | https://github.com/Tramac/awesome-semantic-segmentation-pytorch 5 | ''' 6 | 7 | import logging 8 | import sys 9 | from collections import OrderedDict 10 | import torch 11 | import core.models.encoder as enc 12 | import torch.nn as nn 13 | import numpy as np 14 | 15 | logger = logging.getLogger('pytorch_lightning') 16 | 17 | __all__ = ['SegBaseModel', 'initialize_weights', 'print_network'] 18 | 19 | 20 | def str_to_class(classname): 21 | return getattr(sys.modules[__name__], classname) 22 | 23 | 24 | class SegBaseModel(nn.Module): 25 | def __init__(self, nclass, backbone='ResNet50', pretrained_base=True, dilated=None, **kwargs): 26 | super(SegBaseModel, self).__init__() 27 | 28 | if dilated is None: 29 | dilated = True 30 | 31 | self.nclass = nclass 32 | model_name = backbone.lower() 33 | model = enc.__dict__[model_name.lower()] 34 | self.encoder = model(pretrained=pretrained_base, dilated=dilated, **kwargs) 35 | 36 | def base_forward(self, x, multiscale=False): 37 | """forwarding pre-trained network""" 38 | x = self.encoder.conv1(x) 39 | x = self.encoder.bn1(x) 40 | c0 = self.encoder.relu(x) 41 | x = self.encoder.maxpool(c0) 42 | c1 = self.encoder.layer1(x) 43 | c2 = self.encoder.layer2(c1) 44 | c3 = self.encoder.layer3(c2) 45 | c4 = self.encoder.layer4(c3) 46 | 47 | if multiscale: 48 | return [c0, c1, c2, c3, c4] 49 | else: 50 | return [c1, c2, c3, c4] 51 | 52 | 53 | def get_upsample_filter(size): 54 | """Make a 2D bilinear kernel suitable for upsampling""" 55 | factor = (size + 1) // 2 56 | if size % 2 == 1: 57 | center = factor - 1 58 | else: 59 | center = factor - 0.5 60 | og = np.ogrid[:size, :size] 61 | filter = (1 - abs(og[0] - center) / factor) * \ 62 | (1 - abs(og[1] - center) / factor) 63 | return torch.from_numpy(filter).float() 64 | 65 | 66 | def initialize_weights(module): 67 | if isinstance(module, nn.Conv2d): 68 | nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu') 69 | 70 | elif isinstance(module, nn.BatchNorm2d): 71 | nn.init.constant_(module.weight, 1) 72 | nn.init.constant_(module.bias, 0) 73 | 74 | elif isinstance(module, nn.ConvTranspose2d): 75 | c1, c2, h, w = module.weight.data.size() 76 | weight = get_upsample_filter(h) 77 | module.weight.data = weight.view(1, 1, h, w).repeat(c1, c2, 1, 1) 78 | if module.bias is not None: 79 | module.bias.data.zero_() 80 | 81 | 82 | def check_mismatch(model, pretrained_dict) -> None: 83 | model_dict = model.state_dict() 84 | temp_dict = OrderedDict() 85 | for k, v in pretrained_dict.items(): 86 | if k in model_dict.keys(): 87 | if model_dict[k].shape != pretrained_dict[k].shape: 88 | logger.info(f"Skip loading parameter: {k}, " 89 | f"required shape: {model_dict[k].shape}, " 90 | f"loaded shape: {pretrained_dict[k].shape}") 91 | continue 92 | else: 93 | temp_dict[k] = v 94 | return temp_dict 95 | 96 | 97 | def print_network(net): 98 | num_params = 0 99 | for param in net.parameters(): 100 | num_params += param.numel() 101 | print(net) 102 | print('Total number of parameters: {} M'.format(int(num_params) / 1000**2)) 103 | -------------------------------------------------------------------------------- /Codes/src/bash/Train_and_test_all.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | MODEL=ftnet 3 | BACKBONE=resnext50_32x4d 4 | GPUS=2 5 | FILTERS=128 6 | EDGES='3' 7 | NBLOCKS=2 8 | NODES=1 9 | ALPHA=20 10 | 11 | DATASET=cityscapes_thermal_combine 12 | RUNDIR=./../../../Training_Paper/Lightning/"$MODEL"_"$BACKBONE"_"$FILTERS"_"$EDGES"_"$NBLOCKS"/"$DATASET"/ 13 | mkdir -p "$RUNDIR" 14 | 15 | ## Pretraining 16 | python ./../lightning_scripts/main.py \ 17 | --mode 'train' \ 18 | --no-of-filters "$FILTERS" \ 19 | --edge-extracts "$EDGES" \ 20 | --loss-weight "$ALPHA" \ 21 | --num-blocks "$NBLOCKS" \ 22 | --train_only True \ 23 | --model "$MODEL" \ 24 | --pretrained-base True \ 25 | --backbone "$BACKBONE" \ 26 | --dataset "$DATASET" \ 27 | --dataset-path './../../../../Thermal_Segmentation/Dataset/' \ 28 | --base-size '520' \ 29 | --crop-size '480' \ 30 | --train-batch-size 16 \ 31 | --val-batch-size 16 \ 32 | --epochs 100 \ 33 | --optimizer 'SGD' \ 34 | --lr 0.01 \ 35 | --scheduler-type 'poly_warmstartup' \ 36 | --warmup-iters 0 \ 37 | --warmup-factor 0.3333 \ 38 | --warmup-method 'linear' \ 39 | --save-images-as-subplots False \ 40 | --save-images False \ 41 | --debug False \ 42 | --workers 16 \ 43 | --momentum 0.9 \ 44 | --weight-decay 0.0001 \ 45 | --beta1 0.9 \ 46 | --beta2 0.999 \ 47 | --epsilon 1e-8 \ 48 | --seed 0 \ 49 | --gpus $GPUS \ 50 | --num-nodes $NODES \ 51 | --distributed-backend 'ddp' \ 52 | --wandb-name-ext '' \ 53 | --save-dir "$RUNDIR" >"$RUNDIR"/log_train_cityscape.txt 54 | 55 | 56 | ## FineTuning 57 | for DATASET2 in 'soda' 'mfn' 'scutseg' 58 | do 59 | 60 | RUNDIR2=./../../../Training_Paper/Lightning/"$MODEL"_"$BACKBONE"_"$FILTERS"_"$EDGES"_"$NBLOCKS"/"$DATASET2"/ 61 | mkdir -p "$RUNDIR2" 62 | 63 | python ./../lightning_scripts/main.py \ 64 | --mode 'train' \ 65 | --train_only False \ 66 | --no-of-filters "$FILTERS" \ 67 | --num-blocks "$NBLOCKS" \ 68 | --edge-extracts "$EDGES" \ 69 | --loss-weight "$ALPHA" \ 70 | --model "$MODEL" \ 71 | --pretrain-checkpoint "$RUNDIR"/ckpt/last.ckpt \ 72 | --pretrained-base False \ 73 | --backbone "$BACKBONE" \ 74 | --dataset "$DATASET2" \ 75 | --dataset-path './../../../../Thermal_Segmentation/Dataset/' \ 76 | --base-size '520' \ 77 | --crop-size '480' \ 78 | --train-batch-size 16 \ 79 | --val-batch-size 16 \ 80 | --test-batch-size 1 \ 81 | --epochs 100 \ 82 | --optimizer 'SGD' \ 83 | --lr 0.001 \ 84 | --scheduler-type 'poly_warmstartup' \ 85 | --warmup-iters 0 \ 86 | --warmup-factor 0.3333 \ 87 | --warmup-method 'linear' \ 88 | --save-images-as-subplots False \ 89 | --save-images False \ 90 | --debug False \ 91 | --workers 16 \ 92 | --momentum 0.9 \ 93 | --weight-decay 0.0001 \ 94 | --beta1 0.9 \ 95 | --beta2 0.999 \ 96 | --epsilon 1e-8 \ 97 | --seed 0 \ 98 | --gpus $GPUS \ 99 | --num-nodes $NODES \ 100 | --distributed-backend 'ddp' \ 101 | --wandb-name-ext '' \ 102 | --save-dir "$RUNDIR2" 103 | 104 | 105 | ## Testing 106 | python ./../lightning_scripts/main.py \ 107 | --mode 'test' \ 108 | --model "$MODEL" \ 109 | --edge-extracts "$EDGES" \ 110 | --loss-weight "$ALPHA" \ 111 | --num-blocks "$NBLOCKS" \ 112 | --backbone "$BACKBONE" \ 113 | --no-of-filters "$FILTERS" \ 114 | --test-monitor 'val_mIOU' \ 115 | --test-monitor-path "$RUNDIR2"/ckpt/ \ 116 | --pretrained-base False \ 117 | --dataset "$DATASET2" \ 118 | --dataset-path './../../../../Thermal_Segmentation/Dataset/' \ 119 | --test-batch-size 1 \ 120 | --save-images True \ 121 | --save-images-as-subplots False \ 122 | --debug False \ 123 | --workers 16 \ 124 | --seed 0 \ 125 | --gpus 1 \ 126 | --num-nodes 1 \ 127 | --distributed-backend 'dp' \ 128 | --wandb-name-ext '' \ 129 | --save-dir "$RUNDIR2"/Best_MIOU/ 130 | 131 | done 132 | -------------------------------------------------------------------------------- /Codes/src/core/utils/utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function 2 | 3 | import collections 4 | import logging 5 | import os 6 | import sys 7 | 8 | import matplotlib 9 | import matplotlib.cm 10 | import matplotlib.pyplot as plt 11 | import numpy as np 12 | import torch 13 | import torch.distributed as dist 14 | from torch.autograd import Variable 15 | 16 | __all__ = ['get_rank', 'save_model_summary', 'plot_tensors', 'as_numpy', 17 | 'plot_tensors', 'save_checkpoint', 'print_network', 18 | 'total_gradient', ] 19 | 20 | def th_delete(tensor, indices): 21 | mask = torch.ones(tensor.numel(), dtype=torch.bool) 22 | mask[indices] = False 23 | return tensor[mask] 24 | 25 | 26 | def get_rank() -> int: 27 | if not dist.is_available(): 28 | return 0 29 | if not dist.is_initialized(): 30 | return 0 31 | return dist.get_rank() 32 | 33 | 34 | def save_model_summary(model, dir_): 35 | """ 36 | Print and save the network 37 | """ 38 | path = os.path.join(dir_, 'model.txt') 39 | num_params = 0 40 | for param in model.parameters(): 41 | num_params += param.numel() 42 | confg = repr(model) 43 | confg += '\nTotal number of parameters: {}'.format(num_params) 44 | confg += '\nTotal number of parameters in M: {}M'.format(num_params / (1000**2)) 45 | with open(path, "w") as text_file: 46 | text_file.write(confg) 47 | 48 | 49 | def plot_tensors(img, x=None): 50 | im = img.detach().cpu().numpy() 51 | dim = img.ndim 52 | # (ch, _, _) = im.shape 53 | plt.figure() 54 | if dim == 3: 55 | im = np.moveaxis(im, 0, 2) 56 | plt.imshow(im) 57 | else: 58 | plt.imshow(im, cmap=plt.get_cmap('gray')) 59 | if x is not None: 60 | plt.title(x) 61 | plt.show() 62 | 63 | 64 | def to_python_float(t): 65 | if hasattr(t, 'item'): 66 | if t.numel() > 1: 67 | return t.cpu().detach().numpy() 68 | else: 69 | return t.item() 70 | elif isinstance(t, list): 71 | return t[0] 72 | else: 73 | return t 74 | 75 | 76 | def as_numpy(obj): 77 | if isinstance(obj, collections.Sequence): 78 | return [as_numpy(v) for v in obj] 79 | elif isinstance(obj, collections.Mapping): 80 | return {k: as_numpy(v) for k, v in obj.items()} 81 | elif isinstance(obj, Variable): 82 | return obj.data.cpu().numpy() 83 | elif torch.is_tensor(obj): 84 | return obj.cpu().numpy() 85 | else: 86 | return np.array(obj) 87 | 88 | 89 | def save_checkpoint(states, is_best, output_dir, 90 | filename='checkpoint.pth.tar'): 91 | torch.save(states, os.path.join(output_dir, filename)) 92 | 93 | if is_best and 'state_dict' in states: 94 | torch.save(states, os.path.join(output_dir, 'model_best.pth.tar')) 95 | 96 | 97 | def print_network(net): 98 | num_params = 0 99 | for param in net.parameters(): 100 | num_params += param.numel() 101 | print(net) 102 | print('Total number of parameters: {} M'.format(int(num_params) / 1000**2)) 103 | 104 | 105 | def total_gradient(parameters): 106 | # ============================================================================= 107 | # Computes a gradient clipping coefficient based on gradient norm 108 | # ============================================================================= 109 | parameters = list(filter(lambda p: p.grad is not None, parameters)) 110 | totalnorm = 0 111 | for p in parameters: 112 | modulenorm = p.grad.data.norm() 113 | totalnorm += modulenorm ** 2 114 | totalnorm = totalnorm ** (1. / 2) 115 | return totalnorm 116 | -------------------------------------------------------------------------------- /Codes/src/core/schedulers/OneCycle_LR.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | from torch.optim.lr_scheduler import _LRScheduler 4 | 5 | __all__ = ['OneCycle'] 6 | 7 | 8 | class OneCycle(_LRScheduler): 9 | def __init__(self, optimizer, 10 | num_epochs, 11 | iters_per_epoch=0, 12 | last_epoch=-1, 13 | momentums=(0.85, 0.95), 14 | div_factor=25, 15 | phase1=0.3): 16 | 17 | self.iters_per_epoch = iters_per_epoch 18 | self.cur_iter = 0 19 | self.N = num_epochs * iters_per_epoch 20 | self.phase1_iters = int(self.N * phase1) 21 | self.phase2_iters = (self.N - self.phase1_iters) 22 | self.momentums = momentums 23 | self.mom_diff = momentums[1] - momentums[0] 24 | self.low_lrs = [opt_grp['lr'] / div_factor for opt_grp in optimizer.param_groups] 25 | self.final_lrs = [opt_grp['lr'] / (div_factor * 1e4) for opt_grp in optimizer.param_groups] 26 | super(OneCycle, self).__init__(optimizer, last_epoch) 27 | 28 | def get_lr(self): 29 | T = self.last_epoch * self.iters_per_epoch + self.cur_iter 30 | self.cur_iter %= self.iters_per_epoch 31 | self.cur_iter += 1 32 | 33 | # Going from base_lr / 25 -> base_lr 34 | if T <= self.phase1_iters: 35 | cos_anneling = (1 + math.cos(math.pi * T / self.phase1_iters)) / 2 36 | for i in range(len(self.optimizer.param_groups)): 37 | self.optimizer.param_groups[i]['momentum'] = self.momentums[0] + self.mom_diff * cos_anneling 38 | 39 | return [base_lr - (base_lr - low_lr) * cos_anneling 40 | for base_lr, low_lr in zip(self.base_lrs, self.low_lrs)] 41 | 42 | # Going from base_lr -> base_lr / (25e4) 43 | T -= self.phase1_iters 44 | cos_anneling = (1 + math.cos(math.pi * T / self.phase2_iters)) / 2 45 | 46 | for i in range(len(self.optimizer.param_groups)): 47 | self.optimizer.param_groups[i]['momentum'] = self.momentums[1] - self.mom_diff * cos_anneling 48 | return [final_lr + (base_lr - final_lr) * cos_anneling 49 | for base_lr, final_lr in zip(self.base_lrs, self.final_lrs)] 50 | 51 | 52 | if __name__ == '__main__': 53 | import torch 54 | epochs = 10 55 | 56 | model = torch.nn.Conv2d(16, 16, 3, 1, 1) 57 | optimizer = torch.optim.Adam(model.parameters(), lr=0.01) 58 | 59 | scheduler = OneCycle(optimizer, 60 | num_epochs=epochs, 61 | iters_per_epoch=10, 62 | momentums=(0.85, 0.95), 63 | div_factor=25, 64 | phase1=0.3) 65 | 66 | # scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, 67 | # max_lr = 0.1, 68 | # epochs=10, 69 | # steps_per_epoch=10, 70 | # pct_start=0.3, 71 | # anneal_strategy='cos', 72 | # cycle_momentum=True, 73 | # base_momentum=0.85, 74 | # max_momentum=0.95, 75 | # div_factor=25.0, 76 | # final_div_factor=10000.0, 77 | # last_epoch=-1) 78 | lrs = [] 79 | for e in range(0, epochs): 80 | for _ in range(0, 10): 81 | scheduler.step() 82 | lrs.append(optimizer.param_groups[0]['lr']) 83 | print(e, optimizer.param_groups[0]['lr']) 84 | 85 | import matplotlib.pyplot as plt 86 | import numpy as np 87 | plt.figure() 88 | plt.plot(np.arange(len(lrs)), lrs) 89 | plt.show() 90 | -------------------------------------------------------------------------------- /Codes/src/core/models/encoder/splat.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Split-Attention 3 | Adapted from 4 | https://github.com/zhanghang1989/ResNeSt 5 | ''' 6 | 7 | import torch 8 | from torch import nn 9 | import torch.nn.functional as F 10 | from torch.nn import Conv2d, Module, Linear, BatchNorm2d, ReLU 11 | from torch.nn.modules.utils import _pair 12 | 13 | 14 | __all__ = ['SplAtConv2d'] 15 | 16 | 17 | class SplAtConv2d(Module): 18 | """Split-Attention Conv2d 19 | """ 20 | 21 | def __init__(self, in_channels, channels, kernel_size, stride=(1, 1), padding=(0, 0), 22 | dilation=(1, 1), groups=1, bias=True, 23 | radix=2, reduction_factor=4, 24 | rectify=False, rectify_avg=False, norm_layer=None, 25 | dropblock_prob=0.0, **kwargs): 26 | super(SplAtConv2d, self).__init__() 27 | padding = _pair(padding) 28 | self.rectify = rectify and (padding[0] > 0 or padding[1] > 0) 29 | self.rectify_avg = rectify_avg 30 | inter_channels = max(in_channels * radix // reduction_factor, 32) 31 | self.radix = radix 32 | self.cardinality = groups 33 | self.channels = channels 34 | self.dropblock_prob = dropblock_prob 35 | if self.rectify: 36 | from rfconv import RFConv2d 37 | self.conv = RFConv2d(in_channels, channels * radix, kernel_size, stride, padding, dilation, 38 | groups=groups * radix, bias=bias, average_mode=rectify_avg, **kwargs) 39 | else: 40 | self.conv = Conv2d(in_channels, channels * radix, kernel_size, stride, padding, dilation, 41 | groups=groups * radix, bias=bias, **kwargs) 42 | self.use_bn = norm_layer is not None 43 | if self.use_bn: 44 | self.bn0 = norm_layer(channels * radix) 45 | self.relu = ReLU(inplace=True) 46 | self.fc1 = Conv2d(channels, inter_channels, 1, groups=self.cardinality) 47 | if self.use_bn: 48 | self.bn1 = norm_layer(inter_channels) 49 | self.fc2 = Conv2d(inter_channels, channels * radix, 1, groups=self.cardinality) 50 | if dropblock_prob > 0.0: 51 | self.dropblock = DropBlock2D(dropblock_prob, 3) 52 | self.rsoftmax = rSoftMax(radix, groups) 53 | 54 | def forward(self, x): 55 | x = self.conv(x) 56 | if self.use_bn: 57 | x = self.bn0(x) 58 | if self.dropblock_prob > 0.0: 59 | x = self.dropblock(x) 60 | x = self.relu(x) 61 | 62 | batch, rchannel = x.shape[:2] 63 | if self.radix > 1: 64 | if torch.__version__ < '1.5': 65 | splited = torch.split(x, int(rchannel // self.radix), dim=1) 66 | else: 67 | splited = torch.split(x, rchannel // self.radix, dim=1) 68 | gap = sum(splited) 69 | else: 70 | gap = x 71 | gap = F.adaptive_avg_pool2d(gap, 1) 72 | gap = self.fc1(gap) 73 | 74 | if self.use_bn: 75 | gap = self.bn1(gap) 76 | gap = self.relu(gap) 77 | 78 | atten = self.fc2(gap) 79 | atten = self.rsoftmax(atten).view(batch, -1, 1, 1) 80 | 81 | if self.radix > 1: 82 | if torch.__version__ < '1.5': 83 | attens = torch.split(atten, int(rchannel // self.radix), dim=1) 84 | else: 85 | attens = torch.split(atten, rchannel // self.radix, dim=1) 86 | out = sum([att * split for (att, split) in zip(attens, splited)]) 87 | else: 88 | out = atten * x 89 | return out.contiguous() 90 | 91 | 92 | class rSoftMax(nn.Module): 93 | def __init__(self, radix, cardinality): 94 | super().__init__() 95 | self.radix = radix 96 | self.cardinality = cardinality 97 | 98 | def forward(self, x): 99 | batch = x.size(0) 100 | if self.radix > 1: 101 | x = x.view(batch, self.cardinality, self.radix, -1).transpose(1, 2) 102 | x = F.softmax(x, dim=1) 103 | x = x.reshape(batch, -1) 104 | else: 105 | x = torch.sigmoid(x) 106 | return x 107 | -------------------------------------------------------------------------------- /Codes/src/core/metrics/AverageMeter.py: -------------------------------------------------------------------------------- 1 | # Copyright The PyTorch Lightning team. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from typing import Any, Callable, Optional, Union 15 | 16 | import torch 17 | from torch import Tensor 18 | 19 | from torchmetrics.metric import Metric 20 | 21 | 22 | class AverageMeter(Metric): 23 | """Computes the average of a stream of values. 24 | 25 | Forward accepts 26 | - ``value`` (float tensor): ``(...)`` 27 | - ``weight`` (float tensor): ``(...)`` 28 | 29 | Args: 30 | compute_on_step: 31 | Forward only calls ``update()`` and returns None if this is 32 | set to False. default: True 33 | dist_sync_on_step: 34 | Synchronize metric state across processes at each ``forward()`` 35 | before returning the value at the step. 36 | process_group: 37 | Specify the process group on which synchronization is called. 38 | default: None (which selects the entire world) 39 | dist_sync_fn: 40 | Callback that performs the allgather operation on the metric state. 41 | When `None`, DDP will be used to perform the allgather. 42 | 43 | Example:: 44 | >>> from torchmetrics import AverageMeter 45 | >>> avg = AverageMeter() 46 | >>> avg.update(3) 47 | >>> avg.update(1) 48 | >>> avg.compute() 49 | tensor(2.) 50 | 51 | >>> avg = AverageMeter() 52 | >>> values = torch.tensor([1., 2., 3.]) 53 | >>> avg(values) 54 | tensor(2.) 55 | 56 | >>> avg = AverageMeter() 57 | >>> values = torch.tensor([1., 2.]) 58 | >>> weights = torch.tensor([3., 1.]) 59 | >>> avg(values, weights) 60 | tensor(1.2500) 61 | """ 62 | 63 | def __init__( 64 | self, 65 | compute_on_step: bool = True, 66 | dist_sync_on_step: bool = False, 67 | process_group: Optional[Any] = None, 68 | dist_sync_fn: Callable = None, 69 | ): 70 | super().__init__( 71 | compute_on_step=compute_on_step, 72 | dist_sync_on_step=dist_sync_on_step, 73 | process_group=process_group, 74 | dist_sync_fn=dist_sync_fn, 75 | ) 76 | self.add_state("value", torch.zeros(()), dist_reduce_fx="sum") 77 | self.add_state("weight", torch.zeros(()), dist_reduce_fx="sum") 78 | 79 | # TODO: need to be strings because Unions are not pickleable in Python 3.6 80 | def update( # type: ignore 81 | self, value: "Union[Tensor, float]", weight: "Union[Tensor, float]" = 1.0 82 | ) -> None: 83 | """Updates the average with. 84 | 85 | Args: 86 | value: A tensor of observations (can also be a scalar value) 87 | weight: The weight of each observation (automatically broadcasted 88 | to fit ``value``) 89 | """ 90 | if not isinstance(value, Tensor): 91 | value = torch.as_tensor(value, dtype=torch.float32, device=self.value.device) 92 | if not isinstance(weight, Tensor): 93 | weight = torch.as_tensor(weight, dtype=torch.float32, device=self.weight.device) 94 | 95 | # braodcast_to only supported on PyTorch 1.8+ 96 | if not hasattr(torch, "broadcast_to"): 97 | if weight.shape == (): 98 | weight = torch.ones_like(value) * weight 99 | if weight.shape != value.shape: 100 | raise ValueError("Broadcasting not supported on PyTorch <1.8") 101 | else: 102 | weight = torch.broadcast_to(weight, value.shape) 103 | 104 | self.value += (value * weight).sum() 105 | self.weight += weight.sum() 106 | 107 | def compute(self) -> Tensor: 108 | return self.value / self.weight 109 | -------------------------------------------------------------------------------- /Codes/src/core/optimizers/adamw.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | ''' 3 | Code adapted from 4 | https://github.com/pytorch/pytorch/blob/master/torch/optim/adamw.py 5 | ''' 6 | 7 | 8 | import math 9 | import torch 10 | from torch.optim.optimizer import Optimizer 11 | 12 | __all__ = ['AdamW'] 13 | 14 | 15 | class AdamW(Optimizer): 16 | """Implements Adam algorithm. 17 | Arguments: 18 | params (iterable): iterable of parameters to optimize or dicts defining 19 | parameter groups 20 | lr (float, optional): learning rate (default: 1e-3) 21 | betas (Tuple[float, float], optional): coefficients used for computing 22 | running averages of gradient and its square (default: (0.9, 0.999)) 23 | eps (float, optional): term added to the denominator to improve 24 | numerical stability (default: 1e-8) 25 | weight_decay (float, optional): weight decay (L2 penalty) (default: 0) 26 | amsgrad (boolean, optional): whether to use the AMSGrad variant of this 27 | algorithm from the paper `On the Convergence of Adam and Beyond`_ 28 | """ 29 | 30 | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, 31 | weight_decay=0, amsgrad=False): 32 | if not 0.0 <= betas[0] < 1.0: 33 | raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) 34 | if not 0.0 <= betas[1] < 1.0: 35 | raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) 36 | defaults = dict(lr=lr, betas=betas, eps=eps, 37 | weight_decay=weight_decay, amsgrad=amsgrad) 38 | #super(AdamW, self).__init__(params, defaults) 39 | super().__init__(params, defaults) 40 | 41 | def step(self, closure=None): 42 | """Performs a single optimization step. 43 | Arguments: 44 | closure (callable, optional): A closure that reevaluates the model 45 | and returns the loss. 46 | """ 47 | loss = None 48 | if closure is not None: 49 | loss = closure() 50 | 51 | for group in self.param_groups: 52 | for p in group['params']: 53 | if p.grad is None: 54 | continue 55 | grad = p.grad.data 56 | if grad.is_sparse: 57 | raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead') 58 | amsgrad = group['amsgrad'] 59 | 60 | state = self.state[p] 61 | 62 | # State initialization 63 | if len(state) == 0: 64 | state['step'] = 0 65 | # Exponential moving average of gradient values 66 | state['exp_avg'] = torch.zeros_like(p.data) 67 | # Exponential moving average of squared gradient values 68 | state['exp_avg_sq'] = torch.zeros_like(p.data) 69 | if amsgrad: 70 | # Maintains max of all exp. moving avg. of sq. grad. values 71 | state['max_exp_avg_sq'] = torch.zeros_like(p.data) 72 | 73 | exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] 74 | if amsgrad: 75 | max_exp_avg_sq = state['max_exp_avg_sq'] 76 | beta1, beta2 = group['betas'] 77 | 78 | state['step'] += 1 79 | 80 | # Decay the first and second moment running average coefficient 81 | exp_avg.mul_(beta1).add_(1 - beta1, grad) 82 | exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) 83 | if amsgrad: 84 | # Maintains the maximum of all 2nd moment running avg. till now 85 | torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) 86 | # Use the max. for normalizing running avg. of gradient 87 | denom = max_exp_avg_sq.sqrt().add_(group['eps']) 88 | else: 89 | denom = exp_avg_sq.sqrt().add_(group['eps']) 90 | 91 | bias_correction1 = 1 - beta1 ** state['step'] 92 | bias_correction2 = 1 - beta2 ** state['step'] 93 | step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1 94 | 95 | if group['weight_decay'] != 0: 96 | decayed_weights = torch.mul(p.data, group['weight_decay']) 97 | p.data.addcdiv_(-step_size, exp_avg, denom) 98 | p.data.sub_(decayed_weights) 99 | else: 100 | p.data.addcdiv_(-step_size, exp_avg, denom) 101 | 102 | return loss 103 | -------------------------------------------------------------------------------- /Codes/src/core/data/samplers/multiscale_batch_samplers.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | ''' 3 | Code Adapted from: 4 | https://github.com/CaoWGG/multi-scale-training 5 | 6 | ''' 7 | from typing import Callable 8 | 9 | import numpy as np 10 | from torch.utils.data.sampler import BatchSampler, Sampler 11 | 12 | __all__ = ['MultiscaleBatchSampler', 'IterationBasedMultiscaleBatchSampler'] 13 | 14 | 15 | class MultiscaleBatchSampler(object): 16 | r"""Samples elements with different scales per batch i.e. when scale == 2 then the dataloader 17 | __getitem__ will be a tuple with index and scale 18 | 19 | Usage: 20 | def __getitem__(self, idx): 21 | if type(idx) == list or type(idx) == tuple: 22 | idx, scale = idx 23 | 24 | Arguments: 25 | sampler (Sampler or Iterable): Base sampler. Can be any iterable object 26 | with ``__len__`` implemented. 27 | batch_size (int): Size of mini-batch 28 | drop_last (bool): If ``True``, the sampler will drop the last batch if 29 | its size would be less than ``batch_size`` 30 | multiscale_step(int): Number of times a scale needs to be repeated in a sequential batch 31 | scales(int): Number of scale variations required in the dataloader, for example, is two different crop sizes are required in during training, 32 | then scales =2. 33 | """ 34 | 35 | def __init__(self, sampler: Callable, 36 | batch_size: int, 37 | drop_last: bool, 38 | multiscale_step: int = 1, 39 | scales: int = 1) -> None: 40 | if not isinstance(sampler, Sampler): 41 | raise ValueError("sampler should be an instance of torch.utils.data.Sampler, but got sampler={}".format(sampler)) 42 | if not isinstance(drop_last, bool): 43 | raise ValueError("drop_last should be a boolean value, but got drop_last={}".format(drop_last)) 44 | self.sampler = sampler 45 | self.batch_size = batch_size 46 | self.drop_last = drop_last 47 | if multiscale_step is not None and multiscale_step < 1: 48 | raise ValueError("multiscale_step should be > 0, but got multiscale_step={}".format(multiscale_step)) 49 | 50 | if not isinstance(scales, int): 51 | raise ValueError("scales must a int, got {}".format(scales)) 52 | 53 | self.scales = scales 54 | self.multiscale_step = multiscale_step 55 | 56 | def __iter__(self): 57 | num_batch = 0 58 | batch = [] 59 | scale = np.random.randint(self.scales) 60 | print(scale) 61 | for idx in self.sampler: 62 | batch.append([idx, scale]) 63 | if len(batch) == self.batch_size: 64 | yield batch 65 | num_batch += 1 66 | batch = [] 67 | if self.multiscale_step and num_batch % self.multiscale_step == 0: 68 | scale = np.random.randint(self.scales) 69 | 70 | if len(batch) > 0 and not self.drop_last: 71 | yield batch 72 | 73 | def __len__(self): 74 | if self.drop_last: 75 | return len(self.sampler) // self.batch_size 76 | else: 77 | return (len(self.sampler) + self.batch_size - 1) // self.batch_size 78 | 79 | 80 | class IterationBasedMultiscaleBatchSampler(MultiscaleBatchSampler): 81 | def __init__(self, batch_sampler: Callable, 82 | num_iterations: int, 83 | start_iter: int = 0) -> None: 84 | self.batch_sampler = batch_sampler 85 | self.num_iterations = num_iterations 86 | self.start_iter = start_iter 87 | super(IterationBasedMultiscaleBatchSampler, self).__init__(sampler=batch_sampler.sampler, 88 | batch_size=batch_sampler.batch_size, 89 | drop_last=False, 90 | multiscale_step=batch_sampler.multiscale_step, 91 | scales=batch_sampler.scales) 92 | 93 | def __iter__(self): 94 | iteration = self.start_iter 95 | while iteration <= self.num_iterations: 96 | # if the underlying sampler has a set_epoch method, like 97 | # DistributedSampler, used for making each process see 98 | # a different split of the dataset, then set it 99 | if hasattr(self.batch_sampler.sampler, "set_epoch"): 100 | self.batch_sampler.sampler.set_epoch(iteration) 101 | for batch in self.batch_sampler: 102 | iteration += 1 103 | if iteration > self.num_iterations: 104 | break 105 | yield batch 106 | 107 | def __len__(self): 108 | return self.num_iterations 109 | -------------------------------------------------------------------------------- /Codes/src/datasets/dataloaders/segbase.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Code Adapted from 3 | https://github.com/dmlc/gluon-cv/blob/master/gluoncv/data/segbase.py 4 | ''' 5 | import random 6 | import cv2 7 | import numpy as np 8 | import torch 9 | from PIL import Image, ImageFilter, ImageOps 10 | 11 | __all__ = ['SegmentationDataset'] 12 | 13 | 14 | class SegmentationDataset(object): 15 | """Segmentation Base Dataset""" 16 | 17 | def __init__(self, root, split, mode, base_size=520, crop_size=480, logger=None): 18 | super(SegmentationDataset, self).__init__() 19 | self.root = root 20 | self.logger = logger 21 | self.split = split 22 | self.mode = mode if mode is not None else split 23 | self.base_size = base_size 24 | self.crop_size = crop_size 25 | 26 | edge_radius = 7 27 | self.edge_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, 28 | (edge_radius, edge_radius)) 29 | 30 | if self.mode != 'testval': 31 | self.sizes = dict(zip(crop_size, base_size)) 32 | 33 | def _val_sync_transform(self, img, mask, edge, crop_size=None): 34 | outsize = self.crop_size[0] 35 | short_size = outsize 36 | w, h = img.size 37 | if w > h: 38 | oh = short_size 39 | ow = int(1.0 * w * oh / h) 40 | else: 41 | ow = short_size 42 | oh = int(1.0 * h * ow / w) 43 | img = img.resize((ow, oh), Image.BILINEAR) 44 | mask = mask.resize((ow, oh), Image.NEAREST) 45 | edge = edge.resize((ow, oh), Image.NEAREST) 46 | 47 | # center crop 48 | w, h = img.size 49 | x1 = int(round((w - outsize) / 2.)) 50 | y1 = int(round((h - outsize) / 2.)) 51 | img = img.crop((x1, y1, x1 + outsize, y1 + outsize)) 52 | mask = mask.crop((x1, y1, x1 + outsize, y1 + outsize)) 53 | edge = edge.crop((x1, y1, x1 + outsize, y1 + outsize)) 54 | # final transform 55 | img, mask, edge = self._img_transform(img), self._mask_transform(mask), self._edge_transform(edge) 56 | return img, mask, edge 57 | 58 | def _sync_transform(self, img, mask, edge, crop_size=None): 59 | # random mirror 60 | if random.random() < 0.5: 61 | img = img.transpose(Image.FLIP_LEFT_RIGHT) 62 | mask = mask.transpose(Image.FLIP_LEFT_RIGHT) 63 | edge = edge.transpose(Image.FLIP_LEFT_RIGHT) 64 | 65 | if crop_size is None: 66 | crop_size = self.crop_size[0] 67 | base_size = self.base_size[0] 68 | else: 69 | base_size = self.sizes[crop_size] 70 | 71 | # random scale (short edge) 72 | short_size = random.randint(int(base_size * 0.5), int(base_size * 2.0)) 73 | w, h = img.size 74 | if h > w: 75 | ow = short_size 76 | oh = int(1.0 * h * ow / w) 77 | else: 78 | oh = short_size 79 | ow = int(1.0 * w * oh / h) 80 | img = img.resize((ow, oh), Image.BILINEAR) 81 | mask = mask.resize((ow, oh), Image.NEAREST) 82 | edge = edge.resize((ow, oh), Image.NEAREST) 83 | 84 | # pad crop 85 | if short_size < crop_size: 86 | padh = crop_size - oh if oh < crop_size else 0 87 | padw = crop_size - ow if ow < crop_size else 0 88 | img = ImageOps.expand(img, border=(0, 0, padw, padh), fill=0) 89 | mask = ImageOps.expand(mask, border=(0, 0, padw, padh), fill=0) 90 | edge = ImageOps.expand(edge, border=(0, 0, padw, padh), fill=0) 91 | 92 | # random crop crop_size 93 | w, h = img.size 94 | x1 = random.randint(0, w - crop_size) 95 | y1 = random.randint(0, h - crop_size) 96 | img = img.crop((x1, y1, x1 + crop_size, y1 + crop_size)) 97 | mask = mask.crop((x1, y1, x1 + crop_size, y1 + crop_size)) 98 | edge = edge.crop((x1, y1, x1 + crop_size, y1 + crop_size)) 99 | # gaussian blur as in PSP 100 | if random.random() < 0.5: 101 | img = img.filter(ImageFilter.GaussianBlur(radius=random.random())) 102 | # final transform 103 | img, mask, edge = self._img_transform(img), self._mask_transform(mask), self._edge_transform(edge) 104 | return img, mask, edge 105 | 106 | def im2double(self, im_: np.ndarray) -> np.ndarray: 107 | info = np.iinfo(im_.dtype) # Get the data type of the input image 108 | return im_.astype(np.float) / info.max # Divide all values by the largest possible value in the datatype 109 | 110 | def np2Tensor(self, tensor): 111 | if len(tensor.shape) == 3: 112 | np_transpose = np.ascontiguousarray(tensor.transpose((2, 0, 1))) 113 | else: 114 | np_transpose = np.ascontiguousarray(tensor[np.newaxis, :]) 115 | tensor = torch.from_numpy(np_transpose).float() 116 | return tensor 117 | 118 | def _img_transform(self, img): 119 | return np.array(img) 120 | 121 | def _mask_transform(self, mask): 122 | return np.array(mask).astype('int32') 123 | 124 | def _edge_transform(self, edge): 125 | return np.array(edge).astype('int32') 126 | # [None, :, :] 127 | 128 | @property 129 | def num_class(self): 130 | """Number of categories.""" 131 | return self.NUM_CLASS 132 | 133 | @property 134 | def pred_offset(self): 135 | return 0 136 | -------------------------------------------------------------------------------- /Codes/src/core/utils/filesystem.py: -------------------------------------------------------------------------------- 1 | """Filesystem utility functions.""" 2 | from __future__ import absolute_import 3 | 4 | import argparse 5 | import errno 6 | import os 7 | 8 | __all__ = ['checkpoint', 'makedirs'] 9 | 10 | 11 | class checkpoint(): 12 | def __init__(self, args: argparse.Namespace, test=False): 13 | self.args = args 14 | if test: 15 | 16 | self.args.save_dir = os.path.join(self.args.save_dir, 'Evaluation') 17 | 18 | self.paths = {'save_dir': self.args.save_dir, 19 | 'logs': os.path.join(self.args.save_dir, 'logs')} 20 | if not test: 21 | self.paths['ckpt'] = os.path.join(self.args.save_dir, 'ckpt') 22 | 23 | for _, val in self.paths.items(): 24 | if not os.path.exists(val): 25 | makedirs(val) 26 | 27 | def get_path(self, subdir: str) -> str: 28 | if self.paths.__contains__(subdir): 29 | return self.paths[subdir] 30 | else: 31 | temp = os.path.join(self.args.save_dir, subdir) 32 | self.paths[subdir] = temp 33 | makedirs(temp) 34 | return temp 35 | 36 | 37 | def makedirs(path: str): 38 | """Create directory recursively if not exists. 39 | Similar to `makedir -p`, you can skip checking existence before this function. 40 | Parameters 41 | ---------- 42 | path : str 43 | Path of the desired dir 44 | """ 45 | try: 46 | os.makedirs(path, mode=0o770, exist_ok=True) 47 | except OSError as exc: 48 | if exc.errno != errno.EEXIST: 49 | raise 50 | 51 | 52 | def try_import(package, message=None): 53 | """Try import specified package, with custom message support. 54 | Parameters 55 | ---------- 56 | package : str 57 | The name of the targeting package. 58 | message : str, default is None 59 | If not None, this function will raise customized error message when import error is found. 60 | Returns 61 | ------- 62 | module if found, raise ImportError otherwise 63 | """ 64 | try: 65 | return __import__(package) 66 | except ImportError as e: 67 | if not message: 68 | raise e 69 | raise ImportError(message) 70 | 71 | 72 | def try_import_cv2(): 73 | """Try import cv2 at runtime. 74 | Returns 75 | ------- 76 | cv2 module if found. Raise ImportError otherwise 77 | """ 78 | msg = "cv2 is required, you can install by package manager, e.g. 'apt-get', \ 79 | or `pip install opencv-python --user` (note that this is unofficial PYPI package)." 80 | return try_import('cv2', msg) 81 | 82 | 83 | def import_try_install(package, extern_url=None): 84 | """Try import the specified package. 85 | If the package not installed, try use pip to install and import if success. 86 | Parameters 87 | ---------- 88 | package : str 89 | The name of the package trying to import. 90 | extern_url : str or None, optional 91 | The external url if package is not hosted on PyPI. 92 | For example, you can install a package using: 93 | "pip install git+http://github.com/user/repo/tarball/master/egginfo=xxx". 94 | In this case, you can pass the url to the extern_url. 95 | Returns 96 | ------- 97 | 98 | The imported python module. 99 | """ 100 | try: 101 | return __import__(package) 102 | except ImportError: 103 | try: 104 | from pip import main as pipmain 105 | except ImportError: 106 | from pip._internal import main as pipmain 107 | 108 | # trying to install package 109 | url = package if extern_url is None else extern_url 110 | pipmain(['install', '--user', url]) # will raise SystemExit Error if fails 111 | 112 | # trying to load again 113 | try: 114 | return __import__(package) 115 | except ImportError: 116 | import site 117 | import sys 118 | user_site = site.getusersitepackages() 119 | if user_site not in sys.path: 120 | sys.path.append(user_site) 121 | return __import__(package) 122 | return __import__(package) 123 | 124 | 125 | """Import helper for pycocotools""" 126 | 127 | 128 | # NOTE: for developers 129 | # please do not import any pycocotools in __init__ because we are trying to lazy 130 | # import pycocotools to avoid install it for other users who may not use it. 131 | # only import when you actually use it 132 | 133 | 134 | def try_import_pycocotools(): 135 | """Tricks to optionally install and import pycocotools""" 136 | # first we can try import pycocotools 137 | try: 138 | import pycocotools as _ 139 | except ImportError: 140 | import os 141 | 142 | # we need to install pycootools, which is a bit tricky 143 | # pycocotools sdist requires Cython, numpy(already met) 144 | import_try_install('cython') 145 | # pypi pycocotools is not compatible with windows 146 | win_url = 'git+https://github.com/zhreshold/cocoapi.git#subdirectory=PythonAPI' 147 | try: 148 | if os.name == 'nt': 149 | import_try_install('pycocotools', win_url) 150 | else: 151 | import_try_install('pycocotools') 152 | except ImportError: 153 | faq = 'cocoapi FAQ' 154 | raise ImportError('Cannot import or install pycocotools, please refer to %s.' % faq) 155 | -------------------------------------------------------------------------------- /Codes/src/datasets/utils/Cityscape_folderMap.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # To distribute train, valid, test for SODA 3 | # ============================================================================= 4 | 5 | import os 6 | import pandas as pd 7 | from sklearn.model_selection import train_test_split 8 | from glob import glob 9 | import shutil 10 | import argparse 11 | 12 | 13 | def is_image_file(filename): 14 | return any( 15 | filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg", ".bmp"] 16 | ) 17 | 18 | 19 | def str2bool(v): 20 | if v.lower() in ("yes", "true", "t", "y", "1"): 21 | return True 22 | elif v.lower() in ("no", "false", "f", "n", "0"): 23 | return False 24 | else: 25 | raise argparse.ArgumentTypeError("Boolean value expected.") 26 | 27 | 28 | # %% 29 | # ============================================================================= 30 | # Copying files to different folders 31 | # ============================================================================= 32 | def copying(tiles, path_label, basepath, fileset_path): 33 | tiles.set_index(path_label, inplace=True) 34 | for img_path in tiles.index: 35 | print("Path = {}".format(img_path)) 36 | dst_path = os.path.join(basepath, fileset_path) 37 | 38 | shutil.copy(img_path, dst_path) 39 | 40 | 41 | # %% 42 | # ============================================================================= 43 | # Creating folders 44 | # ============================================================================= 45 | 46 | 47 | def distribute(input_dir, output_dir, reset): 48 | # basepath = './Thermal_Segmentation/Dataset/Cityscapes_thermal/TIR_leftImg8bit/' 49 | basepath = input_dir 50 | 51 | # os.path.abspath(os.path.join(basepath,'..')) 52 | base_dir = os.path.join(output_dir, "CITYSCAPE_5000") 53 | if reset == True and os.path.exists(base_dir): 54 | shutil.rmtree(base_dir) 55 | if not os.path.exists(base_dir): 56 | os.mkdir(base_dir) 57 | 58 | main_dirs_image = ["image/train"] 59 | main_dirs_mask = ["mask/train"] 60 | 61 | for main in main_dirs_image: 62 | 63 | path = os.path.join(base_dir, main) 64 | if not os.path.exists(path): 65 | os.makedirs(path) 66 | 67 | for main in main_dirs_mask: 68 | 69 | path = os.path.join(base_dir, main) 70 | if not os.path.exists(path): 71 | os.makedirs(path) 72 | 73 | # %% 74 | # ============================================================================= 75 | # Creating folders 76 | # ============================================================================= 77 | 78 | imageid_path_dict = { 79 | os.path.splitext(os.path.basename(x))[0]: x 80 | for x in glob(os.path.join(basepath, "**", "**", "*.jpg"), recursive=True) 81 | } 82 | 83 | tile_df = pd.DataFrame( 84 | imageid_path_dict.items(), columns=["Image_Name", "Image_Path"] 85 | ) 86 | tile_df = tile_df.sort_values( 87 | by="Image_Name", axis=0, ascending=True, kind="quicksort" 88 | ).reset_index(drop=True) 89 | tile_df = tile_df.fillna("NA") 90 | tile_df["Mask_Path"] = tile_df["Image_Path"].str.replace(".jpg", ".png") 91 | tile_df["Mask_Path"] = tile_df["Mask_Path"].str.replace( 92 | "TIR_leftImg8bit", "TIR_leftImg8bit/gtFine" 93 | ) 94 | tile_df["Mask_Path"] = tile_df["Mask_Path"].str.replace( 95 | "leftImg8bit_synthesized_image", "gtFine_labelIds" 96 | ) 97 | tile_df["Mask_Name"] = tile_df["Image_Name"].str.replace( 98 | "leftImg8bit_synthesized_image", "gtFine_labelIds" 99 | ) 100 | 101 | # https://stackoverflow.com/questions/28679930/how-to-drop-rows-from-pandas-data-frame-that-contains-a-particular-string-in-a-p 102 | # https://stackoverflow.com/questions/41425945/python-pandas-error-missing-unterminated-subpattern-at-position-2 103 | tile_df = tile_df[ 104 | ~tile_df.Image_Name.str.contains("\(") 105 | ] # there are copies of some files, eg, ABCD.png and ABCD (1).png (I am removing the copies) 106 | # there are copies of some files, eg, ABCD.png and ~temp_ABCD.png (I am removing the copies) 107 | tile_df = tile_df[~tile_df.Image_Name.str.contains("\~")] 108 | 109 | copying( 110 | tiles=tile_df, 111 | path_label="Image_Path", 112 | basepath=base_dir, 113 | fileset_path=main_dirs_image[0], 114 | ) 115 | copying( 116 | tiles=tile_df, 117 | path_label="Mask_Path", 118 | basepath=base_dir, 119 | fileset_path=main_dirs_mask[0], 120 | ) 121 | 122 | 123 | if __name__ == "__main__": 124 | parser = argparse.ArgumentParser(description="Set up Cityscape Dataset") 125 | parser.add_argument( 126 | "--input-image-path", 127 | type=str, 128 | default="/mnt/1842213842211C4E/raw_dataset/SODA-20211127T202136Z-001/SODA/TIR_leftImg8bit/", 129 | help="Path to Cityscape Dataset. Note: This should lead to TIR_leftImg8bit", 130 | ) 131 | parser.add_argument( 132 | "--save-path", 133 | type=str, 134 | default="/mnt/1842213842211C4E/processed_dataset/", 135 | help="Path to Cityscape Dataset", 136 | ) 137 | parser.add_argument( 138 | "--reset", type=bool, default=True, help="Path to Cityscape Dataset" 139 | ) 140 | 141 | args = parser.parse_args() 142 | distribute( 143 | input_dir=args.input_image_path, 144 | output_dir=args.save_path, 145 | reset=args.reset, 146 | ) 147 | -------------------------------------------------------------------------------- /Codes/src/datasets/utils/MFNDataset_folderMap.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # To distribute train, valid, test for MFN 3 | # ============================================================================= 4 | 5 | import os 6 | import pandas as pd 7 | from sklearn.model_selection import train_test_split 8 | from glob import glob 9 | import shutil 10 | import argparse 11 | 12 | 13 | def is_image_file(filename): 14 | return any( 15 | filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg", ".bmp"] 16 | ) 17 | 18 | 19 | def str2bool(v): 20 | if v.lower() in ("yes", "true", "t", "y", "1"): 21 | return True 22 | elif v.lower() in ("no", "false", "f", "n", "0"): 23 | return False 24 | else: 25 | raise argparse.ArgumentTypeError("Boolean value expected.") 26 | 27 | 28 | def copying(tiles, path_label, destpath, fileset_path): 29 | tiles.set_index(path_label, inplace=True) 30 | for img_path in tiles.index: 31 | print("Path = {}".format(img_path)) 32 | dst_path = os.path.join(destpath, fileset_path) 33 | 34 | shutil.copy(img_path, dst_path) 35 | 36 | 37 | def create_df(txt_path, imageid_path_dict, blacklist_df): 38 | df = pd.read_table(txt_path, delim_whitespace=True, names=("Image_Name", "B")) 39 | df = df.sort_values( 40 | by="Image_Name", axis=0, ascending=True, kind="quicksort" 41 | ).reset_index(drop=True) 42 | df = df.fillna("NA") 43 | del df["B"] 44 | df["Image_Path"] = df["Image_Name"].map(imageid_path_dict.get) 45 | df["Mask_Path"] = df["Image_Path"].str.replace("images", "labels") 46 | 47 | df = df[ 48 | ~df.Image_Name.str.contains("\(") 49 | ] # there are copies of some files, eg, ABCD.png and ABCD (1).png (I am removing the copies) 50 | df = df[ 51 | ~df.Image_Name.str.contains("\~") 52 | ] # there are copies of some files, eg, ABCD.png and ~temp_ABCD.png (I am removing the copies) 53 | df = df[ 54 | ~df.Image_Name.str.contains("flip") 55 | ] # there are copies of some files, eg, ABCD.png and ABCD (1).png (I am removing the copies) 56 | cond = df["Image_Name"].isin(blacklist_df["Image_Name"]) 57 | df.drop(df[cond].index, inplace=True) 58 | 59 | return df 60 | 61 | 62 | def copy_files_from_df(df, destpath, main_dirs_image, main_dirs_mask): 63 | 64 | copying( 65 | tiles=df, 66 | path_label="Image_Path", 67 | destpath=destpath, 68 | fileset_path=main_dirs_image, 69 | ) 70 | copying( 71 | tiles=df, path_label="Mask_Path", destpath=destpath, fileset_path=main_dirs_mask 72 | ) 73 | 74 | 75 | def distribute(input_dir, output_dir, reset): 76 | basepath = input_dir 77 | # basepath = "./Thermal_Segmentation/Dataset/ir_seg_dataset-20210510T225620Z-001" 78 | # destpath = os.path.join(basepath, DATASET) 79 | # reset = True 80 | 81 | train_txt_path = basepath + "/train.txt" 82 | validation_txt_path = basepath + "/val.txt" 83 | test_txt_path = basepath + "/test.txt" 84 | blacklist_path = basepath + "/black_list.txt" 85 | 86 | Image_path = basepath + "/images" 87 | Mask_path = basepath + "/labels" 88 | destpath = os.path.join(output_dir, "MFN") 89 | 90 | # os.path.abspath(os.path.join(basepath,'..')) 91 | if reset == True and os.path.exists(destpath): 92 | if os.path.exists(destpath): 93 | shutil.rmtree(destpath) 94 | if not os.path.exists(destpath): 95 | os.mkdir(destpath) 96 | 97 | main_dirs_image = ["image/train", "image/val", "image/test"] 98 | main_dirs_mask = ["mask/train", "mask/val", "mask/test"] 99 | 100 | for main in main_dirs_image: 101 | 102 | path = os.path.join(destpath, main) 103 | if not os.path.exists(path): 104 | os.makedirs(path) 105 | 106 | for main in main_dirs_mask: 107 | 108 | path = os.path.join(destpath, main) 109 | if not os.path.exists(path): 110 | os.makedirs(path) 111 | 112 | imageid_path_dict = { 113 | os.path.splitext(os.path.basename(x))[0]: x 114 | for x in glob(os.path.join(Image_path, "**", "**", "*.png"), recursive=True) 115 | } 116 | 117 | blacklist_df = pd.read_table( 118 | blacklist_path, delim_whitespace=True, names=("Image_Name", "B") 119 | ) 120 | 121 | train_df = create_df(train_txt_path, imageid_path_dict, blacklist_df) 122 | validation_df = create_df(validation_txt_path, imageid_path_dict, blacklist_df) 123 | test_df = create_df(test_txt_path, imageid_path_dict, blacklist_df) 124 | 125 | copy_files_from_df(train_df, destpath, main_dirs_image[0], main_dirs_mask[0]) 126 | copy_files_from_df(validation_df, destpath, main_dirs_image[1], main_dirs_mask[1]) 127 | copy_files_from_df(test_df, destpath, main_dirs_image[2], main_dirs_mask[2]) 128 | 129 | 130 | if __name__ == "__main__": 131 | parser = argparse.ArgumentParser(description="Set up Cityscape Dataset") 132 | parser.add_argument( 133 | "--input-image-path", 134 | type=str, 135 | default="/mnt/1842213842211C4E/raw_dataset/ir_seg_dataset/", 136 | help="Path to MFN Dataset", 137 | ) 138 | parser.add_argument( 139 | "--save-path", 140 | type=str, 141 | default="/mnt/1842213842211C4E/processed_dataset/", 142 | help="Path to save MFN Dataset", 143 | ) 144 | parser.add_argument( 145 | "--reset", type=bool, default=True, help="Path to Cityscape Dataset" 146 | ) 147 | 148 | args = parser.parse_args() 149 | distribute( 150 | input_dir=args.input_image_path, output_dir=args.save_path, reset=args.reset 151 | ) 152 | -------------------------------------------------------------------------------- /Codes/src/core/utils/collect_env.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Adapted from 3 | https://github.com/facebookresearch/detectron2/blob/master/detectron2/utils/collect_env.py 4 | ''' 5 | 6 | import importlib 7 | import os 8 | import re 9 | import subprocess 10 | import sys 11 | from collections import defaultdict 12 | 13 | import numpy as np 14 | import PIL 15 | import torch 16 | import torchvision 17 | from tabulate import tabulate 18 | 19 | __all__ = ["collect_env_info"] 20 | 21 | 22 | def collect_torch_env(): 23 | try: 24 | import torch.__config__ 25 | 26 | return torch.__config__.show() 27 | except ImportError: 28 | # compatible with older versions of pytorch 29 | from torch.utils.collect_env import get_pretty_env_info 30 | 31 | return get_pretty_env_info() 32 | 33 | 34 | def detect_compute_compatibility(CUDA_HOME, so_file): 35 | try: 36 | cuobjdump = os.path.join(CUDA_HOME, "bin", "cuobjdump") 37 | if os.path.isfile(cuobjdump): 38 | output = subprocess.check_output( 39 | "'{}' --list-elf '{}'".format(cuobjdump, so_file), shell=True 40 | ) 41 | output = output.decode("utf-8").strip().split("\n") 42 | sm = [] 43 | for line in output: 44 | line = re.findall(r"\.sm_[0-9]*\.", line)[0] 45 | sm.append(line.strip(".")) 46 | sm = sorted(set(sm)) 47 | return ", ".join(sm) 48 | else: 49 | return so_file + "; cannot find cuobjdump" 50 | except Exception: 51 | # unhandled failure 52 | return so_file 53 | 54 | 55 | def collect_env_info(): 56 | has_gpu = torch.cuda.is_available() # true for both CUDA & ROCM 57 | torch_version = torch.__version__ 58 | 59 | # NOTE: the use of CUDA_HOME and ROCM_HOME requires the CUDA/ROCM build deps, though in 60 | # theory detectron2 should be made runnable with only the corresponding runtimes 61 | from torch.utils.cpp_extension import CUDA_HOME 62 | 63 | has_rocm = False 64 | if tuple(map(int, torch_version.split(".")[:2])) >= (1, 5): 65 | from torch.utils.cpp_extension import ROCM_HOME 66 | 67 | if (getattr(torch.version, "hip", None) is not None) and (ROCM_HOME is not None): 68 | has_rocm = True 69 | has_cuda = has_gpu and (not has_rocm) 70 | 71 | data = [] 72 | data.append(("sys.platform", sys.platform)) 73 | data.append(("Python", sys.version.replace("\n", ""))) 74 | data.append(("numpy", np.__version__)) 75 | 76 | # print system compilers when extension fails to build 77 | if sys.platform != "win32": # don't know what to do for windows 78 | try: 79 | # this is how torch/utils/cpp_extensions.py choose compiler 80 | cxx = os.environ.get("CXX", "c++") 81 | cxx = subprocess.check_output("'{}' --version".format(cxx), shell=True) 82 | cxx = cxx.decode("utf-8").strip().split("\n")[0] 83 | except subprocess.SubprocessError: 84 | cxx = "Not found" 85 | data.append(("Compiler", cxx)) 86 | 87 | if has_cuda and CUDA_HOME is not None: 88 | try: 89 | nvcc = os.path.join(CUDA_HOME, "bin", "nvcc") 90 | nvcc = subprocess.check_output("'{}' -V".format(nvcc), shell=True) 91 | nvcc = nvcc.decode("utf-8").strip().split("\n")[-1] 92 | except subprocess.SubprocessError: 93 | nvcc = "Not found" 94 | data.append(("CUDA compiler", nvcc)) 95 | 96 | data.append(("PyTorch", torch_version + " @" + os.path.dirname(torch.__file__))) 97 | data.append(("PyTorch debug build", torch.version.debug)) 98 | 99 | data.append(("GPU available", has_gpu)) 100 | if has_gpu: 101 | devices = defaultdict(list) 102 | for k in range(torch.cuda.device_count()): 103 | devices[torch.cuda.get_device_name(k)].append(str(k)) 104 | for name, devids in devices.items(): 105 | data.append(("GPU " + ",".join(devids), name)) 106 | 107 | if has_rocm: 108 | data.append(("ROCM_HOME", str(ROCM_HOME))) 109 | else: 110 | data.append(("CUDA_HOME", str(CUDA_HOME))) 111 | 112 | cuda_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST", None) 113 | if cuda_arch_list: 114 | data.append(("TORCH_CUDA_ARCH_LIST", cuda_arch_list)) 115 | data.append(("Pillow", PIL.__version__)) 116 | 117 | try: 118 | data.append( 119 | ( 120 | "torchvision", 121 | str(torchvision.__version__) + " @" + os.path.dirname(torchvision.__file__), 122 | ) 123 | ) 124 | if has_cuda: 125 | try: 126 | torchvision_C = importlib.util.find_spec("torchvision._C").origin 127 | msg = detect_compute_compatibility(CUDA_HOME, torchvision_C) 128 | data.append(("torchvision arch flags", msg)) 129 | except ImportError: 130 | data.append(("torchvision._C", "failed to find")) 131 | except AttributeError: 132 | data.append(("torchvision", "unknown")) 133 | 134 | try: 135 | import fvcore 136 | 137 | data.append(("fvcore", fvcore.__version__)) 138 | except ImportError: 139 | pass 140 | 141 | try: 142 | import cv2 143 | data.append(("cv2", cv2.__version__)) 144 | except ImportError: 145 | pass 146 | 147 | try: 148 | import pytorch_lightning as pl 149 | data.append(("Pyotrch Lighning", pl.__version__)) 150 | except ImportError: 151 | pass 152 | 153 | env_str = tabulate(data) + "\n" 154 | env_str += collect_torch_env() 155 | return env_str 156 | 157 | 158 | if __name__ == "__main__": 159 | 160 | print(collect_env_info()) 161 | -------------------------------------------------------------------------------- /Codes/src/core/schedulers/cyclic_cos_annealing_lr.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | ''' 3 | Code adapted from 4 | https://github.com/bluesky314/Cyclical_LR_Scheduler_With_Decay_Pytorch 5 | ''' 6 | 7 | import math 8 | from bisect import bisect_left, bisect_right 9 | 10 | import numpy as np 11 | import torch 12 | from torch.optim.lr_scheduler import _LRScheduler 13 | from torch.optim.optimizer import Optimizer 14 | 15 | __all__ = ['CyclicCosAnnealingLR', 'CyclicLinearLR'] 16 | 17 | 18 | class CyclicCosAnnealingLR(_LRScheduler): 19 | r""" 20 | 21 | Implements reset on milestones inspired from CosineAnnealingLR pytorch 22 | 23 | Set the learning rate of each parameter group using a cosine annealing 24 | schedule, where :math:`\eta_{max}` is set to the initial lr and 25 | :math:`T_{cur}` is the number of epochs since the last restart in SGDR: 26 | .. math:: 27 | \eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})(1 + 28 | \cos(\frac{T_{cur}}{T_{max}}\pi)) 29 | When last_epoch > last set milestone, lr is automatically set to \eta_{min} 30 | It has been proposed in 31 | `SGDR: Stochastic Gradient Descent with Warm Restarts`_. Note that this only 32 | implements the cosine annealing part of SGDR, and not the restarts. 33 | Args: 34 | optimizer (Optimizer): Wrapped optimizer. 35 | milestones (list of ints): List of epoch indices. Must be increasing. 36 | decay_milestones(list of ints):List of increasing epoch indices. Ideally,decay values should overlap with milestone points 37 | gamma (float): factor by which to decay the max learning rate at each decay milestone 38 | eta_min (float): Minimum learning rate. Default: 1e-6 39 | last_epoch (int): The index of last epoch. Default: -1. 40 | 41 | 42 | .. _SGDR\: Stochastic Gradient Descent with Warm Restarts: 43 | https://arxiv.org/abs/1608.03983 44 | """ 45 | 46 | def __init__(self, optimizer, milestones, decay_milestones=None, gamma=0.5, eta_min=1e-6, last_epoch=-1): 47 | if not list(milestones) == sorted(milestones): 48 | raise ValueError('Milestones should be a list of' 49 | ' increasing integers. Got {}', milestones) 50 | self.eta_min = eta_min 51 | self.milestones = milestones 52 | self.milestones2 = decay_milestones 53 | 54 | self.gamma = gamma 55 | super(CyclicCosAnnealingLR, self).__init__(optimizer, last_epoch) 56 | 57 | def get_lr(self): 58 | 59 | if self.last_epoch >= self.milestones[-1]: 60 | return [self.eta_min for base_lr in self.base_lrs] 61 | 62 | idx = bisect_right(self.milestones, self.last_epoch) 63 | 64 | left_barrier = 0 if idx == 0 else self.milestones[idx - 1] 65 | right_barrier = self.milestones[idx] 66 | 67 | width = right_barrier - left_barrier 68 | curr_pos = self.last_epoch - left_barrier 69 | 70 | if self.milestones2: 71 | return [self.eta_min + (base_lr * self.gamma ** bisect_right(self.milestones2, self.last_epoch) - self.eta_min) * 72 | (1 + math.cos(math.pi * curr_pos / width)) / 2 73 | for base_lr in self.base_lrs] 74 | else: 75 | return [self.eta_min + (base_lr - self.eta_min) * 76 | (1 + math.cos(math.pi * curr_pos / width)) / 2 77 | for base_lr in self.base_lrs] 78 | 79 | 80 | class CyclicLinearLR(_LRScheduler): 81 | r""" 82 | Implements reset on milestones inspired from Linear learning rate decay 83 | 84 | Set the learning rate of each parameter group using a linear decay 85 | schedule, where :math:`\eta_{max}` is set to the initial lr and 86 | :math:`T_{cur}` is the number of epochs since the last restart: 87 | .. math:: 88 | \eta_t = \eta_{min} + (\eta_{max} - \eta_{min})(1 -\frac{T_{cur}}{T_{max}}) 89 | When last_epoch > last set milestone, lr is automatically set to \eta_{min} 90 | 91 | Args: 92 | optimizer (Optimizer): Wrapped optimizer. 93 | milestones (list of ints): List of epoch indices. Must be increasing. 94 | decay_milestones(list of ints):List of increasing epoch indices. Ideally,decay values should overlap with milestone points 95 | gamma (float): factor by which to decay the max learning rate at each decay milestone 96 | eta_min (float): Minimum learning rate. Default: 1e-6 97 | last_epoch (int): The index of last epoch. Default: -1. 98 | .. _SGDR\: Stochastic Gradient Descent with Warm Restarts: 99 | https://arxiv.org/abs/1608.03983 100 | """ 101 | 102 | def __init__(self, optimizer, milestones, decay_milestones=None, gamma=0.5, eta_min=1e-6, last_epoch=-1): 103 | if not list(milestones) == sorted(milestones): 104 | raise ValueError('Milestones should be a list of' 105 | ' increasing integers. Got {}', milestones) 106 | self.eta_min = eta_min 107 | 108 | self.gamma = gamma 109 | self.milestones = milestones 110 | self.milestones2 = decay_milestones 111 | super(CyclicLinearLR, self).__init__(optimizer, last_epoch) 112 | 113 | def get_lr(self): 114 | 115 | if self.last_epoch >= self.milestones[-1]: 116 | return [self.eta_min for base_lr in self.base_lrs] 117 | 118 | idx = bisect_right(self.milestones, self.last_epoch) 119 | 120 | left_barrier = 0 if idx == 0 else self.milestones[idx - 1] 121 | right_barrier = self.milestones[idx] 122 | 123 | width = right_barrier - left_barrier 124 | curr_pos = self.last_epoch - left_barrier 125 | 126 | if self.milestones2: 127 | return [self.eta_min + (base_lr * self.gamma ** bisect_right(self.milestones2, self.last_epoch) - self.eta_min) * 128 | (1. - 1.0 * curr_pos / width) 129 | for base_lr in self.base_lrs] 130 | 131 | else: 132 | return [self.eta_min + (base_lr - self.eta_min) * 133 | (1. - 1.0 * curr_pos / width) 134 | for base_lr in self.base_lrs] 135 | -------------------------------------------------------------------------------- /Codes/src/core/models/segmentation_decoder/ftnet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | ##################################################################################################################################################################### 4 | # FTNet # 5 | # Copyright 2020 Tufts University. # # 6 | # Please see LICENSE file for full terms. # # # 7 | ##################################################################################################################################################################### 8 | 9 | import os 10 | import inspect 11 | import sys 12 | import torch 13 | import torch.nn as nn 14 | import torch.nn.functional as F 15 | from core.models.segbase import SegBaseModel, initialize_weights 16 | from core.models.segmentation_decoder.dcc_net_decoder import FeatureTransverseDecoder, BasicBlock 17 | import logging 18 | import numpy as np 19 | logger = logging.getLogger('pytorch_lightning') 20 | 21 | __all__ = ['FNet', 'get_ftnet'] 22 | 23 | ''' 24 | E0 [1] - 0 25 | E1 [2] - 1 26 | E2 [3] - 2 27 | E3 [4] - 3 28 | ''' 29 | 30 | 31 | class FNet(SegBaseModel): 32 | def __init__(self, nclass, 33 | backbone='resnet50', 34 | pretrained_base=True, 35 | no_of_filters=32, 36 | edge_extracts=None, 37 | num_blocks=None, 38 | dilated=False, 39 | ** kwargs): 40 | 41 | super(FNet, self).__init__(nclass=nclass, 42 | backbone=backbone, 43 | pretrained_base=pretrained_base, 44 | dilated=dilated, 45 | **kwargs) 46 | 47 | num_branches = 4 48 | numblocks = [num_blocks] * num_branches 49 | num_channels = [no_of_filters] * num_branches 50 | edge_extract = [x - 1 for x in edge_extracts] 51 | 52 | if dilated: 53 | dilations = [1, 1, 2, 4] 54 | else: 55 | dilations = [1, 1, 1, 1] 56 | 57 | self.decoder = FeatureTransverseDecoder(num_branches=num_branches, 58 | blocks=BasicBlock, 59 | num_blocks=numblocks, 60 | num_inchannels=self.encoder.feature_list[1:], 61 | num_channels=num_channels, 62 | dilation=dilations, 63 | edge_extract=edge_extract) 64 | 65 | self.last_layer = nn.Sequential(nn.Conv2d(in_channels=np.sum(num_channels) + 1, 66 | out_channels=np.sum(num_channels) // 2, 67 | kernel_size=1, 68 | stride=1, 69 | padding=0), 70 | nn.BatchNorm2d(np.sum(num_channels) // 2, momentum=0.01), 71 | nn.ReLU(inplace=True), 72 | nn.Conv2d(in_channels=np.sum(num_channels) // 2, 73 | out_channels=nclass, 74 | kernel_size=1, 75 | stride=1, 76 | padding=0) 77 | ) 78 | 79 | self.__setattr__('exclusive', ['decoder', 'last_layer']) 80 | 81 | if pretrained_base: 82 | for name, mod in self.named_children(): 83 | if name != 'encoder': 84 | for m in mod.modules(): 85 | initialize_weights(m) 86 | else: 87 | for m in self.modules(): 88 | initialize_weights(m) 89 | 90 | def forward(self, x): 91 | 92 | _, _, h, w = x.shape 93 | out, edge = self.decoder(self.base_forward(x, multiscale=False), h, w) 94 | out.append(edge) 95 | 96 | class_maps = self.last_layer(torch.cat(out, 1)) 97 | return (class_maps, edge) 98 | 99 | 100 | def get_ftnet(dataset='soda', 101 | backbone='resnet50', 102 | root=None, 103 | pretrained_base=False, 104 | **kwargs): 105 | 106 | from datasets import datasets 107 | model = FNet(nclass=datasets[dataset.lower()].NUM_CLASS, 108 | backbone=backbone, 109 | pretrained_base=pretrained_base, 110 | **kwargs) 111 | return model 112 | 113 | 114 | if __name__ == '__main__': 115 | from core.models.segbase import print_network 116 | img = torch.randint(3, 5, (2, 3, 512, 512)).float() 117 | 118 | model = get_ftnet(pretrained_base=False, edge_extracts=[3], num_blocks=4, no_of_filters=32, backbone='resnext101_32x8d') 119 | print_network(model) 120 | 121 | outputs = model(img) 122 | from ptflops import get_model_complexity_info 123 | 124 | macs, params = get_model_complexity_info(model, (3, 640, 480), as_strings=True, 125 | print_per_layer_stat=False, verbose=True) 126 | print('{:<30} {:<8}'.format('Computational complexity: ', macs)) 127 | print('{:<30} {:<8}'.format('Number of parameters: ', params)) 128 | -------------------------------------------------------------------------------- /Codes/src/core/utils/logger.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Adapted from 3 | https://detectron2.readthedocs.io/en/latest/_modules/detectron2/utils/logger.html 4 | ''' 5 | 6 | import functools 7 | import logging 8 | import os 9 | import sys 10 | import time 11 | from collections import Counter 12 | from logging.handlers import RotatingFileHandler 13 | 14 | from tabulate import tabulate 15 | from termcolor import colored 16 | from pytorch_lightning.utilities import rank_zero_only, rank_zero_warn 17 | 18 | __all__ = ['setup_logger'] 19 | 20 | 21 | class _ColorfulFormatter(logging.Formatter): 22 | def __init__(self, *args, **kwargs): 23 | self._root_name = kwargs.pop("root_name") + "." 24 | self._abbrev_name = kwargs.pop("abbrev_name", "") 25 | if len(self._abbrev_name): 26 | self._abbrev_name = self._abbrev_name + "." 27 | super(_ColorfulFormatter, self).__init__(*args, **kwargs) 28 | 29 | def formatMessage(self, record): 30 | record.name = record.name.replace(self._root_name, self._abbrev_name) 31 | log = super(_ColorfulFormatter, self).formatMessage(record) 32 | if record.levelno == logging.WARNING: 33 | prefix = colored("WARNING", "red", attrs=["blink"]) 34 | elif record.levelno == logging.ERROR: 35 | prefix = colored("ERROR", "red", attrs=["blink", "underline"]) 36 | elif record.levelno == logging.CRITICAL: 37 | prefix = colored("CRITICAL", "blue", attrs=["blink", "underline"]) 38 | else: 39 | return log 40 | return prefix + " " + log 41 | 42 | 43 | @rank_zero_only 44 | @functools.lru_cache() # so that calling setup_logger multiple times won't add many handlers 45 | def setup_logger(save_dir: str, 46 | name: str = 'iDINE', 47 | distributed_rank: int = 0, 48 | filename: str = "log.txt", 49 | mode: str = 'w', 50 | color: bool = True, 51 | abbrev_name: str = None, 52 | print_to_console: bool = True): 53 | 54 | logger = logging.getLogger(name) 55 | logger.setLevel(logging.DEBUG) 56 | logger.propagate = False 57 | # don't log results for the non-master process 58 | 59 | if distributed_rank > 0: 60 | return logger 61 | 62 | if abbrev_name is None: 63 | abbrev_name = "SG" if name == "segmentation" else name 64 | 65 | plain_formatter = logging.Formatter( 66 | "[%(asctime)s] %(name)s %(levelname)s: %(message)s", datefmt="%m/%d %H:%M:%S") 67 | 68 | if print_to_console: 69 | ch = logging.StreamHandler(stream=sys.stdout) 70 | ch.setLevel(logging.DEBUG) 71 | if color: 72 | formatter = _ColorfulFormatter( 73 | colored("[%(asctime)s %(name)s]: ", "green") + "%(message)s", 74 | datefmt="%m/%d %H:%M:%S", 75 | root_name=name, 76 | abbrev_name=str(abbrev_name), 77 | ) 78 | else: 79 | formatter = plain_formatter 80 | ch.setFormatter(formatter) 81 | logger.addHandler(ch) 82 | 83 | if save_dir: 84 | if not os.path.exists(save_dir): 85 | os.makedirs(save_dir) 86 | 87 | fh = logging.FileHandler(os.path.join(save_dir, filename), mode=mode) # 'a+' for add, 'w' for overwrite 88 | fh.setLevel(logging.DEBUG) 89 | fh.setFormatter(plain_formatter) 90 | logger.addHandler(fh) 91 | 92 | return logger 93 | 94 | 95 | # @functools.lru_cache() # so that calling setup_logger multiple times won't add many handlers 96 | # def setup_logger(save_dir: str, 97 | # name: str = 'iDINE', 98 | # distributed_rank: int = 0, 99 | # level=logging.DEBUG, 100 | # filename: str = "log.txt", 101 | # mode: str = 'w', 102 | # color: bool = True, 103 | # log_size_mb=1000, 104 | # abbrev_name: str = None, 105 | # num_log_archives=10, 106 | # stdout: bool = True): 107 | 108 | # # Get the Lightning logger and add handlers/formatter 109 | # if not stdout and filename is None: 110 | # raise ValueError('ConsoleLogger will have no handlers if stdout=False and file=None') 111 | 112 | # logger = logging.getLogger(name) 113 | # logger.setLevel(level) 114 | # logger.propagate = False 115 | 116 | # if distributed_rank > 0: 117 | # return logger 118 | 119 | # if abbrev_name is None: 120 | # abbrev_name = "SG" if name == "segmentation" else name 121 | 122 | # plain_formatter = logging.Formatter( 123 | # "[%(asctime)s] %(name)s %(levelname)s: %(message)s", datefmt="%m/%d %H:%M:%S") 124 | 125 | # if stdout: 126 | # ch = logging.StreamHandler(stream=sys.stdout) 127 | # ch.setLevel(logging.DEBUG) 128 | # if color: 129 | # formatter = _ColorfulFormatter( 130 | # colored("[%(asctime)s %(name)s]: ", "green") + "%(message)s", 131 | # datefmt="%m/%d %H:%M:%S", 132 | # root_name=name, 133 | # abbrev_name=str(abbrev_name), 134 | # ) 135 | # else: 136 | # formatter = plain_formatter 137 | # ch.setFormatter(formatter) 138 | # logger.addHandler(ch) 139 | 140 | # if save_dir: 141 | # if not os.path.exists(save_dir): 142 | # os.makedirs(save_dir) 143 | 144 | # # should_roll_over = os.path.isfile(os.path.join(save_dir, filename)) 145 | # # if should_roll_over: 146 | # # fh = RotatingFileHandler(os.path.join(save_dir, filename), mode='w', encoding='utf-8', backupCount=5) 147 | # # fh.doRollover() 148 | # # else: 149 | # file_handler = RotatingFileHandler(os.path.join(save_dir, filename), 150 | # maxBytes=log_size_mb * 1024 * 1024, 151 | # backupCount=num_log_archives) 152 | # file_handler.setFormatter(logging.Formatter( 153 | # "[%(asctime)s] %(name)s %(levelname)s: %(message)s", datefmt="%m/%d %H:%M:%S") 154 | # ) 155 | # logger.addHandler(file_handler) 156 | # logger.debug('Initialized ConsoleLogger') 157 | 158 | # return logger 159 | -------------------------------------------------------------------------------- /Codes/src/datasets/utils/scutseg_foldermap.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # To distribute train, valid, test for SODA 3 | # ============================================================================= 4 | 5 | import os 6 | import pandas as pd 7 | from sklearn.model_selection import train_test_split 8 | from glob import glob 9 | import shutil 10 | import argparse 11 | 12 | 13 | def is_image_file(filename): 14 | return any( 15 | filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg", ".bmp"] 16 | ) 17 | 18 | 19 | def str2bool(v): 20 | if v.lower() in ("yes", "true", "t", "y", "1"): 21 | return True 22 | elif v.lower() in ("no", "false", "f", "n", "0"): 23 | return False 24 | else: 25 | raise argparse.ArgumentTypeError("Boolean value expected.") 26 | 27 | 28 | # %% 29 | # ============================================================================= 30 | # Copying files to different folders 31 | # ============================================================================= 32 | def copying(tiles, path_label, basepath, fileset_path): 33 | tiles.set_index(path_label, inplace=True) 34 | # count=0 35 | for img_path in tiles.index: 36 | print("Path = {}".format(img_path)) 37 | dst_path = os.path.join(basepath, fileset_path) 38 | 39 | shutil.copy(img_path, dst_path) 40 | 41 | 42 | # %% 43 | # ============================================================================= 44 | # Creating folders 45 | # ============================================================================= 46 | 47 | 48 | def distribute(input_dir, output_dir, reset): 49 | # basepath = './Dataset/SCUT/' 50 | basepath = input_dir 51 | 52 | train_Name_path = basepath + "scut_train.txt" 53 | val_Name_path = basepath + "scut_val.txt" 54 | Image_path = basepath + "images/" 55 | Mask_path = basepath + "gt_instance/" 56 | 57 | # os.path.abspath(os.path.join(basepath,'..')) 58 | base_dir = os.path.join(output_dir, "SCUTSEG") 59 | if reset == True and os.path.exists(base_dir): 60 | if os.path.exists(base_dir): 61 | shutil.rmtree(base_dir) 62 | if not os.path.exists(base_dir): 63 | os.mkdir(base_dir) 64 | 65 | main_dirs_image = ["image/train", "image/val"] 66 | main_dirs_mask = ["mask/train", "mask/val"] 67 | 68 | for main in main_dirs_image: 69 | 70 | path = os.path.join(base_dir, main) 71 | if not os.path.exists(path): 72 | os.makedirs(path) 73 | 74 | for main in main_dirs_mask: 75 | 76 | path = os.path.join(base_dir, main) 77 | if not os.path.exists(path): 78 | os.makedirs(path) 79 | # %% 80 | # ============================================================================= 81 | # Creating folders 82 | # ============================================================================= 83 | 84 | imageid_path_dict = { 85 | os.path.splitext(os.path.basename(x))[0]: x 86 | for x in glob(os.path.join(Image_path, "**", "**", "*.jpg"), recursive=True) 87 | } 88 | maskid_path_dict = { 89 | os.path.splitext(os.path.basename(x))[0]: x 90 | for x in glob(os.path.join(Mask_path, "**", "**", "*.jpg"), recursive=True) 91 | } 92 | 93 | train_df = pd.read_table( 94 | train_Name_path, delim_whitespace=True, names=("Image_subpath", "Mask_subpath") 95 | ) 96 | train_df = train_df.sort_values( 97 | by="Image_subpath", axis=0, ascending=True, kind="quicksort" 98 | ).reset_index(drop=True) 99 | train_df = train_df.fillna("NA") 100 | 101 | train_df["Image_Name"] = train_df["Image_subpath"].apply( 102 | lambda x: os.path.splitext(os.path.basename(x))[0] 103 | ) 104 | train_df["Image_Path"] = train_df["Image_Name"].map(imageid_path_dict.get) 105 | train_df["Mask_Path"] = train_df["Image_Name"].map(maskid_path_dict.get) 106 | 107 | train_df["Mask_Path"] = train_df["Mask_Path"].apply( 108 | lambda x: x.split(".")[0] + "_labelIds.png" 109 | ) 110 | train_df1 = train_df.copy(deep=True) 111 | 112 | copying( 113 | tiles=train_df, 114 | path_label="Image_Path", 115 | basepath=base_dir, 116 | fileset_path=main_dirs_image[0], 117 | ) 118 | copying( 119 | tiles=train_df1, 120 | path_label="Mask_Path", 121 | basepath=base_dir, 122 | fileset_path=main_dirs_mask[0], 123 | ) 124 | 125 | validation_df = pd.read_table( 126 | val_Name_path, delim_whitespace=True, names=("Image_subpath", "Mask_subpath") 127 | ) 128 | validation_df = validation_df.sort_values( 129 | by="Image_subpath", axis=0, ascending=True, kind="quicksort" 130 | ).reset_index(drop=True) 131 | validation_df = validation_df.fillna("NA") 132 | 133 | validation_df["Image_Name"] = validation_df["Image_subpath"].apply( 134 | lambda x: os.path.splitext(os.path.basename(x))[0] 135 | ) 136 | validation_df["Image_Path"] = validation_df["Image_Name"].map(imageid_path_dict.get) 137 | validation_df["Mask_Path"] = validation_df["Image_Name"].map(maskid_path_dict.get) 138 | # validation_df['Mask_Path'] = validation_df['Mask_Path'].apply(lambda x: x+'_labelIds') 139 | validation_df["Mask_Path"] = validation_df["Mask_Path"].apply( 140 | lambda x: x.split(".")[0] + "_labelIds.png" 141 | ) 142 | 143 | validation_df1 = validation_df.copy(deep=True) 144 | 145 | copying( 146 | tiles=validation_df, 147 | path_label="Image_Path", 148 | basepath=base_dir, 149 | fileset_path=main_dirs_image[1], 150 | ) 151 | copying( 152 | tiles=validation_df1, 153 | path_label="Mask_Path", 154 | basepath=base_dir, 155 | fileset_path=main_dirs_mask[1], 156 | ) 157 | 158 | 159 | if __name__ == "__main__": 160 | parser = argparse.ArgumentParser(description="Set up Cityscape Dataset") 161 | parser.add_argument( 162 | "--input-image-path", 163 | type=str, 164 | default="/mnt/1842213842211C4E/raw_dataset/SCUT-SEG/", 165 | help="Path to Scut-Seg Dataset", 166 | ) 167 | parser.add_argument( 168 | "--save-path", 169 | type=str, 170 | default="/mnt/1842213842211C4E/processed_dataset/", 171 | help="Path to save Scut-Seg Dataset", 172 | ) 173 | parser.add_argument( 174 | "--reset", type=bool, default=True, help="Path to Cityscape Dataset" 175 | ) 176 | 177 | args = parser.parse_args() 178 | distribute( 179 | input_dir=args.input_image_path, output_dir=args.save_path, reset=args.reset 180 | ) 181 | -------------------------------------------------------------------------------- /Codes/src/core/utils/visualize.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | from PIL import Image 5 | 6 | __all__ = ['get_color_pallete', 'print_iou', 'set_img_color', 7 | 'show_prediction', 'show_colorful_images', 'save_colorful_images'] 8 | 9 | cityspallete = [ 10 | 128, 64, 128, 11 | 244, 35, 232, 12 | 70, 70, 70, 13 | 102, 102, 156, 14 | 190, 153, 153, 15 | 153, 153, 153, 16 | 250, 170, 30, 17 | 220, 220, 0, 18 | 107, 142, 35, 19 | 152, 251, 152, 20 | 0, 130, 180, 21 | 220, 20, 60, 22 | 255, 0, 0, 23 | 0, 0, 142, 24 | 0, 0, 70, 25 | 0, 60, 100, 26 | 0, 80, 100, 27 | 0, 0, 230, 28 | 119, 11, 32, 29 | ] 30 | 31 | 32 | scutseg = {0: np.array([0, 0, 0]), 33 | 1: np.array([128, 64, 128]), 34 | 2: np.array([60, 20, 220]), 35 | 3: np.array([0, 0, 255]), 36 | 4: np.array([142, 0, 0]), 37 | 5: np.array([70, 0, 0]), 38 | 6: np.array([153, 153, 190]), 39 | 7: np.array([35, 142, 107]), 40 | 8: np.array([100, 60, 0]), 41 | 9: np.array([153, 153, 153])} 42 | 43 | 44 | soda = {0: np.array([0, 0, 0]), 45 | 1: np.array([128, 64, 128]), 46 | 2: np.array([244, 35, 232]), 47 | 3: np.array([70, 70, 70]), 48 | 4: np.array([102, 102, 156]), 49 | 5: np.array([190, 153, 153]), 50 | 6: np.array([153, 153, 153]), 51 | 7: np.array([250, 170, 30]), 52 | 8: np.array([220, 220, 0]), 53 | 9: np.array([107, 142, 35]), 54 | 10: np.array([0, 130, 180]), 55 | 11: np.array([220, 20, 60]), 56 | 12: np.array([60, 20, 220]), 57 | 13: np.array([0, 0, 255]), 58 | 14: np.array([142, 0, 0]), 59 | 15: np.array([70, 0, 0]), 60 | 16: np.array([153, 153, 190]), 61 | 17: np.array([35, 142, 107]), 62 | 18: np.array([100, 60, 0]), 63 | 19: np.array([255, 0, 0]), 64 | 20: np.array([0, 64, 128]) 65 | } 66 | 67 | 68 | mfn = {0: np.array([0, 0, 0]), 69 | 1: np.array([128, 64, 128]), 70 | 2: np.array([244, 35, 232]), 71 | 3: np.array([70, 70, 70]), 72 | 4: np.array([102, 102, 156]), 73 | 5: np.array([190, 153, 153]), 74 | 6: np.array([153, 153, 153]), 75 | 7: np.array([250, 170, 30]), 76 | 8: np.array([220, 220, 0]), } 77 | 78 | 79 | def print_iou(iu, mean_pixel_acc, class_names=None, show_no_back=False): 80 | n = iu.size 81 | lines = [] 82 | for i in range(n): 83 | if class_names is None: 84 | cls = 'Class %d:' % (i + 1) 85 | else: 86 | cls = '%d %s' % (i + 1, class_names[i]) 87 | # lines.append('%-8s: %.3f%%' % (cls, iu[i] * 100)) 88 | mean_IU = np.nanmean(iu) 89 | mean_IU_no_back = np.nanmean(iu[1:]) 90 | if show_no_back: 91 | lines.append('mean_IU: %.3f%% || mean_IU_no_back: %.3f%% || mean_pixel_acc: %.3f%%' % ( 92 | mean_IU * 100, mean_IU_no_back * 100, mean_pixel_acc * 100)) 93 | else: 94 | lines.append('mean_IU: %.3f%% || mean_pixel_acc: %.3f%%' % (mean_IU * 100, mean_pixel_acc * 100)) 95 | lines.append('=================================================') 96 | line = "\n".join(lines) 97 | 98 | print(line) 99 | 100 | 101 | def set_img_color(img, label, colors, background=0, show255=False): 102 | for i in range(len(colors)): 103 | if i != background: 104 | img[np.where(label == i)] = colors[i] 105 | if show255: 106 | img[np.where(label == 255)] = 255 107 | 108 | return img 109 | 110 | 111 | def show_prediction(img, pred, colors, background=0): 112 | im = np.array(img, np.uint8) 113 | set_img_color(im, pred, colors, background) 114 | out = np.array(im) 115 | 116 | return out 117 | 118 | 119 | def show_colorful_images(prediction, palettes): 120 | im = Image.fromarray(palettes[prediction.astype('uint8').squeeze()]) 121 | im.show() 122 | 123 | 124 | def save_colorful_images(prediction, filename, output_dir, palettes): 125 | ''' 126 | :param prediction: [B, H, W, C] 127 | ''' 128 | im = Image.fromarray(palettes[prediction.astype('uint8').squeeze()]) 129 | fn = os.path.join(output_dir, filename) 130 | out_dir = os.path.split(fn)[0] 131 | if not os.path.exists(out_dir): 132 | os.mkdir(out_dir) 133 | im.save(fn) 134 | 135 | 136 | def decode_segmap(label_mask, label_colours, n_classes): 137 | # label_colours = self.dataset.get_class_colors() 138 | r = label_mask.copy() 139 | g = label_mask.copy() 140 | b = label_mask.copy() 141 | for ll in range(0, n_classes): 142 | r[label_mask == ll] = label_colours[ll][0] 143 | g[label_mask == ll] = label_colours[ll][1] 144 | b[label_mask == ll] = label_colours[ll][2] 145 | rgb = np.zeros((label_mask.shape[0], label_mask.shape[1], 3)) 146 | rgb[:, :, 0] = r / 255.0 147 | rgb[:, :, 1] = g / 255.0 148 | rgb[:, :, 2] = b / 255.0 149 | return rgb * 255 150 | 151 | 152 | def get_color_pallete(npimg, dataset='idine'): 153 | """Visualize image. 154 | 155 | Parameters 156 | ---------- 157 | npimg : numpy.ndarray 158 | Single channel image with shape `H, W, 1`. 159 | dataset : str, default: 'idine' 160 | Returns 161 | ------- 162 | out_img : PIL.Image 163 | Image with color pallete 164 | """ 165 | if 'cityscapes' in dataset.lower(): 166 | out_img = Image.fromarray(npimg.astype('uint8')) 167 | out_img.putpalette(cityspallete) 168 | return out_img 169 | 170 | if dataset.lower() == 'soda': 171 | out_img = decode_segmap(npimg, label_colours=soda, n_classes=len(soda)) 172 | return out_img.astype(np.uint8) 173 | 174 | if dataset.lower() == 'mfn': 175 | out_img = decode_segmap(npimg, label_colours=mfn, n_classes=len(mfn)) 176 | return out_img.astype(np.uint8) 177 | 178 | if dataset.lower() == 'scutseg': 179 | out_img = decode_segmap(npimg, label_colours=scutseg, n_classes=len(scutseg)) 180 | return out_img.astype(np.uint8) 181 | 182 | out_img = Image.fromarray(npimg.astype('uint8')) 183 | out_img.putpalette(generic) 184 | return out_img 185 | 186 | 187 | def _generatepallete(num_cls): 188 | n = num_cls 189 | pallete = [0] * (n * 3) 190 | for j in range(0, n): 191 | lab = j 192 | pallete[j * 3 + 0] = 0 193 | pallete[j * 3 + 1] = 0 194 | pallete[j * 3 + 2] = 0 195 | i = 0 196 | while (lab > 0): 197 | pallete[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i)) 198 | pallete[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i)) 199 | pallete[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i)) 200 | i = i + 1 201 | lab >>= 3 202 | return pallete 203 | 204 | 205 | # mfn1 = _generatepallete(9) 206 | # soda1 = _generatepallete(21) 207 | generic = _generatepallete(1000) 208 | -------------------------------------------------------------------------------- /Codes/src/core/models/encoder/resnext.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | ''' 3 | Adapted from 4 | https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py 5 | ''' 6 | import os 7 | import torch.nn as nn 8 | import torch.utils.model_zoo as model_zoo 9 | 10 | os.environ['TORCH_HOME'] = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..', '..', 'pretrained_models')) 11 | os.makedirs(os.environ['TORCH_HOME'], exist_ok=True) 12 | 13 | root_pretrained_path = os.environ['TORCH_HOME'] 14 | __all__ = ['ResNext', 'resnext50_32x4d', 'resnext101_32x8d'] 15 | 16 | model_urls = { 17 | 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth', 18 | 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth', 19 | } 20 | 21 | 22 | class Bottleneck(nn.Module): 23 | expansion = 4 24 | 25 | def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, 26 | base_width=64, dilation=1, norm_layer=None, **kwargs): 27 | super(Bottleneck, self).__init__() 28 | width = int(planes * (base_width / 64.)) * groups 29 | 30 | self.conv1 = nn.Conv2d(inplanes, width, 1, bias=False) 31 | self.bn1 = norm_layer(width) 32 | self.conv2 = nn.Conv2d(width, width, 3, stride, dilation, dilation, groups, bias=False) 33 | self.bn2 = norm_layer(width) 34 | self.conv3 = nn.Conv2d(width, planes * self.expansion, 1, bias=False) 35 | self.bn3 = norm_layer(planes * self.expansion) 36 | self.relu = nn.ReLU(True) 37 | self.downsample = downsample 38 | self.stride = stride 39 | 40 | def forward(self, x): 41 | identity = x 42 | 43 | out = self.conv1(x) 44 | out = self.bn1(out) 45 | out = self.relu(out) 46 | 47 | out = self.conv2(out) 48 | out = self.bn2(out) 49 | out = self.relu(out) 50 | 51 | out = self.conv3(out) 52 | out = self.bn3(out) 53 | 54 | if self.downsample is not None: 55 | identity = self.downsample(x) 56 | 57 | out += identity 58 | out = self.relu(out) 59 | 60 | return out 61 | 62 | 63 | class ResNext(nn.Module): 64 | 65 | def __init__(self, block, layers, num_classes=1000, zero_init_residual=False, groups=1, 66 | width_per_group=64, dilated=False, norm_layer=nn.BatchNorm2d, **kwargs): 67 | super(ResNext, self).__init__() 68 | self.feature_list = [] 69 | self.inplanes = 64 70 | self.groups = groups 71 | self.base_width = width_per_group 72 | 73 | self.conv1 = nn.Conv2d(3, self.inplanes, 7, 2, 3, bias=False) 74 | self.feature_list.append(self.inplanes) 75 | self.bn1 = norm_layer(self.inplanes) 76 | self.relu = nn.ReLU(True) 77 | self.maxpool = nn.MaxPool2d(3, 2, 1) 78 | 79 | self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer) 80 | self.feature_list.append(self.inplanes) 81 | 82 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer) 83 | self.feature_list.append(self.inplanes) 84 | if dilated: 85 | self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2, norm_layer=norm_layer) 86 | self.feature_list.append(self.inplanes) 87 | self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4, norm_layer=norm_layer) 88 | self.feature_list.append(self.inplanes) 89 | else: 90 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2, norm_layer=norm_layer) 91 | self.feature_list.append(self.inplanes) 92 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2, norm_layer=norm_layer) 93 | self.feature_list.append(self.inplanes) 94 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 95 | self.fc = nn.Linear(512 * block.expansion, num_classes) 96 | 97 | for m in self.modules(): 98 | if isinstance(m, nn.Conv2d): 99 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 100 | elif isinstance(m, nn.BatchNorm2d): 101 | nn.init.constant_(m.weight, 1) 102 | nn.init.constant_(m.bias, 0) 103 | 104 | if zero_init_residual: 105 | for m in self.modules(): 106 | if isinstance(m, Bottleneck): 107 | nn.init.constant_(m.bn3.weight, 0) 108 | 109 | def _make_layer(self, block, planes, blocks, stride=1, dilation=1, norm_layer=nn.BatchNorm2d): 110 | downsample = None 111 | if stride != 1 or self.inplanes != planes * block.expansion: 112 | downsample = nn.Sequential( 113 | nn.Conv2d(self.inplanes, planes * block.expansion, 1, stride, bias=False), 114 | norm_layer(planes * block.expansion) 115 | ) 116 | 117 | layers = list() 118 | if dilation in (1, 2): 119 | layers.append(block(self.inplanes, planes, stride, downsample, self.groups, 120 | self.base_width, norm_layer=norm_layer)) 121 | elif dilation == 4: 122 | layers.append(block(self.inplanes, planes, stride, downsample, self.groups, 123 | self.base_width, dilation=2, norm_layer=norm_layer)) 124 | else: 125 | raise RuntimeError("=> unknown dilation size: {}".format(dilation)) 126 | self.inplanes = planes * block.expansion 127 | for _ in range(1, blocks): 128 | layers.append(block(self.inplanes, planes, groups=self.groups, base_width=self.base_width, 129 | dilation=dilation, norm_layer=norm_layer)) 130 | 131 | return nn.Sequential(*layers) 132 | 133 | def forward(self, x): 134 | x = self.conv1(x) 135 | x = self.bn1(x) 136 | x = self.relu(x) 137 | x = self.maxpool(x) 138 | 139 | x = self.layer1(x) 140 | x = self.layer2(x) 141 | x = self.layer3(x) 142 | x = self.layer4(x) 143 | 144 | x = self.avgpool(x) 145 | x = x.view(x.size(0), -1) 146 | x = self.fc(x) 147 | 148 | return x 149 | 150 | 151 | def resnext50_32x4d(pretrained=False, **kwargs): 152 | kwargs['groups'] = 32 153 | kwargs['width_per_group'] = 4 154 | model = ResNext(Bottleneck, [3, 4, 6, 3], **kwargs) 155 | if pretrained: 156 | state_dict = model_zoo.load_url(model_urls['resnext50_32x4d']) 157 | model.load_state_dict(state_dict) 158 | return model 159 | 160 | 161 | def resnext101_32x8d(pretrained=False, **kwargs): 162 | kwargs['groups'] = 32 163 | kwargs['width_per_group'] = 8 164 | model = ResNext(Bottleneck, [3, 4, 23, 3], **kwargs) 165 | if pretrained: 166 | state_dict = model_zoo.load_url(model_urls['resnext101_32x8d']) 167 | model.load_state_dict(state_dict) 168 | return model 169 | 170 | 171 | if __name__ == '__main__': 172 | model = resnext101_32x8d() 173 | -------------------------------------------------------------------------------- /Codes/src/core/utils/optimizer_scheduler_helper.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | import math 3 | import torch.optim as optim 4 | from core.optimizers.adabound import AdaBound 5 | from core.optimizers.adamw import AdamW 6 | from core.schedulers.cyclic_warm_restart import CyclicLRWithRestarts 7 | from core.schedulers.iteration_polyLR import IterationPolyLR 8 | from core.schedulers.lr_scheduler import WarmupMultiStepLR, WarmupPolyLR 9 | from core.schedulers.OneCycle_LR import OneCycle 10 | from core.schedulers.Poly_LR import Poly 11 | from torch.optim import lr_scheduler 12 | from torch.optim.lr_scheduler import _LRScheduler 13 | 14 | __all__ = ['make_optimizer', 'make_scheduler'] 15 | 16 | 17 | def make_optimizer(args, params, logger=None): 18 | 19 | if args.optimizer.lower() == 'sgd': 20 | optimizer_function = optim.SGD 21 | kwargs = {'momentum': args.momentum, 22 | 'nesterov': args.nesterov} 23 | if logger is not None: 24 | logger.info('Optimizer SGD is being used') 25 | 26 | elif args.optimizer.lower() == 'adam': 27 | optimizer_function = optim.Adam 28 | kwargs = { 29 | 'betas': (args.beta1, args.beta2), 30 | 'eps': args.epsilon 31 | } 32 | if logger is not None: 33 | logger.info('Optimizer ADAM is being used') 34 | 35 | elif args.optimizer.lower() == 'rmsprop': 36 | optimizer_function = optim.RMSprop 37 | kwargs = {'eps': args.epsilon} 38 | if logger is not None: 39 | logger.info('Optimizer RMSprop is being used') 40 | 41 | elif args.optimizer.lower() == 'adabound': 42 | optimizer_function = AdaBound 43 | kwargs = {'eps': args.epsilon, 44 | 'betas': (args.beta1, args.beta2)} 45 | if logger is not None: 46 | logger.info('Optimizer AdaBound is being used') 47 | 48 | elif args.optimizer.lower() == 'adamw': 49 | optimizer_function = optim.AdamW 50 | kwargs = {'eps': args.epsilon, 51 | 'betas': (args.beta1, args.beta2)} 52 | if logger is not None: 53 | logger.info('Optimizer AdamW is being used') 54 | 55 | if isinstance(params, list): 56 | if not any("weight_decay" in s for s in params): 57 | kwargs['weight_decay'] = args.weight_decay 58 | else: 59 | kwargs['lr'] = args.lr 60 | 61 | return optimizer_function(params, **kwargs) 62 | 63 | 64 | class ConstantLR(_LRScheduler): 65 | def __init__(self, optimizer, last_epoch=-1): 66 | super(ConstantLR, self).__init__(optimizer, last_epoch) 67 | 68 | def get_lr(self): 69 | return [base_lr for base_lr in self.base_lrs] 70 | 71 | 72 | def make_scheduler(args, optimizer, 73 | iters_per_epoch=None, 74 | logger=None, 75 | last_epoch=-1): 76 | 77 | if args.scheduler_type.lower() == 'step': 78 | scheduler = lr_scheduler.StepLR(optimizer, 79 | step_size=args.lr_decay, 80 | gamma=args.gamma, 81 | last_epoch=last_epoch 82 | ) 83 | scheduler.__setattr__('__interval__', 'epoch') 84 | if logger is not None: 85 | logger.info('Loading Step scheduler') 86 | 87 | elif args.scheduler_type.find('step') >= 0: 88 | milestones = args.scheduler_type.split('_') 89 | milestones.pop(0) 90 | milestones = list(map(lambda x: int(x), milestones)) 91 | scheduler = lr_scheduler.MultiStepLR( 92 | optimizer, 93 | milestones=milestones, 94 | gamma=args.gamma, 95 | last_epoch=last_epoch 96 | ) 97 | 98 | scheduler.__setattr__('__interval__', 'epoch') 99 | if logger is not None: 100 | logger.info('Loading Multi step scheduler ') 101 | 102 | # elif args.scheduler_type.lower() == 'reduce_on_plateau': 103 | # scheduler = lr_scheduler.ReduceLROnPlateau( 104 | # optimizer, 105 | # ) 106 | # scheduler.__setattr__('__interval__', 'epoch') 107 | 108 | elif args.scheduler_type.lower() == 'poly_warmstartup': 109 | # https://github.com/PyTorchLightning/pytorch-lightning/issues/1038 110 | scheduler = WarmupPolyLR(optimizer, 111 | power=0.9, 112 | epochs=args.epochs, 113 | steps_per_epoch=iters_per_epoch, 114 | warmup_factor=args.warmup_factor, 115 | warmup_iters=args.warmup_iters, 116 | warmup_method=args.warmup_method) 117 | 118 | scheduler.__setattr__('__interval__', 'step') 119 | if logger is not None: 120 | logger.info('Loading Warm Startup scheduler') 121 | 122 | elif args.scheduler_type.lower() == 'multistep_warmstartup': 123 | # ============================================================================= 124 | # https://github.com/Tony-Y/pytorch_warmup 125 | # ============================================================================= 126 | scheduler = WarmupMultiStepLR(optimizer, 127 | milestones=args.milestones, 128 | gamma=0.1, 129 | warmup_factor=args.warmup_factor, 130 | warmup_iters=args.warmup_iters, 131 | warmup_method=args.warmup_method) 132 | 133 | if logger is not None: 134 | logger.info('Loading multi step Warm Startup scheduler') 135 | warnings.warn('This is not in compliance with pytorch lightning') 136 | scheduler.__setattr__('__interval__', 'step') 137 | 138 | elif args.scheduler_type.lower() == 'onecycle': 139 | scheduler = lr_scheduler.OneCycleLR(optimizer, 140 | max_lr=args.lr, 141 | epochs=args.epochs, 142 | steps_per_epoch=iters_per_epoch, 143 | # pct_start=0.3, 144 | # anneal_strategy='cos', 145 | # cycle_momentum=True, 146 | # base_momentum=0.85, 147 | # max_momentum=0.95, 148 | # div_factor=25.0, 149 | # final_div_factor=10000.0, 150 | last_epoch=last_epoch) 151 | scheduler.__setattr__('__interval__', 'step') 152 | if logger is not None: 153 | logger.info('Loading OneCycle scheduler') 154 | 155 | else: 156 | scheduler = ConstantLR(optimizer) 157 | scheduler.__setattr__('__interval__', 'epoch') 158 | logger.info('Loading Constant scheduler') 159 | 160 | return scheduler 161 | -------------------------------------------------------------------------------- /Codes/src/datasets/dataloaders/mfn.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Code Adapted from 3 | https://github.com/dmlc/gluon-cv/blob/master/gluoncv/data/cityscapes.py 4 | ''' 5 | import os 6 | import cv2 7 | import torch 8 | import numpy as np 9 | from skimage.exposure import equalize_adapthist 10 | 11 | from PIL import Image 12 | from PIL import ImageFilter 13 | from torchvision import transforms 14 | 15 | 16 | try: 17 | from segbase import SegmentationDataset 18 | except: 19 | from datasets.dataloaders.segbase import SegmentationDataset 20 | 21 | 22 | class MFNDataset(SegmentationDataset): 23 | NUM_CLASS = 9 24 | IGNORE_INDEX = 255 25 | NAME = "MFNDataset" 26 | BASE_FOLDER = 'MFNDataset' 27 | 28 | def __init__(self, root='./Dataset/', 29 | split='train', 30 | base_size=1024, 31 | crop_size=720, 32 | mode=None, 33 | logger=None, 34 | transform=None, 35 | sobel_edges=False): 36 | 37 | root = os.path.join(root, self.BASE_FOLDER) 38 | super(MFNDataset, self).__init__(root, split, mode, base_size, crop_size, logger) 39 | 40 | self.sobel_edges = sobel_edges 41 | assert os.path.exists(self.root), "Error: data root path is wrong!" 42 | 43 | self.mean = [0.3954101, 0.3954101, 0.3954101] 44 | self.std = [0.07577764, 0.07577764, 0.07577764] 45 | 46 | self.images, self.mask_paths, self.edge_paths = _get_mfn_pairs(self.root, self.split, logger) 47 | assert (len(self.images) == len(self.mask_paths)) 48 | 49 | if len(self.images) == 0: 50 | raise RuntimeError("Found 0 images in subfolders of:" + root + "\n") 51 | 52 | def __getitem__(self, index): 53 | if type(index) == list or type(index) == tuple: 54 | index, scale = index 55 | input_size = self.crop_size[scale] 56 | else: 57 | input_size = None 58 | 59 | img = np.asarray(Image.open(self.images[index]))[:, :, 3] 60 | img = Image.fromarray(img).convert('RGB') 61 | 62 | if self.mode == 'test': 63 | img = self.normalize(img) 64 | return img, os.path.basename(self.images[index]) 65 | 66 | ''' For Training and validation purposes''' 67 | mask = Image.open(self.mask_paths[index]) 68 | 69 | if self.sobel_edges: 70 | id255 = np.where(mask == 255) 71 | no255_gt = np.array(mask) 72 | no255_gt[id255] = 0 73 | edge = cv2.Canny(no255_gt, 5, 5, apertureSize=7) 74 | edge = cv2.dilate(edge, self.edge_kernel) 75 | edge[edge == 255] = 1 76 | edge = Image.fromarray(edge) 77 | else: 78 | edge = Image.open(self.edge_paths[index]) 79 | 80 | # synchrosized transform 81 | if self.mode == 'train': 82 | img, mask, edge = self._sync_transform(img=img, mask=mask, edge=edge, crop_size=input_size) 83 | elif self.mode == 'val': 84 | img, mask, edge = self._val_sync_transform(img=img, mask=mask, edge=edge, crop_size=input_size) 85 | else: 86 | assert self.mode == 'testval' 87 | img, mask, edge = self._img_transform(img), self._mask_transform(mask), self._edge_transform(edge) 88 | 89 | # general resize, normalize and to Tensor 90 | img = self.normalize(img) 91 | return img, mask, edge, os.path.basename(self.images[index]) 92 | 93 | def normalize(self, img): 94 | img = self.im2double(np.array(img)) 95 | img = (img - self.mean) * np.reciprocal(self.std) 96 | img = self.np2Tensor(img).float() 97 | return img 98 | 99 | def _mask_transform(self, mask): 100 | return torch.LongTensor(np.array(mask).astype("int32")) 101 | 102 | def __len__(self): 103 | return len(self.images) 104 | 105 | @property 106 | def pred_offset(self): 107 | return 0 108 | 109 | @property 110 | def class_names(self,): 111 | return ["unlabeled", 112 | "car", 113 | "person", 114 | "bike", 115 | "curve", 116 | "car_stop", 117 | "guardrail", 118 | "color_cone", 119 | "bump"] 120 | 121 | 122 | def _get_mfn_pairs(folder, split='train', logger=None): 123 | def get_path_pairs(img_folder, mask_folder, edge_folder): 124 | img_folder = os.path.join(img_folder, split) 125 | mask_folder = os.path.join(mask_folder, split) 126 | edge_folder = os.path.join(edge_folder, split) 127 | img_paths = [] 128 | mask_paths = [] 129 | edge_paths = [] 130 | 131 | for root, _, files in os.walk(img_folder): 132 | for filename in files: 133 | if filename.endswith('.png'): 134 | 135 | imgpath = os.path.join(root, filename) 136 | maskpath = os.path.join(mask_folder, filename) 137 | edgepath = os.path.join(edge_folder, filename) 138 | 139 | if os.path.isfile(imgpath) and os.path.isfile(maskpath) and os.path.isfile(edgepath): 140 | img_paths.append(imgpath) 141 | mask_paths.append(maskpath) 142 | edge_paths.append(edgepath) 143 | else: 144 | print('cannot find the mask or image:', imgpath, maskpath) 145 | if logger is not None: 146 | logger.info('Found {} images in the folder {}'.format(len(img_paths), img_folder)) 147 | return img_paths, mask_paths, edge_paths 148 | 149 | if split in ('train', 'val', 'test'): 150 | img_folder = os.path.join(folder, 'image') 151 | mask_folder = os.path.join(folder, 'mask') 152 | edge_folder = os.path.join(folder, 'edges') 153 | 154 | img_paths, mask_paths, edge_paths = get_path_pairs(img_folder, mask_folder, edge_folder) 155 | return img_paths, mask_paths, edge_paths 156 | 157 | 158 | if __name__ == '__main__': 159 | from os import path, sys 160 | import matplotlib.pyplot as plt 161 | sys.path.append(path.dirname(path.dirname(path.abspath(__file__)))) 162 | sys.path.append(os.path.join(path.dirname(path.dirname(path.abspath(__file__))), '..')) 163 | from core.data.samplers import (make_batch_data_sampler, make_data_sampler, 164 | make_multiscale_batch_data_sampler) 165 | from core.utils.utils import plot_tensors 166 | from datasets import * 167 | from torch.utils.data import DataLoader 168 | from skimage import color, io 169 | 170 | data_kwargs = {"base_size": [520], "crop_size": [480]} 171 | 172 | train_dataset = MFNDataset( 173 | root="/mnt/ACF29FC2F29F8F68/Work/Deep_Learning/Thermal_Segmentation/Dataset/", 174 | split='test', 175 | mode='testval', 176 | **data_kwargs 177 | ) 178 | 179 | train_sampler = make_data_sampler( 180 | dataset=train_dataset, shuffle=True, distributed=False 181 | ) 182 | 183 | train_batch_sampler = make_multiscale_batch_data_sampler( 184 | sampler=train_sampler, 185 | batch_size=1, 186 | multiscale_step=1, 187 | scales=1) 188 | 189 | loader_random_sample = DataLoader( 190 | dataset=train_dataset, 191 | batch_sampler=train_batch_sampler, 192 | num_workers=0, 193 | pin_memory=True, 194 | ) 195 | 196 | x = train_dataset.__getitem__(200) 197 | -------------------------------------------------------------------------------- /Codes/src/core/callbacks/progressbar.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | ''' 3 | Adapted from 4 | https://github.com/zhutmost/neuralzip/blob/master/apputil/progressbar.py 5 | ''' 6 | from time import time 7 | 8 | import pytorch_lightning as pl 9 | 10 | # ============================================================================= 11 | # https://github.com/zhutmost/neuralzip/blob/master/apputil/progressbar.py 12 | # https://github.com/PyTorchLightning/pytorch-lightning/issues/4876 13 | # ============================================================================= 14 | 15 | 16 | class ProgressBar(pl.callbacks.ProgressBarBase): 17 | def __init__(self, logger, refresh_rate: int = 1): 18 | super().__init__() 19 | self._logger = logger 20 | self._refresh_rate = refresh_rate 21 | self._enabled = True 22 | 23 | # a time flag to indicate the beginning of an epoch 24 | self._time = 0 25 | 26 | @property 27 | def refresh_rate(self) -> int: 28 | return self._refresh_rate 29 | 30 | @property 31 | def is_enabled(self) -> bool: 32 | return self._enabled 33 | 34 | @property 35 | def is_disabled(self) -> bool: 36 | return not self.is_enabled 37 | 38 | def disable(self) -> None: 39 | # No need to disable the ProgressBar on processes with LOCAL_RANK != 1, because the 40 | # StreamHandler of logging is disabled on these processes. 41 | self._enabled = True 42 | 43 | def enable(self) -> None: 44 | self._enabled = True 45 | 46 | @staticmethod 47 | def _serialize_metrics(progressbar_log_dict, filter_fn=None): 48 | if filter_fn: 49 | progressbar_log_dict = {k: v for k, v in progressbar_log_dict.items() if filter_fn(k)} 50 | msg = '' 51 | for metric, value in progressbar_log_dict.items(): 52 | if 'v_num' not in metric: 53 | if type(value) is str: 54 | msg += f'[{metric}: {value}]' 55 | elif 'acc' in metric: 56 | msg += f'[{metric}: {value:.3%}]' 57 | else: 58 | msg += f'[{metric}: {value:f}]' 59 | return msg 60 | 61 | def on_train_start(self, trainer, pl_module): 62 | super().on_train_start(trainer, pl_module) 63 | self._logger.info(f'Trainer fit begins ... ' 64 | f'Current epoch: {trainer.current_epoch}, batch: {self.train_batch_idx}') 65 | 66 | def on_train_epoch_start(self, trainer, pl_module): 67 | super().on_train_epoch_start(trainer, pl_module) 68 | total_train_batches = self.total_train_batches 69 | total_val_batches = self.total_val_batches 70 | if total_train_batches != float('inf') and total_val_batches != float('inf'): 71 | # val can be checked multiple times per epoch 72 | val_checks_per_epoch = total_train_batches // trainer.val_check_batch 73 | total_val_batches = total_val_batches * val_checks_per_epoch 74 | total_batches = total_train_batches + total_val_batches 75 | self._logger.info(f'\n ' 76 | f'>>> >>> >>> >>> Epoch {trainer.current_epoch}, including {total_batches} batches ' 77 | f'[train: {total_train_batches} & val: {total_val_batches}] ' 78 | f'<<< <<< <<< <<<') 79 | self._time = time() 80 | 81 | def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): 82 | super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) 83 | if self.is_enabled and self.train_batch_idx % self.refresh_rate == 0: 84 | batch_time = (time() - self._time) / self.train_batch_idx 85 | msg = f'Train [Epoch {trainer.current_epoch}, ' \ 86 | f'Batch {self.train_batch_idx} / {self.total_train_batches}, {batch_time:.2f}s/it] ' 87 | msg += self._serialize_metrics(trainer.progress_bar_dict, 88 | filter_fn=lambda x: not x.startswith('val_') and not x.startswith('test_') and not x.startswith('train_')) 89 | self._logger.info(msg) 90 | 91 | def on_train_end(self, trainer, pl_module): 92 | super().on_train_end(trainer, pl_module) 93 | self._logger.info(f'Trainer fit ends.') 94 | 95 | def on_validation_start(self, trainer, pl_module): 96 | super().on_validation_start(trainer, pl_module) 97 | if not trainer.running_sanity_check: 98 | self._logger.info(f'\n ' 99 | f'>>> Validate step begins ... Epoch {trainer.current_epoch}, ' 100 | f'including {self.total_val_batches} batches') 101 | self._time = time() 102 | 103 | def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): 104 | super().on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) 105 | if self.is_enabled and self.val_batch_idx % self.refresh_rate == 0: 106 | batch_time = (time() - self._time) / self.val_batch_idx 107 | msg = f'Validate [Epoch {trainer.current_epoch}, ' \ 108 | f'Batch {self.val_batch_idx} / {self.total_val_batches}, {batch_time:.2f}s/it]' 109 | msg += self._serialize_metrics(trainer.progress_bar_dict, 110 | filter_fn=lambda x: x.startswith('val_') and x.endswith('_step')) 111 | self._logger.info(msg) 112 | 113 | def on_validation_end(self, trainer, pl_module): 114 | super().on_validation_end(trainer, pl_module) 115 | if not trainer.running_sanity_check: 116 | msg = 'Validate ends' 117 | msg += self._serialize_metrics(trainer.progress_bar_dict, 118 | filter_fn=lambda x: x.startswith('val_') and x.endswith('_epoch')) 119 | self._logger.info(msg) 120 | 121 | def on_test_start(self, trainer, pl_module): 122 | super().on_test_start(trainer, pl_module) 123 | self._logger.info(f'\n >>> >>> >>> >>> Test, ' 124 | f'including {self.total_test_batches} batches ' 125 | f'<<< <<< <<< <<<') 126 | self._time = time() 127 | 128 | def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): 129 | super().on_test_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) 130 | if self.is_enabled and self.test_batch_idx % self.refresh_rate == 0: 131 | batch_time = (time() - self._time) / self.test_batch_idx 132 | msg = f'Test [Batch {self.test_batch_idx} / {self.total_test_batches}, {batch_time:.2f}s/it]' 133 | msg += self._serialize_metrics(trainer.progress_bar_dict, 134 | filter_fn=lambda x: x.startswith('test_') and x.endswith('_step')) 135 | self._logger.info(msg) 136 | 137 | def on_test_end(self, trainer, pl_module): 138 | super().on_test_end(trainer, pl_module) 139 | msg = '>>> Test ends' 140 | msg += self._serialize_metrics(trainer.progress_bar_dict, 141 | filter_fn=lambda x: x.startswith('test_') and x.endswith('_epoch')) 142 | self._logger.info(msg + '\n') 143 | 144 | def on_sanity_check_start(self, trainer, pl_module): 145 | super().on_sanity_check_start(trainer, pl_module) 146 | self._logger.info('Validate set sanity check begins.') 147 | 148 | def on_sanity_check_end(self, trainer, pl_module): 149 | super().on_sanity_check_end(trainer, pl_module) 150 | self._logger.info('Validate set sanity check ends.') 151 | -------------------------------------------------------------------------------- /Codes/src/core/metrics/pl_segmentation_score.py: -------------------------------------------------------------------------------- 1 | """ 2 | Code Adapted from 3 | https://github.com/mseg-dataset/mseg-semantic 4 | 5 | Evaluation Metrics for Semantic Segmentation 6 | """ 7 | from typing import Tuple 8 | 9 | import numpy as np 10 | import torch 11 | import torch.nn.functional as F 12 | from torchmetrics import Metric 13 | from typing import Any, Optional 14 | 15 | 16 | __all__ = ['pl_IOU', 'intersectionAndUnionGPU', 'intersectionAndUnion'] 17 | 18 | 19 | class pl_IOU(Metric): 20 | def __init__( 21 | self, 22 | num_classes: int, 23 | ignore_index: Optional[int] = None, 24 | compute_on_step: bool = True, 25 | dist_sync_on_step: bool = False, 26 | process_group: Optional[Any] = None, 27 | ): 28 | super().__init__(compute_on_step=compute_on_step, 29 | dist_sync_on_step=dist_sync_on_step, 30 | process_group=process_group, 31 | ) 32 | 33 | self.ignore_index = ignore_index 34 | self.num_classes = num_classes 35 | self.add_state("area_intersection", default=torch.zeros(num_classes), dist_reduce_fx="sum") 36 | self.add_state("area_union", default=torch.zeros(num_classes), dist_reduce_fx="sum") 37 | self.add_state("area_target", default=torch.zeros(num_classes), dist_reduce_fx="sum") 38 | 39 | def update(self, preds: torch.Tensor, target: torch.Tensor): 40 | 41 | preds = torch.argmax(preds.long(), 1) 42 | assert (preds.dim() in [1, 2, 3]) 43 | preds = preds.view(-1) 44 | target = target.view(-1) 45 | assert preds.shape == target.shape 46 | preds[target == self.ignore_index] = self.ignore_index 47 | intersection = preds[preds == target] 48 | 49 | area_intersection = torch.histc(intersection, bins=self.num_classes, min=0, max=self.num_classes - 1) 50 | area_output = torch.histc(preds, bins=self.num_classes, min=0, max=self.num_classes - 1) 51 | area_target = torch.histc(target, bins=self.num_classes, min=0, max=self.num_classes - 1) 52 | 53 | self.area_union += area_output + area_target - area_intersection 54 | self.area_intersection += area_intersection 55 | self.area_target += area_target 56 | 57 | def compute_mean(self): 58 | mean_iou = torch.mean(self.area_intersection / (self.area_union + 1e-10)) 59 | mean_accuracy = torch.mean(self.area_intersection / (self.area_target + 1e-10)) 60 | all_accuracy = torch.sum(self.area_intersection) / torch.sum(self.area_target + 1e-10) 61 | return mean_iou, mean_accuracy, all_accuracy 62 | 63 | def compute(self): 64 | iou = self.area_intersection / (self.area_union + 1e-10) 65 | accuracy = self.area_intersection / (self.area_target + 1e-10) 66 | all_accuracy = torch.sum(self.area_intersection) / torch.sum(self.area_target + 1e-10) 67 | return iou, accuracy, all_accuracy 68 | 69 | 70 | def intersectionAndUnionGPU(output: torch.Tensor, 71 | target: torch.Tensor, 72 | K: int, 73 | ignore_index: int = -1, 74 | test: bool = False, 75 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 76 | """ 77 | Note output and target sizes are N or N * L or N * H * W 78 | Args: 79 | - output: Pytorch tensor represeting predicted label map, 80 | each value in range 0 to K - 1. 81 | - target: Pytorch tensor representing ground truth label map, 82 | each value in range 0 to K - 1. 83 | - K: integer number of possible classes 84 | - ignore_index: integer representing class index to ignore 85 | Returns: 86 | - area_intersection: 1d Pytorch tensor of length (K,) with counts 87 | for each of K classes, where pred & target matched 88 | - area_union: 1d Pytorch tensor of length (K,) with counts 89 | - area_target: 1d Pytorch tensor of length (K,) with bin counts 90 | for each of K classes, present in this GT label map. 91 | """ 92 | 93 | if test: 94 | if output.shape != target.shape: 95 | output = F.interpolate(output, 96 | size=(target.shape[1], target.shape[2]), 97 | mode='bilinear', 98 | align_corners=True) 99 | 100 | output = torch.argmax(output.long(), 1) 101 | assert (output.dim() in [1, 2, 3]) 102 | output = output.view(-1) 103 | target = target.view(-1) 104 | assert output.shape == target.shape 105 | output[target == ignore_index] = ignore_index 106 | intersection = output[output == target] 107 | area_intersection = torch.histc(intersection, bins=K, min=0, max=K - 1) 108 | area_output = torch.histc(output, bins=K, min=0, max=K - 1) 109 | area_target = torch.histc(target, bins=K, min=0, max=K - 1) 110 | area_union = area_output + area_target - area_intersection 111 | return area_intersection, area_union, area_target 112 | 113 | 114 | def intersectionAndUnion(output: np.ndarray, 115 | target: np.ndarray, 116 | K: int, 117 | ignore_index: int = 255 118 | ) -> Tuple[np.array, np.array, np.array]: 119 | """ 120 | Compute IoU on Numpy arrays on CPU. We will be reasoning about each 121 | matrix cell individually, so we can reshape (flatten) these arrays 122 | into column vectors and the evaluation result won’t change. Compare 123 | horizontally-corresponding cells. Wherever ground truth (target) 124 | pixels should be ignored, set prediction also to the ignore label. 125 | `intersection` represents values (class indices) in cells where 126 | output and target are identical. We bin such correct class indices. 127 | Note output and target sizes are N or N * L or N * H * W 128 | Args: 129 | - output: Numpy array represeting predicted label map, 130 | each value in range 0 to K - 1. 131 | - target: Numpy array representing ground truth label map, 132 | each value in range 0 to K - 1. 133 | - K: integer number of possible classes 134 | - ignore_index: integer representing class index to ignore 135 | Returns: 136 | - area_intersection: 1d Numpy array of length (K,) with counts 137 | for each of K classes, where pred & target matched 138 | - area_union: 1d Numpy array of length (K,) with counts 139 | - area_target: 1d Numpy array of length (K,) with bin counts 140 | for each of K classes, present in this GT label map. 141 | """ 142 | assert (output.ndim in [1, 2, 3]) 143 | assert output.shape == target.shape 144 | # flatten the tensors to 1d arrays 145 | output = output.reshape(output.size).copy() 146 | target = target.reshape(target.size) 147 | output[np.where(target == ignore_index)[0]] = 255 148 | intersection = output[np.where(output == target)[0]] 149 | # contains the number of samples in each bin. 150 | area_intersection, _ = np.histogram(intersection, bins=np.arange(K + 1)) 151 | area_output, _ = np.histogram(output, bins=np.arange(K + 1)) 152 | area_target, _ = np.histogram(target, bins=np.arange(K + 1)) 153 | area_union = area_output + area_target - area_intersection 154 | return area_intersection, area_union, area_target 155 | 156 | 157 | if __name__ == '__main__': 158 | 159 | target = torch.randint(-1, 1, (10, 25, 25)).unsqueeze(0) 160 | pred = torch.tensor(target.clone()) 161 | pred[2:5, 7:13, 9:15] = 1 - pred[2:5, 7:13, 9:15] 162 | 163 | # metric = SegmentationMetric(nclass=2) 164 | # pixAcc, miou = metric(pred, target) 165 | -------------------------------------------------------------------------------- /Codes/src/datasets/dataloaders/scutseg.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Code Adapted from 3 | https://github.com/dmlc/gluon-cv/blob/master/gluoncv/data/cityscapes.py 4 | ''' 5 | import os 6 | import torch 7 | import numpy as np 8 | from skimage.exposure import equalize_adapthist 9 | 10 | from PIL import Image 11 | from PIL import ImageFilter 12 | from torchvision import transforms 13 | try: 14 | from segbase import SegmentationDataset 15 | except: 16 | from datasets.dataloaders.segbase import SegmentationDataset 17 | 18 | import cv2 19 | 20 | 21 | class SCUTSEGDataset(SegmentationDataset): 22 | NUM_CLASS = 10 23 | IGNORE_INDEX = 255 24 | NAME = "SCUTSEG" 25 | BASE_FOLDER = 'SCUTSEG' 26 | 27 | def __init__(self, root='./', 28 | split='train', 29 | base_size=1024, 30 | crop_size=720, 31 | mode=None, 32 | logger=None, 33 | transform=None, 34 | 35 | sobel_edges=False): 36 | 37 | root = os.path.join(root, self.BASE_FOLDER) 38 | super(SCUTSEGDataset, self).__init__(root, split, mode, base_size, crop_size, logger) 39 | self.sobel_edges = sobel_edges 40 | 41 | assert os.path.exists(self.root), "Error: data root path is wrong!" 42 | 43 | # FOR same channel 44 | self.mean = [0.41213047, 0.42389206, 0.416051] 45 | self.std = [0.13611181, 0.13612076, 0.13611817] 46 | 47 | self.images, self.mask_paths, self.edge_paths = _get_scutseg_pairs(self.root, self.split, logger) 48 | assert (len(self.images) == len(self.mask_paths)) 49 | 50 | if len(self.images) == 0: 51 | raise RuntimeError("Found 0 images in subfolders of:" + root + "\n") 52 | 53 | self.valid_classes = [0, 7, 24, 25, 26, 27, 13, 21, 28, 17] 54 | self.class_map = dict(zip(self.valid_classes, range(10))) 55 | 56 | def __getitem__(self, index): 57 | if type(index) == list or type(index) == tuple: 58 | index, scale = index 59 | input_size = self.crop_size[scale] 60 | else: 61 | input_size = None 62 | 63 | img = Image.open(self.images[index]) 64 | if self.mode == 'test': 65 | img = self.normalize(img) 66 | return img, os.path.basename(self.images[index]) 67 | 68 | ''' For Training and validation purposes''' 69 | mask = Image.open(self.mask_paths[index]) 70 | 71 | if self.sobel_edges: 72 | id255 = np.where(mask == 255) 73 | no255_gt = np.array(mask) 74 | no255_gt[id255] = 0 75 | edge = cv2.Canny(no255_gt, 5, 5, apertureSize=7) 76 | edge = cv2.dilate(edge, self.edge_kernel) 77 | edge[edge == 255] = 1 78 | edge = Image.fromarray(edge) 79 | else: 80 | edge = Image.open(self.edge_paths[index]) 81 | 82 | # synchrosized transform 83 | if self.mode == 'train': 84 | img, mask, edge = self._sync_transform(img=img, mask=mask, edge=edge, crop_size=input_size) 85 | elif self.mode == 'val': 86 | img, mask, edge = self._val_sync_transform(img=img, mask=mask, edge=edge, crop_size=input_size) 87 | else: 88 | assert self.mode == 'testval' 89 | img, mask, edge = self._img_transform(img), self._mask_transform(mask), self._edge_transform(edge) 90 | 91 | # general resize, normalize and to Tensor 92 | img = self.normalize(img) 93 | return img, mask, edge, os.path.basename(self.images[index]) 94 | 95 | def encode_segmap(self, mask): 96 | for _validc in self.valid_classes: 97 | mask[mask == _validc] = self.class_map[_validc] 98 | return mask 99 | 100 | def normalize(self, img): 101 | img = self.im2double(np.array(img)) 102 | img = (img - self.mean) * np.reciprocal(self.std) 103 | img = self.np2Tensor(img).float() 104 | return img 105 | 106 | def _mask_transform(self, mask): 107 | mask = self.encode_segmap(np.array(mask)) 108 | return torch.LongTensor(mask.astype("int32")) 109 | 110 | def __len__(self): 111 | return len(self.images) 112 | 113 | @property 114 | def pred_offset(self): 115 | return 0 116 | 117 | @property 118 | def class_names(self,): 119 | return ["background", 120 | "road", 121 | "person", 122 | "rider", 123 | "car", 124 | "truck", 125 | "fence", 126 | "tree", 127 | "bus", 128 | "pole"] 129 | 130 | 131 | def _get_scutseg_pairs(folder, split='train', logger=None): 132 | def get_path_pairs(img_folder, mask_folder, edge_folder): 133 | img_folder = os.path.join(img_folder, split) 134 | mask_folder = os.path.join(mask_folder, split) 135 | edge_folder = os.path.join(edge_folder, split) 136 | img_paths = [] 137 | mask_paths = [] 138 | edge_paths = [] 139 | for root, _, files in os.walk(img_folder): 140 | for filename in files: 141 | if filename.endswith('.jpg'): 142 | imgpath = os.path.join(root, filename) 143 | maskname = filename.replace('.jpg', '_labelIds.png') 144 | maskpath = os.path.join(mask_folder, maskname) 145 | edgepath = os.path.join(edge_folder, maskname) 146 | 147 | if os.path.isfile(imgpath) and os.path.isfile(maskpath) and os.path.isfile(edgepath): 148 | img_paths.append(imgpath) 149 | mask_paths.append(maskpath) 150 | edge_paths.append(edgepath) 151 | else: 152 | print('cannot find the image, mask, or edge:', imgpath, maskpath, edgepath) 153 | if logger is not None: 154 | logger.info('Found {} images in the folder {}'.format(len(img_paths), img_folder)) 155 | return img_paths, mask_paths, edge_paths 156 | 157 | if split in ('train', 'val', 'test'): 158 | img_folder = os.path.join(folder, 'image') 159 | mask_folder = os.path.join(folder, 'mask') 160 | edge_folder = os.path.join(folder, 'edges') 161 | img_paths, mask_paths, edge_paths = get_path_pairs(img_folder, mask_folder, edge_folder) 162 | return img_paths, mask_paths, edge_paths 163 | 164 | 165 | if __name__ == '__main__': 166 | from os import path, sys 167 | import matplotlib.pyplot as plt 168 | sys.path.append(path.dirname(path.dirname(path.abspath(__file__)))) 169 | sys.path.append(os.path.join(path.dirname(path.dirname(path.abspath(__file__))), '..')) 170 | from core.data.samplers import (make_batch_data_sampler, make_data_sampler, 171 | make_multiscale_batch_data_sampler) 172 | from core.utils.utils import plot_tensors 173 | from datasets import * 174 | from torch.utils.data import DataLoader 175 | from skimage import color, io 176 | 177 | data_kwargs = {"base_size": [520], "crop_size": [480]} 178 | 179 | train_dataset = SCUTSEGDataset( 180 | root="/mnt/ACF29FC2F29F8F68/Work/Deep_Learning/Thermal_Segmentation/Dataset/", 181 | split='train', 182 | mode='train', 183 | **data_kwargs 184 | ) 185 | 186 | train_sampler = make_data_sampler( 187 | dataset=train_dataset, shuffle=False, distributed=False 188 | ) 189 | 190 | train_batch_sampler = make_multiscale_batch_data_sampler( 191 | sampler=train_sampler, 192 | batch_size=1, 193 | multiscale_step=1, 194 | scales=1) 195 | 196 | loader_random_sample = DataLoader( 197 | dataset=train_dataset, 198 | batch_sampler=train_batch_sampler, 199 | num_workers=0, 200 | pin_memory=True, 201 | ) 202 | 203 | x = train_dataset.__getitem__(200) 204 | -------------------------------------------------------------------------------- /Codes/src/datasets/dataloaders/soda.py: -------------------------------------------------------------------------------- 1 | """Prepare SODA dataset""" 2 | ''' 3 | Code Adapted from 4 | https://github.com/dmlc/gluon-cv/blob/master/gluoncv/data/cityscapes.py 5 | ''' 6 | import os 7 | 8 | import numpy as np 9 | import torch 10 | from PIL import Image, ImageFilter 11 | from skimage.exposure import equalize_adapthist 12 | from torchvision import transforms 13 | 14 | try: 15 | from segbase import SegmentationDataset 16 | except: 17 | from datasets.dataloaders.segbase import SegmentationDataset 18 | 19 | import cv2 20 | 21 | 22 | class SODADataset(SegmentationDataset): 23 | NUM_CLASS = 21 24 | IGNORE_INDEX = 255 25 | NAME = "SODA" 26 | BASE_FOLDER = 'InfraredSemanticLabel-20210430T150555Z-001/SODA' 27 | 28 | def __init__(self, root='./datasets/InfraredSemanticLabel-20210430T150555Z-001/SODA', 29 | split='train', 30 | base_size=1024, 31 | crop_size=720, 32 | mode=None, 33 | logger=None, 34 | transform=None, 35 | sobel_edges=False): 36 | 37 | root = os.path.join(root, self.BASE_FOLDER) 38 | super(SODADataset, self).__init__(root, split, mode, base_size, crop_size, logger) 39 | assert os.path.exists(self.root), "Error: data root path is wrong!" 40 | self.sobel_edges = sobel_edges 41 | 42 | self.mean = [0.41079543, 0.41079543, 0.41079543] 43 | self.std = [0.18772296, 0.18772296, 0.18772296] 44 | 45 | self.images, self.mask_paths, self.edge_paths = _get_soda_pairs(self.root, self.split, logger) 46 | assert (len(self.images) == len(self.mask_paths)) 47 | 48 | if len(self.images) == 0: 49 | raise RuntimeError("Found 0 images in subfolders of:" + root + "\n") 50 | 51 | def __getitem__(self, index): 52 | if type(index) == list or type(index) == tuple: 53 | index, scale = index 54 | input_size = self.crop_size[scale] 55 | else: 56 | input_size = None 57 | 58 | img = Image.open(self.images[index]).convert('RGB') 59 | 60 | if self.mode == 'test': 61 | img = self.normalize(img) 62 | return img, os.path.basename(self.images[index]) 63 | 64 | ''' For Training and validation purposes''' 65 | mask = Image.open(self.mask_paths[index]) 66 | 67 | if self.sobel_edges: 68 | id255 = np.where(mask == 255) 69 | no255_gt = np.array(mask) 70 | no255_gt[id255] = 0 71 | edge = cv2.Canny(no255_gt, 5, 5, apertureSize=7) 72 | edge = cv2.dilate(edge, self.edge_kernel) 73 | edge[edge == 255] = 1 74 | edge = Image.fromarray(edge) 75 | else: 76 | edge = Image.open(self.edge_paths[index]) 77 | 78 | # synchrosized transform 79 | if self.mode == 'train': 80 | img, mask, edge = self._sync_transform(img=img, mask=mask, edge=edge, crop_size=input_size) 81 | elif self.mode == 'val': 82 | img, mask, edge = self._val_sync_transform(img=img, mask=mask, edge=edge, crop_size=input_size) 83 | else: 84 | assert self.mode == 'testval' 85 | img, mask, edge = self._img_transform(img), self._mask_transform(mask), self._edge_transform(edge) 86 | 87 | # general resize, normalize and to Tensor 88 | img = self.normalize(img) 89 | return img, mask, edge, os.path.basename(self.images[index]) 90 | 91 | def normalize(self, img): 92 | img = self.im2double(np.array(img)) 93 | img = (img - self.mean) * np.reciprocal(self.std) 94 | img = self.np2Tensor(img).float() 95 | return img 96 | 97 | def _mask_transform(self, mask): 98 | return torch.LongTensor(np.array(mask).astype("int32")) 99 | 100 | def __len__(self): 101 | return len(self.images) 102 | 103 | @property 104 | def pred_offset(self): 105 | return 0 106 | 107 | @property 108 | def class_names(self,): 109 | return ["background", 110 | "person", 111 | "building", 112 | "tree", 113 | "road", 114 | "pole", 115 | "grass", 116 | "door", 117 | "table", 118 | "chair", 119 | "car", 120 | "bicycle", 121 | "lamp", 122 | "monitor", 123 | "trafficCone", 124 | "trash can", 125 | "animal", 126 | "fence", 127 | "sky", 128 | "river", 129 | "sidewalk"] 130 | 131 | 132 | def _get_soda_pairs(folder, split='train', logger=None): 133 | def get_path_pairs(img_folder, mask_folder, edge_folder): 134 | img_folder = os.path.join(img_folder, split) 135 | mask_folder = os.path.join(mask_folder, split) 136 | edge_folder = os.path.join(edge_folder, split) 137 | img_paths = [] 138 | mask_paths = [] 139 | edge_paths = [] 140 | for root, _, files in os.walk(img_folder): 141 | for filename in files: 142 | if filename.endswith('.jpg'): 143 | imgpath = os.path.join(root, filename) 144 | maskname = filename.replace('.jpg', '.png') 145 | maskpath = os.path.join(mask_folder, maskname) 146 | edgepath = os.path.join(edge_folder, maskname) 147 | 148 | if os.path.isfile(imgpath) and os.path.isfile(maskpath) and os.path.isfile(edgepath): 149 | img_paths.append(imgpath) 150 | mask_paths.append(maskpath) 151 | edge_paths.append(edgepath) 152 | else: 153 | print('cannot find the image, mask, or edge:', imgpath, maskpath, edgepath) 154 | if logger is not None: 155 | logger.info('Found {} images in the folder {}'.format(len(img_paths), img_folder)) 156 | return img_paths, mask_paths, edge_paths 157 | 158 | if split in ('train', 'val', 'test'): 159 | img_folder = os.path.join(folder, 'image') 160 | mask_folder = os.path.join(folder, 'mask') 161 | edge_folder = os.path.join(folder, 'edges') 162 | img_paths, mask_paths, edge_paths = get_path_pairs(img_folder, mask_folder, edge_folder) 163 | return img_paths, mask_paths, edge_paths 164 | 165 | 166 | if __name__ == '__main__': 167 | from os import path, sys 168 | 169 | import matplotlib.pyplot as plt 170 | sys.path.append(path.dirname(path.dirname(path.abspath(__file__)))) 171 | sys.path.append(os.path.join(path.dirname(path.dirname(path.abspath(__file__))), '..')) 172 | from core.data.samplers import (make_batch_data_sampler, make_data_sampler, 173 | make_multiscale_batch_data_sampler) 174 | from core.utils.utils import plot_tensors 175 | from datasets import * 176 | from skimage import color, io 177 | from torch.utils.data import DataLoader 178 | 179 | data_kwargs = {"base_size": [520], "crop_size": [480]} 180 | 181 | train_dataset = SODADataset( 182 | root="/mnt/ACF29FC2F29F8F68/Work/Deep_Learning/Thermal_Segmentation/Dataset/", 183 | split='test', 184 | mode='testval', 185 | **data_kwargs 186 | ) 187 | 188 | train_sampler = make_data_sampler( 189 | dataset=train_dataset, shuffle=True, distributed=False 190 | ) 191 | 192 | train_batch_sampler = make_multiscale_batch_data_sampler( 193 | sampler=train_sampler, 194 | batch_size=1, 195 | multiscale_step=1, 196 | scales=1) 197 | 198 | loader_random_sample = DataLoader( 199 | dataset=train_dataset, 200 | batch_sampler=train_batch_sampler, 201 | num_workers=0, 202 | pin_memory=True, 203 | ) 204 | 205 | x = train_dataset.__getitem__(10) 206 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Overview of Agreement: 2 | 3 | Permissions: Private use, non-commercial use, reproduction & dissemination 4 | Conditions: License & copyright notice 5 | Limitation: Limitation of liability and does not include any warranty 6 | 7 | 8 | 9 | Terms and Conditions of Use 10 | 11 | Please read these Terms & Conditions of Use ("Terms") and Tufts University’s Privacy Statement (including its European Economic Area Privacy Statement) carefully before accessing Software developed in the Panetta Lab (“The Panetta Lab”) at Tufts University (“Tufts”) and provided hereunder (“Licensed Software”). By using the Licensed Software, you signify that you have read these terms and agree to be bound by and comply with them. If you do not agree to be bound by these terms, please promptly cease use of the Licensed Software. 12 | 13 | Tufts University reserves the right to modify these Terms at any time and any such changes will be posted on this page. 14 | 15 | “Software” means any and all (a) computer programs, including any and all software implementation of algorithms, models and methodologies, whether in source code, object code, human readable form or other form, (b) databases and compilations, including any and all data and collections of data, whether machine readable or otherwise, (c) descriptions, flow charts and other work products used to design, plan, organize and develop any of the foregoing, (d) screens, user interfaces, report formats, firmware, development tools, templates, menus, buttons and icons and (e) documentation, including user manuals and other training documentation, relating to any of the foregoing. 16 | 17 | 18 | Terms of Use 19 | 20 | Downloading or Accessing Software 21 | 22 | By downloading or accessing the Licensed Software, you agree to comply with Terms, and, to the extent applicable, comply with the terms of use of any third party owners identified in the Licensed Software. In consideration for your agreement to the Terms, Tufts grants you a personal, non-exclusive, non-transferable, non-commercial license to access and use the Licensed Software. 23 | You agree to use the Licensed Software solely for private study, research and teaching purposes, and non-commercial uses. You agree to use the Licensed Software in compliance with all applicable laws, policies and regulations (including any third party restrictions indicated in the Licensed Software). 24 | To the extent you desire to request a commercial license to access and use the Licensed Software for commercial purposes, please contact Tufts at OVPR@tufts.edu. 25 | 26 | 27 | Acknowledgement 28 | 29 | Provide proper attribution to The Panetta Lab and Tufts University, and any third party Software owners (if applicable) as the source of the Licensed Software, and name the Software that you have used. Whenever Licensed Software or Derivatives (as such term is defined below) are used, there needs to be a specific sentence, as a footnote, stating that data was provided by The Panetta Lab and Tufts University (in association with the third party Software owners, if applicable). Any publication relating to the Software or Derivatives (as such term is defined below) should name the source of the Software and should state that The Panetta Lab’s and Tufts University’s endorsement of users' views, products or services is not stated or implied in any way. The following format is to be used when providing attribution: 30 | 31 | @ARTICLE{9585453, 32 | author={Panetta, Karen and Shreyas Kamath, K. M. and Rajeev, Srijith and Agaian, Sos S.}, 33 | journal={IEEE Access}, 34 | title={FTNet: Feature Transverse Network for Thermal Image Semantic Segmentation}, 35 | year={2021}, 36 | volume={9}, 37 | number={}, 38 | pages={145212-145227}, 39 | doi={10.1109/ACCESS.2021.3123066}} 40 | 41 | 42 | Reproduction or Sharing of Software 43 | 44 | You agree to not reproduce, share, or further disseminate the original Licensed Software to others without written permission from Tufts, and not use the Licensed Software for commercial purposes without written permission from Tufts. If further distribution or dissemination is given, you warrant that you will not remove or export any part of the Licensed Software from the United States except in full compliance with all applicable laws, including export regulations. 45 | 46 | Title and copyright to the Licensed Software and any associated documentation will remain with Tufts University. You may make modifications to the Licensed Software and/or integrate the Licensed Software into your own software; provided that you agree that any thing created based on or through the use of the Licensed Software will be considered a derivative of the Licensed Software (“Derivative”). 47 | 48 | Under no circumstances may you sell the Licensed Software or Derivatives without a commercial use license from Tufts to the Licensed Software. Some Software contained in the Licensed Software may have been shared by Tufts courtesy of third party Software owners, which are identified in the Licensed Software. If you wish to redistribute or reuse any portion of any Software that has been provided by a third party data owner, it is your responsibility to obtain permission from the Software owner. 49 | 50 | 51 | Choice of Law 52 | 53 | These Terms shall be governed in all respects by the Commonwealth of Massachusetts without giving effect to its conflicts of law provisions. Both parties submit to the personal jurisdiction of and venue in, the state and federal courts sitting in the judicial district that includes Middlesex County, Massachusetts. The parties further agree that any cause of action arising under these Terms shall exclusively be brought in such courts. If any provision of these Terms is held to be invalid or unenforceable, such provision shall be struck and the remaining provisions shall be enforced. Headings are for reference purposes only and in no way define, limit, construe, or describe the scope or extent of such section. Tufts University’s failure to act with respect to a breach by you or others does not waive its right to act with respect to subsequent or similar breaches. 54 | 55 | 56 | Disclaimer of Warranty/Limitation of Liability 57 | 58 | THE LICENSED SOFTWARE AND ANY INFORMATION SUPPLIED HEREUNDER ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION, THE IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR USE OF A PARTICULAR PURPOSE, OR NONINFRINGEMENT OF THIRD PARTY RIGHTS, AND WITHOUT ANY ACCOMPANYING SERVICES, SUPPORT OR IMPROVEMENTS FROM TUFTS UNIVERSITY OR THE PANETTA LAB. 59 | 60 | Tufts University does not warrant, and hereby expressly disclaims, any warranties of any kind, either express or implied, including with respect to the accuracy, adequacy or completeness of any of the Licensed Software or any information supplied hereunder. 61 | 62 | YOU AGREE THAT TUFTS UNIVERSITY WILL NOT BE LIABLE TO YOU FOR ANY LOSS OR DAMAGES, EITHER ACTUAL OR CONSEQUENTIAL, ARISING OUT OF OR RELATING TO THESE TERMS, OR TO YOUR (OR ANY THIRD PARTY'S) USE OR INABILITY TO USE THE LICENSED SOFTWARE OR ANY INFORMATION SUPPLIED HEREUNDER. IN PARTICULAR, TUFTS WILL HAVE NO LIABILTY FOR ANY CONSEQUENTIAL, INDIRECT, PUNITIVE, SPECIAL OR INCIDENTAL DAMAGES, WHETHER FORESEEABLE OR UNFORESEEABLE, ARISING OUT OF OR RELATING TO THESE TERMS, YOUR USE OR INABILITY TO USE THE LICENSED SOFTWARE OR ANY INFORMATION SUPPLIED HEREUNDER, WHETHER BASED IN CONTRACT, TORT, STATUTORY OR OTHER LAW, EXCEPT ONLY IN THE CASE OF DEATH OR PERSONAL INJURY WHERE AND ONLY TO THE EXTENT THAT APPLICABLE LAW REQUIRES SUCH LIABILITY. 63 | 64 | Trademarks 65 | The logo, name and all graphics of Tufts University or any of its schools, clinics or affiliates, are trademarks of Tufts or its affiliates. Use, reproduction, copying or redistribution of trademarks, without the written permission of Tufts or its affiliates is prohibited. 66 | 67 | Indemnification 68 | You agree to indemnify and hold Tufts harmless from any claims, losses or damages, including legal fees, arising out of or related to the exercise of any rights granted under this Agreement or your breach of these Terms and to fully cooperate in Tufts’ defense against any such claims. 69 | 70 | 71 | 72 | -------------------------------------------------------------------------------- /Codes/src/lightning_scripts/trainers/thermal_edge_trainer.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import os 3 | from collections import OrderedDict 4 | 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from .base_trainer import BaseTrainer 10 | from core.metrics import AverageMeter 11 | from core.utils.json_extension import save_to_json_pretty 12 | from core.utils.utils import as_numpy 13 | from torchmetrics import ConfusionMatrix as pl_ConfusionMatrix 14 | from torchmetrics.classification import IoU 15 | 16 | BatchNorm2d = nn.BatchNorm2d 17 | 18 | 19 | # https://github.com/hszhao/semseg/blob/master/tool/train.py 20 | 21 | 22 | class SegmentationLightningModel(BaseTrainer): 23 | def __init__(self, *args, **kwargs): 24 | super(SegmentationLightningModel, self).__init__(*args, **kwargs) 25 | 26 | def load_metrics(self, mode, num_class, ignore_index): 27 | setattr(self, f'{mode}_confmat', pl_ConfusionMatrix(num_classes=num_class).to(self.device)) 28 | setattr(self, f'{mode}_IOU', IoU(num_classes=num_class, ignore_index=ignore_index, reduction='none').to(self.device)) 29 | setattr(self, f'{mode}_edge_accuracy', AverageMeter().to(self.device)) 30 | 31 | def training_step(self, batch, batch_idx): 32 | images, target, edges, _ = batch 33 | output = self.forward(images) 34 | loss_val = self.criterion(output, (target, edges)) 35 | 36 | class_map, edge_map = output 37 | class_map = class_map[0] if isinstance(class_map, tuple) or isinstance(class_map, list) else class_map 38 | class_map = torch.argmax(class_map.long(), 1) 39 | edge_pred = torch.mean( 40 | ((edge_map > 0) == edges).float(), dim=[1, 2, 3]) 41 | 42 | self.train_edge_accuracy.update(edge_pred) 43 | self.train_confmat.update(preds=class_map, target=target) 44 | self.train_IOU.update(preds=class_map, target=target) 45 | 46 | log_dict = OrderedDict({'loss': loss_val}) 47 | self.log('loss', loss_val) 48 | 49 | return log_dict 50 | 51 | def training_epoch_end(self, outputs): 52 | confusion_matrix = self.train_confmat.compute() 53 | iou = self.train_IOU.compute() 54 | accuracy = self.accuracy_(confusion_matrix) 55 | 56 | log_dict = {"train_mIOU": torch.mean(iou), 57 | "train_mAcc": torch.mean(accuracy), 58 | "train_avg_mIOU_Acc": (torch.mean(iou) + torch.mean(accuracy)) / 2, 59 | "train_edge_accuracy": self.train_edge_accuracy.compute()} 60 | 61 | log_dict["loss"] = torch.stack( 62 | [output["loss"] for output in outputs]).mean() 63 | 64 | self.log_dict(log_dict, prog_bar=True) 65 | self.train_confmat.reset() 66 | self.train_IOU.reset() 67 | self.train_edge_accuracy.reset() 68 | 69 | def validation_step(self, batch, batch_idx): 70 | images, target, edges, filename = batch 71 | output = self.forward(images) 72 | loss_val = self.criterion(output, (target, edges)) 73 | 74 | class_map, edge_map = output 75 | class_map = class_map[0] if isinstance(class_map, tuple) or isinstance(class_map, list) else class_map 76 | class_map = torch.argmax(class_map.long(), 1) 77 | edge_pred = torch.mean(((edge_map > 0) == edges).float(), dim=[1, 2, 3]) 78 | 79 | self.val_edge_accuracy.update(edge_pred) 80 | self.val_confmat.update(preds=class_map, target=target) 81 | self.val_IOU.update(preds=class_map, target=target) 82 | 83 | if self.hparams.args.save_images: 84 | image_dict = {'original': as_numpy(images), 85 | 'groundtruth': as_numpy(target), 86 | 'prediction': as_numpy(class_map), 87 | 'edge_map': as_numpy(edge_map), 88 | 'filename': filename} 89 | 90 | self.save_edge_images(**image_dict) 91 | 92 | log_dict = OrderedDict({'val_loss': loss_val}) 93 | 94 | return log_dict 95 | 96 | def validation_epoch_end(self, outputs): 97 | confusion_matrix = self.val_confmat.compute() 98 | iou = self.val_IOU.compute() 99 | accuracy = self.accuracy_(confusion_matrix) 100 | 101 | log_dict = {"val_mIOU": torch.mean(iou), 102 | "val_mAcc": torch.mean(accuracy), 103 | "val_avg_mIOU_Acc": (torch.mean(iou) + torch.mean(accuracy)) / 2, 104 | "val_edge_accuracy": self.val_edge_accuracy.compute()} 105 | 106 | log_dict["val_loss"] = torch.stack( 107 | [output["val_loss"] for output in outputs]).mean() 108 | self.log_dict(log_dict) 109 | 110 | self.val_edge_accuracy.reset() 111 | self.val_confmat.reset() 112 | self.val_IOU.reset() 113 | 114 | def test_step(self, batch, batch_idx): 115 | 116 | def upsample_output(output, target): 117 | if output.shape != target.shape: 118 | return F.interpolate(output, 119 | size=(target.shape[1], target.shape[2]), 120 | mode='bilinear', 121 | align_corners=True) 122 | else: 123 | output 124 | 125 | self.seg_dir = self.ckp.get_path('Segmented_images/test/{}'.format(self.trainer.test_name) 126 | if hasattr(self.trainer, 'test_name') 127 | else 'Segmented_images/test/') 128 | 129 | images, target, edges, filename = batch 130 | output = self.forward(images) 131 | 132 | class_map, edge_map = output 133 | class_map = class_map[0] if isinstance(class_map, tuple) or isinstance(class_map, list) else class_map 134 | class_map = upsample_output(class_map, target) 135 | edge_map = upsample_output(edge_map, target) 136 | 137 | class_map = torch.argmax(class_map.long(), 1) 138 | edge_pred = torch.mean(((edge_map > 0) == edges).float(), dim=[1, 2, 3]) 139 | 140 | self.test_confmat.update(preds=class_map, target=target) 141 | self.test_IOU.update(preds=class_map, target=target) 142 | self.test_edge_accuracy.update(edge_pred) 143 | 144 | if self.hparams.args.save_images: 145 | image_dict = {'original': as_numpy(images), 146 | 'groundtruth': as_numpy(target), 147 | 'prediction': as_numpy(class_map), 148 | 'edge_map': as_numpy(edge_map), 149 | 'filename': filename} 150 | 151 | self.save_edge_images(**image_dict) 152 | 153 | return 154 | 155 | def test_epoch_end(self, outputs): 156 | # https://stackoverflow.com/questions/20927368/how-to-normalize-a-confusion-matrix 157 | 158 | def accuracy_(confusion_matrix): 159 | acc = confusion_matrix.diag() / confusion_matrix.sum(1) 160 | acc[torch.isnan(acc)] = 0 161 | return acc 162 | 163 | confusion_matrix = self.test_confmat.compute() 164 | iou = self.test_IOU.compute() 165 | accuracy = accuracy_(confusion_matrix) 166 | 167 | self.plot_confusion_matrix(confusion_matrix.cpu().numpy()) 168 | 169 | log_dict = {"test_mIOU": torch.mean(iou), 170 | "test_mAcc": torch.mean(accuracy), 171 | "test_avg_mIOU_Acc": (torch.mean(iou) + torch.mean(accuracy)) / 2, 172 | "test_edge_accuracy": self.test_edge_accuracy.compute(), 173 | } 174 | 175 | per_class_IOU = iou.cpu().numpy() * 100 176 | save_to_json_pretty(log_dict, path=os.path.join(self.seg_dir, 'Average.txt')) 177 | np.savetxt(os.path.join(self.seg_dir, 'per_class_iou.txt'), per_class_IOU) 178 | 179 | self.log_dict(log_dict) 180 | self.test_confmat.reset() 181 | self.test_IOU.reset() 182 | self.test_edge_accuracy.reset() 183 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/ftnet-feature-transverse-network-for-thermal/thermal-image-segmentation-on-soda-dataset)](https://paperswithcode.com/sota/thermal-image-segmentation-on-soda-dataset?p=ftnet-feature-transverse-network-for-thermal) 6 | 7 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/ftnet-feature-transverse-network-for-thermal/thermal-image-segmentation-on-scut-seg)](https://paperswithcode.com/sota/thermal-image-segmentation-on-scut-seg?p=ftnet-feature-transverse-network-for-thermal) 8 | 9 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/ftnet-feature-transverse-network-for-thermal/thermal-image-segmentation-on-mfn-dataset)](https://paperswithcode.com/sota/thermal-image-segmentation-on-mfn-dataset?p=ftnet-feature-transverse-network-for-thermal) 10 | 11 | # FTNet 12 | 13 | This repository is an official PyTorch implementation of the paper **" [FTNet: Feature Transverse Network for Thermal Semantic Segmentation](https://ieeexplore.ieee.org/abstract/document/9585453) "** 14 | 15 | ![Main Figure](https://github.com/shreyaskamathkm/FTNet/blob/main/images/Main%20Figure.png) 16 | We provide scripts for the models from our paper. You can train your own model from scratch, or use pretrained models for testing. 17 | 18 | ## FTNet Model Weights 19 | 20 | Model weights are provided for ResNeXt50 and ResNeXt101. 21 | 22 | For user convenience, the Thermal Cityscape pretrained model and weights for all datasets are provided [here](https://tufts.box.com/s/19pdu8w3ruerbzee5lt1wkgo5y0serum). 23 | 24 | This link also provides the semantic maps generated during testing phase. 25 | 26 | ## Highlight: 27 | - Completely Built on Pytorch Lightning with well designed code structures. This comes with built in DistributedDataParallel, DataParallel support. 28 | - All initialization models, trained models and predictions are available. 29 | - Can be easily used to plug in new models with minimal changes. 30 | 31 | ## Requirements 32 | * Hardware: 1 - 2 GPUs (better with >=11G GPU memory) 33 | * Python 3.8 34 | * Pytorch >=1.6 (Code tested on 1.6) 35 | 36 | ## Code 37 | Clone this repository into any place you want. 38 | ```bash 39 | git clone https://github.com/shreyaskamathkm/FTNet.git 40 | cd FTNet 41 | ``` 42 | ## Dependencies 43 | Please run the following to meet the requirements of the model 44 | ``` 45 | pip install -r requirements.txt 46 | ``` 47 | ## Setting up the environment for training and testing 48 | We train and test the models on three dataset: 49 | - [SODA Dataset](https://arxiv.org/abs/1907.10303) which can be downloaded from [here](https://drive.google.com/drive/folders/1ZF2vDk9j69kP5U0zcp-liOBk-atWcw-5). 50 | - [MFN Dataset](https://ieeexplore.ieee.org/document/8206396) which can be downloaded from [here](https://www.mi.t.u-tokyo.ac.jp/static/projects/mil_multispectral/). Their github repo can be found [here](https://github.com/haqishen/MFNet-pytorch) 51 | - [SCUT-Seg Dataset](https://www.sciencedirect.com/science/article/abs/pii/S1350449520306769) which can be downloaded from [here](https://drive.google.com/drive/folders/1soPrrx2_AXNzbrlOE89i5aYb3TxbmcB5). Their github repo can be found [here](https://github.com/haitaobiyao/MCNet) 52 | 53 | ### Extracting Dataset 54 | Please download all the datasets from the link provided above. Once downloaded, run the following commands to get the dataset into the following data structure. 55 | 56 | For simplicity sake, consider all the images are downloaded to a folder name `Raw_Dataset`. The rest of the steps are as follows 57 | 58 | For Cityscapes thermal dataset 59 | ``` 60 | cd Codes/src/datasets/utils/ # You are now in */src/datasets/utils/ 61 | python Cityscape_folderMap.py --input-image-path /raw_dataset/SODA-20211127T202136Z-001/SODA/TIR_leftImg8bit/ --save-path /Processed_dataset/ 62 | ``` 63 | For SODA thermal dataset 64 | ``` 65 | cd Codes/src/datasets/utils/ # You are now in */src/datasets/utils/ 66 | python SODA_folderMap.py --input-image-path /raw_dataset/SODA-20211127T202136Z-001/SODA/InfraredSemanticLabel/ --save-path /Processed_dataset/ 67 | ``` 68 | For SCUTSeg thermal dataset 69 | ``` 70 | cd Codes/src/datasets/utils/ # You are now in */src/datasets/utils/ 71 | python scutseg_foldermap.py --input-image-path /raw_dataset/SCUT-SEG/ --save-path /Processed_dataset/ 72 | ``` 73 | For MFN thermal dataset 74 | ``` 75 | cd Codes/src/datasets/utils/ # You are now in */src/datasets/utils/ 76 | python MFNDataset_folderMap.py --input-image-path /raw_dataset/ir_seg_dataset/ --save-path /Processed_dataset/ 77 | ``` 78 | ### Generating Edges 79 | Please Note: Current implementation requires MATLAB to generate edges. 80 | ``` 81 | cd Codes/src/datasets/edge_generation/ 82 | Change the path in the 'main.m' file and run it to generate edges 83 | ``` 84 | 85 | #### Dataset Structure 86 | Once the extracting and edge generation is completed, the dataset looks similar to the structure provided below: 87 | 88 | ├── ... 89 | ├── Processed_dataset # Dataset Folder 90 | │ ├── Cityscapes_thermal 91 | │ ├── CITYSCAPE_5000 92 | │ ├── edges 93 | │ └── train 94 | │ ├── image 95 | │ └── train 96 | │ └── mask 97 | │ └── train 98 | │ ├── SODA 99 | │ ├── edges 100 | │ ├── train 101 | │ ├── val 102 | │ └── test 103 | │ ├── image 104 | │ ├── train 105 | │ ├── val 106 | │ └── test 107 | │ └── mask 108 | │ ├── train 109 | │ ├── val 110 | │ └── test 111 | │ ├── MFNDataset 112 | │ ├── edges 113 | │ ├── train 114 | │ ├── val 115 | │ └── test 116 | │ ├── image 117 | │ ├── train 118 | │ ├── val 119 | │ └── test 120 | │ └── mask 121 | │ ├── train 122 | │ ├── val 123 | │ └── test 124 | │ ├── SCUTSEG 125 | │ ├── edges 126 | │ ├── train 127 | │ └── val 128 | │ ├── image 129 | │ ├── train 130 | │ └── val 131 | │ └── mask 132 | │ ├── train 133 | │ └── val 134 | └── ... 135 | 136 | The new processed dataset will be used for training purposes. You can now train FTNet by yourself. Training and testing script is provided in the ``*/FTNet/Codes/src/bash`` folder. Before you run them, please fill in the appropriate details in the **.sh** file before you execute. 137 | 138 | ```bash 139 | cd /Codes/src/bash # You are now in */src/bash/ 140 | bash Train_and_test.sh # To train and test one dataset. eg: SODA 141 | ``` 142 | ```bash 143 | cd /Codes/src/bash # You are now in */src/bash/ 144 | bash Train_and_test_all.sh # To train and test more than one dataset. eg: SODA, MFN, SCUT-Seg 145 | ``` 146 | 147 | 148 | ## License 149 | Please read the LICENSE file in the repository 150 | 151 | ## Citation 152 | If you find the code or trained models useful, please consider citing: 153 | 154 | ``` 155 | @ARTICLE{9585453, 156 | author={Panetta, Karen and Shreyas Kamath, K. M. and Rajeev, Srijith and Agaian, Sos S.}, 157 | journal={IEEE Access}, 158 | title={FTNet: Feature Transverse Network for Thermal Image Semantic Segmentation}, 159 | year={2021}, 160 | volume={9}, 161 | number={}, 162 | pages={145212-145227}, 163 | doi={10.1109/ACCESS.2021.3123066}} 164 | ``` 165 | 166 | ## References 167 | * [Pytorch Lightning](https://www.pytorchlightning.ai/) 168 | * [Semantic Segmentation on PyTorch](https://github.com/Tramac/awesome-semantic-segmentation-pytorch) 169 | * [Edge Detection](https://github.com/Lavender105/DFF/blob/152397cec4a3dac2aa86e92a65cc27e6c8016ab9/lib/matlab/modules/data/seg2edge.m) 170 | * [Progress Bar](https://github.com/zhutmost/neuralzip/blob/master/apputil/progressbar.py) 171 | * [Multi Scale Training](https://github.com/CaoWGG/multi-scale-training) 172 | * [Metrics](https://github.com/mseg-dataset/mseg-semantic) 173 | * [ResNet variants](https://github.com/zhanghang1989/ResNeSt) 174 | * [Logger](https://detectron2.readthedocs.io/en/latest/_modules/detectron2/utils/logger.html) 175 | -------------------------------------------------------------------------------- /Codes/src/datasets/utils/SODA_folderMap.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # To distribute train, valid, test for SODA 3 | # ============================================================================= 4 | 5 | import os 6 | import pandas as pd 7 | from sklearn.model_selection import train_test_split 8 | from glob import glob 9 | import shutil 10 | import argparse 11 | 12 | 13 | def is_image_file(filename): 14 | return any( 15 | filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg", ".bmp"] 16 | ) 17 | 18 | 19 | def str2bool(v): 20 | if v.lower() in ("yes", "true", "t", "y", "1"): 21 | return True 22 | elif v.lower() in ("no", "false", "f", "n", "0"): 23 | return False 24 | else: 25 | raise argparse.ArgumentTypeError("Boolean value expected.") 26 | 27 | 28 | # %% 29 | # ============================================================================= 30 | # Copying files to different folders 31 | # ============================================================================= 32 | def copying(tiles, path_label, basepath, fileset_path): 33 | tiles.set_index(path_label, inplace=True) 34 | for img_path in tiles.index: 35 | print("Path = {}".format(img_path)) 36 | dst_path = os.path.join(basepath, fileset_path) 37 | 38 | shutil.copy(img_path, dst_path) 39 | 40 | 41 | # %% 42 | # ============================================================================= 43 | # Creating folders 44 | # ============================================================================= 45 | 46 | 47 | def distribute(input_dir, output_dir, reset): 48 | # basepath = './Thermal_Segmentation/Dataset/InfraredSemanticLabel-20210430T150555Z-001/' 49 | basepath = input_dir 50 | 51 | train_Name_path = basepath + "/train_infrared.txt" 52 | test_Name_path = basepath + "/test_infrared.txt" 53 | Image_path = basepath + "/JPEGImages" 54 | # Mask_path = basepath + '/InfraredSemanticLabel/SegmentationClassOne' 55 | 56 | # os.path.abspath(os.path.join(basepath,'..')) 57 | base_dir = os.path.join(output_dir, "SODA") 58 | if reset == True and os.path.exists(base_dir): 59 | if os.path.exists(base_dir): 60 | shutil.rmtree(base_dir) 61 | if not os.path.exists(base_dir): 62 | os.mkdir(base_dir) 63 | 64 | main_dirs_image = ["image/train", "image/val", "image/test"] 65 | main_dirs_mask = ["mask/train", "mask/val", "mask/test"] 66 | 67 | for main in main_dirs_image: 68 | 69 | path = os.path.join(base_dir, main) 70 | if not os.path.exists(path): 71 | os.makedirs(path) 72 | 73 | for main in main_dirs_mask: 74 | 75 | path = os.path.join(base_dir, main) 76 | if not os.path.exists(path): 77 | os.makedirs(path) 78 | # %% 79 | # ============================================================================= 80 | # Creating folders 81 | # ============================================================================= 82 | 83 | imageid_path_dict = { 84 | os.path.splitext(os.path.basename(x))[0]: x 85 | for x in glob(os.path.join(Image_path, "**", "**", "*.jpg"), recursive=True) 86 | } 87 | 88 | name_df = pd.read_table( 89 | test_Name_path, delim_whitespace=True, names=("Image_Name", "B") 90 | ) 91 | name_df = name_df.sort_values( 92 | by="Image_Name", axis=0, ascending=True, kind="quicksort" 93 | ).reset_index(drop=True) 94 | name_df = name_df.fillna("NA") 95 | del name_df["B"] 96 | name_df["Image_Path"] = name_df["Image_Name"].map(imageid_path_dict.get) 97 | name_df["Mask_Path"] = name_df["Image_Path"].str.replace(".jpg", ".png") 98 | name_df["Mask_Path"] = name_df["Mask_Path"].str.replace( 99 | "JPEGImages", "SegmentationClassOne" 100 | ) 101 | 102 | test_df, validation_df = train_test_split(name_df, test_size=0.1, random_state=2) 103 | 104 | test_df = test_df[ 105 | ~test_df.Image_Name.str.contains("\(") 106 | ] # there are copies of some files, eg, ABCD.png and ABCD (1).png (I am removing the copies) 107 | # there are copies of some files, eg, ABCD.png and ~temp_ABCD.png (I am removing the copies) 108 | test_df = test_df[~test_df.Image_Name.str.contains("\~")] 109 | # there are copies of some files, eg, ABCD.png and ABCD (1).png (I am removing the copies) 110 | validation_df = validation_df[~validation_df.Image_Name.str.contains("\(")] 111 | # there are copies of some files, eg, ABCD.png and ~temp_ABCD.png (I am removing the copies) 112 | validation_df = validation_df[~validation_df.Image_Name.str.contains("\~")] 113 | 114 | validation_df1 = validation_df.copy(deep=True) 115 | 116 | copying( 117 | tiles=validation_df, 118 | path_label="Image_Path", 119 | basepath=base_dir, 120 | fileset_path=main_dirs_image[1], 121 | ) 122 | copying( 123 | tiles=validation_df, 124 | path_label="Mask_Path", 125 | basepath=base_dir, 126 | fileset_path=main_dirs_mask[1], 127 | ) 128 | 129 | copying( 130 | tiles=test_df, 131 | path_label="Image_Path", 132 | basepath=base_dir, 133 | fileset_path=main_dirs_image[2], 134 | ) 135 | copying( 136 | tiles=test_df, 137 | path_label="Mask_Path", 138 | basepath=base_dir, 139 | fileset_path=main_dirs_mask[2], 140 | ) 141 | 142 | # Include validation set in test as well 143 | copying( 144 | tiles=validation_df1, 145 | path_label="Image_Path", 146 | basepath=base_dir, 147 | fileset_path=main_dirs_image[2], 148 | ) 149 | copying( 150 | tiles=validation_df1, 151 | path_label="Mask_Path", 152 | basepath=base_dir, 153 | fileset_path=main_dirs_mask[2], 154 | ) 155 | # 156 | 157 | train_df = pd.read_table( 158 | train_Name_path, delim_whitespace=True, names=("Image_Name", "B") 159 | ) 160 | train_df = train_df.sort_values( 161 | by="Image_Name", axis=0, ascending=True, kind="quicksort" 162 | ).reset_index(drop=True) 163 | train_df = train_df.fillna("NA") 164 | del train_df["B"] 165 | train_df["Image_Path"] = train_df["Image_Name"].map(imageid_path_dict.get) 166 | train_df["Mask_Path"] = train_df["Image_Path"].str.replace(".jpg", ".png") 167 | train_df["Mask_Path"] = train_df["Mask_Path"].str.replace( 168 | "JPEGImages", "SegmentationClassOne" 169 | ) 170 | # https://stackoverflow.com/questions/28679930/how-to-drop-rows-from-pandas-data-frame-that-contains-a-particular-string-in-a-p 171 | # https://stackoverflow.com/questions/41425945/python-pandas-error-missing-unterminated-subpattern-at-position-2 172 | # there are copies of some files, eg, ABCD.png and ABCD (1).png (I am removing the copies) 173 | train_df = train_df[~train_df.Image_Name.str.contains("\(")] 174 | # there are copies of some files, eg, ABCD.png and ~temp_ABCD.png (I am removing the copies) 175 | train_df = train_df[~train_df.Image_Name.str.contains("\~")] 176 | 177 | copying( 178 | tiles=train_df, 179 | path_label="Image_Path", 180 | basepath=base_dir, 181 | fileset_path=main_dirs_image[0], 182 | ) 183 | copying( 184 | tiles=train_df, 185 | path_label="Mask_Path", 186 | basepath=base_dir, 187 | fileset_path=main_dirs_mask[0], 188 | ) 189 | 190 | 191 | if __name__ == "__main__": 192 | parser = argparse.ArgumentParser(description="Set up Cityscape Dataset") 193 | parser.add_argument( 194 | "--input-image-path", 195 | type=str, 196 | default="/mnt/1842213842211C4E/raw_dataset/SODA-20211127T202136Z-001/SODA/InfraredSemanticLabel/", 197 | help="Path to Soda Dataset", 198 | ) 199 | parser.add_argument( 200 | "--save-path", 201 | type=str, 202 | default="/mnt/1842213842211C4E/processed_dataset/", 203 | help="Path to Soda Dataset", 204 | ) 205 | parser.add_argument( 206 | "--reset", type=bool, default=True, help="Path to Cityscape Dataset" 207 | ) 208 | 209 | args = parser.parse_args() 210 | distribute( 211 | input_dir=args.input_image_path, 212 | output_dir=args.save_path, 213 | reset=args.reset, 214 | ) 215 | --------------------------------------------------------------------------------