Пример #1
0
def main(flags):
    nn_utils.set_gpu(GPU)

    # define network
    model = pspnet.PSPNet(flags.num_classes, flags.tile_size, suffix=flags.model_suffix, learn_rate=flags.learning_rate,
                          decay_step=flags.decay_step, decay_rate=flags.decay_rate, epochs=flags.epochs,
                          batch_size=flags.batch_size, weight_decay=flags.weight_decay, momentum=flags.momentum)

    cm_train = cityscapes_reader.CollectionMakerCityscapes(flags.data_dir, flags.rgb_type, flags.gt_type, 'train',
                                                           flags.rgb_ext, flags.gt_ext, ['png', 'png'],
                                                           clc_name='{}_train'.format(flags.ds_name),
                                                           force_run=flags.force_run)
    cm_valid = cityscapes_reader.CollectionMakerCityscapes(flags.data_dir, flags.rgb_type, flags.gt_type, 'val',
                                                           flags.rgb_ext, flags.gt_ext, ['png', 'png'],
                                                           clc_name='{}_valid'.format(flags.ds_name),
                                                           force_run=flags.force_run)
    cm_train.print_meta_data()
    cm_valid.print_meta_data()

    resize_func = lambda img: resize_image(img, flags.tile_size)
    train_init_op, valid_init_op, reader_op = dataReaderSegmentation.DataReaderSegmentationTrainValid(
            flags.tile_size, cm_train.meta_data['file_list'], cm_valid.meta_data['file_list'],
            flags.batch_size, cm_train.meta_data['chan_mean'], aug_func=[reader_utils.image_flipping_hori,
                                                                         reader_utils.image_scaling_with_label],
            random=True, has_gt=True, gt_dim=1, include_gt=True, valid_mult=flags.val_mult, global_func=resize_func)\
        .read_op()
    feature, label = reader_op

    model.create_graph(feature)
    model.load_resnet(flags.res_dir)
    model.compile(feature, label, flags.n_train, flags.n_valid, flags.tile_size, ersaPath.PATH['model'],
                  par_dir=flags.model_par_dir, val_mult=flags.val_mult, loss_type='xent')
    train_hook = hook.ValueSummaryHook(flags.verb_step, [model.loss, model.lr_op],
                                       value_names=['train_loss', 'learning_rate'], print_val=[0])
    model_save_hook = hook.ModelSaveHook(model.get_epoch_step()*flags.save_epoch, model.ckdir)
    valid_loss_hook = hook.ValueSummaryHook(model.get_epoch_step(), [model.loss, model.loss_iou],
                                            value_names=['valid_loss', 'valid_mIoU'], log_time=True,
                                            run_time=model.n_valid, iou_pos=1)
    image_hook = hook.ImageValidSummaryHook(model.input_size, model.get_epoch_step(), feature, label, model.output,
                                            cityscapes_labels.image_summary, img_mean=cm_train.meta_data['chan_mean'])
    start_time = time.time()
    model.train(train_hooks=[train_hook, model_save_hook], valid_hooks=[valid_loss_hook, image_hook],
                train_init=train_init_op, valid_init=valid_init_op)
    print('Duration: {:.3f}'.format((time.time() - start_time)/3600))
