Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from dataclasses import asdict, dataclass | |
| from pathlib import Path | |
| import random | |
| from PIL import Image | |
| class AugmentationSpec: | |
| width: int | |
| height: int | |
| crop_box: tuple[int, int, int, int] | |
| rotation: int = 0 | |
| flip: bool = False | |
| def _fit_crop_box( | |
| image_size: tuple[int, int], | |
| aspect_ratio: float, | |
| scale: float, | |
| required_box: tuple[int, int, int, int] | None, | |
| rng: random.Random, | |
| ) -> tuple[int, int, int, int]: | |
| width, height = image_size | |
| crop_area = width * height * scale | |
| crop_w = int((crop_area * aspect_ratio) ** 0.5) | |
| crop_h = int(crop_w / aspect_ratio) | |
| if crop_w > width: | |
| crop_w = width | |
| crop_h = int(crop_w / aspect_ratio) | |
| if crop_h > height: | |
| crop_h = height | |
| crop_w = int(crop_h * aspect_ratio) | |
| crop_w = max(16, min(width, crop_w)) | |
| crop_h = max(16, min(height, crop_h)) | |
| if required_box is None: | |
| left = rng.randint(0, max(0, width - crop_w)) | |
| top = rng.randint(0, max(0, height - crop_h)) | |
| return left, top, left + crop_w, top + crop_h | |
| rx, ry, rw, rh = required_box | |
| min_left = max(0, rx + rw - crop_w) | |
| max_left = min(rx, width - crop_w) | |
| min_top = max(0, ry + rh - crop_h) | |
| max_top = min(ry, height - crop_h) | |
| if min_left > max_left or min_top > max_top: | |
| return 0, 0, width, height | |
| left = rng.randint(int(min_left), int(max_left)) | |
| top = rng.randint(int(min_top), int(max_top)) | |
| return left, top, left + crop_w, top + crop_h | |
| def generate_specs( | |
| image_size: tuple[int, int], | |
| count: int, | |
| seed: int = 0, | |
| output_area: int = 1024 * 1024, | |
| aspect_ratio_range: tuple[float, float] = (1 / 3, 3), | |
| required_box: tuple[int, int, int, int] | None = None, | |
| ) -> list[AugmentationSpec]: | |
| rng = random.Random(seed) | |
| specs = [] | |
| angles = [-15, -10, -5, 0, 5, 10, 15] | |
| for _ in range(count): | |
| ar = rng.uniform(*aspect_ratio_range) | |
| scale = rng.uniform(0.55, 1.0) | |
| crop_box = _fit_crop_box(image_size, ar, scale, required_box, rng) | |
| target_w = int((output_area * ar) ** 0.5) | |
| target_h = int(target_w / ar) | |
| target_w = max(16, (target_w // 16) * 16) | |
| target_h = max(16, (target_h // 16) * 16) | |
| specs.append( | |
| AugmentationSpec( | |
| width=target_w, | |
| height=target_h, | |
| crop_box=crop_box, | |
| rotation=rng.choice(angles), | |
| flip=bool(rng.getrandbits(1)), | |
| ) | |
| ) | |
| return specs | |
| def apply_spec(image: Image.Image, spec: AugmentationSpec) -> Image.Image: | |
| out = image.convert("RGB") | |
| if spec.flip: | |
| out = out.transpose(Image.Transpose.FLIP_LEFT_RIGHT) | |
| out = out.crop(spec.crop_box) | |
| if spec.rotation: | |
| out = out.rotate(spec.rotation, resample=Image.Resampling.BICUBIC, expand=False) | |
| return out.resize((spec.width, spec.height), Image.Resampling.LANCZOS) | |
| def save_augmented_pair( | |
| source: Image.Image, | |
| target: Image.Image, | |
| output_dir: str | Path, | |
| count: int, | |
| seed: int = 0, | |
| prefix: str = "aug", | |
| required_box: tuple[int, int, int, int] | None = None, | |
| ) -> list[dict[str, object]]: | |
| output_dir = Path(output_dir) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| specs = generate_specs(source.size, count=count, seed=seed, required_box=required_box) | |
| pairs = [] | |
| for idx, spec in enumerate(specs): | |
| source_out = output_dir / f"{prefix}_{idx:03d}_source.png" | |
| target_out = output_dir / f"{prefix}_{idx:03d}_target.png" | |
| apply_spec(source, spec).save(source_out) | |
| apply_spec(target, spec).save(target_out) | |
| pairs.append( | |
| { | |
| "source": str(source_out), | |
| "target": str(target_out), | |
| "augmentation": asdict(spec), | |
| } | |
| ) | |
| return pairs | |