| from __future__ import annotations | |
| import numpy as np | |
| from PIL import Image | |
| def pad_image_to_square(image: Image.Image) -> Image.Image: | |
| """Pad an image with black pixels to create a square image. | |
| This keeps the localization coordinate space consistent with the MedGemma | |
| localization notebook while avoiding a second uint8 rescaling step. | |
| """ | |
| rgb_image = image.convert("RGB") | |
| image_array = np.asarray(rgb_image, dtype=np.uint8) | |
| height, width = image_array.shape[:2] | |
| max_dim = max(height, width) | |
| pad_top = (max_dim - height) // 2 | |
| pad_bottom = max_dim - height - pad_top | |
| pad_left = (max_dim - width) // 2 | |
| pad_right = max_dim - width - pad_left | |
| padded = np.pad( | |
| image_array, | |
| ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), | |
| mode="constant", | |
| constant_values=0, | |
| ) | |
| return Image.fromarray(padded) | |