labeling_module.py
2.12 KB
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
# plaidml
# import plaidml.keras
# plaidml.keras.install_backend()
# packages
from keras.models import load_model
from keras.preprocessing import image
# import queue
import numpy as np
from queue import Full, Empty
from multiprocessing import Process, Queue
class LabelingModule:
def __init__(self):
# self.model1 = load_model('svhn_model.h5')
self.model2 = load_model('svhn_model.h5')
self.image_queue = Queue(maxsize=3000)
self.label_queue = Queue(maxsize=10)
self.signal_queue = Queue()
self.predict_process = Process(target=_predict, \
args=(self.model2, self.image_queue, self.label_queue, self.signal_queue))
def run(self):
self.predict_process.start()
def close(self):
self.image_queue.close()
self.label_queue.close()
def new_tensor(self, tensor):
try:
self.image_queue.put(tensor)
except Full:
print('[LabelingModule] image_queue is full')
def new_image(self, filename):
tensor = self._img_to_tensor(filename)
try:
self.image_queue.put(tensor)
except Full:
print('[LabelingModule] image_queue is full')
def _img_to_tensor(self, filename):
img = image.load_img(filename, target_size=(48, 48))
img_tensor = image.img_to_array(img)
img_tensor = np.squeeze(img_tensor)
img_tensor /= 255.
img_tensor = img_tensor - img_tensor.mean()
return img_tensor
def _predict(model, input_queue, output_queue, signal_queue):
print('predict process started.')
while True:
try:
signal = signal_queue.get_nowait()
if signal == 'stop':
break
except Empty:
pass
tensor = input_queue.get(timeout=-1)
dat = model.predict(np.array([tensor]))
o1 = np.argmax(dat[0])
o2 = np.argmax(dat[1])
o3 = np.argmax(dat[2])
o4 = np.argmax(dat[3])
o5 = np.argmax(dat[4])
o6 = np.argmax(dat[5])
output = [o1, o2, o3, o4, o5, o6]
print('[LabelingModule] predict result :', output)