diff --git a/README.md b/README.md index 896b174..23eb376 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,7 @@ The performance of each model variant using the pre-trained weights converted fr * `Keras >= 2.2.0` / `TensorFlow >= 1.12.0` * `keras_applications >= 1.0.7` -* `scikit-image` +* `opencv >= 3.4.2` or `scikit-image` (for resizing images; `opencv` seems to be faster) ### Installing from the source diff --git a/efficientnet/preprocessing.py b/efficientnet/preprocessing.py index 079b4b4..0f5b878 100644 --- a/efficientnet/preprocessing.py +++ b/efficientnet/preprocessing.py @@ -13,37 +13,75 @@ # limitations under the License. # ============================================================================== import numpy as np -from skimage.transform import resize +try: + import cv2 +except ModuleNotFoundError: + import skimage.transform -MAP_INTERPOLATION_TO_ORDER = { - "nearest": 0, - "bilinear": 1, - "biquadratic": 2, - "bicubic": 3, -} + +try: + # OpenCV: Map interpolation string to OpenCV interpolation enum value + INTERPOLATION_DICT = { + "nearest": cv2.INTER_NEAREST, + "bilinear": cv2.INTER_LINEAR, + "bicubic": cv2.INTER_CUBIC, + "laconzos": cv2.INTER_LANCZOS4, + } +except NameError: + # scikit-image: Map interpolation string to interpolation order + INTERPOLATION_DICT = { + "nearest": 0, + "bilinear": 1, + "biquadratic": 2, + "bicubic": 3, + } def center_crop_and_resize(image, image_size, crop_padding=32, interpolation="bicubic"): assert image.ndim in {2, 3} - assert interpolation in MAP_INTERPOLATION_TO_ORDER.keys() + assert interpolation in INTERPOLATION_DICT.keys() + + in_h, in_w = image.shape[:2] + + if isinstance(image_size, (int, float)): + out_h = out_w = image_size + else: + out_h, out_w = image_size + + if isinstance(crop_padding, (int, float)): + crop_padding_h = crop_padding_w = crop_padding + else: + crop_padding_h, crop_padding_w = crop_padding - h, w = image.shape[:2] + padded_center_crop_shape_post_scaling = (out_h + crop_padding_h, + out_w + crop_padding_w) - padded_center_crop_size = int( - (image_size / (image_size + crop_padding)) * min(h, w) - ) - offset_height = ((h - padded_center_crop_size) + 1) // 2 - offset_width = ((w - padded_center_crop_size) + 1) // 2 + inv_scale = min(in_h / padded_center_crop_shape_post_scaling[0], + in_w / padded_center_crop_shape_post_scaling[1]) + + unpadded_center_crop_size_pre_scaling = (round(out_h * inv_scale), + round(out_w * inv_scale)) + + offset_h = ((in_h - unpadded_center_crop_size_pre_scaling[0]) + 1) // 2 + offset_w = ((in_w - unpadded_center_crop_size_pre_scaling[1]) + 1) // 2 image_crop = image[ - offset_height: padded_center_crop_size + offset_height, - offset_width: padded_center_crop_size + offset_width, + offset_h : unpadded_center_crop_size_pre_scaling[0] + offset_h, + offset_w : unpadded_center_crop_size_pre_scaling[1] + offset_w, ] - resized_image = resize( - image_crop, - (image_size, image_size), - order=MAP_INTERPOLATION_TO_ORDER[interpolation], - preserve_range=True, - ) + + try: + resized_image = cv2.resize( + image_crop, + (out_w, out_h), + interpolation=INTERPOLATION_DICT[interpolation] if inv_scale < 1 else cv2.INTER_AREA, + ) + except NameError: + resized_image = skimage.transform.resize( + image_crop, + (out_h, out_w), + order=INTERPOLATION_DICT[interpolation], + preserve_range=True, + ) return resized_image