test_directory.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from watermark3d.attacks import add_noise, affine_transform, laplacian_smooth, vertex_decimate
from watermark3d.mesh import Mesh, load_mesh, save_mesh
from watermark3d.watermark import WatermarkConfig, bit_error_rate, embed_watermark, extract_watermark

SUPPORTED = {".obj", ".ply"}


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Embed/extract a 3D interval watermark for every OBJ/PLY in a directory."
    )
    parser.add_argument("--input-dir", required=True, help="directory containing .obj or ASCII .ply meshes")
    parser.add_argument("--output-dir", required=True, help="directory for watermarked and attacked meshes")
    parser.add_argument("--bits", default="1011001110001011", help="binary watermark payload")
    parser.add_argument("--repeats", type=int, default=5, help="interval-pair repetitions per bit")
    parser.add_argument("--make-sample-if-empty", action="store_true", help="write a sample OBJ when input-dir has no meshes")
    args = parser.parse_args()

    input_dir = Path(args.input_dir)
    output_dir = Path(args.output_dir)
    input_dir.mkdir(parents=True, exist_ok=True)
    output_dir.mkdir(parents=True, exist_ok=True)

    meshes = sorted(path for path in input_dir.iterdir() if path.suffix.lower() in SUPPORTED)
    if not meshes and args.make_sample_if_empty:
        sample_path = input_dir / "sample_ellipsoid.obj"
        save_mesh(make_sample_ellipsoid(), sample_path)
        meshes = [sample_path]
    if not meshes:
        raise SystemExit("No .obj or .ply meshes found")

    cfg = WatermarkConfig(repeats=args.repeats)
    attacks = {
        "none": lambda mesh: mesh,
        "affine": affine_transform,
        "noise": add_noise,
        "smooth": laplacian_smooth,
        "simplify": vertex_decimate,
    }

    report = []
    for mesh_path in meshes:
        mesh = load_mesh(mesh_path)
        watermarked = embed_watermark(mesh, args.bits, cfg)
        base = output_dir / mesh_path.stem
        marked_path = base.with_name(base.name + "_watermarked" + mesh_path.suffix)
        save_mesh(watermarked, marked_path)

        for attack_name, attack in attacks.items():
            attacked = attack(watermarked.copy())
            attacked_path = base.with_name(base.name + f"_{attack_name}" + mesh_path.suffix)
            save_mesh(attacked, attacked_path)
            observed = extract_watermark(attacked, len(args.bits), cfg)
            report.append(
                {
                    "mesh": str(mesh_path),
                    "attack": attack_name,
                    "observed": observed,
                    "expected": args.bits,
                    "ber": bit_error_rate(args.bits, observed),
                    "attacked_path": str(attacked_path),
                }
            )

    report_path = output_dir / "watermark_report.json"
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report, indent=2))


def make_sample_ellipsoid(rows: int = 26, cols: int = 52) -> Mesh:
    import math

    vertices = []
    for r in range(rows + 1):
        theta = math.pi * r / rows
        for c in range(cols):
            phi = 2.0 * math.pi * c / cols
            ripple = 1.0 + 0.045 * math.sin(3.0 * phi + 0.7 * r) + 0.025 * math.cos(5.0 * theta + phi)
            x = 2.4 * ripple * math.sin(theta) * math.cos(phi)
            y = 1.0 * ripple * math.sin(theta) * math.sin(phi) + 0.05 * math.sin(2.0 * theta)
            z = 0.7 * ripple * math.cos(theta) + 0.18 * x + 0.08 * x * x
            vertices.append((x, y, z))
    faces = []
    for r in range(rows):
        for c in range(cols):
            a = r * cols + c
            b = r * cols + (c + 1) % cols
            d = (r + 1) * cols + c
            e = (r + 1) * cols + (c + 1) % cols
            faces.append([a, d, e, b])
    return Mesh(vertices, faces, "obj")


if __name__ == "__main__":
    main()