alice.py

"""Reproduce Alice's full public HILL dataset."""

import hmac
import os
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

from steganography.common import read_grayscale_image, validate_stego_image
from steganography.hill import hill

COVER_DIR = Path("data/BOSSbase_1.01/cover")
STEGO_DIR = Path("data/BOSSbase_1.01/public_stego")
KEY_ENV_VAR = "AUTOSTEGO_EMBED_KEY"
PAYLOAD = 0.4
TRAIN_COUNT = 7_000
DEV_COUNT = 1_000


def _image_seed(key: bytes, filename: str) -> int:
    digest = hmac.digest(key, filename.encode(), "sha256")
    return int.from_bytes(digest[:16], byteorder="big", signed=False)


def _embed_one(job: tuple[Path, Path, int]) -> None:
    cover_path, stego_path, seed = job
    hill(cover_path, stego_path, PAYLOAD, seed=seed)
    validate_stego_image(
        read_grayscale_image(cover_path),
        read_grayscale_image(stego_path),
    )


def _select_public_covers() -> list[Path]:
    covers = sorted(COVER_DIR.glob("*.pgm"))
    train = [path for path in covers if path.stem.endswith(tuple("0123456"))]
    dev = [path for path in covers if path.stem.endswith("7")]
    if len(train) != TRAIN_COUNT or len(dev) != DEV_COUNT:
        raise RuntimeError(
            f"Expected {TRAIN_COUNT} training and {DEV_COUNT} development covers; "
            f"found {len(train)} and {len(dev)}."
        )
    return train + dev


def main() -> None:
    key = os.environ.get(KEY_ENV_VAR, "").encode()
    if not key:
        raise RuntimeError(f"Set {KEY_ENV_VAR} before embedding.")

    covers = _select_public_covers()
    jobs = [
        (
            cover_path,
            STEGO_DIR / cover_path.name,
            _image_seed(key, cover_path.name),
        )
        for cover_path in covers
    ]
    if len({seed for _, _, seed in jobs}) != len(jobs):
        raise RuntimeError("Per-image embedding seeds must be distinct.")

    STEGO_DIR.mkdir(parents=True)
    with ProcessPoolExecutor() as executor:
        list(executor.map(_embed_one, jobs))

    print(f"Wrote {len(jobs)} validated HILL stegos to {STEGO_DIR}")


if __name__ == "__main__":
    main()