# -*- coding: utf-8 -*- import math import cv2 import numpy as np def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA): """Rezise the sample to ensure the given size. Keeps aspect ratio. Args: sample (dict): sample size (tuple): image size Returns: tuple: new size """ shape = list(sample['disparity'].shape) if shape[0] >= size[0] and shape[1] >= size[1]: return sample scale = [0, 0] scale[0] = size[0] / shape[0] scale[1] = size[1] / shape[1] scale = max(scale) shape[0] = math.ceil(scale * shape[0]) shape[1] = math.ceil(scale * shape[1]) # resize sample['image'] = cv2.resize(sample['image'], tuple(shape[::-1]), interpolation=image_interpolation_method) sample['disparity'] = cv2.resize(sample['disparity'], tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST) sample['mask'] = cv2.resize( sample['mask'].astype(np.float32), tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST, ) sample['mask'] = sample['mask'].astype(bool) return tuple(shape) class Resize(object): """Resize sample to given size (width, height). """ def __init__( self, width, height, resize_target=True, keep_aspect_ratio=False, ensure_multiple_of=1, resize_method='lower_bound', image_interpolation_method=cv2.INTER_AREA, ): """Init. Args: width (int): desired output width height (int): desired output height resize_target (bool, optional): True: Resize the full sample (image, mask, target). False: Resize image only. Defaults to True. keep_aspect_ratio (bool, optional): True: Keep the aspect ratio of the input sample. Output sample might not have the given width and height, and resize behaviour depends on the parameter 'resize_method'. Defaults to False. ensure_multiple_of (int, optional): Output width and height is constrained to be multiple of this parameter. Defaults to 1. resize_method (str, optional): "lower_bound": Output will be at least as large as the given size. "upper_bound": Output will be at max as large as the given size. " "(Output size might be smaller than given size.)" "minimal": Scale as least as possible. (Output size might be smaller than given size.) Defaults to "lower_bound". """ self.__width = width self.__height = height self.__resize_target = resize_target self.__keep_aspect_ratio = keep_aspect_ratio self.__multiple_of = ensure_multiple_of self.__resize_method = resize_method self.__image_interpolation_method = image_interpolation_method def constrain_to_multiple_of(self, x, min_val=0, max_val=None): y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int) if max_val is not None and y > max_val: y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int) if y < min_val: y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int) return y def get_size(self, width, height): # determine new height and width scale_height = self.__height / height scale_width = self.__width / width if self.__keep_aspect_ratio: if self.__resize_method == 'lower_bound': # scale such that output size is lower bound if scale_width > scale_height: # fit width scale_height = scale_width else: # fit height scale_width = scale_height elif self.__resize_method == 'upper_bound': # scale such that output size is upper bound if scale_width < scale_height: # fit width scale_height = scale_width else: # fit height scale_width = scale_height elif self.__resize_method == 'minimal': # scale as least as possbile if abs(1 - scale_width) < abs(1 - scale_height): # fit width scale_height = scale_width else: # fit height scale_width = scale_height else: raise ValueError( f'resize_method {self.__resize_method} not implemented') if self.__resize_method == 'lower_bound': new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height) new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width) elif self.__resize_method == 'upper_bound': new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height) new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width) elif self.__resize_method == 'minimal': new_height = self.constrain_to_multiple_of(scale_height * height) new_width = self.constrain_to_multiple_of(scale_width * width) else: raise ValueError( f'resize_method {self.__resize_method} not implemented') return (new_width, new_height) def __call__(self, sample): width, height = self.get_size(sample['image'].shape[1], sample['image'].shape[0]) # resize sample sample['image'] = cv2.resize( sample['image'], (width, height), interpolation=self.__image_interpolation_method, ) if self.__resize_target: if 'disparity' in sample: sample['disparity'] = cv2.resize( sample['disparity'], (width, height), interpolation=cv2.INTER_NEAREST, ) if 'depth' in sample: sample['depth'] = cv2.resize(sample['depth'], (width, height), interpolation=cv2.INTER_NEAREST) sample['mask'] = cv2.resize( sample['mask'].astype(np.float32), (width, height), interpolation=cv2.INTER_NEAREST, ) sample['mask'] = sample['mask'].astype(bool) return sample class NormalizeImage(object): """Normlize image by given mean and std. """ def __init__(self, mean, std): self.__mean = mean self.__std = std def __call__(self, sample): sample['image'] = (sample['image'] - self.__mean) / self.__std return sample class PrepareForNet(object): """Prepare sample for usage as network input. """ def __init__(self): pass def __call__(self, sample): image = np.transpose(sample['image'], (2, 0, 1)) sample['image'] = np.ascontiguousarray(image).astype(np.float32) if 'mask' in sample: sample['mask'] = sample['mask'].astype(np.float32) sample['mask'] = np.ascontiguousarray(sample['mask']) if 'disparity' in sample: disparity = sample['disparity'].astype(np.float32) sample['disparity'] = np.ascontiguousarray(disparity) if 'depth' in sample: depth = sample['depth'].astype(np.float32) sample['depth'] = np.ascontiguousarray(depth) return sample