sactharaka/ID-Background-Edit

1import cv2
2from matplotlib import pyplot as plt
3import PIL.Image as Image
4import numpy as np
5
6
7def crop_for_filling_pre(image: np.array, mask: np.array, crop_size: int = 512):
8    # Calculate the aspect ratio of the image
9    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 proportionally
13    if min(height, width) < crop_size:
14        if height < width:
15            new_height = crop_size
16            new_width = int(new_height * aspect_ratio)
17        else:
18            new_width = crop_size
19            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 mask
25    x, y, w, h = cv2.boundingRect(mask)
26
27    # Update the height and width of the resized image
28    height, width = image.shape[:2]
29
30    # # If the 512x512 square cannot cover the entire mask, resize the image accordingly
31    if w > crop_size or h > crop_size:
32        # padding to square at first
33        if height < width:
34            padding = width - height
35            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 - width
39            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 coordinates
48    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 image
52    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_mask
56    
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 image
67    height, width = image.shape[:2]
68    height_ori, width_ori = height, width
69    aspect_ratio = float(width) / float(height)
70
71    # If the shorter side is less than 512, resize the image proportionally
72    if min(height, width) < crop_size:
73        if height < width:
74            new_height = crop_size
75            new_width = int(new_height * aspect_ratio)
76        else:
77            new_width = crop_size
78            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 mask
84    x, y, w, h = cv2.boundingRect(mask)
85
86    # Update the height and width of the resized image
87    height, width = image.shape[:2]
88
89    # # If the 512x512 square cannot cover the entire mask, resize the image accordingly
90    if w > crop_size or h > crop_size:
91        flag_padding = True
92        # padding to square at first
93        if height < width:
94            padding = width - height
95            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 - width
100            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 = False
110
111    # Calculate the crop coordinates
112    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 image
116    image[crop_y:crop_y + crop_size, crop_x:crop_x + crop_size] = filled_image
117    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_copy
128
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 mask
146    # ^ ------------------------------------------------------------------------------------
147
148    # e.g.
149    # cropped_image[cropped_mask==255] = 0
150    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