# pylint: skip-file import argparse import cPickle import os import os.path as osp import re import sys import time #added line by Tets based on comments on github issue "Run errors #2" (https://github.com/itijyou/ademxapp/issues/2) sys.path.append('/home/tets/ademxapp') from functools import partial from PIL import Image from multiprocessing import Pool import numpy as np import mxnet as mx from util import mxutil from util import transformer as ts from util import util from util.lr_scheduler import FixedScheduler, LinearScheduler from data import FileIter, make_divisible, parse_split_file def parse_model_label(args): assert args.model is not None fields = [_.strip() for _ in osp.basename(args.model).split('_')] fieldtets= [_.strip() for _ in osp.basename(args.data_root).split('/')] #print(args) print("len(fieldtets)=" + str(len(fieldtets)) + " : field n=" + fieldtets[len(fieldtets)-1]) # parse fields i = 0 num_fields = len(fields) # database #print("len(fields)=" + str(num_fields) + " : dataset=" + fields[i]) dataset = fields[i] if args.dataset is None else args.dataset dataset = fieldtets[len(fieldtets)-1] i += 1 # network structure assert fields[i].startswith('rn') net_type = re.compile('rn[a-z]*').findall(fields[i])[0] net_name = fields[i][len(net_type):].strip('-') i += 1 # number of classes assert fields[i].startswith('cls') classes = int(fields[i][len('cls'):]) i += 1 # feature resolution feat_stride = 32 if i < num_fields and fields[i].startswith('s'): feat_stride = int(fields[i][len('s'):]) i += 1 # learning rate lr_params = { 'type': 'fixed', 'base': 0.1, 'args': None, } if args.base_lr is not None: lr_params['base'] = args.base_lr # linear if args.lr_type in ('linear',): lr_params['type'] = args.lr_type elif args.lr_type == 'step': lr_params['args'] = {'step': [int(_) for _ in args.lr_steps.split(',')], 'factor': 0.1} model_specs = { # model 'lr_params': lr_params, 'net_type': net_type, 'net_name': net_name, 'classes': classes, 'feat_stride': feat_stride, # data 'dataset': dataset, } return model_specs def parse_args(): parser = argparse.ArgumentParser(description='Tune FCRNs from ResNets.') parser.add_argument('--gpus', default='0', help='The devices to use, e.g. 0,1,2,3') parser.add_argument('--dataset', default=None, help='The dataset to use, e.g. cityscapes, voc.') parser.add_argument('--split', default='train', help='The split to use, e.g. train, trainval.') parser.add_argument('--data-root', dest='data_root', help='The root data dir.', default=None, type=str) parser.add_argument('--output', default=None, help='The output dir.') parser.add_argument('--model', default=None, help='The unique label of this model.') parser.add_argument('--batch-images', dest='batch_images', help='The number of images per batch.', default=None, type=int) parser.add_argument('--crop-size', dest='crop_size', help='The size of network input during training.', default=None, type=int) parser.add_argument('--origin-size', dest='origin_size', help='The size of images to crop from.', default=None, type=int) parser.add_argument('--scale-rate-range', dest='scale_rate_range', help='The range of rescaling', default='0.7,1.3', type=str) parser.add_argument('--weights', default=None, help='The path of a pretrained model.') # parser.add_argument('--lr-type', dest='lr_type', help='The learning rate scheduler, e.g., fixed(default)/step/linear', default=None, type=str) parser.add_argument('--base-lr', dest='base_lr', help='The lr to start from.', default=None, type=float) parser.add_argument('--lr-steps', dest='lr_steps', help='The steps when to reduce lr.', default=None, type=str) parser.add_argument('--weight-decay', dest='weight_decay', help='The weight decay in sgd.', default=0.0005, type=float) # parser.add_argument('--from-epoch', dest='from_epoch', help='The epoch to start from.', default=None, type=int) parser.add_argument('--stop-epoch', dest='stop_epoch', help='The index of epoch to stop.', default=None, type=int) parser.add_argument('--to-epoch', dest='to_epoch', help='The number of epochs to run.', default=None, type=int) # parser.add_argument('--phase', help='Phase of this call, e.g., train/val.', default='train', type=str) # for testing parser.add_argument('--test-scales', dest='test_scales', help='Lengths of the longer side to resize an image into, e.g., 224,256.', default=None, type=str) parser.add_argument('--test-flipping', dest='test_flipping', help='If average predictions of original and flipped images.', default=False, action='store_true') parser.add_argument('--test-steps', dest='test_steps', help='The number of steps to take, for predictions at a higher resolution.', default=1, type=int) parser.add_argument('--save-predictions', dest='save_predictions', help='If save the predicted score maps.', default=False, action='store_true') parser.add_argument('--no-save-results', dest='save_results', help='If save the predicted pixel-wise labels.', default=True, action='store_false') # parser.add_argument('--kvstore', dest='kvstore', help='The type of kvstore, e.g., local/device.', default='device', type=str) parser.add_argument('--prefetch-threads', dest='prefetch_threads', help='The number of threads to fetch data.', default=1, type=int) parser.add_argument('--prefetcher', dest='prefetcher', help='The type of prefetercher, e.g., process/thread.', default='thread', type=str) parser.add_argument('--cache-images', dest='cache_images', help='If cache images, e.g., 0/1', default=None, type=int) parser.add_argument('--log-file', dest='log_file', default=None, type=str) parser.add_argument('--check-start', dest='check_start', help='The first epoch to snapshot.', default=1, type=int) parser.add_argument('--check-step', dest='check_step', help='The steps between adjacent snapshots.', default=4, type=int) parser.add_argument('--debug', help='True means logging debug info.', default=False, action='store_true') parser.add_argument('--backward-do-mirror', dest='backward_do_mirror', help='True means less gpu memory usage.', default=False, action='store_true') if len(sys.argv) == 1: parser.print_help() sys.exit(1) args = parser.parse_args() if args.debug: os.environ['MXNET_ENGINE_TYPE'] = 'NaiveEngine' if args.backward_do_mirror: os.environ['MXNET_BACKWARD_DO_MIRROR'] = '1' if args.output is None: if args.phase == 'val': args.output = osp.dirname(args.weights) else: args.output = 'output' if args.weights is not None: # if args.model is None: assert '_ep-' in args.weights parts = osp.basename(args.weights).split('_ep-') args.model = '_'.join(parts[:-1]) # if args.phase == 'train': if args.from_epoch is None: assert '_ep-' in args.weights parts = os.path.basename(args.weights).split('_ep-') assert len(parts) == 2 from_model = parts[0] if from_model == args.model: parts = os.path.splitext(os.path.basename(args.weights))[0].split('-') args.from_epoch = int(parts[-1]) if args.model is None: raise NotImplementedError('Missing argument: args.model') if args.from_epoch is None: args.from_epoch = 0 if args.log_file is None: if args.phase == 'train': args.log_file = '{}.log'.format(args.model) elif args.phase == 'val': suffix = '' if args.split != 'val': suffix = '_{}'.format(args.split) args.log_file = '{}{}.log'.format(osp.splitext(osp.basename(args.weights))[0], suffix) else: raise NotImplementedError('Unknown phase: {}'.format(args.phase)) model_specs = parse_model_label(args) if args.data_root is None: args.data_root = osp.join('data', model_specs['dataset']) return args, model_specs def get_dataset_specs(args, model_specs): dataset = model_specs['dataset'] meta = {} meta_path = osp.join('issegm/data', dataset, 'meta.pkl') if osp.isfile(meta_path): with open(meta_path) as f: meta = cPickle.load(f) label_2_id = None id_2_label = None cmap = None cmap_path = 'data/shared/cmap.pkl' ident_size = False cache_images = args.phase == 'train' mx_workspace = 1650 if dataset == 'ade20k': num_classes = model_specs.get('classes', 150) label_2_id = np.arange(-1, 150) label_2_id[0] = 255 id_2_label = np.arange(1, 256+1) id_2_label[255] = 0 valid_labels = range(1, 150+1) # if args.split == 'test': cmap_path = None # max_shape = np.array((2100, 2100)) if model_specs.get('balanced', False) and args.split == 'trainval': meta['image_classes']['trainval'] = meta['image_classes']['train'] + meta['image_classes']['val'] elif dataset == 'cityscapes': sys.path.insert(0, 'data/cityscapesScripts/cityscapesscripts/helpers') from labels import id2label, trainId2label # num_classes = model_specs.get('classes', 19) label_2_id = 255 * np.ones((256,)) for l in id2label: if l in (-1, 255): continue label_2_id[l] = id2label[l].trainId id_2_label = np.array([trainId2label[_].id for _ in trainId2label if _ not in (-1, 255)]) valid_labels = sorted(set(id_2_label.ravel())) # cmap = np.zeros((256,3), dtype=np.uint8) for i in id2label.keys(): cmap[i] = id2label[i].color # ident_size = True # max_shape = np.array((1024, 2048)) # if args.split in ('train+', 'trainval+'): cache_images = False # if args.phase in ('val',): mx_workspace = 8000 elif dataset == 'coco': sys.path.insert(0, osp.join(args.data_root, 'PythonAPI')) from pycocotools.coco import COCO coco = COCO(osp.join(args.data_root, 'annotations', 'instances_minival2014.json')) # id_2_label = np.array([0] + sorted(coco.getCatIds())) assert len(id_2_label) == 81 valid_labels = id_2_label.tolist() num_classes = model_specs.get('classes', 81) label_2_id = 255 * np.ones((256,)) for i, l in enumerate(id_2_label): label_2_id[l] = i # max_shape = np.array((640, 640)) elif dataset == 'pascal-context': num_classes = model_specs.get('classes', 60) valid_labels = range(num_classes) # max_shape = np.array((500, 500)) elif dataset == 'voc': num_classes = model_specs.get('classes', 21) valid_labels = range(num_classes) # if args.split in ('train++',): max_shape = np.array((640, 640)) else: max_shape = np.array((500, 500)) else: #default to voc type -- Tets modif num_classes = model_specs.get('classes', 21) valid_labels = range(num_classes) # if args.split in ('train++',): max_shape = np.array((640, 640)) else: max_shape = np.array((500, 500)) #else: # old defaut -- Tets modif # raise NotImplementedError('Unknow dataset: {}'.format(dataset)) if cmap is None and cmap_path is not None: if osp.isfile(cmap_path): with open(cmap_path) as f: cmap = cPickle.load(f) meta['label_2_id'] = label_2_id meta['id_2_label'] = id_2_label meta['valid_labels'] = valid_labels meta['cmap'] = cmap meta['ident_size'] = ident_size meta['max_shape'] = meta.get('max_shape', max_shape) meta['cache_images'] = args.cache_images if args.cache_images is not None else cache_images meta['mx_workspace'] = mx_workspace return meta def _get_metric(): def _eval_func(label, pred): gt_label = label.ravel() valid_flag = gt_label != 255 gt_label = gt_label[valid_flag] pred_label = pred.argmax(1).ravel()[valid_flag] sum_metric = (gt_label == pred_label).sum() num_inst = valid_flag.sum() return (sum_metric, num_inst + (num_inst == 0)) return mx.metric.CustomMetric(_eval_func, 'fcn_valid') def _get_scalemeanstd(): if model_specs['net_type'] == 'rn': return -1, np.array([123.68, 116.779, 103.939]).reshape((1, 1, 3)), None if model_specs['net_type'] in ('rna',): return (1.0/255, np.array([0.485, 0.456, 0.406]).reshape((1, 1, 3)), np.array([0.229, 0.224, 0.225]).reshape((1, 1, 3))) return None, None, None def _get_transformer_image(): scale, mean_, std_ = _get_scalemeanstd() transformers = [] if scale > 0: transformers.append(ts.ColorScale(np.single(scale))) transformers.append(ts.ColorNormalize(mean_, std_)) return transformers def _get_module(args, margs, dargs, net=None): if net is None: # the following lines show how to create symbols for our networks if model_specs['net_type'] == 'rna': from util.symbol.symbol import cfg as symcfg symcfg['lr_type'] = 'alex' symcfg['workspace'] = dargs.mx_workspace symcfg['bn_use_global_stats'] = True if model_specs['net_name'] == 'a1': from util.symbol.resnet_v2 import fcrna_model_a1 net = fcrna_model_a1(margs.classes, margs.feat_stride, bootstrapping=True) if net is None: raise NotImplementedError('Unknown network: {}'.format(vars(margs))) contexts = [mx.gpu(int(_)) for _ in args.gpus.split(',')] mod = mx.mod.Module(net, context=contexts) return mod def _make_dirs(path): if not osp.isdir(path): os.makedirs(path) def _train_impl(args, model_specs, logger): if len(args.output) > 0: _make_dirs(args.output) # dataiter dataset_specs = get_dataset_specs(args, model_specs) scale, mean_, _ = _get_scalemeanstd() if scale > 0: mean_ /= scale margs = argparse.Namespace(**model_specs) dargs = argparse.Namespace(**dataset_specs) dataiter = FileIter(dataset=margs.dataset, split=args.split, data_root=args.data_root, sampler='random', batch_images=args.batch_images, meta=dataset_specs, rgb_mean=mean_, feat_stride=margs.feat_stride, label_stride=margs.feat_stride, origin_size=args.origin_size, crop_size=args.crop_size, scale_rate_range=[float(_) for _ in args.scale_rate_range.split(',')], transformer=None, transformer_image=ts.Compose(_get_transformer_image()), prefetch_threads=args.prefetch_threads, prefetcher_type=args.prefetcher,) dataiter.reset() # optimizer assert args.to_epoch is not None if args.stop_epoch is not None: assert args.stop_epoch > args.from_epoch and args.stop_epoch <= args.to_epoch else: args.stop_epoch = args.to_epoch from_iter = args.from_epoch * dataiter.batches_per_epoch to_iter = args.to_epoch * dataiter.batches_per_epoch lr_params = model_specs['lr_params'] base_lr = lr_params['base'] if lr_params['type'] == 'fixed': scheduler = FixedScheduler() elif lr_params['type'] == 'step': left_step = [] for step in lr_params['args']['step']: if from_iter > step: base_lr *= lr_params['args']['factor'] continue left_step.append(step - from_iter) model_specs['lr_params']['step'] = left_step scheduler = mx.lr_scheduler.MultiFactorScheduler(**lr_params['args']) elif lr_params['type'] == 'linear': scheduler = LinearScheduler(updates=to_iter+1, frequency=50, stop_lr=min(base_lr/100., 1e-6), offset=from_iter) optimizer_params = { 'learning_rate': base_lr, 'momentum': 0.9, 'wd': args.weight_decay, 'lr_scheduler': scheduler, 'rescale_grad': 1.0/len(args.gpus.split(',')), } # initializer net_args = None net_auxs = None if args.weights is not None: net_args, net_auxs = mxutil.load_params_from_file(args.weights) initializer = mx.init.Xavier(rnd_type='gaussian', factor_type='in', magnitude=2) # to_model = osp.join(args.output, '{}_ep'.format(args.model)) mod = _get_module(args, margs, dargs) mod.fit( dataiter, eval_metric=_get_metric(), batch_end_callback=mx.callback.Speedometer(dataiter.batch_size, 1), epoch_end_callback=mx.callback.do_checkpoint(to_model), kvstore=args.kvstore, optimizer='sgd', optimizer_params=optimizer_params, initializer=initializer, arg_params=net_args, aux_params=net_auxs, allow_missing=args.from_epoch == 0, begin_epoch=args.from_epoch, num_epoch=args.stop_epoch, ) def _interp_preds_as_impl(num_classes, im_size, pred_stride, imh, imw, pred): imh0, imw0 = im_size pred = pred.astype(np.single, copy=False) input_h, input_w = pred.shape[0] * pred_stride, pred.shape[1] * pred_stride assert pred_stride >= 1. this_interp_pred = np.array(Image.fromarray(pred).resize((input_w, input_h), Image.CUBIC)) if imh0 == imh: interp_pred = this_interp_pred[:imh, :imw] else: interp_method = util.get_interp_method(imh, imw, imh0, imw0) interp_pred = np.array(Image.fromarray(this_interp_pred[:imh, :imw]).resize((imw0, imh0), interp_method)) return interp_pred def interp_preds_as(im_size, net_preds, pred_stride, imh, imw, threads=4): num_classes = net_preds.shape[0] worker = partial(_interp_preds_as_impl, num_classes, im_size, pred_stride, imh, imw) if threads == 1: ret = [worker(_) for _ in net_preds] else: pool = Pool(threads) ret = pool.map(worker, net_preds) pool.close() return np.array(ret) class ScoreUpdater(object): def __init__(self, valid_labels, c_num, x_num, logger=None, label=None, info=None): self._valid_labels = valid_labels self._confs = np.zeros((c_num, c_num, x_num)) self._pixels = np.zeros((c_num, x_num)) self._logger = logger self._label = label self._info = info @property def info(self): return self._info def reset(self): self._start = time.time() self._computed = np.zeros((self._pixels.shape[1],)) self._confs[:] = 0 self._pixels[:] = 0 @staticmethod def calc_updates(valid_labels, pred_label, label): num_classes = len(valid_labels) pred_flags = [set(np.where((pred_label == _).ravel())[0]) for _ in valid_labels] class_flags = [set(np.where((label == _).ravel())[0]) for _ in valid_labels] conf = [len(class_flags[j].intersection(pred_flags[k])) for j in xrange(num_classes) for k in xrange(num_classes)] pixel = [len(class_flags[j]) for j in xrange(num_classes)] return np.single(conf).reshape((num_classes, num_classes)), np.single(pixel) def do_updates(self, conf, pixel, i, computed=True): if computed: self._computed[i] = 1 self._confs[:, :, i] = conf self._pixels[:, i] = pixel def update(self, pred_label, label, i, computed=True): conf, pixel = ScoreUpdater.calc_updates(self._valid_labels, pred_label, label) self.do_updates(conf, pixel, i, computed) self.scores(i) def scores(self, i=None, logger=None): confs = self._confs pixels = self._pixels num_classes = pixels.shape[0] x_num = pixels.shape[1] class_pixels = pixels.sum(1) class_pixels += class_pixels == 0 scores = confs[xrange(num_classes), xrange(num_classes), :].sum(1) acc = scores.sum() / pixels.sum() cls_accs = scores / class_pixels class_preds = confs.sum(0).sum(1) ious = scores / (class_pixels + class_preds - scores) logger = self._logger if logger is None else logger if logger is not None: if i is not None: speed = 1.*self._computed.sum() / (time.time() - self._start) logger.info('Done {}/{} with speed: {:.2f}/s'.format(i+1, x_num, speed)) name = '' if self._label is None else '{}, '.format(self._label) logger.info('{}pixel acc: {:.2f}%, mean acc: {:.2f}%, mean iou: {:.2f}%'.\ format(name, acc*100, cls_accs.mean()*100, ious.mean()*100)) with util.np_print_options(formatter={'float': '{:5.2f}'.format}): logger.info('\n{}'.format(cls_accs*100)) logger.info('\n{}'.format(ious*100)) return acc, cls_accs, ious def overall_scores(self, logger=None): acc, cls_accs, ious = self.scores(None, logger) return acc, cls_accs.mean(), ious.mean() #@profile def _val_impl(args, model_specs, logger): assert args.prefetch_threads == 1 assert args.weights is not None margs = argparse.Namespace(**model_specs) dargs = argparse.Namespace(**get_dataset_specs(args, model_specs)) image_list, label_list = parse_split_file(margs.dataset, args.split) net_args, net_auxs = mxutil.load_params_from_file(args.weights) net = None mod = _get_module(args, margs, dargs, net) has_gt = args.split in ('train', 'val',) crop_sizes = sorted([int(_) for _ in args.test_scales.split(',')])[::-1] # TODO: multi-scale testing assert len(crop_sizes) == 1, 'multi-scale testing not implemented' label_stride = margs.feat_stride crop_size = crop_sizes[0] save_dir = osp.join(args.output, osp.splitext(args.log_file)[0]) _make_dirs(save_dir) x_num = len(image_list) do_forward = True if do_forward: batch = None transformers = [ts.Scale(crop_size, Image.CUBIC, False)] transformers += _get_transformer_image() transformer = ts.Compose(transformers) scorer = ScoreUpdater(dargs.valid_labels, margs.classes, x_num, logger) scorer.reset() start = time.time() done_count = 0 for i in xrange(x_num): sample_name = osp.splitext(osp.basename(image_list[i]))[0] # skip computed images if args.save_predictions: pred_save_path = osp.join(save_dir, 'predictions', '{}.h5'.format(sample_name)) if osp.isfile(pred_save_path): logger.info('Skipped {} {}/{}'.format(sample_name, i+1, x_num)) continue im_path = osp.join(args.data_root, image_list[i]) rim = np.array(Image.open(im_path).convert('RGB'), np.uint8) if do_forward: im = transformer(rim) imh, imw = im.shape[:2] # init if batch is None: if dargs.ident_size: input_h = make_divisible(imh, margs.feat_stride) input_w = make_divisible(imw, margs.feat_stride) else: input_h = input_w = make_divisible(crop_size, margs.feat_stride) label_h, label_w = input_h / label_stride, input_w / label_stride test_steps = args.test_steps pred_stride = label_stride / test_steps pred_h, pred_w = label_h * test_steps, label_w * test_steps input_data = np.zeros((1, 3, input_h, input_w), np.single) input_label = 255 * np.ones((1, label_h * label_w), np.single) dataiter = mx.io.NDArrayIter(input_data, input_label) batch = dataiter.next() mod.bind(dataiter.provide_data, dataiter.provide_label, for_training=False, force_rebind=True) if not mod.params_initialized: mod.init_params(arg_params=net_args, aux_params=net_auxs) nim = np.zeros((3, imh+label_stride, imw+label_stride), np.single) sy = sx = label_stride // 2 nim[:, sy:sy+imh, sx:sx+imw] = im.transpose(2, 0, 1) net_preds = np.zeros((margs.classes, pred_h, pred_w), np.single) sy = sx = pred_stride // 2 + np.arange(test_steps) * pred_stride for ix in xrange(test_steps): for iy in xrange(test_steps): input_data = np.zeros((1, 3, input_h, input_w), np.single) input_data[0, :, :imh, :imw] = nim[:, sy[iy]:sy[iy]+imh, sx[ix]:sx[ix]+imw] batch.data[0] = mx.nd.array(input_data) mod.forward(batch, is_train=False) this_call_preds = mod.get_outputs()[0].asnumpy()[0] if args.test_flipping: batch.data[0] = mx.nd.array(input_data[:, :, :, ::-1]) mod.forward(batch, is_train=False) this_call_preds = 0.5 * (this_call_preds + mod.get_outputs()[0].asnumpy()[0][:, :, ::-1]) net_preds[:, iy:iy+pred_h:test_steps, ix:ix+pred_w:test_steps] = this_call_preds # save predicted probabilities if args.save_predictions: _make_dirs(osp.dirname(pred_save_path)) tmp = (rim.shape[:2], net_preds.astype(np.float16), pred_stride, imh, imw) util.h5py_save(pred_save_path, *tmp) if args.save_results: # compute pixel-wise predictions interp_preds = interp_preds_as(rim.shape[:2], net_preds, pred_stride, imh, imw) pred_label = interp_preds.argmax(0) if dargs.id_2_label is not None: pred_label = dargs.id_2_label[pred_label] # save predicted labels into an image out_path = osp.join(save_dir, '{}.png'.format(sample_name)) im_to_save = Image.fromarray(pred_label.astype(np.uint8)) if dargs.cmap is not None: im_to_save.putpalette(dargs.cmap.ravel()) im_to_save.save(out_path) else: assert not has_gt done_count += 1 if not has_gt: logger.info('Done {}/{} with speed: {:.2f}/s'.format(i+1, x_num, 1.*done_count / (time.time() - start))) continue # Tets modif - deactivate the save of labels and perf scoring since I want to test here, not validate. #label_path = osp.join(args.data_root, label_list[i]) #label = np.array(Image.open(label_path), np.uint8) # save correctly labeled pixels into an image #out_path = osp.join(save_dir, 'correct', '{}.png'.format(sample_name)) #_make_dirs(osp.dirname(out_path)) #invalid_mask = np.logical_not(np.in1d(label, dargs.valid_labels)).reshape(label.shape) #Image.fromarray((invalid_mask*255 + (label == pred_label)*127).astype(np.uint8)).save(out_path) #scorer.update(pred_label, label, i) logger.info('Done in %.2f s.', time.time() - start) if __name__ == "__main__": util.cfg['choose_interpolation_method'] = True print("Hello Tets!") args, model_specs = parse_args() if len(args.output) > 0: _make_dirs(args.output) logger = util.set_logger(args.output, args.log_file, args.debug) logger.info('start with arguments %s', args) logger.info('and model specs %s', model_specs) if args.phase == 'train': _train_impl(args, model_specs, logger) elif args.phase == 'val': _val_impl(args, model_specs, logger) else: raise NotImplementedError('Unknown phase: {}'.format(args.phase))