Spaces:
Sleeping
Sleeping
| import PIL | |
| import numpy as np | |
| from PIL import Image | |
| class Croper: | |
| def __init__( | |
| self, | |
| input_image: PIL.Image, | |
| target_mask: np.ndarray, | |
| ): | |
| self.input_image = input_image | |
| self.target_mask = target_mask | |
| def corp_mask_image(self): | |
| target_mask = self.target_mask | |
| input_image = self.input_image | |
| crop_length = 512 | |
| expand_size = 40 | |
| original_width, original_height = input_image.size | |
| mask_indices = np.where(target_mask) | |
| start_y = np.min(mask_indices[0]) - expand_size | |
| if start_y < 0: | |
| start_y = 0 | |
| end_y = np.max(mask_indices[0]) + expand_size | |
| if end_y > original_height: | |
| end_y = original_height | |
| start_x = np.min(mask_indices[1]) - expand_size | |
| if start_x < 0: | |
| start_x = 0 | |
| end_x = np.max(mask_indices[1]) + expand_size | |
| if end_x > original_width: | |
| end_x = original_width | |
| mask_height = end_y - start_y | |
| mask_width = end_x - start_x | |
| # choose the max side length | |
| max_side_length = max(mask_height, mask_width) | |
| # calculate the crop area | |
| crop_mask = target_mask[start_y:end_y, start_x:end_x] | |
| crop_mask_start_y = (max_side_length - mask_height) // 2 | |
| crop_mask_end_y = crop_mask_start_y + mask_height | |
| crop_mask_start_x = (max_side_length - mask_width) // 2 | |
| crop_mask_end_x = crop_mask_start_x + mask_width | |
| # create a square mask | |
| square_mask = np.zeros((max_side_length, max_side_length), dtype=target_mask.dtype) | |
| square_mask[crop_mask_start_y:crop_mask_end_y, crop_mask_start_x:crop_mask_end_x] = crop_mask | |
| square_mask_image = Image.fromarray((square_mask * 255).astype(np.uint8)) | |
| crop_image = input_image.crop((start_x, start_y, end_x, end_y)) | |
| square_image = Image.new("RGB", (max_side_length, max_side_length)) | |
| square_image.paste(crop_image, (crop_mask_start_x, crop_mask_start_y)) | |
| self.origin_start_x = start_x | |
| self.origin_start_y = start_y | |
| self.origin_end_x = end_x | |
| self.origin_end_y = end_y | |
| self.square_start_x = crop_mask_start_x | |
| self.square_start_y = crop_mask_start_y | |
| self.square_end_x = crop_mask_end_x | |
| self.square_end_y = crop_mask_end_y | |
| self.square_length = max_side_length | |
| self.square_mask_image = square_mask_image | |
| self.square_image = square_image | |
| self.resized_square_mask_image = square_mask_image.resize((crop_length, crop_length)) | |
| self.resized_square_image = square_image.resize((crop_length, crop_length)) | |
| return self.square_image, self.resized_square_image |