forked from dhlab-epfl/dhSegment
-
Notifications
You must be signed in to change notification settings - Fork 0
/
evaluate.py
143 lines (105 loc) · 5.27 KB
/
evaluate.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
139
140
141
142
import cv2 as cv
import os
import re
from glob import glob
import tensorflow as tf
import numpy as np
import mask_unused_gpus
import sys
from tqdm import tqdm
from dh_segment.inference import LoadedModel
if __name__ == '__main__':
# Determine the model name:
if len(sys.argv) < 2:
print('Please provide a model name as argument')
sys.exit()
model_name = str(sys.argv[1])
# Is the model trained on MS images`:
# print(sys.argv[2])
if (len(sys.argv) == 3):
use_ms = eval(sys.argv[2])
else:
use_ms = True
mask_unused_gpus.mask_unused_gpus(2)
use_ms=True
if use_ms:
input_path = 'binarization/msbin/gray/test/images'
# input_path = 'binarization/mstex/data/test/images'
model_dir = model_name + '/export/'
# input_path = 'binarization/test_ms_tex/msi/'
# model_dir = 'binarization/model_ms_tex/msi/' + model_name + '/export/'
else:
input_path = 'binarization/test_ms_tex/single_channel/'
model_dir = model_name
output_dir = 'binarization/ms_tex_binar_test/' + model_name + '/'
if use_ms:
input_image_filenames = glob(os.path.join(input_path, '*.jpg'),
recursive=False) + \
glob(os.path.join(input_path, '*.png'),
recursive=False)
input_image_filenames = [re.sub(r'_\d\d?', '', f) for f in input_image_filenames]
# input_image_filenames = [re.sub(r'_\d', '', f) for f in input_image_filenames]
input_files = set(input_image_filenames)
else:
input_image_filenames = glob(os.path.join(input_path, '*.jpg'),
recursive=False) + \
glob(os.path.join(input_path, '*.png'),
recursive=False)
input_files = input_image_filenames
print('Found {} images'.format(len(input_image_filenames)))
os.makedirs(output_dir, exist_ok=True)
os.environ['CUDA_VISIBLE_DEVICES'] = "3"
with tf.Session(): # Start a tensorflow session
# Load the model
m = LoadedModel(model_dir, predict_mode='filename')
total_parameters = 0
for variable in tf.trainable_variables():
# shape is an array of tf.Dimension
shape = variable.get_shape()
variable_parameters = 1
for dim in shape:
variable_parameters *= dim.value
total_parameters += variable_parameters
print('number of network parameters: ' + str(total_parameters))
# Iterate over the images:
for filename in tqdm(input_files, desc='Processed files'):
# For each image, predict each pixel's label
prediction_outputs = m.predict(filename)
probs = prediction_outputs['probs'][0]
original_shape = prediction_outputs['original_shape']
# probs = probs[:, :, 1] # Take only class '1' (class 0 is the background, class 1 is the page)
# probs = probs / np.max(probs) # Normalize to be in [0, 1]
# img = cv.imread(filename, 0)
# cv.imwrite(output_dir + os.path.basename(filename), img)
if use_ms:
filename = re.sub(r'.png', '_2.png', filename)
img = cv.imread(filename, cv.IMREAD_COLOR)
p = probs[:, :, 1] * 255
# Upscale to have full resolution image (cv2 uses (w,h) and not (h,w) for giving shapes)
bin_upscaled = cv.resize(p.astype(np.uint8, copy=False), tuple(
original_shape[::-1]), interpolation=cv.INTER_NEAREST)
pred = probs[:, :, 2] * 255
# Upscale to have full resolution image (cv2 uses (w,h) and not (h,w) for giving shapes)
bin_upscaled_red = cv.resize(pred.astype(np.uint8, copy=False), tuple(
original_shape[::-1]), interpolation=cv.INTER_NEAREST)
img[:, :, 1] = bin_upscaled
img[:, :, 2] = bin_upscaled_red
pseudo_filename = re.sub(r'.png', 'pseudo.png', filename)
filename = re.sub(r'_\d.png', '.png', filename)
b = (bin_upscaled > 128) * 255
b = np.array(b, dtype=np.uint8)
b_red = (bin_upscaled_red > 128) * 255
b_red = np.array(b_red, dtype=np.uint8)
b_multiclass = np.zeros(img.shape)
b_multiclass[:,:,1] = b
b_multiclass[:,:,2] = b_red
cv.imwrite(output_dir + os.path.basename(pseudo_filename), img)
cv.imwrite(output_dir + os.path.basename(filename), b_multiclass)
# ===============================================================================================================
# \\moe\ResultTray\holl\ms\binarization\ms_tex_binar_test\ms_resnet_large\
# F: 0.56551 R: 0.47877 P: 0.83511
# F: 0.76683 0.53423 0.49183 0.55508 0.74797 0.64981 0.84789 0.63763 0.41311 0.010691
# ===============================================================================================================
# \\moe\ResultTray\holl\ms\binarization\ms_tex_binar_test\rn\ (single channel!)
# F: 0.62753 R: 0.61162 P: 0.72608
# F: 0.76141 0.51172 0.59566 0.60379 0.78367 0.74422 0.8069 0.65014 0.63192 0.18591