scripts/dataset/augment_images.py

"""
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()