| import cv2 | |
| from PIL import Image | |
| import numpy as np | |
| def get_mask_bbox(mask_org, threshold=0): | |
| ''' | |
| Returns the minimum bounding box of the mask. If the shadow area is all 0, a single-pixel image is symbolically returned. | |
| ''' | |
| mask = mask_org.copy() | |
| mask[mask > threshold] = 255 | |
| mask[mask <= threshold] = 0 | |
| h, w = mask.shape[:2] | |
| if len(mask.shape) == 3: | |
| coords = np.where(np.mean(mask, axis=-1, keepdims=False) > threshold) | |
| else: | |
| coords = np.where(mask > 20) | |
| if len(coords[0]) > 0: | |
| ymin, ymax = coords[0].min(), coords[0].max() | |
| else: | |
| ymin, ymax = h // 2, h // 2 + 1 | |
| if len(coords[1]) > 0: | |
| xmin, xmax = coords[1].min(), coords[1].max() | |
| else: | |
| xmin, xmax = w // 2, w // 2 + 1 | |
| return (xmin, xmax, ymin, ymax) | |
| def padding_image(img, padding_rate, canve_reso=1024, padding_color=(0, 0, 0)): | |
| ''' | |
| img:PIL.Image | |
| Scale the image proportionally and place it on a white square canvas, | |
| making sure the long side of the original image is filled with padding_rate on the edge of the canvas. | |
| ''' | |
| long_size = int(canve_reso*(1-padding_rate)) | |
| img = np.array(img) | |
| h, w = img.shape[:2] | |
| if h > w: | |
| new_h = long_size | |
| new_w = int(w / h * long_size) | |
| else: | |
| new_w = long_size | |
| new_h = int(h / w * long_size) | |
| resized_img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA) | |
| padding_h = (canve_reso - new_h) // 2 | |
| padding_w = (canve_reso - new_w) // 2 | |
| padding_output = np.ones((canve_reso, canve_reso, resized_img.shape[-1])) * \ | |
| np.array(padding_color)[np.newaxis, np.newaxis, :] | |
| padding_output[padding_h:padding_h + new_h, padding_w:padding_w + new_w, :] = resized_img | |
| return Image.fromarray(padding_output.astype(np.uint8)) |