ICML2022/resefa
4
1# python3.72"""Contains utility functions for image processing.3 4The module is primarily built on `cv2`. But, differently, we assume all colorful5images are with `RGB` channel order by default. Also, we assume all gray-scale6images to be with shape [height, width, 1].7"""8 9import os10import cv211import numpy as np12 13from .misc import IMAGE_EXTENSIONS14from .misc import check_file_ext15 16__all__ = [17 'get_blank_image', 'load_image', 'save_image', 'resize_image',18 'add_text_to_image', 'preprocess_image', 'postprocess_image',19 'parse_image_size', 'get_grid_shape', 'list_images_from_dir'20]21 22 23def _check_2d_image(image):24 """Checks whether a given image is valid.25 26 A valid image is expected to be with dtype `uint8`. Also, it should have27 shape like:28 29 (1) (height, width, 1) # gray-scale image.30 (2) (height, width, 3) # colorful image.31 (3) (height, width, 4) # colorful image with transparency (RGBA)32 """33 assert isinstance(image, np.ndarray)34 assert image.dtype == np.uint835 assert image.ndim == 3 and image.shape[2] in [1, 3, 4]36 37 38def get_blank_image(height, width, channels=3, use_black=True):39 """Gets a blank image, either white of black.40 41 NOTE: This function will always return an image with `RGB` channel order for42 color image and pixel range [0, 255].43 44 Args:45 height: Height of the returned image.46 width: Width of the returned image.47 channels: Number of channels. (default: 3)48 use_black: Whether to return a black image. (default: True)49 """50 shape = (height, width, channels)51 if use_black:52 return np.zeros(shape, dtype=np.uint8)53 return np.ones(shape, dtype=np.uint8) * 25554 55 56def load_image(path):57 """Loads an image from disk.58 59 NOTE: This function will always return an image with `RGB` channel order for60 color image and pixel range [0, 255].61 62 Args:63 path: Path to load the image from.64 65 Returns:66 An image with dtype `np.ndarray`, or `None` if `path` does not exist.67 """68 image = cv2.imread(path, cv2.IMREAD_UNCHANGED)69 if image is None:70 return None71 72 if image.ndim == 2:73 image = image[:, :, np.newaxis]74 _check_2d_image(image)75 if image.shape[2] == 3:76 return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)77 if image.shape[2] == 4:78 return cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)79 return image80 81 82def save_image(path, image):83 """Saves an image to disk.84 85 NOTE: The input image (if colorful) is assumed to be with `RGB` channel86 order and pixel range [0, 255].87 88 Args:89 path: Path to save the image to.90 image: Image to save.91 """92 if image is None:93 return94 95 _check_2d_image(image)96 if image.shape[2] == 1:97 cv2.imwrite(path, image)98 elif image.shape[2] == 3:99 cv2.imwrite(path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))100 elif image.shape[2] == 4:101 cv2.imwrite(path, cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA))102 103 104def resize_image(image, *args, **kwargs):105 """Resizes image.106 107 This is a wrap of `cv2.resize()`.108 109 NOTE: The channel order of the input image will not be changed.110 111 Args:112 image: Image to resize.113 *args: Additional positional arguments.114 **kwargs: Additional keyword arguments.115 116 Returns:117 An image with dtype `np.ndarray`, or `None` if `image` is empty.118 """119 if image is None:120 return None121 122 _check_2d_image(image)123 if image.shape[2] == 1: # Re-expand the squeezed dim of gray-scale image.124 return cv2.resize(image, *args, **kwargs)[:, :, np.newaxis]125 return cv2.resize(image, *args, **kwargs)126 127 128def add_text_to_image(image,129 text='',130 position=None,131 font=cv2.FONT_HERSHEY_TRIPLEX,132 font_size=1.0,133 line_type=cv2.LINE_8,134 line_width=1,135 color=(255, 255, 255)):136 """Overlays text on given image.137 138 NOTE: The input image is assumed to be with `RGB` channel order.139 140 Args:141 image: The image to overlay text on.142 text: Text content to overlay on the image. (default: empty)143 position: Target position (bottom-left corner) to add text. If not set,144 center of the image will be used by default. (default: None)145 font: Font of the text added. (default: cv2.FONT_HERSHEY_TRIPLEX)146 font_size: Font size of the text added. (default: 1.0)147 line_type: Line type used to depict the text. (default: cv2.LINE_8)148 line_width: Line width used to depict the text. (default: 1)149 color: Color of the text added in `RGB` channel order. (default:150 (255, 255, 255))151 152 Returns:153 An image with target text overlaid on.154 """155 if image is None or not text:156 return image157 158 _check_2d_image(image)159 cv2.putText(img=image,160 text=text,161 org=position,162 fontFace=font,163 fontScale=font_size,164 color=color,165 thickness=line_width,166 lineType=line_type,167 bottomLeftOrigin=False)168 return image169 170 171def preprocess_image(image, min_val=-1.0, max_val=1.0):172 """Pre-processes image by adjusting the pixel range and to dtype `float32`.173 174 This function is particularly used to convert an image or a batch of images175 to `NCHW` format, which matches the data type commonly used in deep models.176 177 NOTE: The input image is assumed to be with pixel range [0, 255] and with178 format `HWC` or `NHWC`. The returned image will be always be with format179 `NCHW`.180 181 Args:182 image: The input image for pre-processing.183 min_val: Minimum value of the output image.184 max_val: Maximum value of the output image.185 186 Returns:187 The pre-processed image.188 """189 assert isinstance(image, np.ndarray)190 191 image = image.astype(np.float64)192 image = image / 255.0 * (max_val - min_val) + min_val193 194 if image.ndim == 3:195 image = image[np.newaxis]196 assert image.ndim == 4 and image.shape[3] in [1, 3, 4]197 return image.transpose(0, 3, 1, 2)198 199 200def postprocess_image(image, min_val=-1.0, max_val=1.0):201 """Post-processes image to pixel range [0, 255] with dtype `uint8`.202 203 This function is particularly used to handle the results produced by deep204 models.205 206 NOTE: The input image is assumed to be with format `NCHW`, and the returned207 image will always be with format `NHWC`.208 209 Args:210 image: The input image for post-processing.211 min_val: Expected minimum value of the input image.212 max_val: Expected maximum value of the input image.213 214 Returns:215 The post-processed image.216 """217 assert isinstance(image, np.ndarray)218 219 image = image.astype(np.float64)220 image = (image - min_val) / (max_val - min_val) * 255221 image = np.clip(image + 0.5, 0, 255).astype(np.uint8)222 223 assert image.ndim == 4 and image.shape[1] in [1, 3, 4]224 return image.transpose(0, 2, 3, 1)225 226 227def parse_image_size(obj):228 """Parses an object to a pair of image size, i.e., (height, width).229 230 Args:231 obj: The input object to parse image size from.232 233 Returns:234 A two-element tuple, indicating image height and width respectively.235 236 Raises:237 If the input is invalid, i.e., neither a list or tuple, nor a string.238 """239 if obj is None or obj == '':240 height = 0241 width = 0242 elif isinstance(obj, int):243 height = obj244 width = obj245 elif isinstance(obj, (list, tuple, str, np.ndarray)):246 if isinstance(obj, str):247 splits = obj.replace(' ', '').split(',')248 numbers = tuple(map(int, splits))249 else:250 numbers = tuple(obj)251 if len(numbers) == 0:252 height = 0253 width = 0254 elif len(numbers) == 1:255 height = int(numbers[0])256 width = int(numbers[0])257 elif len(numbers) == 2:258 height = int(numbers[0])259 width = int(numbers[1])260 else:261 raise ValueError('At most two elements for image size.')262 else:263 raise ValueError(f'Invalid type of input: `{type(obj)}`!')264 265 return (max(0, height), max(0, width))266 267 268def get_grid_shape(size, height=0, width=0, is_portrait=False):269 """Gets the shape of a grid based on the size.270 271 This function makes greatest effort on making the output grid square if272 neither `height` nor `width` is set. If `is_portrait` is set as `False`, the273 height will always be equal to or smaller than the width. For example, if274 input `size = 16`, output shape will be `(4, 4)`; if input `size = 15`,275 output shape will be (3, 5). Otherwise, the height will always be equal to276 or larger than the width.277 278 Args:279 size: Size (height * width) of the target grid.280 height: Expected height. If `size % height != 0`, this field will be281 ignored. (default: 0)282 width: Expected width. If `size % width != 0`, this field will be283 ignored. (default: 0)284 is_portrait: Whether to return a portrait size of a landscape size.285 (default: False)286 287 Returns:288 A two-element tuple, representing height and width respectively.289 """290 assert isinstance(size, int)291 assert isinstance(height, int)292 assert isinstance(width, int)293 if size <= 0:294 return (0, 0)295 296 if height > 0 and width > 0 and height * width != size:297 height = 0298 width = 0299 300 if height > 0 and width > 0 and height * width == size:301 return (height, width)302 if height > 0 and size % height == 0:303 return (height, size // height)304 if width > 0 and size % width == 0:305 return (size // width, width)306 307 height = int(np.sqrt(size))308 while height > 0:309 if size % height == 0:310 width = size // height311 break312 height = height - 1313 314 return (width, height) if is_portrait else (height, width)315 316 317def list_images_from_dir(directory):318 """Lists all images from the given directory.319 320 NOTE: Do NOT support finding images recursively.321 322 Args:323 directory: The directory to find images from.324 325 Returns:326 A list of sorted filenames, with the directory as prefix.327 """328 image_list = []329 for filename in os.listdir(directory):330 if check_file_ext(filename, *IMAGE_EXTENSIONS):331 image_list.append(os.path.join(directory, filename))332 return sorted(image_list)333