# MIT License # # Copyright (c) 2019-2024 Iván de Paz Centeno # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in all # copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. import tensorflow as tf import numpy as np from tensorflow.python.framework.errors_impl import NotFoundError def build_scale_pyramid(width, height, min_face_size, scale_factor, min_size=12): """ Builds a scale pyramid for detecting objects (e.g., faces) at different sizes in an image. Args: width (int): The width of the input image, in pixels. height (int): The height of the input image, in pixels. min_face_size (int): The minimum size (in pixels) of the object (e.g., face) that should be detectable. scale_factor (float): The factor by which the image is downscaled at each level (e.g., 0.7 means 70% of the previous scale). min_size (int, optional): The smallest size (in pixels) to which the image can be downscaled. Default is 12 pixels. Returns: np.ndarray: Array of scales to apply to the image at each level of the pyramid. """ # Find the smallest dimension of the image min_dim = min(width, height) # Calculate how many scales are needed based on the smallest dimension and the scale factor scales_count = round(-((np.log(min_dim / min_size) / np.log(scale_factor)) + 1)) # Calculate the base scale value (based on the smallest detectable face size) m = min_size / min_face_size # Generate an array of scales for the pyramid return m * scale_factor ** np.arange(scales_count) def scale_images(images, scale: float=None, new_shape: tuple=None): """ Scales the input images either by a given factor or to a specified new shape. Args: images (np.ndarray or tf.Tensor): A batch of images or a single image. The expected input should have the shape (..., height, width, channels), where the last three dimensions represent the height, width, and color channels of the images. scale (float, optional): A scaling factor to resize the images. For example, a value of 0.5 will reduce the image to 50% of its original size, while 2.0 will double its size. This parameter is ignored if `new_shape` is provided. new_shape (tuple, optional): A tuple specifying the new shape (height, width) to resize the images to. If provided, this will override the scaling factor. Returns: tf.Tensor: The scaled images as a tensor with resized dimensions, determined either by the scaling factor or the new shape provided. """ # Extract the shape from the images shape = np.asarray(images.shape[-3:-1]) if scale is None and new_shape is None: new_shape = shape new_shape = shape * scale if new_shape is None else new_shape # Resize the images using the specified scaling factor images_scaled = tf.image.resize(images, new_shape, method=tf.image.ResizeMethod.AREA) return images_scaled def normalize_images(images): """ Normalizes the input images to the range [-1, 1]. Args: images (np.ndarray or tf.Tensor): A batch of images or a single image. The expected input should have pixel values in the range [0, 255]. Returns: np.ndarray or tf.Tensor: The normalized images, where pixel values are rescaled to the range [-1, 1]. """ # Normalize the images to the range [-1, 1] return (images - 127.5) / 128 def pad_stack_np(images, justification="center"): """ Pads and stacks a list of images to ensure they all have the same size, based on the specified justification. Args: images (list of np.ndarray): A list of images with varying shapes. Each image is expected to be a NumPy array. justification (str, optional): Specifies the justification (alignment) for padding. Available options are: "center", "top", "topleft", "topright", "bottom", "bottomleft", "bottomright", "left", "right". Default is "center". Returns: np.ndarray: A stacked NumPy array where all images have been padded to the same size, based on the chosen justification. np.ndarray: A stacked NumPy array of each original shape. np.ndarray: A NumPy array containing the padding parameters applied to each image. """ # Stack the shapes of all images into an array sizes_stack = np.stack([img.shape for img in images], axis=0) # Find the maximum shape along each dimension sizes_max = sizes_stack.max(axis=0, keepdims=True) # Calculate the difference in size for padding sizes_diff = sizes_max - sizes_stack # Calculate if any padding size is odd, to adjust padding sizes_mod = sizes_diff % 2 sizes_diff = sizes_diff - sizes_mod # Justification masks for padding alignment justification_mask = { "top": np.asarray([[[0, 1], [0.5, 0.5], [0, 0]]]), "topleft": np.asarray([[[0, 1], [0, 1], [0, 0]]]), "topright": np.asarray([[[0, 1], [1, 0], [0, 0]]]), "bottom": np.asarray([[[1, 0], [0.5, 0.5], [0, 0]]]), "bottomleft": np.asarray([[[1, 0], [0, 1], [0, 0]]]), "bottomright": np.asarray([[[1, 0], [1, 0], [0, 0]]]), "left": np.asarray([[[0.5, 0.5], [0, 1], [0, 0]]]), "right": np.asarray([[[0.5, 0.5], [1, 0], [0, 0]]]), "center": np.asarray([[[0.5, 0.5], [0.5, 0.5], [0, 0]]]), } # Justification adjustments for padding if needed justification_pad_mask = { "top": "topleft", "bottom": "bottomleft", "left": "topleft", "right": "topright", "center": "topleft" } # Get the correct padding mask based on justification pad_mask = justification_mask[justification] mod_mask = justification_mask[justification_pad_mask.get(justification, justification)] # Calculate the exact padding parameters pad_param = (pad_mask * sizes_diff[:,:,None] + mod_mask * sizes_mod[:,:,None]).astype(int) # Apply the calculated padding to each image and stack them into a single array images_padded = np.stack([np.pad(img, pad) for img, pad in zip(images, pad_param)], axis=0) # We keep the original faces to return as extra info original_shapes = np.stack([img.shape for img in images], axis=0) return images_padded, original_shapes, pad_param def pad_stack_tf(images, justification="center"): """ Pads and stacks a list of images to ensure they all have the same size, based on the specified justification. Args: images (list of tf.Tensor): A list of images with varying shapes. Each image is expected to be a TensorFlow tensor. justification (str, optional): Specifies the justification (alignment) for padding. Available options are: "center", "top", "topleft", "topright", "bottom", "bottomleft", "bottomright", "left", "right". Default is "center". Returns: tf.Tensor: A stacked TensorFlow tensor where all images have been padded to the same size, based on the chosen justification. tf.Tensor: A TensorFlow tensor of the original shapes of each image, for reference. tf.Tensor: A TensorFlow tensor containing the padding parameters applied to each image. """ # Stack the shapes of all images into a tensor sizes_stack = tf.stack([tf.shape(img) for img in images], axis=0) # Find the maximum shape along each dimension sizes_max = tf.reduce_max(sizes_stack, axis=0, keepdims=True) # Calculate the difference in size for padding sizes_diff = sizes_max - sizes_stack # Calculate if any padding size is odd, to adjust padding sizes_mod = tf.cast(sizes_diff % 2, tf.float32) sizes_diff = tf.cast(sizes_diff, tf.float32) - sizes_mod # Justification masks for padding alignment justification_mask = { "top": tf.constant([[[0, 1.], [0.5, 0.5], [0, 0]]]), "topleft": tf.constant([[[0, 1.], [0, 1.], [0, 0]]]), "topright": tf.constant([[[0, 1.], [1., 0], [0, 0]]]), "bottom": tf.constant([[[1., 0], [0.5, 0.5], [0, 0]]]), "bottomleft": tf.constant([[[1., 0], [0, 1.], [0, 0]]]), "bottomright": tf.constant([[[1., 0], [1., 0], [0, 0]]]), "left": tf.constant([[[0.5, 0.5], [0, 1.], [0, 0]]]), "right": tf.constant([[[0.5, 0.5], [1., 0], [0, 0]]]), "center": tf.constant([[[0.5, 0.5], [0.5, 0.5], [0, 0]]]), } # Justification adjustments for padding if needed justification_pad_mask = { "top": "topleft", "bottom": "bottomleft", "left": "topleft", "right": "topright", "center": "topleft" } # Get the correct padding mask based on justification pad_mask = justification_mask[justification] mod_mask = justification_mask[justification_pad_mask.get(justification, justification)] # Calculate the exact padding parameters pad_param = (pad_mask * sizes_diff[:,:,None] + mod_mask * sizes_mod[:,:,None]) pad_param = tf.cast(pad_param, tf.int32) # Apply the calculated padding to each image and stack them into a single tensor images_padded = tf.stack([tf.pad(img, paddings=pad) for img, pad in zip(images, pad_param)], axis=0) # We keep the original faces to return as extra info original_shapes = tf.stack([tf.shape(img) for img in images], axis=0) return images_padded, original_shapes, pad_param def ensure_stack(images): """ Ensures that the input is a properly stacked array of images. This function should be called to format the input for a given model. Args: images (list or np.ndarray): A list of images or a NumPy array of images. If it's a list, images will be padded and stacked. If it is a single image, the image's dimension will be expanded as if it were a list of a single image. Returns: np.ndarray: A properly stacked array of images, ensuring they have the same shape. """ # If images is a list, pad and stack them if isinstance(images, list): images = pad_stack_np(images) # Broadcast to ensure the images have a consistent shape (batch dimension) return np.broadcast_to(images, [(len(images.shape) < 4) + (len(images.shape) >= 4) * images.shape[0],] + list(images.shape[len(images.shape) >= 4:])) def load_image(image, dtype=tf.float32, device="CPU:0"): """ Loads an image and decodes it into a TensorFlow tensor, optionally normalizing it. Args: image (str, np.ndarray, or tf.Tensor): The input image. It can be: - A file path to the image as a string. - A TensorFlow tensor or a NumPy array representing an image. dtype (tf.DType, optional): The desired data type for the decoded image. Default is tf.float32. device (str, optional): the target device for the operation. Using CPU most of the times should be fine. Returns: tf.Tensor: The decoded image tensor, with shape (height, width, channels) and dtype `dtype`. If `normalize=True`, the image values will be scaled to the range [0, 1]. """ with tf.device(device): is_tensor = tf.is_tensor(image) or isinstance(image, np.ndarray) if is_tensor: decoded_image = image else: try: if isinstance(image, str): image_data = tf.io.read_file(image) # Read image from file else: image_data = image # Assume image data is provided directly except NotFoundError: image_data = image # If file not found, use the input directly # Decode the image with 3 channels (RGB) decoded_image = tf.image.decode_image(image_data, channels=3, dtype=dtype).numpy() # If dtype is float, adjust the image scale if dtype in [tf.float16, tf.float32]: decoded_image *= 255 # Convert pixel values to [0, 255] if using float data type return decoded_image def load_images_batch(images, dtype=tf.float32, device="CPU:0"): """ Loads a batch of images into memory. If the images are not already in tensor or NumPy array format, they are loaded from their file paths. Args: images (list of str, np.ndarray, or tf.Tensor): A list of images, where each image can either be a TensorFlow tensor, NumPy array, or a file path (string). dtype (tf.DType, optional): The data type for the loaded images. Default is tf.float32. device (str, optional): the target device for the operation. Using CPU most of the times should be fine. Returns: list of tf.Tensor: A list of TensorFlow tensors representing the raw images. """ is_tensor = tf.is_tensor(images[0]) or isinstance(images[0], np.ndarray) images_raw = images if is_tensor else [load_image(img, dtype=dtype, device=device) for img in images] return images_raw def standarize_batch(images_raw, normalize=True, justification="center"): """ Pads and stacks a batch of images to ensure they all have the same size, with an option to normalize them. Args: images_raw (list of tf.Tensor or np.ndarray): A list of raw images, each either as a TensorFlow tensor or a NumPy array. normalize (bool, optional): Whether to normalize the images after stacking. Default is True. justification (str, optional): The alignment for padding the images. Available options are: "center", "top", "topleft", "topright", "bottom", "bottomleft", "bottomright", "left", "right". Default is "center". Returns: np.ndarray: A stacked array of images, padded to the same shape based on the chosen justification. np.ndarray: An array containing the original shapes of each image before padding. np.ndarray: The padding parameters applied to each image. """ images_result, images_oshapes, pad_param = pad_stack_np(images_raw, justification=justification) if normalize: images_result = normalize_images(images_result) return images_result, images_oshapes, pad_param def apply_scales(images_normalized, scales_groups): """ Applies scales to the normalized images based on the largest group of scales. Args: images_normalized (np.ndarray): A normalized image or batch of images. scales_groups (list of np.ndarray): A list of different scale groups, where each group contains multiple possible scales. Returns: tuple: - list of np.ndarray: A list of images scaled according to the largest group of scales. - np.ndarray: The largest group of scales that was applied to the images. """ # Select the scale group with the largest number of elements selected_scaleset_as_index = np.argmax([x.shape[0] for x in scales_groups]) largest_scale_group_set = scales_groups[selected_scaleset_as_index] # Apply the scales from the largest scale group to the normalized images result = [scale_images(images_normalized, scale) for scale in largest_scale_group_set] return result, largest_scale_group_set def extract_patches(images_normalized, bboxes_batch, expected_size=(24, 24)): """ Extracts patches from a batch of normalized images based on bounding boxes, and resizes each patch to a specified size. Args: images_normalized (tf.Tensor): A batch of images or a single image, normalized, with the shape (batch_size, height, width, channels). bboxes_batch (np.ndarray): An array of bounding boxes of shape (n, 5), where each row represents [batch_index, x1, y1, x2, y2]. The coordinates are in pixel format, and the first column indicates the image index in the batch. expected_size (tuple, optional): A tuple specifying the size (height, width) to resize each extracted patch. Defaults to (24, 24). Returns: tf.Tensor: A tensor of extracted patches with shape (n, height, width, channels), where each patch corresponds to a bounding box in `bboxes_batch`, resized to `expected_size`. """ # Get the shape of the input images shape = images_normalized.shape # Normalize the bounding box coordinates to be within [0, 1] relative to image dimensions selector = [2, 1, 4, 3] bboxes_batch_coords = bboxes_batch[:, selector] / np.asarray([[shape[selector[1]], shape[selector[0]], shape[selector[1]], shape[selector[0]]]]) # Extract patches from the images using the bounding boxes, resizing them to `expected_size` result = tf.image.crop_and_resize( images_normalized, # Input image tensor bboxes_batch_coords, # Bounding boxes in format [y1, x1, y2, x2], normalized to [0.0, 1.0] bboxes_batch[:, 0].astype(int), # Indices of the images in the batch corresponding to the bounding boxes expected_size # Size to resize the cropped patches (height, width) ) return result