WompUniversity/Inpaint-Anything-no-errors

1import cv22from matplotlib import pyplot as plt3import PIL.Image as Image4import numpy as np5 6 7def crop_for_filling_pre(image: np.array, mask: np.array, crop_size: int = 512):8    # Calculate the aspect ratio of the image9    height, width = image.shape[:2]10    aspect_ratio = float(width) / float(height)11 12    # If the shorter side is less than 512, resize the image proportionally13    if min(height, width) < crop_size:14        if height < width:15            new_height = crop_size16            new_width = int(new_height * aspect_ratio)17        else:18            new_width = crop_size19            new_height = int(new_width / aspect_ratio)20 21        image = cv2.resize(image, (new_width, new_height))22        mask = cv2.resize(mask, (new_width, new_height))23 24    # Find the bounding box of the mask25    x, y, w, h = cv2.boundingRect(mask)26 27    # Update the height and width of the resized image28    height, width = image.shape[:2]29 30    # # If the 512x512 square cannot cover the entire mask, resize the image accordingly31    if w > crop_size or h > crop_size:32        # padding to square at first33        if height < width:34            padding = width - height35            image = np.pad(image, ((padding // 2, padding - padding // 2), (0, 0), (0, 0)), 'constant')36            mask = np.pad(mask, ((padding // 2, padding - padding // 2), (0, 0)), 'constant')37        else:38            padding = height - width39            image = np.pad(image, ((0, 0), (padding // 2, padding - padding // 2), (0, 0)), 'constant')40            mask = np.pad(mask, ((0, 0), (padding // 2, padding - padding // 2)), 'constant')41 42        resize_factor = crop_size / max(w, h)43        image = cv2.resize(image, (0, 0), fx=resize_factor, fy=resize_factor)44        mask = cv2.resize(mask, (0, 0), fx=resize_factor, fy=resize_factor)45        x, y, w, h = cv2.boundingRect(mask)46 47    # Calculate the crop coordinates48    crop_x = min(max(x + w // 2 - crop_size // 2, 0), width - crop_size)49    crop_y = min(max(y + h // 2 - crop_size // 2, 0), height - crop_size)50 51    # Crop the image52    cropped_image = image[crop_y:crop_y + crop_size, crop_x:crop_x + crop_size]53    cropped_mask = mask[crop_y:crop_y + crop_size, crop_x:crop_x + crop_size]54 55    return cropped_image, cropped_mask56    57    58def crop_for_filling_post(59        image: np.array,60        mask: np.array,61        filled_image: np.array, 62        crop_size: int = 512,63        ):64    image_copy = image.copy()65    mask_copy = mask.copy()66    # Calculate the aspect ratio of the image67    height, width = image.shape[:2]68    height_ori, width_ori = height, width69    aspect_ratio = float(width) / float(height)70 71    # If the shorter side is less than 512, resize the image proportionally72    if min(height, width) < crop_size:73        if height < width:74            new_height = crop_size75            new_width = int(new_height * aspect_ratio)76        else:77            new_width = crop_size78            new_height = int(new_width / aspect_ratio)79 80        image = cv2.resize(image, (new_width, new_height))81        mask = cv2.resize(mask, (new_width, new_height))82 83    # Find the bounding box of the mask84    x, y, w, h = cv2.boundingRect(mask)85 86    # Update the height and width of the resized image87    height, width = image.shape[:2]88 89    # # If the 512x512 square cannot cover the entire mask, resize the image accordingly90    if w > crop_size or h > crop_size:91        flag_padding = True92        # padding to square at first93        if height < width:94            padding = width - height95            image = np.pad(image, ((padding // 2, padding - padding // 2), (0, 0), (0, 0)), 'constant')96            mask = np.pad(mask, ((padding // 2, padding - padding // 2), (0, 0)), 'constant')97            padding_side = 'h'98        else:99            padding = height - width100            image = np.pad(image, ((0, 0), (padding // 2, padding - padding // 2), (0, 0)), 'constant')101            mask = np.pad(mask, ((0, 0), (padding // 2, padding - padding // 2)), 'constant')102            padding_side = 'w'103 104        resize_factor = crop_size / max(w, h)105        image = cv2.resize(image, (0, 0), fx=resize_factor, fy=resize_factor)106        mask = cv2.resize(mask, (0, 0), fx=resize_factor, fy=resize_factor)107        x, y, w, h = cv2.boundingRect(mask)108    else:109        flag_padding = False110 111    # Calculate the crop coordinates112    crop_x = min(max(x + w // 2 - crop_size // 2, 0), width - crop_size)113    crop_y = min(max(y + h // 2 - crop_size // 2, 0), height - crop_size)114 115    # Fill the image116    image[crop_y:crop_y + crop_size, crop_x:crop_x + crop_size] = filled_image117    if flag_padding:118        image = cv2.resize(image, (0, 0), fx=1/resize_factor, fy=1/resize_factor)119        if padding_side == 'h':120            image = image[padding // 2:padding // 2 + height_ori, :]121        else:122            image = image[:, padding // 2:padding // 2 + width_ori]123 124    image = cv2.resize(image, (width_ori, height_ori))125 126    image_copy[mask_copy==255] = image[mask_copy==255]127    return image_copy128 129 130if __name__ == '__main__':131 132    # image = cv2.imread('example/boat.jpg')133    # mask = cv2.imread('example/boat_mask_2.png', cv2.IMREAD_GRAYSCALE)134    image = cv2.imread('./example/groceries.jpg')135    mask = cv2.imread('example/groceries_mask_2.png', cv2.IMREAD_GRAYSCALE)136    # image = cv2.imread('example/bridge.jpg')137    # mask = cv2.imread('example/bridge_mask_2.png', cv2.IMREAD_GRAYSCALE)138    # image = cv2.imread('example/person_umbrella.jpg')139    # mask = cv2.imread('example/person_umbrella_mask_2.png', cv2.IMREAD_GRAYSCALE)140    # image = cv2.imread('example/hippopotamus.jpg')141    # mask = cv2.imread('example/hippopotamus_mask_1.png', cv2.IMREAD_GRAYSCALE)142 143    cropped_image, cropped_mask = crop_for_filling_pre(image, mask)144    # ^ ------------------------------------------------------------------------------------145    # ^ Please conduct inpainting or filling here on the cropped image with the cropped mask146    # ^ ------------------------------------------------------------------------------------147 148    # e.g.149    # cropped_image[cropped_mask==255] = 0150    cv2.imwrite('cropped_image.jpg', cropped_image)151    cv2.imwrite('cropped_mask.jpg', cropped_mask)152    print(cropped_image.shape)153    print(cropped_mask.shape)154 155    image = crop_for_filling_post(image, mask, cropped_image)156    cv2.imwrite('filled_image.jpg', image)157    print(image.shape)158 159 160