From 18a786cdfa0f9344afc7f359794319cd780321fb Mon Sep 17 00:00:00 2001 From: Rauf Yagfarov Date: Wed, 19 Feb 2020 22:25:43 +0300 Subject: [PATCH] add embeddings loading --- embedding_net/backbones.py | 5 +- embedding_net/model_new.py | 98 +++++++++++++++----------------------- 2 files changed, 41 insertions(+), 62 deletions(-) diff --git a/embedding_net/backbones.py b/embedding_net/backbones.py index 2fb7873..03361bc 100644 --- a/embedding_net/backbones.py +++ b/embedding_net/backbones.py @@ -1,4 +1,4 @@ -from tensorflow.keras.layers import Dense, Input, Lambda, Dropout, Flatten +from tensorflow.keras.layers import Dense, Input, Lambda, Dropout, Flatten, GlobalAveragePooling2D from tensorflow.keras.layers import Conv2D, MaxPool2D, BatchNormalization, concatenate from tensorflow.keras.models import Model from tensorflow.keras.regularizers import l2 @@ -105,7 +105,8 @@ def get_backbone(input_shape, layer.trainable = False after_backbone = backbone_model.output - x = Flatten()(after_backbone) + x = GlobalAveragePooling2D()(after_backbone) + # x = Flatten()(after_backbone) x = Dense(encodings_len//2, activation="relu")(x) diff --git a/embedding_net/model_new.py b/embedding_net/model_new.py index 3e33cb7..65276ba 100644 --- a/embedding_net/model_new.py +++ b/embedding_net/model_new.py @@ -103,8 +103,7 @@ def pretrain_backbone_softmax(self): ModelCheckpoint(filepath=weights_save_file, verbose=1, monitor=checkpoint_callback_monitor, - save_best_only=True) - ] + save_best_only=True)] history = model.fit_generator(train_generator, steps_per_epoch=steps_per_epoch, @@ -135,65 +134,48 @@ def train(self, callbacks=[], verbose=1): validation_data=self.val_generator, epochs=self.params_train['n_epochs'], callbacks=callbacks, - verbose=verbose) + verbose=verbose, + use_multiprocessing=True) return history - def generate_encodings(self, save_file_name='encodings.pkl', - only_centers=False, - max_num_samples_of_each_class=10, - knn_k=1, + def train_embeddings_classifier(self, + classification_model, + max_n_samples=10, + shuffle=True): + encodings = self.generate_encodings(max_n_samples=max_n_samples, + shuffle=shuffle) + classification_model.fit(encodings['encodings'], + encodings['labels']) + + def generate_encodings(self, max_n_samples=10, shuffle=True): data_paths, data_labels, data_encodings = [], [], [] - classes_counter = {} - classes_encodings = {} - k_val = 1 if only_centers else knn_k - - if shuffle: - c = list(zip( - self.data_loader.images_paths['train'], self.data_loader.images_labels['train'])) - random.shuffle(c) - self.data_loader.images_paths['train'], self.data_loader.images_labels['train'] = zip( - *c) - - for img_path, img_label in zip(self.data_loader.images_paths['train'], - self.data_loader.images_labels['train']): - if only_centers: - if img_label not in classes_encodings: - classes_encodings[img_label] = [] - else: - if img_label not in classes_counter: - classes_counter[img_label] = 0 - if classes_counter[img_label] < max_num_samples_of_each_class: + encoded_training_data = {} + + for class_name in self.data_loader.class_names: + data_list = self.data_loader.train_data[class_name] + if len(data_list)>max_n_samples: + if shuffle: + random.shuffle(data_list) + data_list = data_list[:max_n_samples] + + data_paths += data_list + data_labels += class_name * len(data_list) + for img_path in data_list: encod = self._generate_encoding(img_path) - - if encod is not None: - if only_centers: - classes_encodings[img_label].append(encod) - else: - data_paths.append(img_path) - data_labels.append(img_label) - data_encodings.append(encod) - classes_counter[img_label] += 1 - if only_centers: - for class_i, encodings_i in classes_encodings.items(): - encodings_i_np = np.array(encodings_i) - class_encoding = np.mean(encodings_i_np, axis = 0) - data_encodings.append(class_encoding) - data_labels.append(class_i) - self.encoded_training_data['paths'] = data_paths - self.encoded_training_data['labels'] = data_labels - self.encoded_training_data['encodings'] = np.squeeze( - np.array(data_encodings)) - self.encoded_training_data['knn_classifier'] = KNeighborsClassifier( - n_neighbors=k_val) - self.encoded_training_data['knn_classifier'].fit(self.encoded_training_data['encodings'], - self.encoded_training_data['labels']) - with open(save_file_name, "wb") as f: - pickle.dump(self.encoded_training_data, f) + data_encodings.append(encod) - def load_encodings(self, path_to_encodings): - self.encoded_training_data = load_encodings(path_to_encodings) + encoded_training_data['paths'] = data_paths + encoded_training_data['labels'] = data_labels + encoded_training_data['encodings'] = np.squeeze(np.array(data_encodings)) + + return encoded_training_data + + def save_encodings(self, encoded_training_data, + save_file_name='encodings.pkl'): + with open(save_file_name, "wb") as f: + pickle.dump(encoded_training_data, f) def load_model(self, file_path): from keras_radam import RAdam @@ -207,17 +189,13 @@ def load_model(self, file_path): outputs=[self.model.layers[3].layers[-1].output]) self.base_model._make_predict_function() - def calculate_distances(self, encoding): - training_encodings = self.encoded_training_data['encodings'] - return np.sqrt( - np.sum((training_encodings - np.array(encoding))**2, axis=1)) - def predict(self, image): if type(image) is str: img = cv2.imread(image) else: img = image - img = cv2.resize(img, (self.input_shape[0], self.input_shape[1])) + img = cv2.resize(img, (self.params_model['input_shape'][0], + self.params_model['input_shape'][1])) encoding = self.base_model.predict(np.expand_dims(img, axis=0)) distances = self.calculate_distances(encoding) max_element = np.argmin(distances)