-
Notifications
You must be signed in to change notification settings - Fork 29
/
Copy pathTestModule.py
138 lines (122 loc) · 7.14 KB
/
TestModule.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
from PIL import Image
import torch
from torchvision import transforms
from math import log10
from torchvision.utils import save_image as imwrite
import os
from utility import to_variable
class Middlebury_eval:
def __init__(self, input_dir='./evaluation'):
self.im_list = ['Backyard', 'Basketball', 'Dumptruck', 'Evergreen', 'Mequon', 'Schefflera', 'Teddy', 'Urban']
self.transform = transforms.Compose([transforms.ToTensor()])
self.input0_list = []
self.input1_list = []
for item in self.im_list:
self.input0_list.append(to_variable(self.transform(Image.open(input_dir + '/input/' + item + '/frame10.png')).unsqueeze(0)))
self.input1_list.append(to_variable(self.transform(Image.open(input_dir + '/input/' + item + '/frame11.png')).unsqueeze(0)))
def Test(self, model, output_dir='./evaluation/output', output_name='frame10i11.png'):
model.eval()
for idx in range(len(self.im_list)):
if not os.path.exists(output_dir + '/' + self.im_list[idx]):
os.makedirs(output_dir + '/' + self.im_list[idx])
frame_out = model(self.input0_list[idx], self.input1_list[idx])
imwrite(frame_out, output_dir + '/' + self.im_list[idx] + '/' + output_name, range=(0, 1))
class Middlebury_other:
def __init__(self, input_dir, gt_dir):
self.im_list = ['Beanbags', 'Dimetrodon', 'DogDance', 'Grove2', 'Grove3', 'Hydrangea', 'MiniCooper', 'RubberWhale', 'Urban2', 'Urban3', 'Venus', 'Walking']
self.transform = transforms.Compose([transforms.ToTensor()])
self.input0_list = []
self.input1_list = []
self.gt_list = []
for item in self.im_list:
self.input0_list.append(to_variable(self.transform(Image.open(input_dir + '/' + item + '/frame10.png')).unsqueeze(0)))
self.input1_list.append(to_variable(self.transform(Image.open(input_dir + '/' + item + '/frame11.png')).unsqueeze(0)))
self.gt_list.append(to_variable(self.transform(Image.open(gt_dir + '/' + item + '/frame10i11.png')).unsqueeze(0)))
def Test(self, model, output_dir, current_epoch, logfile=None, output_name='output.png'):
model.eval()
av_psnr = 0
if logfile is not None:
logfile.write('{:<7s}{:<3d}'.format('Epoch: ', current_epoch) + '\n')
for idx in range(len(self.im_list)):
if not os.path.exists(output_dir + '/' + self.im_list[idx]):
os.makedirs(output_dir + '/' + self.im_list[idx])
frame_out = model(self.input0_list[idx], self.input1_list[idx])
gt = self.gt_list[idx]
psnr = -10 * log10(torch.mean((gt - frame_out) * (gt - frame_out)).item())
av_psnr += psnr
imwrite(frame_out, output_dir + '/' + self.im_list[idx] + '/' + output_name, range=(0, 1))
msg = '{:<15s}{:<20.16f}'.format(self.im_list[idx] + ': ', psnr) + '\n'
print(msg, end='')
if logfile is not None:
logfile.write(msg)
av_psnr /= len(self.im_list)
msg = '{:<15s}{:<20.16f}'.format('Average: ', av_psnr) + '\n'
print(msg, end='')
if logfile is not None:
logfile.write(msg)
class Davis:
def __init__(self, input_dir, gt_dir):
self.im_list = ['bike-trial', 'boxing', 'burnout', 'choreography', 'demolition', 'dive-in', 'dolphins', 'e-bike', 'grass-chopper', 'hurdles', 'inflatable', 'juggle', 'kart-turn', 'kids-turning', 'lions', 'mbike-santa', 'monkeys', 'ocean-birds', 'pole-vault', 'running', 'selfie', 'skydive', 'speed-skating', 'swing-boy', 'tackle', 'turtle', 'varanus-tree', 'vietnam', 'wings-turn']
self.transform = transforms.Compose([transforms.ToTensor()])
self.input0_list = []
self.input1_list = []
self.gt_list = []
for item in self.im_list:
self.input0_list.append(to_variable(self.transform(Image.open(input_dir + '/' + item + '/frame10.png')).unsqueeze(0)))
self.input1_list.append(to_variable(self.transform(Image.open(input_dir + '/' + item + '/frame11.png')).unsqueeze(0)))
self.gt_list.append(to_variable(self.transform(Image.open(gt_dir + '/' + item + '/frame10i11.png')).unsqueeze(0)))
def Test(self, model, output_dir, logfile=None, output_name='output.png'):
model.eval()
av_psnr = 0
if logfile is not None:
logfile.write('{:<7s}{:<3d}'.format('Epoch: ', model.get_epoch()) + '\n')
for idx in range(len(self.im_list)):
if not os.path.exists(output_dir + '/' + self.im_list[idx]):
os.makedirs(output_dir + '/' + self.im_list[idx])
frame_out = model(self.input0_list[idx], self.input1_list[idx])
gt = self.gt_list[idx]
psnr = -10 * log10(torch.mean((gt - frame_out) * (gt - frame_out)).item())
av_psnr += psnr
imwrite(frame_out, output_dir + '/' + self.im_list[idx] + '/' + output_name, range=(0, 1))
msg = '{:<15s}{:<20.16f}'.format(self.im_list[idx] + ': ', psnr) + '\n'
print(msg, end='')
if logfile is not None:
logfile.write(msg)
av_psnr /= len(self.im_list)
msg = '{:<15s}{:<20.16f}'.format('Average: ', av_psnr) + '\n'
print(msg, end='')
if logfile is not None:
logfile.write(msg)
class ucf:
def __init__(self, input_dir):
self.im_list = ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21']
self.transform = transforms.Compose([transforms.ToTensor()])
self.input0_list = []
self.input1_list = []
self.gt_list = []
for item in self.im_list:
self.input0_list.append(to_variable(self.transform(Image.open(input_dir + '/' + item + '/frame0.png')).unsqueeze(0)))
self.input1_list.append(to_variable(self.transform(Image.open(input_dir + '/' + item + '/frame2.png')).unsqueeze(0)))
self.gt_list.append(to_variable(self.transform(Image.open(input_dir + '/' + item + '/frame1.png')).unsqueeze(0)))
def Test(self, model, output_dir, logfile=None, output_name='output.png'):
model.eval()
av_psnr = 0
if logfile is not None:
logfile.write('{:<7s}{:<3d}'.format('Epoch: ', model.get_epoch()) + '\n')
for idx in range(len(self.im_list)):
if not os.path.exists(output_dir + '/' + self.im_list[idx]):
os.makedirs(output_dir + '/' + self.im_list[idx])
frame_out = model(self.input0_list[idx], self.input1_list[idx])
gt = self.gt_list[idx]
psnr = -10 * log10(torch.mean((gt - frame_out) * (gt - frame_out)).item())
av_psnr += psnr
imwrite(frame_out, output_dir + '/' + self.im_list[idx] + '/' + output_name, range=(0, 1))
msg = '{:<15s}{:<20.16f}'.format(self.im_list[idx] + ': ', psnr) + '\n'
print(msg, end='')
if logfile is not None:
logfile.write(msg)
av_psnr /= len(self.im_list)
msg = '{:<15s}{:<20.16f}'.format('Average: ', av_psnr) + '\n'
print(msg, end='')
if logfile is not None:
logfile.write(msg)