From 329fac23f67b2a46fe3c00b0340bd3a56143bed2 Mon Sep 17 00:00:00 2001 From: Robert Sachunsky Date: Wed, 4 Dec 2024 18:29:49 +0000 Subject: [PATCH] do not reload enhancement model in dir_in mode, simplify --- src/eynollah/eynollah.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/eynollah/eynollah.py b/src/eynollah/eynollah.py index c1e0f4d..145f722 100644 --- a/src/eynollah/eynollah.py +++ b/src/eynollah/eynollah.py @@ -363,10 +363,11 @@ class Eynollah: def predict_enhancement(self, img): self.logger.debug("enter predict_enhancement") - model_enhancement, session_enhancement = self.start_new_session_and_model(self.model_dir_of_enhancement) + if not self.dir_in: + self.model_enhancement, _ = self.start_new_session_and_model(self.model_dir_of_enhancement) - img_height_model = model_enhancement.layers[len(model_enhancement.layers) - 1].output_shape[1] - img_width_model = model_enhancement.layers[len(model_enhancement.layers) - 1].output_shape[2] + img_height_model = self.model_enhancement.layers[len(self.model_enhancement.layers) - 1].output_shape[1] + img_width_model = self.model_enhancement.layers[len(self.model_enhancement.layers) - 1].output_shape[2] if img.shape[0] < img_height_model: img = cv2.resize(img, (img.shape[1], img_width_model), interpolation=cv2.INTER_NEAREST) @@ -409,9 +410,8 @@ class Eynollah: index_y_u = img_h index_y_d = img_h - img_height_model - img_patch = img[index_y_d:index_y_u, index_x_d:index_x_u, :] - label_p_pred = model_enhancement.predict(img_patch.reshape(1, img_patch.shape[0], img_patch.shape[1], img_patch.shape[2]), - verbose=0) + img_patch = img[np.newaxis, index_y_d:index_y_u, index_x_d:index_x_u, :] + label_p_pred = self.model_enhancement.predict(img_patch, verbose=0) seg = label_p_pred[0, :, :, :] seg = seg * 255