Ejemplo n.º 1
0
def eval_dir_func(true_dir, test_dir, mask_dir=None, metrics=MetricType.ALL):
    logger = get_logger(__name__)
    avg_scores = deepcopy(default_score_dict)
    flist = os.listdir(true_dir)
    file_num = len(flist)
    logger.info(
        'Eval image num: {}\nTrue img: [ {} ]\nTest img: [ {} ]\nMask img: [ {} ]'
        .format(file_num, true_dir, test_dir, mask_dir))
    pixel_types = ['all']
    if mask_dir:
        pixel_types += ['shadow', 'shadow_free']
    for fname in flist:
        ftrue = os.path.join(true_dir, fname)
        ftest = os.path.join(test_dir, fname)
        if mask_dir:
            fmask = os.path.join(mask_dir, fname)
            scores = eval_func(ftrue, ftest, fmask, metrics)
        else:
            scores = eval_func(ftrue, ftest, None, metrics)
        # if scores['rmse']['all'] < 3:
        visualize_result(true_dir, test_dir, mask_dir, fname, scores)
        # print(scores)
        for pixel_type in pixel_types:
            if metrics & MetricType.MSE:
                avg_scores['mse'][
                    pixel_type] += scores['mse'][pixel_type] / file_num
            if metrics & MetricType.RMSE:
                avg_scores['rmse'][
                    pixel_type] += scores['rmse'][pixel_type] / file_num
    logger.info(avg_scores)
    return avg_scores
Ejemplo n.º 2
0
def print_netowrk(net):
    logger = get_logger(__name__)
    num_params = 0
    for param in net.parameters():
        num_params += param.numel()
    logger.info(net)
    logger.info('Total number of parameters: {}'.format(num_params))
Ejemplo n.º 3
0
    def __init__(self, args, path):
        self.logger = get_logger(__name__)
        self.mdl_name = 'STCGAN'
        self.epoch = args.epochs
        self.batch_size = args.batch_size
        self.gpu_mode = args.gpu_mode
        self.mdl_dir = path.mdl_dir
        self.train_hist = {
            'G1_loss': [],
            'G2_loss': [],
            'D1_loss': [],
            'D2_loss': [],
        }
        # data_loader
        # split data
        train_size, test_size = len(os.listdir(path.train_shadow_dir)), len(os.listdir(path.test_shadow_dir))
        train_img_list, test_img_list = list(range(train_size)), list(range(test_size))
        
        # TODO: add augmentation here
        training_augmentation = get_composed_augmentation()
        if args.valid_ratio:
            split_size = int((1 - args.valid_ratio) * total_size)
            train_img_list, valid_img_list = train_img_list[:split_size], train_img_list[split_size:]
            train_dataset = ShadowRemovalDataset(path, 'training', train_img_list, training_augmentation)
            valid_dataset = ShadowRemovalDataset(path, 'validation', valid_img_list)
            self.train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=8) 
            self.valid_loader = DataLoader(valid_dataset, batch_size=self.batch_size, shuffle=False, num_workers=8)
        else:
            train_dataset = ShadowRemovalDataset(path, 'training', train_img_list, training_augmentation)
            self.train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=8)
        test_dataset = ShadowRemovalDataset(path, 'testing', test_img_list)
        self.test_loader = DataLoader(test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=8)

        # model
        self.G1 = STCGAN_G1()
        self.G2 = STCGAN_G2()
        self.D1 = STCGAN_D1()
        self.D2 = STCGAN_D2()

        self.G_opt = optim.Adam(list(self.G1.parameters()) + list(self.G2.parameters()), lr=args.lrG, betas=(args.beta1, args.beta2))
        self.D_opt = optim.Adam(list(self.D1.parameters()) + list(self.D2.parameters()), lr=args.lrD, betas=(args.beta1, args.beta2))

        if self.gpu_mode:
            self.G1.cuda()
            self.G2.cuda()
            self.D1.cuda()
            self.D2.cuda()
            self.l1_loss = nn.L1Loss().cuda()
            self.adversial_loss = nn.CrossEntropyLoss().cuda()
        else:
            self.l1_loss = nn.L1Loss()
            self.adversial_loss = nn.CrossEntropyLoss()
        
        self.logger.info('-' * 10 + ' Networks Architecture ' + '-' * 10)
        utils.print_netowrk(self.G1)
        utils.print_netowrk(self.G2)
        utils.print_netowrk(self.D1)
        utils.print_netowrk(self.D2)
        self.logger.info('-' * 43)
Ejemplo n.º 4
0
 def __init__(self,
              path=None,
              data_type='training',
              img_list=None,
              transform=None):
     self.logger = get_logger(__name__)
     if not path:
         from pathHandler import PathHandler
         self.path = PathHandler()
     else:
         self.path = path
     self.data_type = data_type
     self.img_list = img_list
     self.transform = transform
Ejemplo n.º 5
0
    fig.suptitle(
        'Image: {}\n[ rmse score ] all: {:.3f}, shadow: {:.3f}, non-shadow: {:.3f}'
        .format(fname, scores['rmse']['all'], scores['rmse']['shadow'],
                scores['rmse']['shadow_free']))
    for pos, title, filename, kwargs in zip(pos_list, title_list,
                                            filename_list, kwargs_list):
        ax = fig.add_subplot(pos)
        ax.set_title(title)
        plt.imshow(mpimg.imread(filename), **kwargs)
    plt.show()


