ReflectionEraser/ReflectionEraserApp

1import os
2from os.path import join
3
4import data.sirs_dataset as datasets # Custom dataset module for SIRS dataset
5import util.util as util
6from data.image_folder import read_fns
7from engine import Engine
8from options.net_options.train_options import TrainOptions
9from tools import mutils
10
11opt = TrainOptions().parse()
12print(opt)
13
14opt.display_freq = 10
15
16if opt.debug:
17    opt.display_id = 1
18    opt.display_freq = 1
19    opt.print_freq = 20
20    opt.nEpochs = 40
21    opt.max_dataset_size = 100 #was 9999 (mingcv)
22    opt.no_log = False
23    opt.nThreads = 0
24    opt.decay_iter = 0
25    opt.serial_batches = True
26    opt.no_flip = True
27
28# modify the following code to
29datadir = os.path.join(opt.base_dir)
30datadir_syn = join(datadir, 'train/VOCdevkit/VOC2012/PNGImages')
31datadir_real = join(datadir, 'train/real')
32
33train_dataset = datasets.DSRDataset(
34    datadir_syn, read_fns('data/VOC2012_224_train_png.txt'), size=opt.max_dataset_size, enable_transforms=True)
35
36train_dataset_real = datasets.DSRTestDataset(datadir_real, enable_transforms=True, if_align=opt.if_align)
37
38train_dataset_fusion = datasets.FusionDataset([train_dataset,
39                                               train_dataset_real], [0.7, 0.3])
40
41train_dataloader_fusion = datasets.DataLoader(
42    train_dataset_fusion, batch_size=opt.batchSize, shuffle=not opt.serial_batches,
43    pin_memory=True, prefetch_factor=32, num_workers=32)
44
45eval_dataset_real = datasets.DSRTestDataset(join(datadir, f'test/real20_{opt.real20_size}'),
46                                            fns=read_fns('data/real_test.txt'), if_align=opt.if_align)
47eval_dataset_solidobject = datasets.DSRTestDataset(join(datadir, 'test/SIR2/SolidObjectDataset'),
48                                                   if_align=opt.if_align)
49eval_dataset_postcard = datasets.DSRTestDataset(join(datadir, 'test/SIR2/PostcardDataset'), if_align=opt.if_align)
50eval_dataset_wild = datasets.DSRTestDataset(join(datadir, 'test/SIR2/WildSceneDataset'), if_align=opt.if_align)
51
52eval_dataloader_real = datasets.DataLoader(
53    eval_dataset_real, batch_size=1, shuffle=False,
54    pin_memory=True, prefetch_factor=32, num_workers=32)
55
56eval_dataloader_solidobject = datasets.DataLoader(
57    eval_dataset_solidobject, batch_size=1, shuffle=False,
58    pin_memory=True, prefetch_factor=32, num_workers=32)
59eval_dataloader_postcard = datasets.DataLoader(
60    eval_dataset_postcard, batch_size=1, shuffle=False,
61    pin_memory=True, prefetch_factor=32, num_workers=32)
62
63eval_dataloader_wild = datasets.DataLoader(
64    eval_dataset_wild, batch_size=1, shuffle=False,
65    pin_memory=True, prefetch_factor=32, num_workers=32)
66
67"""Main Loop"""
68engine = Engine(opt)
69result_dir = os.path.join(f'./checkpoints/{opt.name}/results',
70                          mutils.get_formatted_time())
71
72
73def set_learning_rate(lr):
74    for optimizer in engine.model.optimizers:
75        print('[i] set learning rate to {}'.format(lr))
76        util.set_opt_param(optimizer, 'lr', lr)
77
78
79
80if opt.resume or opt.debug_eval:
81    save_dir = os.path.join(result_dir, '%03d' % engine.epoch)
82    os.makedirs(save_dir, exist_ok=True)
83    engine.save_model()
84
85    engine.eval(eval_dataloader_real, dataset_name='testdata_real20',
86                savedir=save_dir, suffix='real20', max_save_size=10)
87    engine.eval(eval_dataloader_solidobject, dataset_name='testdata_solidobject',
88                savedir=save_dir, suffix='solidobject', max_save_size=10)
89    engine.eval(eval_dataloader_postcard, dataset_name='testdata_postcard',
90                savedir=save_dir, suffix='postcard', max_save_size=10)
91    engine.eval(eval_dataloader_wild, dataset_name='testdata_wild',
92                savedir=save_dir, suffix='wild', max_save_size=10)
93# define training strategy
94set_learning_rate(opt.lr)
95while engine.epoch < 50:
96    print('random_seed: ', opt.seed)
97    engine.train(train_dataloader_fusion)
98
99    if engine.epoch % 1 == 0:
100        save_dir = os.path.join(result_dir, '%03d' % engine.epoch)
101        os.makedirs(save_dir, exist_ok=True)
102        engine.eval(eval_dataloader_real, dataset_name='testdata_real20',
103                    savedir=save_dir, suffix='real20', max_save_size=10)
104        engine.eval(eval_dataloader_solidobject, dataset_name='testdata_solidobject',
105                    savedir=save_dir, suffix='solidobject', max_save_size=10)
106        engine.eval(eval_dataloader_postcard, dataset_name='testdata_postcard',
107                    savedir=save_dir, suffix='postcard', max_save_size=10)
108        engine.eval(eval_dataloader_wild, dataset_name='testdata_wild',
109                    savedir=save_dir, suffix='wild', max_save_size=10)
110