| import numpy as np |
| from PIL import Image |
|
|
| def expand2square(pil_img, background_color): |
| width, height = pil_img.size |
| if width == height: |
| return pil_img |
| elif width > height: |
| result = Image.new(pil_img.mode, (width, width), background_color) |
| result.paste(pil_img, (0, (width - height) // 2)) |
| return result |
| else: |
| result = Image.new(pil_img.mode, (height, height), background_color) |
| result.paste(pil_img, ((height - width) // 2, 0)) |
| return result |
|
|
| def expand2square_mask(mask): |
| |
| n_mask, width, height = mask.shape |
| if width == height: |
| return mask |
| elif width > height: |
| n_pad = width - height |
| n_pad_1 = n_pad // 2 |
| n_pad_2 = n_pad - n_pad_1 |
| pad_mask_1 = np.zeros((n_mask, width, n_pad_1), dtype=np.uint8) |
| pad_mask_2 = np.zeros((n_mask, width, n_pad_2), dtype=np.uint8) |
| result = np.concatenate([pad_mask_1, mask, pad_mask_2], axis=2) |
| return result |
| else: |
| n_pad = height - width |
| n_pad_1 = n_pad // 2 |
| n_pad_2 = n_pad - n_pad_1 |
| pad_mask_1 = np.zeros((n_mask, n_pad_1, height), dtype=np.uint8) |
| pad_mask_2 = np.zeros((n_mask, n_pad_2, height), dtype=np.uint8) |
| result = np.concatenate([pad_mask_1, mask, pad_mask_2], axis=1) |
| return result |
|
|
| def expand2square_bbox(bboxes, width, height): |
| bboxes = np.array(bboxes) |
| if width == height: |
| return bboxes |
| elif width > height: |
| n_pad = width - height |
| n_pad_1 = n_pad // 2 |
| n_pad_2 = n_pad - n_pad_1 |
| bboxes[:, 1] += n_pad_1 |
| return bboxes |
| else: |
| n_pad = height - width |
| n_pad_1 = n_pad // 2 |
| n_pad_2 = n_pad - n_pad_1 |
| bboxes[:, 0] += n_pad_1 |
| return bboxes |
|
|
| def expand2square_points(points, width, height): |
| if width == height: |
| return points |
| elif width > height: |
| n_pad = width - height |
| n_pad_1 = n_pad // 2 |
| n_pad_2 = n_pad - n_pad_1 |
| points[:, 1] += n_pad_1 |
| return points |
| else: |
| n_pad = height - width |
| n_pad_1 = n_pad // 2 |
| n_pad_2 = n_pad - n_pad_1 |
| points[:, 0] += n_pad_1 |
| return points |
|
|
|
|
