WompUniversity/Inpaint-Anything-no-errors

1import cv22import numpy as np3from typing import Tuple4 5def resize_and_pad(image: np.ndarray, mask: np.ndarray, target_size: int = 512) -> Tuple[np.ndarray, np.ndarray]:6    """7    Resizes an image and its corresponding mask to have the longer side equal to `target_size` and pads them to make them8    both have the same size. The resulting image and mask have dimensions (target_size, target_size).9 10    Args:11        image: A numpy array representing the image to resize and pad.12        mask: A numpy array representing the mask to resize and pad.13        target_size: An integer specifying the desired size of the longer side after resizing.14 15    Returns:16        A tuple containing two numpy arrays - the resized and padded image and the resized and padded mask.17    """18    height, width, _ = image.shape19    max_dim = max(height, width)20    scale = target_size / max_dim21    new_height = int(height * scale)22    new_width = int(width * scale)23    image_resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LINEAR)24    mask_resized = cv2.resize(mask, (new_width, new_height), interpolation=cv2.INTER_LINEAR)25    pad_height = target_size - new_height26    pad_width = target_size - new_width27    top_pad = pad_height // 228    bottom_pad = pad_height - top_pad29    left_pad = pad_width // 230    right_pad = pad_width - left_pad31    image_padded = np.pad(image_resized, ((top_pad, bottom_pad), (left_pad, right_pad), (0, 0)), mode='constant')32    mask_padded = np.pad(mask_resized, ((top_pad, bottom_pad), (left_pad, right_pad)), mode='constant')33    return image_padded, mask_padded, (top_pad, bottom_pad, left_pad, right_pad)34 35def recover_size(image_padded: np.ndarray, mask_padded: np.ndarray, orig_size: Tuple[int, int], 36                 padding_factors: Tuple[int, int, int, int]) -> Tuple[np.ndarray, np.ndarray]:37    """38    Resizes a padded and resized image and mask to the original size.39 40    Args:41        image_padded: A numpy array representing the padded and resized image.42        mask_padded: A numpy array representing the padded and resized mask.43        orig_size: A tuple containing two integers - the original height and width of the image before resizing and padding.44 45    Returns:46        A tuple containing two numpy arrays - the recovered image and the recovered mask with dimensions `orig_size`.47    """48    h,w,c = image_padded.shape49    top_pad, bottom_pad, left_pad, right_pad = padding_factors50    image = image_padded[top_pad:h-bottom_pad, left_pad:w-right_pad, :]51    mask = mask_padded[top_pad:h-bottom_pad, left_pad:w-right_pad]52    image_resized = cv2.resize(image, orig_size[::-1], interpolation=cv2.INTER_LINEAR)53    mask_resized = cv2.resize(mask, orig_size[::-1], interpolation=cv2.INTER_LINEAR)54    return image_resized, mask_resized55 56 57 58 59if __name__ == '__main__':60 61    # image = cv2.imread('example/boat.jpg')62    # mask = cv2.imread('example/boat_mask_2.png', cv2.IMREAD_GRAYSCALE)63    # image = cv2.imread('example/groceries.jpg')64    # mask = cv2.imread('example/groceries_mask_2.png', cv2.IMREAD_GRAYSCALE)65    # image = cv2.imread('example/bridge.jpg')66    # mask = cv2.imread('example/bridge_mask_2.png', cv2.IMREAD_GRAYSCALE)67    # image = cv2.imread('example/person_umbrella.jpg')68    # mask = cv2.imread('example/person_umbrella_mask_2.png', cv2.IMREAD_GRAYSCALE)69    # image = cv2.imread('example/hippopotamus.jpg')70    # mask = cv2.imread('example/hippopotamus_mask_1.png', cv2.IMREAD_GRAYSCALE)71    image = cv2.imread('/data1/yutao/projects/IAM/Inpaint-Anything/example/fill-anything/sample5.jpeg')72    mask = cv2.imread('/data1/yutao/projects/IAM/Inpaint-Anything/example/fill-anything/sample5/mask.png', cv2.IMREAD_GRAYSCALE)73    print(image.shape)74    print(mask.shape)75    cv2.imwrite('original_image.jpg', image)76    cv2.imwrite('original_mask.jpg', mask)77    image_padded, mask_padded, padding_factors = resize_and_pad(image, mask)78    cv2.imwrite('padded_image.png', image_padded)79    cv2.imwrite('padded_mask.png', mask_padded)80    print(image_padded.shape, mask_padded.shape, padding_factors)81 82    # ^ ------------------------------------------------------------------------------------83    # ^ Please conduct inpainting or filling here on the cropped image with the cropped mask84    # ^ ------------------------------------------------------------------------------------85 86    # resize and pad the image and mask87 88    # perform some operation on the 512x512 image and mask89    # ...90 91    # recover the image and mask to the original size92    height, width, _ = image.shape93    image_resized, mask_resized = recover_size(image_padded, mask_padded, (height, width), padding_factors)94 95    # save the resized and recovered image and mask96    cv2.imwrite('resized_and_padded_image.png', image_padded)97    cv2.imwrite('resized_and_padded_mask.png', mask_padded)98    cv2.imwrite('recovered_image.png', image_resized)99    cv2.imwrite('recovered_mask.png', mask_resized)100 101