forked from killthekitten/kaggle-carvana-2017
-
Notifications
You must be signed in to change notification settings - Fork 0
/
predict_multithreaded.py
121 lines (94 loc) · 3.79 KB
/
predict_multithreaded.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
import os
import numpy as np
import tensorflow as tf
import pandas as pd
from datasets import generate_filenames
from keras.applications.imagenet_utils import preprocess_input
from keras.preprocessing.image import array_to_img, load_img, img_to_array, flip_axis
from tensorflow.python.client import device_lib
from models import make_model
from params import args
import threading
import queue
from tqdm import tqdm
gpus = [x.name for x in device_lib.list_local_devices() if x.name[:4] == '/gpu']
prediction_dir = args.pred_mask_dir
output_dir = args.pred_mask_dir
batch_size = args.pred_batch_size
if args.predict_on_val:
folds_df = pd.read_csv(os.path.join(args.dataset_dir, args.folds_source))
ids = generate_filenames(folds_df['id'])
else:
test_images = set(os.listdir(args.test_data_dir))
already_tested_images = set(map(lambda x: x.replace('.png', '.jpg'), os.listdir(output_dir)))
ids = sorted(list(test_images - already_tested_images))
filenames = [os.path.join(args.test_data_dir, f) for f in ids]
q_size = 10
def do_tta(x, tta_type):
if tta_type == 'hflip':
# batch, img_col = 2
return flip_axis(x, 2)
else:
return x
def undo_tta(pred, tta_type):
if tta_type == 'hflip':
# batch, img_col = 2
return flip_axis(pred, 2)
else:
return pred
def create_model(gpu):
with tf.device(gpu):
model = make_model((None, None, args.stacked_channels + 3))
model.load_weights(args.weights, by_name=True)
return model
def data_loader(q, ):
for start in tqdm(range(0, len(filenames), batch_size)):
x_batch = []
end = min(start + batch_size, len(filenames))
filenames_batch = filenames[start:end]
for filename in filenames_batch:
img = load_img(filename)
stacked_channels = []
for i in range(args.stacked_channels):
channel_path = os.path.join(args.stacked_channels_dir,
str(i),
filename.split('/')[-1].replace('.jpg', '.png'))
stacked_channel = load_img(channel_path, grayscale=True)
stacked_channels.append(stacked_channel)
stacked_img = np.dstack((img, *stacked_channels))
x_batch.append(img_to_array(stacked_img))
x_batch = preprocess_input(np.array(x_batch, np.float32), mode=args.preprocessing_function)
if args.pred_tta:
x_batch = do_tta(x_batch, args.pred_tta)
padded_x = np.zeros((batch_size, 1280, 1920, args.stacked_channels + 3))
padded_x[:, :, 1:-1, :] = x_batch
q.put((filenames_batch, padded_x))
for gpu in gpus:
q.put((None, None))
def predictor(q, gpu):
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
with sess.as_default():
model = create_model(gpu)
while True:
batch_fnames, x_batch = q.get()
if x_batch is None:
break
preds = model.predict_on_batch(x_batch)
if args.pred_tta:
preds = undo_tta(preds, args.pred_tta)
for i, pred in enumerate(preds):
filename = batch_fnames[i]
prediction = pred[:, 1:-1, :]
array_to_img(prediction * 255).save(os.path.join(output_dir, filename.split('/')[-1][:-4] + ".png"))
print('Predicting on {} samples with batch_size = {}...'.format(len(filenames), batch_size))
q = queue.Queue(maxsize=q_size)
threads = []
threads.append(threading.Thread(target=data_loader, name='DataLoader', args=(q,)))
threads[0].start()
for gpu in gpus:
print("Starting predictor at device " + gpu)
t = threading.Thread(target=predictor, name='Predictor', args=(q, gpu))
threads.append(t)
t.start()