| """Image preprocessing for QuPAINT. |
| |
| The released checkpoint expects the same pipeline used at tuning time: |
| |
| 1. the micrograph is resized to a fixed 1792x1344 canvas, |
| 2. that canvas is tiled into 448x448 patches by aspect-ratio matching |
| (InternVL-style dynamic tiling, up to `max_num` tiles), |
| 3. a downscaled thumbnail of the whole image is appended as a final tile, |
| 4. every tile is normalised with ImageNet statistics. |
| |
| Because boxes are predicted in *percent* of image size, the fixed canvas does |
| not distort results -- but changing `DEFAULT_CANVAS` moves the model away from |
| the resolution it was tuned at, so keep it unless you know why you are changing it. |
| """ |
|
|
| from typing import List, Tuple |
|
|
| import torch |
| import torchvision.transforms as T |
| from PIL import Image |
| from torchvision.transforms.functional import InterpolationMode |
|
|
| IMAGENET_MEAN = (0.485, 0.456, 0.406) |
| IMAGENET_STD = (0.229, 0.224, 0.225) |
|
|
| DEFAULT_CANVAS: Tuple[int, int] = (1792, 1344) |
| DEFAULT_TILE_SIZE: int = 448 |
| DEFAULT_MAX_TILES: int = 12 |
|
|
|
|
| def build_transform(input_size: int = DEFAULT_TILE_SIZE) -> T.Compose: |
| """Tile transform: RGB, bicubic resize to `input_size`, ImageNet normalisation.""" |
| return T.Compose( |
| [ |
| T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img), |
| T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC), |
| T.ToTensor(), |
| T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), |
| ] |
| ) |
|
|
|
|
| def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size): |
| best_ratio_diff = float("inf") |
| best_ratio = (1, 1) |
| area = width * height |
| for ratio in target_ratios: |
| target_aspect_ratio = ratio[0] / ratio[1] |
| ratio_diff = abs(aspect_ratio - target_aspect_ratio) |
| if ratio_diff < best_ratio_diff: |
| best_ratio_diff = ratio_diff |
| best_ratio = ratio |
| elif ratio_diff == best_ratio_diff: |
| if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]: |
| best_ratio = ratio |
| return best_ratio |
|
|
|
|
| def dynamic_preprocess( |
| image: Image.Image, |
| min_num: int = 1, |
| max_num: int = DEFAULT_MAX_TILES, |
| image_size: int = DEFAULT_TILE_SIZE, |
| use_thumbnail: bool = False, |
| ) -> List[Image.Image]: |
| """Split `image` into `image_size` tiles using the closest allowed aspect ratio.""" |
| orig_width, orig_height = image.size |
| aspect_ratio = orig_width / orig_height |
|
|
| target_ratios = sorted( |
| { |
| (i, j) |
| for n in range(min_num, max_num + 1) |
| for i in range(1, n + 1) |
| for j in range(1, n + 1) |
| if min_num <= i * j <= max_num |
| }, |
| key=lambda x: x[0] * x[1], |
| ) |
|
|
| target_aspect_ratio = find_closest_aspect_ratio( |
| aspect_ratio, target_ratios, orig_width, orig_height, image_size |
| ) |
| target_width = image_size * target_aspect_ratio[0] |
| target_height = image_size * target_aspect_ratio[1] |
| blocks = target_aspect_ratio[0] * target_aspect_ratio[1] |
|
|
| resized_img = image.resize((target_width, target_height)) |
| cols = target_width // image_size |
| processed_images = [] |
| for i in range(blocks): |
| box = ( |
| (i % cols) * image_size, |
| (i // cols) * image_size, |
| ((i % cols) + 1) * image_size, |
| ((i // cols) + 1) * image_size, |
| ) |
| processed_images.append(resized_img.crop(box)) |
| assert len(processed_images) == blocks |
|
|
| if use_thumbnail and len(processed_images) != 1: |
| processed_images.append(image.resize((image_size, image_size))) |
| return processed_images |
|
|
|
|
| def load_image( |
| image_file, |
| input_size: int = DEFAULT_TILE_SIZE, |
| max_num: int = DEFAULT_MAX_TILES, |
| canvas: Tuple[int, int] = DEFAULT_CANVAS, |
| ) -> torch.Tensor: |
| """Load a micrograph and return the stacked tile tensor of shape (tiles, 3, S, S).""" |
| image = image_file |
| if not isinstance(image, Image.Image): |
| image = Image.open(image_file) |
| image = image.convert("RGB").resize(canvas) |
|
|
| transform = build_transform(input_size=input_size) |
| tiles = dynamic_preprocess( |
| image, image_size=input_size, use_thumbnail=True, max_num=max_num |
| ) |
| return torch.stack([transform(tile) for tile in tiles]) |
|
|
