def create_model(args):
"""
Create the model based on configuration file
:param args: the configuration parameter class defined in config.py
:return:
"""
args['encoder_name'] = misc_utils.stem_string(args['encoder_name'])
args['decoder_name'] = misc_utils.stem_string(args['decoder_name'])
if args['optimizer']['aux_loss']:
aux_loss = True
else:
aux_loss = False
# TODO this is for compatible issue only, we might want to get rid of this later
if 'imagenet' not in args:
args['imagenet'] = 'True'
if args['decoder_name'] == 'unet':
if args['encoder_name'] == 'base':
model = unet.UNet(sfn=args['sfn'],
n_class=args['dataset']['class_num'],
encoder_name=args['encoder_name'],
aux_loss=aux_loss,
use_emau=args['use_emau'],
use_ocr=args['use_ocr'])
else:
model = unet.UNet(n_class=args['dataset']['class_num'],
encoder_name=args['encoder_name'],
pretrained=eval(args['imagenet']),
aux_loss=aux_loss,
use_emau=args['use_emau'],
use_ocr=args['use_ocr'])
elif args['decoder_name'] in ['psp', 'pspnet']:
model = pspnet.PSPNet(n_class=args['dataset']['class_num'],
encoder_name=args['encoder_name'],
pretrained=eval(args['imagenet']),
aux_loss=aux_loss,
use_emau=args['use_emau'],
use_ocr=args['use_ocr'])
elif args['decoder_name'] == 'dlinknet':
model = dlinknet.DLinkNet(n_class=args['dataset']['class_num'],
encoder_name=args['encoder_name'],
pretrained=eval(args['imagenet']),
aux_loss=aux_loss,
use_emau=args['use_emau'],
use_ocr=args['use_ocr'])
elif args['decoder_name'] == 'deeplabv3':
model = deeplabv3.DeepLabV3(n_class=args['dataset']['class_num'],
encoder_name=args['encoder_name'],
pretrained=eval(args['imagenet']),
aux_loss=aux_loss,
use_emau=args['use_emau'],
use_ocr=args['use_ocr'])
else:
raise NotImplementedError(
'Decoder structure {} is not supported'.format(
args['decoder_name']))
return model
def create_loss(args, **kwargs):
"""
Create loss based on configuration
:param args: the configuration parameter class defined in config.py
:return:
"""
criterions = []
for c_name in misc_utils.stem_string(
args['trainer']['criterion_name']).split(','):
if c_name == 'xent':
criterions.append(
metric_utils.CrossEntropyLoss(
eval(args['trainer']['class_weight'])))
elif c_name == 'iou':
# this metric is non-differentiable
criterions.append(metric_utils.IoU())
elif c_name == 'softiou':
criterions.append(metric_utils.SoftIoULoss(kwargs['device']))
elif c_name == 'focal':
criterions.append(
metric_utils.FocalLoss(kwargs['device'],
gamma=args['trainer']['gamma'],
alpha=args['trainer']['alpha']))
elif c_name == 'lovasz':
criterions.append(metric_utils.LovaszSoftmax())
else:
raise NotImplementedError(
'Criterion type {} is not supported'.format(
args['trainer']['criterion_name']))
return criterions
def create_preproc_layer(preproc_name):
preproc_name = misc_utils.stem_string(preproc_name)
if preproc_name == 'gamma':
preproc_layer = GammaAdjustTransform()
elif preproc_name == 'affine':
preproc_layer = AffineTransform()
elif preproc_name == 'colormap':
preproc_layer = ColorMap()
else:
raise NotImplementedError(
'Preproc layer {} not supported'.format(preproc_name))
return preproc_layer
def create_optimizer(optm_name, train_params, lr):
"""
Create optimizer based on configuration
:param optm_name: the optimizer name defined in config.py
:param train_params: learning rate arrangement for the training parameters
:param lr: learning rate
:return: corresponding torch optim class
"""
o_name = misc_utils.stem_string(optm_name)
if o_name == 'sgd':
optm = optim.SGD(train_params, lr=lr, momentum=0.9, weight_decay=5e-4)
elif o_name == 'adam':
optm = optim.Adam(train_params, lr=lr)
else:
raise NotImplementedError(
'Optimizer name {} is not supported'.format(optm_name))
return optm
def __init__(self,
n_class,
sfn=32,
encoder_name='base',
pretrained=True,
aux_loss=False,
use_emau=False,
use_ocr=False):
"""
Initialize the Unet model
:param sfn: the start filter number, following blocks have n*sfn number of filters
:param n_class: the number of class
:param encoder_name: name of the encoder, could be 'base', 'vgg16'
:param pretrained: if True, load the weights from pretrained model
:param aux_loss: if True, will create a classification branch for extracted features
:param use_emau: if True or int, an EMAU will be appended at the end of the encoder
:param use_ocr: if True, the OCR module will be appended at the end of the encoder
"""
super(UNet, self).__init__()
self.n_class = n_class
self.encoder_name = misc_utils.stem_string(encoder_name)
self.aux_loss = aux_loss
self.use_emau = use_emau
self.use_ocr = use_ocr
if self.encoder_name == 'base':
self.sfn = sfn
self.encoder = UnetBaseEncoder(self.sfn)
self.margins = [4, 16, 40, 88]
filter_nums = [self.sfn * (2**a) for a in range(4, -1, -1)]
self.decode_in_chans = filter_nums[:-1]
self.decode_out_chans = filter_nums[1:]
self.lbl_margin = 92
conv_chan = None
pad = 0
up_sample = 0
else:
self.encoder = encoders.models(self.encoder_name, pretrained,
(2, 2, 2, 2, 2), True)
self.margins = [0, 0, 0, 0]
filter_nums = self.encoder.chans
self.decode_in_chans = filter_nums[:-1]
self.decode_out_chans = filter_nums[1:]
self.lbl_margin = 0
conv_chan = [
d_in // 2 + d_out
for (d_in,
d_out) in zip(self.decode_in_chans, self.decode_out_chans)
]
pad = 1
up_sample = 0 if 'vgg' in self.encoder_name else 2
if self.aux_loss:
self.cls = nn.Sequential(nn.Linear(self.decode_in_chans[0], 256),
nn.ReLU(inplace=True),
nn.Linear(256, self.n_class))
else:
self.cls = None
if self.use_emau:
if isinstance(self.use_emau, int):
c = self.use_emau
else:
c = 64
self.encoder.emau = emau.EMAU(self.decode_in_chans[0], c)
if self.use_ocr:
self.encoder.ocr = ocr.OCRModule(self.n_class,
*self.encoder.chans[:2][::-1],
self.encoder.chans[0])
self.decoder = UnetDecoder(self.decode_in_chans, self.decode_out_chans,
self.margins, self.n_class, conv_chan, pad,
up_sample)
def __init__(self,
ds_name,
data_dir,
tsfm,
device,
load_func=None,
infer=False,
ensembler=None,
**kwargs):
ds_name = misc_utils.stem_string(ds_name)
self.tsfm = tsfm
self.device = device
if ensembler is None:
self.ensembler = BaseEnsemble()
else:
self.ensembler = ensembler
if ds_name == 'inria':
from data.inria import preprocess
self.rgb_files, self.lbl_files = preprocess.get_images(
data_dir, **kwargs)
assert len(self.rgb_files) == len(self.lbl_files)
self.truth_val = 255
self.decode_func = None
self.encode_func = None
self.class_names = [
'building',
]
elif ds_name == 'deepglobe':
from data.deepglobe import preprocess
self.rgb_files, self.lbl_files = preprocess.get_images(data_dir)
assert len(self.rgb_files) == len(self.lbl_files)
self.truth_val = 1
self.decode_func = None
self.encode_func = lambda x: x * 255
self.class_names = [
'building',
]
elif ds_name == 'deepgloberoad':
from data.deepgloberoad import preprocess
self.rgb_files, self.lbl_files = preprocess.get_images(
data_dir, **kwargs)
assert len(self.rgb_files) == len(self.lbl_files)
self.truth_val = 255
self.decode_func = preprocess.decode_map
self.encode_func = None
self.class_names = [
'road',
]
elif ds_name == 'deepglobeland':
from data.deepglobeland import preprocess
if not infer:
self.rgb_files, self.lbl_files = preprocess.get_images(
data_dir, **kwargs)
else:
self.rgb_files, self.lbl_files = preprocess.get_test_images(
data_dir, **kwargs)
assert len(self.rgb_files) == len(self.lbl_files)
self.truth_val = 1
self.decode_func = preprocess.decode_map
self.encode_func = preprocess.encode_map
self.class_names = preprocess.CLASS_NAMES[:6]
elif ds_name == 'mnih':
from data.mnih import preprocess
self.rgb_files, self.lbl_files = preprocess.get_images(
data_dir, **kwargs)
assert len(self.rgb_files) == len(self.lbl_files)
self.truth_val = 255
self.decode_func = None
self.encode_func = None
self.class_names = [
'road',
]
elif ds_name == 'spca':
from data.spca import preprocess
self.rgb_files, self.lbl_files = preprocess.get_images(
data_dir, **kwargs)
assert len(self.rgb_files) == len(self.lbl_files)
self.truth_val = 1
self.decode_func = None
self.encode_func = None
self.class_names = [
'panel',
]
elif load_func:
self.truth_val = kwargs.pop('truth_val', 1)
self.decode_func = kwargs.pop('decode_func', None)
self.encode_func = kwargs.pop('encode_func', None)
self.class_names = kwargs.pop('class_names', [
'building',
])
self.rgb_files, self.lbl_files = load_func(data_dir, **kwargs)
assert len(self.rgb_files) == len(self.lbl_files)
else:
raise NotImplementedError('Dataset {} is not supported')