if __name__ == '__main__':
    log_file = os.path.join('log', os.path.basename(__file__) + '.log')
    set_logger(log_file)
    logger = get_logger(__name__)

    test_metrics = MetricType.MSE | MetricType.RMSE

    true_dir = os.path.join('processed_dataset', 'ISTD', 'test', 'non_shadow')
    mask_dir = os.path.join('processed_dataset', 'ISTD', 'test', 'mask')
    guo_dir = os.path.join('processed_dataset', 'ISTD', 'result', 'Guo')
    yang_dir = os.path.join('processed_dataset', 'ISTD', 'result', 'Yang')
    gong_dir = os.path.join('processed_dataset', 'ISTD', 'result', 'Gong')
    stcgan_dir = os.path.join('processed_dataset', 'ISTD', 'result', 'ST-CGAN')

    avg_eval_scores = eval_dir_func(true_dir, guo_dir, mask_dir, test_metrics)
    avg_eval_scores = eval_dir_func(true_dir, yang_dir, mask_dir, test_metrics)
    avg_eval_scores = eval_dir_func(true_dir, gong_dir, mask_dir, test_metrics)
    avg_eval_scores = eval_dir_func(true_dir, stcgan_dir, mask_dir,
                                    test_metrics)
Ejemplo n.º 6
0
    def __init__(self, args, path):
        self.logger = get_logger(__name__)
        self.mdl_name = 'STCGAN'
        self.epoch = args.epochs
        self.batch_size = args.batch_size
        self.valid_ratio = args.valid_ratio

        self.lambda1 = args.lambda1
        self.lambda2 = args.lambda2
        self.lambda3 = args.lambda3

        self.gpu_mode = args.gpu_mode
        # self.gpu_id = args.gpu_id
        self.path = path
        self.mdl_dir = path.mdl_dir
        self.train_hist = {
            'G_loss': [],
            'D_loss': [],
            'G1_loss': [],
            'G2_loss': [],
            'D1_loss': [],
            'D2_loss': [],
        }
        self.device = torch.device('cuda:{}'.format(
            args.gpu_id)) if args.gpu_id else torch.device('cpu')
        # data_loader
        # split data
        train_size, test_size = len(os.listdir(path.train_shadow_dir)), len(
            os.listdir(path.test_shadow_dir))
        train_img_list, test_img_list = list(range(train_size)), list(
            range(test_size))

        # TODO: add augmentation here
        training_augmentation = get_composed_augmentation()
        if self.valid_ratio:
            split_size = int((1 - args.valid_ratio) * train_size)
            train_img_list, valid_img_list = train_img_list[:
                                                            split_size], train_img_list[
                                                                split_size:]
            train_dataset = ShadowRemovalDataset(path, 'training',
                                                 train_img_list,
                                                 training_augmentation)
            valid_dataset = ShadowRemovalDataset(path, 'validation',
                                                 valid_img_list)
            self.train_loader = DataLoader(train_dataset,
                                           batch_size=self.batch_size,
                                           shuffle=True,
                                           num_workers=8)
            self.valid_loader = DataLoader(valid_dataset,
                                           batch_size=2,
                                           shuffle=False,
                                           num_workers=8)
            self.logger.info('Training size: {} Validation size: {}'.format(
                split_size, train_size - split_size))
        else:
            train_dataset = ShadowRemovalDataset(path, 'training',
                                                 train_img_list,
                                                 training_augmentation)
            self.train_loader = DataLoader(train_dataset,
                                           batch_size=self.batch_size,
                                           shuffle=True,
                                           num_workers=8)
        test_dataset = ShadowRemovalDataset(path, 'testing', test_img_list)
        self.test_loader = DataLoader(test_dataset,
                                      batch_size=self.batch_size,
                                      shuffle=False,
                                      num_workers=8)

        # model
        self.G1 = networks.define_G(input_nc=3,
                                    output_nc=1,
                                    ngf=64,
                                    netG='unet_256',
                                    gpu_ids=[args.gpu_id])
        self.G2 = networks.define_G(input_nc=4,
                                    output_nc=3,
                                    ngf=64,
                                    netG='unet_256',
                                    gpu_ids=[args.gpu_id])
        self.D1 = networks.define_D(input_nc=3 + 1,
                                    ndf=64,
                                    netD='pixel',
                                    use_sigmoid=True,
                                    gpu_ids=[args.gpu_id])
        self.D2 = networks.define_D(input_nc=3 + 3 + 1,
                                    ndf=64,
                                    netD='pixel',
                                    use_sigmoid=True,
                                    gpu_ids=[args.gpu_id])

        self.G_opt = optim.Adam(list(self.G1.parameters()) +
                                list(self.G2.parameters()),
                                lr=args.lrG,
                                betas=(args.beta1, args.beta2))
        self.D_opt = optim.Adam(list(self.D1.parameters()) +
                                list(self.D2.parameters()),
                                lr=args.lrD,
                                betas=(args.beta1, args.beta2))

        self.l1_loss = nn.L1Loss()
        self.mse_loss = nn.MSELoss()  # for validation
        # self.bce_loss = nn.BCELoss()
        self.gan_loss = networks.GANLoss(use_lsgan=False).to(self.device)