""" Create deterministic Stage 3A image augmentations from a manifest or split CSV. The script writes new image files plus a manifest-compatible CSV fragment. It does not modify the source images or source CSV, so validation/test splits stay clean unless you explicitly run this script for a separate augmented eval set. Augmentations ------------- - JPEG recompression - Downscale then upscale - Deterministic crop - Mild blur or sharpen - Mild color/contrast/brightness jitter Usage ----- python scripts/dataset/augment_images.py \\ --manifest data/train.csv \\ --data-root data \\ --out-dir data/augmented/train \\ --out-manifest data/train_augmented.csv \\ --copies 1 \\ --seed 0 For a separate robustness eval set, point --manifest at val.csv or test.csv and write to a distinct output manifest such as data/test_augmented.csv. """ from __future__ import annotations import argparse import csv import hashlib import json import random from io import BytesIO from pathlib import Path from typing import Any from PIL import Image, ImageEnhance, ImageFilter AUGMENTATION_FIELDS = [ "width", "height", "original_path", "augmentation", "augmentation_seed", "augmentation_params_json", ] def _stable_seed(global_seed: int, path: str, copy_index: int) -> int: payload = f"{global_seed}|{path}|{copy_index}".encode("utf-8") return int(hashlib.sha256(payload).hexdigest()[:16], 16) % (2**31) def _sha256_file(path: Path) -> str: h = hashlib.sha256() with path.open("rb") as fh: for chunk in iter(lambda: fh.read(1024 * 1024), b""): h.update(chunk) return h.hexdigest() def _params_json(params: dict[str, Any]) -> str: return json.dumps(params, sort_keys=True, separators=(",", ":")) def _relative_to_data_root(path: Path, data_root: Path) -> str: try: return path.relative_to(data_root).as_posix() except ValueError as exc: raise ValueError( f"Output path {path} is not under data root {data_root}; " "choose an --out-dir inside --data-root" ) from exc def _resize_roundtrip( image: Image.Image, rng: random.Random, ) -> tuple[Image.Image, dict]: width, height = image.size scale = rng.uniform(0.65, 0.95) small_size = ( max(1, int(round(width * scale))), max(1, int(round(height * scale))), ) resample_down = Image.Resampling.BICUBIC resample_up = Image.Resampling.BILINEAR resized = image.resize(small_size, resample_down).resize( (width, height), resample_up, ) return resized, {"resize_scale": round(scale, 4)} def _crop_and_restore( image: Image.Image, rng: random.Random, ) -> tuple[Image.Image, dict]: width, height = image.size crop_scale = rng.uniform(0.9, 1.0) crop_w = max(1, int(round(width * crop_scale))) crop_h = max(1, int(round(height * crop_scale))) max_left = max(0, width - crop_w) max_top = max(0, height - crop_h) left = rng.randint(0, max_left) if max_left else 0 top = rng.randint(0, max_top) if max_top else 0 cropped = image.crop((left, top, left + crop_w, top + crop_h)) restored = cropped.resize((width, height), Image.Resampling.BICUBIC) return restored, { "crop_scale": round(crop_scale, 4), "crop_left": left, "crop_top": top, } def _filter(image: Image.Image, rng: random.Random) -> tuple[Image.Image, dict]: mode = rng.choice(["none", "blur", "sharpen"]) if mode == "blur": radius = rng.uniform(0.15, 0.45) return image.filter(ImageFilter.GaussianBlur(radius=radius)), { "filter": mode, "blur_radius": round(radius, 4), } if mode == "sharpen": return image.filter(ImageFilter.SHARPEN), {"filter": mode} return image, {"filter": mode} def _enhance(image: Image.Image, rng: random.Random) -> tuple[Image.Image, dict]: brightness = rng.uniform(0.92, 1.08) contrast = rng.uniform(0.9, 1.1) color = rng.uniform(0.9, 1.1) image = ImageEnhance.Brightness(image).enhance(brightness) image = ImageEnhance.Contrast(image).enhance(contrast) image = ImageEnhance.Color(image).enhance(color) return image, { "brightness": round(brightness, 4), "contrast": round(contrast, 4), "color": round(color, 4), } def _jpeg_roundtrip(image: Image.Image, rng: random.Random) -> tuple[Image.Image, dict]: quality = rng.randint(55, 95) buffer = BytesIO() image.save(buffer, format="JPEG", quality=quality, optimize=False) buffer.seek(0) with Image.open(buffer) as jpeg: jpeg.load() result = jpeg.convert("RGB") return result, {"jpeg_quality": quality} def augment_image(image: Image.Image, seed: int) -> tuple[Image.Image, dict[str, Any]]: """Return a deterministic augmented image and its parameter record.""" rng = random.Random(seed) augmented = image.convert("RGB") params: dict[str, Any] = {"version": "stage3a-v1"} augmented, resize_params = _resize_roundtrip(augmented, rng) params.update(resize_params) augmented, crop_params = _crop_and_restore(augmented, rng) params.update(crop_params) augmented, filter_params = _filter(augmented, rng) params.update(filter_params) augmented, enhance_params = _enhance(augmented, rng) params.update(enhance_params) augmented, jpeg_params = _jpeg_roundtrip(augmented, rng) params.update(jpeg_params) return augmented, params def _output_path(out_dir: Path, seed: int, original_path: str, copy_index: int) -> Path: key = hashlib.sha256(f"{original_path}|{seed}|{copy_index}".encode()).hexdigest() return out_dir / f"{key[:24]}.jpg" def _augmented_row( row: dict[str, str], *, original_path: str, output_path: Path, data_root: Path, seed: int, params: dict[str, Any], ) -> dict[str, str]: updated = dict(row) updated["path"] = _relative_to_data_root(output_path, data_root) updated["sha256"] = _sha256_file(output_path) updated["width"] = str(params["width"]) updated["height"] = str(params["height"]) updated["original_path"] = original_path updated["augmentation"] = params["version"] updated["augmentation_seed"] = str(seed) updated["augmentation_params_json"] = _params_json(params) return updated def main() -> None: parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--manifest", type=Path, required=True) parser.add_argument("--data-root", type=Path, required=True) parser.add_argument("--out-dir", type=Path, required=True) parser.add_argument("--out-manifest", type=Path, required=True) parser.add_argument("--copies", type=int, default=1) parser.add_argument("--seed", type=int, default=0) args = parser.parse_args() if args.copies < 1: raise ValueError("--copies must be >= 1") data_root = args.data_root.resolve() out_dir = args.out_dir.resolve() out_manifest = args.out_manifest.resolve() out_dir.mkdir(parents=True, exist_ok=True) out_manifest.parent.mkdir(parents=True, exist_ok=True) with args.manifest.open() as fh: reader = csv.DictReader(fh) source_rows = list(reader) source_fieldnames = reader.fieldnames or [] rows: list[dict[str, str]] = [] for row in source_rows: original_path = row["path"] src = data_root / original_path for copy_index in range(args.copies): seed = _stable_seed(args.seed, original_path, copy_index) dst = _output_path(out_dir, seed, original_path, copy_index) with Image.open(src) as image: augmented, params = augment_image(image, seed) params["width"], params["height"] = augmented.size augmented.save(dst, format="JPEG", quality=params["jpeg_quality"]) rows.append( _augmented_row( row, original_path=original_path, output_path=dst, data_root=data_root, seed=seed, params=params, ) ) fieldnames = [ *source_fieldnames, *[field for field in AUGMENTATION_FIELDS if field not in source_fieldnames], ] with out_manifest.open("w", newline="") as fh: writer = csv.DictWriter(fh, fieldnames=fieldnames) writer.writeheader() writer.writerows(rows) print(f"Augmented rows: {len(rows)}") print(f"Images: {out_dir}") print(f"Manifest: {out_manifest}") if __name__ == "__main__": main()