forked from Boyiliee/EVD-Net
/
test.py
107 lines (87 loc) · 3.81 KB
/
test.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
import os
import numpy as np
import caffe
import sys
import re
import random
import time
import copy
import cv2
import scipy
import shutil
import csv
from PIL import Image
import datetime
def EditFcnProto(templateFile, height, width):
with open(templateFile, 'r') as ft:
template = ft.read()
print templateFile
outFile = 'DeployT.prototxt'
with open(outFile, 'w') as fd:
fd.write(template.format(height=height,width=width))
def test():
caffe.set_mode_gpu()
caffe.set_device(0)
#caffe.set_mode_cpu();
info = os.listdir(r'VIDEO_test_img');
model = r'EVD-Net.caffemodel'
net = caffe.Net('test.prototxt', model, caffe.TEST);
imagesnum=0;
for line in info:
reg = re.compile(r'(.*?).jpg');
all = reg.findall(line)
if (all != []):
imagename = str(all[0]);
line=imagename
reg = re.compile(r'ILSVRC2015_train_00124006_([0-9]{6})_1_3');
all = reg.findall(line)
labelnum = int(all[0]);
if (os.path.isfile(r'VIDEO_test_img\ILSVRC2015_train_00124006_%s_1_3.jpg' % str(labelnum-1).zfill(6)) == False or
os.path.isfile(r'VIDEO_test_img\ILSVRC2015_train_00124006_%s_1_3.jpg' % str(labelnum-2).zfill(6)) == False or
os.path.isfile(r'VIDEO_test_img\ILSVRC2015_train_00124006_%s_1_3.jpg' % str(labelnum+1).zfill(6)) == False or
os.path.isfile(r'VIDEO_test_img\ILSVRC2015_train_00124006_%s_1_3.jpg' % str(labelnum+2).zfill(6)) == False):
continue;
else:
imagesnum = imagesnum + 1;
npstore_1 = caffe.io.load_image(r'VIDEO_test_img\ILSVRC2015_train_00124006_%s_1_3.jpg' % str(labelnum-2).zfill(6))
npstore_2 = caffe.io.load_image(r'VIDEO_test_img\ILSVRC2015_train_00124006_%s_1_3.jpg' % str(labelnum-1).zfill(6))
npstore = caffe.io.load_image(r'VIDEO_test_img\ILSVRC2015_train_00124006_%s_1_3.jpg' % str(labelnum).zfill(6))
npstore_3 = caffe.io.load_image(r'VIDEO_test_img\ILSVRC2015_train_00124006_%s_1_3.jpg' % str(labelnum+1).zfill(6))
npstore_4 = caffe.io.load_image(r'VIDEO_test_img\ILSVRC2015_train_00124006_%s_1_3.jpg' % str(labelnum+2).zfill(6))
batchdata = []
data = npstore_1
data = data.transpose((2, 0, 1))
batchdata.append(data)
net.blobs['img_1'].data[...] = batchdata;
batchdata = []
data = npstore_2
data = data.transpose((2, 0, 1))
batchdata.append(data)
net.blobs['img_2'].data[...] = batchdata;
batchdata = []
data = npstore
data = data.transpose((2, 0, 1))
batchdata.append(data)
net.blobs['img'].data[...] = batchdata;
batchdata = []
data = npstore_3
data = data.transpose((2, 0, 1))
batchdata.append(data)
net.blobs['img_3'].data[...] = batchdata;
batchdata = []
data = npstore_4
data = data.transpose((2, 0, 1))
batchdata.append(data)
net.blobs['img_4'].data[...] = batchdata;
net.forward()
data = net.blobs['sum'].data[0];
data = data.transpose((1, 2, 0));
data = data[:, :, ::-1]
savepath = 'result\\' + imagename + '_EVD-Net.jpg'
cv2.imwrite(savepath, data * 255.0,[cv2.IMWRITE_JPEG_QUALITY, 100])
print imagename
print 'image numbers:',imagesnum;
def main():
test()
if __name__ == '__main__':
main();