ReflectionEraser/ReflectionEraserApp

1import os
2from os.path import join
3
4import torch.backends.cudnn as cudnn
5
6import data.sirs_dataset as datasets
7from data.image_folder import read_fns
8from engine import Engine
9from options.net_options.train_options import TrainOptions
10from tools import mutils
11
12opt = TrainOptions().parse() #parses training options from command-line arguments using the TrainOptions class and stores them in opt
13#evaluation settings:
14opt.isTrain = False 
15cudnn.benchmark = True
16opt.no_log = True
17opt.display_id = 0
18opt.verbose = False
19
20datadir = os.path.join(opt.base_dir, 'test')# test data directory
21
22#(mingcv) Define evaluation/test dataset
23eval_dataset_real = datasets.DSRTestDataset(join(datadir, f'real20_420'),
24                                             fns=read_fns('data/real_test.txt'),
25                                             if_align=opt.if_align)
26eval_dataset_postcard = datasets.DSRTestDataset(join(datadir, 'SIR2/PostcardDataset'),
27                                                 if_align=opt.if_align)
28eval_dataset_solidobject = datasets.DSRTestDataset(join(datadir, 'SIR2/SolidObjectDataset'),
29                                                    if_align=opt.if_align)
30eval_dataset_wild = datasets.DSRTestDataset(join(datadir, 'SIR2/WildSceneDataset'),
31                                             if_align=opt.if_align)
32
33eval_dataloader_real = datasets.DataLoader(
34    eval_dataset_real, batch_size=1, shuffle=True,
35    num_workers=opt.nThreads, pin_memory=True)
36
37eval_dataloader_solidobject = datasets.DataLoader(
38    eval_dataset_solidobject, batch_size=1, shuffle=False,
39    num_workers=opt.nThreads, pin_memory=True)
40
41eval_dataloader_postcard = datasets.DataLoader(
42    eval_dataset_postcard, batch_size=1, shuffle=False,
43    num_workers=opt.nThreads, pin_memory=True)
44
45eval_dataloader_wild = datasets.DataLoader(
46    eval_dataset_wild, batch_size=1, shuffle=False,
47    num_workers=opt.nThreads, pin_memory=True)
48
49engine = Engine(opt)
50
51"""Main Loop"""
52result_dir = os.path.join('./checkpoints', opt.name, mutils.get_formatted_time())
53
54#
55res = engine.eval(eval_dataloader_real, dataset_name='testdata_real',
56                  savedir=join(result_dir, 'real20'), suffix='real20')
57print(res)
58
59res = engine.eval(eval_dataloader_solidobject, dataset_name='testdata_solidobject',
60                  savedir=join(result_dir, 'solidobject'), suffix='solidobject')
61print(res)
62res = engine.eval(eval_dataloader_postcard, dataset_name='testdata_postcard',
63                  savedir=join(result_dir, 'postcard'), suffix='postcard')
64print(res)
65res = engine.eval(eval_dataloader_wild, dataset_name='testdata_wild',
66                  savedir=join(result_dir, 'wild'), suffix='wild')
67print(res)
68