Пример #2
0
def main(flags):
    nn_utils.set_gpu(GPU)

    # define network
    model = stn.STN(flags.num_classes,
                    flags.tile_size,
                    suffix=flags.model_suffix,
                    learn_rate=flags.learning_rate,
                    decay_step=flags.decay_step,
                    decay_rate=flags.decay_rate,
                    epochs=flags.epochs,
                    batch_size=flags.batch_size)

    cm_train = cityscapes_reader.CollectionMakerCityscapes(
        flags.data_dir,
        flags.rgb_type,
        flags.gt_type,
        'train',
        flags.rgb_ext,
        flags.gt_ext, ['png', 'png'],
        clc_name='{}_train'.format(flags.ds_name),
        force_run=flags.force_run)

    cm_valid = cityscapes_reader.CollectionMakerCityscapes(
        flags.data_dir,
        flags.rgb_type,
        flags.gt_type,
        'val',
        flags.rgb_ext,
        flags.gt_ext, ['png', 'png'],
        clc_name='{}_valid'.format(flags.ds_name),
        force_run=flags.force_run)
    cm_train.print_meta_data()
    cm_valid.print_meta_data()

    train_pred_dir = r'/home/lab/Documents/bohao/data/deeplab_model/vis_train/raw_segmentation_results'
    valid_pred_dir = r'/home/lab/Documents/bohao/data/deeplab_model/vis/raw_segmentation_results'
    file_list_train = get_image_list(flags.data_dir, train_pred_dir, 'train')
    file_list_valid = get_image_list(flags.data_dir, valid_pred_dir, 'val')

    resize_func = lambda img: resize_image(img, flags.tile_size)
    train_init_op, valid_init_op, reader_op = DataReaderSegmentationTrainValid(
            flags.tile_size, file_list_train, file_list_valid,
            flags.batch_size, cm_train.meta_data['chan_mean'], aug_func=[reader_utils.image_flipping_hori],
            random=True, has_gt=True, gt_dim=1, include_gt=True, valid_mult=flags.val_mult, global_func=resize_func)\
        .read_op()
    feature, label, pred = reader_op
    train_init_op_valid, _, reader_op = DataReaderSegmentationTrainValid(
        flags.tile_size, file_list_valid, file_list_train,
        flags.batch_size, cm_valid.meta_data['chan_mean'], aug_func=[reader_utils.image_flipping_hori],
        random=True, has_gt=True, gt_dim=1, include_gt=True, valid_mult=flags.val_mult, global_func=resize_func) \
        .read_op()
    _, _, pred_valid = reader_op

    model.create_graph(pred, feature_valid=pred_valid, rgb=feature)

    model.compile(pred,
                  label,
                  flags.n_train,
                  flags.n_valid,
                  flags.tile_size,
                  ersaPath.PATH['model'],
                  par_dir=flags.model_par_dir,
                  val_mult=flags.val_mult,
                  loss_type='xent')
    train_hook = hook.ValueSummaryHook(
        flags.verb_step, [
            model.loss, model.g_loss, model.d_loss, model.lr_op[0],
            model.lr_op[1], model.lr_op[2]
        ],
        value_names=['seg_loss', 'g_loss', 'd_loss', 'lr_seg', 'lr_g', 'lr_d'],
        print_val=[0, 1, 2])
    model_save_hook = hook.ModelSaveHook(
        model.get_epoch_step() * flags.save_epoch, model.ckdir)
    valid_loss_hook = hook.ValueSummaryHookIters(
        model.get_epoch_step(), [model.loss_xent, model.loss_iou],
        value_names=['valid_loss', 'valid_mIoU'],
        log_time=True,
        run_time=model.n_valid)
    image_hook = hook.ImageValidSummaryHook(
        model.input_size,
        model.get_epoch_step(),
        feature,
        label,
        model.refine,
        cityscapes_labels.image_summary,
        img_mean=cm_train.meta_data['chan_mean'])
    start_time = time.time()
    model.train(train_hooks=[train_hook, model_save_hook],
                valid_hooks=[valid_loss_hook, image_hook],
                train_init=[train_init_op, train_init_op_valid],
                valid_init=valid_init_op)
    print('Duration: {:.3f}'.format((time.time() - start_time) / 3600))
Пример #3
0
def main(flags):
    nn_utils.set_gpu(GPU)

    # define network
    model = deeplab.DeepLab(flags.num_classes,
                            flags.tile_size,
                            batch_size=flags.batch_size)

    cm_train = cityscapes_reader.CollectionMakerCityscapes(
        flags.data_dir,
        flags.rgb_type,
        flags.gt_type,
        'train',
        flags.rgb_ext,
        flags.gt_ext, ['png', 'png'],
        clc_name='{}_train'.format(flags.ds_name),
        force_run=False)
    cm_test = cityscapes_reader.CollectionMakerCityscapes(
        flags.data_dir,
        flags.rgb_type,
        flags.gt_type,
        'val',
        flags.rgb_ext,
        flags.gt_ext, ['png', 'png'],
        clc_name='{}_valid'.format(flags.ds_name),
        force_run=False)
    cm_test.print_meta_data()
    resize_func_train = lambda img: skimage.transform.resize(
        img, flags.tile_size, mode='reflect')
    resize_func_test = lambda img: skimage.transform.resize(
        img,
        cm_test.meta_data['tile_dim'],
        order=0,
        preserve_range=True,
        mode='reflect')

    init_op, reader_op = dataReaderSegmentation.DataReaderSegmentation(
        flags.tile_size,
        cm_test.meta_data['file_list'],
        batch_size=flags.batch_size,
        random=False,
        chan_mean=cm_train.meta_data['chan_mean'],
        is_train=False,
        has_gt=True,
        gt_dim=1,
        include_gt=True,
        global_func=resize_func_train).read_op()
    estimator = nn_processor.NNEstimatorSegmentScene(
        model,
        cm_test.meta_data['file_list'],
        flags.res_dir,
        init_op,
        reader_op,
        ds_name='city_scapes',
        save_result_parent_dir='Cityscapes',
        gpu=flags.GPU,
        score_result=True,
        split_char='.',
        post_func=resize_func_test,
        save_func=make_general_id_map,
        ignore_label=(-1, 255))
    estimator.run(force_run=flags.force_run)