diff --git a/image_prediction/redai_adapter/model_wrapper.py b/image_prediction/redai_adapter/model_wrapper.py index 5ad93a9..657d20d 100644 --- a/image_prediction/redai_adapter/model_wrapper.py +++ b/image_prediction/redai_adapter/model_wrapper.py @@ -35,9 +35,7 @@ class ModelWrapper(abc.ABC): # RED-5170: fails if image is 'broken' try: image, valid = self.__resize_and_convert(image), True - except OSError as err: - image, valid = self.__handle_resize_exception(err) - except Exception as err: + except (OSError, Exception) as err: image, valid = self.__handle_resize_exception(err) return image, valid @@ -46,7 +44,7 @@ class ModelWrapper(abc.ABC): return image.resize(self.input_shape[:-1]).convert("RGB") def __handle_resize_exception(self, err): - logger.warn(f"{err}: couldn't resize image, replace and passthrough.") + logger.warn(f"{err}: Couldn't resize image, replace with blank image and passthrough.") image = Image.new("RGB", self.input_shape[:-1]) valid = False return image, valid diff --git a/image_prediction/stitching/merging.py b/image_prediction/stitching/merging.py index fb96665..6286066 100644 --- a/image_prediction/stitching/merging.py +++ b/image_prediction/stitching/merging.py @@ -191,6 +191,8 @@ def concat_images(im1: Image, im2: Image, metadata: dict, axis): try: # RED-5170: fails if image is 'broken' im_aggr.paste(im, box=box) except Exception as err: - logger.warn(f"{err}: couldn't merge image, replace and passthrough. (page: {metadata[Info.PAGE_IDX]})") + logger.warn( + f"{err}: Couldn't merge image, replace broken part by blank image and passthrough. (page: {metadata[Info.PAGE_IDX]})" + ) return im_aggr diff --git a/test/fixtures/model.py b/test/fixtures/model.py index 812da4b..3cceec6 100644 --- a/test/fixtures/model.py +++ b/test/fixtures/model.py @@ -102,7 +102,7 @@ def model_handle_mock(estimator_mock): self.model = estimator_mock def prep_images(self, batch): - return [True for _ in batch], [None for _ in batch] + return [None for _ in batch], [True for _ in batch] def predict(self, batch): return [None for _ in batch]