Anubis - Bachelor thesis project

computer_vision.py
Go to the documentation of this file.
1 # Anubis is a program to control industrial cameras and train and use artificial neural networks
2 # Copyright (C) 2021 Lukaszczyk Jakub
3 #
4 # This program is free software: you can redistribute it and/or modify
5 # it under the terms of the GNU General Public License as published by
6 # the Free Software Foundation, either version 3 of the License, or
7 # any later version.
8 #
9 # This program is distributed in the hope that it will be useful,
10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12 # GNU General Public License for more details.
13 #
14 # You should have received a copy of the GNU General Public License
15 # along with this program. If not, see <https://www.gnu.org/licenses/>.
16
17 from tensorflow import keras
18 from tensorflow.python.keras.utils import tf_utils
19 import numpy as np
20 import cv2
21 import os
22 import random
23
24 class Computer_vision():
25 """!brief Implements Tensorflow communication"""
26 def __init__(self,plot):
27 """!@brief Initialize object and its variables
28 @param[in] plot The class needs to have a reference to the GUI object the
29 classification results will be shown in.
30 """
31 self.model = None
32 self.width = 1
33 self.height = 1
34 self.categories = []
35 self.plot = plot
36 self.split = 0
37 self.busy = False
38
39 self.x = []
40 self.y = []
41
42 self.training_data = []
43
44 def load_model(self,path):
45 """!@brief Load Keras model.
46 @details Model is loaded from the path and variables defining input
47 dimensions of the model are set.
48 @param[in] path Path to the folder with saved model.
49 @return True if loading is successful else False.
50 """
51 if(self.model):
52 keras.backend.clear_session()
53 self.plot.__init__()
54 self.__init__(self.plot)
55
56
57 try:
58 self.model = keras.models.load_model(path)
59 input_dim = list(self.model.get_layer(index=0).input_shape)
60 self.width = input_dim[1]
61 self.height = input_dim[2]
62
63 return [self.width, self.height]
64 except:
65 self.model = None
66 return False
67
68 def save_model(self,path):
69 """!@brief Save model to defined path
70 @param[in] path Path to the folder where model will be saved.
71 @return True if loading is successful else False.
72 """
73 if(self.model):
74 self.model.save(path)
75 return True
76 else:
77 return False
78
79 def train(self, epochs, train_vals, split, progress_flag, training_flag):
80 """!@brief Start training loaded model
81 @details If the object is not busy, the tensorflow/Keras training will
82 start. During training the callback object will push current values of
83 training variavles int the references passed in train_vals.
84 @param[in] epochs How many times should the dataset be introduced to
85 the network.
86 @param[in] train_vals Reference to variables used in callback to inform
87 calling functions about training progress
88 @param[in] split Fraction of dataset to be used as validation data (float).
89 @param[in] progress_flag Signal to tell that it is time to update
90 progress values
91 @param[in] training_flag Used to announce finished training
92 """
93 if(self.model != None and not self.busy):
94 self.busy = True
95 if(split):
96 self.split = split
97 else:
98 self.split = 0
99 callback = Gui_callback(train_vals,progress_flag)
100
101 self.model.fit(self.x,
102 self.y,
103 batch_size=32,
104 validation_split=self.split,epochs=epochs,
105 callbacks=[callback])
106 self.busy = False
107
108 self.train_vals['progress'] = 100
109 self.progress_flag.set()
110
111 training_flag.clear()
112
113 def classify(self,frame,prediction_flag):
114 """!@brief Classify image using trained model
115 @details This method does not change the behavior of the network but
116 only uses it to make prediction about what is in the picture. The
117 result is shown in the plot passed by reference in the __init__ method.
118 @param[in] frame OpenCV image that the prediction will be done for.
119 @param[in] prediction_flag Announces that the prediction was completed
120 """
121
122 #get current frame
123 if self.model != None:
124 input_dim = list(self.model.get_layer(index=0).input_shape)
125
126 if(self.model.get_layer(index=-1).output_shape[1] != len(self.categories)):
127 #Add reading categories
128 for ctg in range(0,self.model.get_layer(index=-1).output_shape[1]):
129 self.categories.append(f'ctg {ctg}')
130 self.plot.add_categories(self.categories)
131
132
133 if(input_dim[0] == None):
134 input_dim[0] = -1
135
136 frame_data = cv2.resize(frame,(input_dim[1],input_dim[2]))
137 #maybe change dimensions according to model
138 frame_data = np.array(frame_data)
139 frame_data = frame_data.reshape(input_dim)
140 #transform it to numpy array
141
142 #run .predict
143 out_data = self.model.predict([frame_data])
144 self.plot.write_probability(out_data[0].tolist())
145 #save output to variable and write it to the plot
146
147
148 prediction_flag.clear()
149
150 def process_dataset(self, path, process_perc , categories=[], process_flag = None, callback_flag = None):
151 """!@brief Prepare dataset to be used as input and desired output.
152 @details The method will proceed only if a model is loaded and path is valid.
153 dimensions of the model are set.
154 @param[in] path Path to the dataset
155 @param[in] process_perc Reference to the progress variable in the
156 calling function. Value should be read during callbacks
157 @param[in] categories Names of the individual folders containing data
158 for different categories.
159 @param[in] process_flag When processing finishes the flag is raised.
160 @param[in] callback_flag Threading event telling the calling function
161 to update progress based on the valu in process_perc
162 @return True if loading is successful else False.
163 """
164
165 if(path != None and self.model != None and not self.busy):
166 self.busy = True
167 del self.x
168 del self.y
169
170 self.x = []
171 self.y = []
172
173 if(self._create_training_data(path,categories, callback_flag, process_perc)):
174
175 for features,label in self.training_data:
176 self.x.append(features)
177 self.y.append(label)
178
179 self.x = np.array(self.x).reshape(-1, self.width, self.height, 1)
180 self.y = np.array(self.y)
181
182 self.busy = False
183 process_flag.set()
184 return True
185 else:
186 self.busy = False
187 process_flag.set()
188 return False
189
190 def _create_training_data(self, path, categories, callback_flag, process_perc):
191 """!@brief Auxiliary method to resize all input data to the size of input
192 of the loaded model.
193 dimensions of the model are set.
194 @param[in] path Path to the dataset
195 @param[in] categories Names of the individual folders containing data
196 for different categories.
197 @param[in] callback_flag Threading event telling the calling function
198 to update progress based on the valu in process_perc
199 @param[in] process_perc Reference to the progress variable in the
200 calling function. Value should be read during callbacks
201 @return True if loading is successful else False.
202 """
203
204 self.training_data.clear()
205 items = 0
206 process_perc[0] = 0
207 callback_flag.set()
208
209 try:
210 for category in categories:
211 path_cat = os.path.join(path, category)
212 items += len(os.listdir(path_cat))
213 except:
214 return False
215
216 percent = int(items/100)
217 items_done = 0
218
219 for category in categories:
220 path_cat = os.path.join(path, category)
221 class_num = categories.index(category)
222 for img in os.listdir(path_cat):
223 try:
224 img_array = cv2.imread(os.path.join(path_cat,img),cv2.IMREAD_GRAYSCALE)
225 new_array = cv2.resize(img_array, (self.width, self.height))
226 self.training_data.append([new_array,class_num])
227 items_done += 1
228 if(items_done % percent == 0):
229 process_perc[0] = items_done/percent
230 callback_flag.set()
231 except Exception as e:
232 pass
233 process_perc[0] = 100
234 callback_flag.set()
235 random.shuffle(self.training_data)
236
237 return True
238
239
240 class Gui_callback(keras.callbacks.Callback):
241 """!@brief callback for training progress to GUI"""
242 def __init__(self, train_vals, progress_flag, count_mode='samples', stateful_metrics=None):
243 super(Gui_callback, self).__init__()
244
245 self.train_vals = train_vals
246 self.progress_flag = progress_flag
247
248 self._supports_tf_logs = True
249 if count_mode == 'samples':
250 self.use_steps = False
251 elif count_mode == 'steps':
252 self.use_steps = True
253 else:
254 raise ValueError('Unknown `count_mode`: ' + str(count_mode))
255 # Defaults to all Model's metrics except for loss.
256 self.stateful_metrics = set(stateful_metrics) if stateful_metrics else None
257
258 self.seen = 0
259 self.active_epoch = 0
260 self.progbar = None
261 self.target = None
262 self.verbose = 1
263 self.epochs = 1
264
265 self._called_in_fit = False
266
267 def set_params(self, params):
268 self.verbose = params['verbose']
269 self.epochs = params['epochs']
270 self.target = params['steps']
271
272
273 def on_train_begin(self, logs=None):
274 # When this logger is called inside `fit`, validation is silent.
275 self._called_in_fit = True
276
277 #Pass in label for epochs
278 def on_epoch_begin(self, epoch, logs=None):
279 self.active_epoch = epoch
280
281 def on_train_batch_end(self, batch, logs=None):
282 self._update_progbar(batch, logs)
283
284 def on_test_batch_end(self, batch, logs=None):
285 if not self._called_in_fit:
286 self._update_progbar(batch, logs)
287
288 def on_predict_batch_end(self, batch, logs=None):
289 # Don't pass prediction results.
290 self._update_progbar(batch, None)
291
292 def on_epoch_end(self, epoch, logs=None):
293 self._update_progbar(logs=logs)
294
295 def on_test_end(self, logs=None):
296 if not self._called_in_fit:
297 self._update_progbar(logs=logs)
298
299 def on_predict_end(self, logs=None):
300 self._update_progbar(logs=logs)
301
302 def _reset_progbar(self):
303 self.train_vals['progress'] = 0
304 self.progress_flag.set()
305 self.seen = 0
306
307 def _maybe_init_progbar(self):
308 if self.stateful_metrics is None:
309 if self.model:
310 self.stateful_metrics = (set(m.name for m in self.model.metrics))
311 else:
312 self.stateful_metrics = set()
313
314
315 def _update_progbar(self, batch=None, logs=None):
316 """!@brief Updates the progress bar"""
317 logs = logs or {}
318 self._maybe_init_progbar()
319 self.seen += 1
320 logs = tf_utils.to_numpy_or_python_type(logs)
321
322 # Only block async when verbose = 1.
323 if(self.target == None):
324 self.target = 0
325
326 percentage = int( (self.seen)/((self.target*self.epochs)/100))
327 try:
328 self.train_vals['loss'] = format(float(logs['loss']),'.3f')
329 except:
330 pass
331 try:
332 self.train_vals['val_loss'] = format(float(logs['val_loss']),'.3f')
333 except:
334 pass
335 try:
336 self.train_vals['acc'] = format(float(logs['accuracy']),'.3f')
337 except:
338 pass
339 try:
340 self.train_vals['val_acc'] = format(float(logs['val_accuracy']),'.3f')
341 except:
342 pass
343 self.train_vals['progress'] = percentage
344 self.train_vals['max_epoch'] = self.epochs
345 self.train_vals['epoch'] = self.active_epoch + 1
346 self.progress_flag.set()
computer_vision.Gui_callback.progress_flag
progress_flag
Definition: computer_vision.py:246
computer_vision.Gui_callback.train_vals
train_vals
Definition: computer_vision.py:245
computer_vision.Computer_vision.height
height
Definition: computer_vision.py:33
computer_vision.Gui_callback.on_test_batch_end
def on_test_batch_end(self, batch, logs=None)
Definition: computer_vision.py:284
computer_vision.Computer_vision.split
split
Definition: computer_vision.py:36
computer_vision.Computer_vision.plot
plot
Definition: computer_vision.py:35
computer_vision.Computer_vision.__init__
def __init__(self, plot)
Initialize object and its variables.
Definition: computer_vision.py:26
computer_vision.Gui_callback.__init__
def __init__(self, train_vals, progress_flag, count_mode='samples', stateful_metrics=None)
Definition: computer_vision.py:242
computer_vision.Gui_callback.target
target
Definition: computer_vision.py:261
computer_vision.Computer_vision.model
model
Definition: computer_vision.py:31
computer_vision.Gui_callback.on_epoch_end
def on_epoch_end(self, epoch, logs=None)
Definition: computer_vision.py:292
computer_vision.Gui_callback.verbose
verbose
Definition: computer_vision.py:262
computer_vision.Computer_vision.y
y
Definition: computer_vision.py:40
computer_vision.Computer_vision.training_data
training_data
Definition: computer_vision.py:42
computer_vision.Computer_vision.width
width
Definition: computer_vision.py:32
computer_vision.Gui_callback.use_steps
use_steps
Definition: computer_vision.py:250
computer_vision.Computer_vision.train
def train(self, epochs, train_vals, split, progress_flag, training_flag)
Start training loaded model.
Definition: computer_vision.py:79
computer_vision.Gui_callback.on_train_batch_end
def on_train_batch_end(self, batch, logs=None)
Definition: computer_vision.py:281
computer_vision.Computer_vision
brief Implements Tensorflow communication
Definition: computer_vision.py:24
computer_vision.Computer_vision.categories
categories
Definition: computer_vision.py:34
computer_vision.Gui_callback.on_epoch_begin
def on_epoch_begin(self, epoch, logs=None)
Definition: computer_vision.py:278
computer_vision.Gui_callback.on_test_end
def on_test_end(self, logs=None)
Definition: computer_vision.py:295
computer_vision.Gui_callback._called_in_fit
_called_in_fit
Definition: computer_vision.py:265
computer_vision.Gui_callback._maybe_init_progbar
def _maybe_init_progbar(self)
Definition: computer_vision.py:307
computer_vision.Gui_callback._update_progbar
def _update_progbar(self, batch=None, logs=None)
Updates the progress bar.
Definition: computer_vision.py:315
computer_vision.Computer_vision.busy
busy
Definition: computer_vision.py:37
computer_vision.Computer_vision.x
x
Definition: computer_vision.py:39
computer_vision.Gui_callback.set_params
def set_params(self, params)
Definition: computer_vision.py:267
computer_vision.Computer_vision.classify
def classify(self, frame, prediction_flag)
Classify image using trained model.
Definition: computer_vision.py:113
computer_vision.Gui_callback.seen
seen
Definition: computer_vision.py:258
computer_vision.Gui_callback.stateful_metrics
stateful_metrics
Definition: computer_vision.py:256
computer_vision.Computer_vision.load_model
def load_model(self, path)
Load Keras model.
Definition: computer_vision.py:44
computer_vision.Gui_callback.epochs
epochs
Definition: computer_vision.py:263
computer_vision.Computer_vision.process_dataset
def process_dataset(self, path, process_perc, categories=[], process_flag=None, callback_flag=None)
Prepare dataset to be used as input and desired output.
Definition: computer_vision.py:150
computer_vision.Gui_callback.on_train_begin
def on_train_begin(self, logs=None)
Definition: computer_vision.py:273
computer_vision.Gui_callback._supports_tf_logs
_supports_tf_logs
Definition: computer_vision.py:248
computer_vision.Gui_callback
callback for training progress to GUI
Definition: computer_vision.py:240
computer_vision.Gui_callback.progbar
progbar
Definition: computer_vision.py:260
computer_vision.Computer_vision.save_model
def save_model(self, path)
Save model to defined path.
Definition: computer_vision.py:68
computer_vision.Gui_callback.active_epoch
active_epoch
Definition: computer_vision.py:259
computer_vision.Gui_callback._reset_progbar
def _reset_progbar(self)
Definition: computer_vision.py:302
computer_vision.Gui_callback.on_predict_batch_end
def on_predict_batch_end(self, batch, logs=None)
Definition: computer_vision.py:288
computer_vision.Computer_vision._create_training_data
def _create_training_data(self, path, categories, callback_flag, process_perc)
Auxiliary method to resize all input data to the size of input of the loaded model.
Definition: computer_vision.py:190
computer_vision.Gui_callback.on_predict_end
def on_predict_end(self, logs=None)
Definition: computer_vision.py:299