mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-31 09:40:38 +02:00
120 lines
4.0 KiB
Python
120 lines
4.0 KiB
Python
"""Export fixed, large, and scale-registered SynthID observations for images.
|
|
|
|
The output is an input manifest for ``synthid_conformal_cascade.py``. All
|
|
experts consume decoded RGB pixels only. Unsupported geometry is recorded
|
|
explicitly and never represented by a synthetic score.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, TypedDict
|
|
|
|
import click
|
|
import numpy as np
|
|
|
|
if TYPE_CHECKING:
|
|
from numpy.typing import NDArray
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(PROJECT_ROOT / "scripts"))
|
|
sys.path.insert(0, str(PROJECT_ROOT / "src"))
|
|
|
|
from synthid_pixel_attack import load_rgb # noqa: E402
|
|
from synthid_research_manifest import artifact_sha256 # noqa: E402
|
|
from synthid_runtime import synthid_detector # noqa: E402
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
FIXED_EXPERT_NAME = synthid_detector.DETECTOR_ID
|
|
REGISTERED_EXPERT_NAME = synthid_detector.REGISTERED_DETECTOR_ID
|
|
LARGE_EXPERT_NAME = synthid_detector.LARGE_DETECTOR_ID
|
|
|
|
|
|
class ExpertScore(TypedDict):
|
|
"""One JSON-safe runtime expert observation."""
|
|
|
|
name: str
|
|
supported: bool
|
|
score: float | None
|
|
|
|
|
|
class ScoredImage(TypedDict):
|
|
"""One hash-pinned image with every runtime expert observation."""
|
|
|
|
id: str
|
|
path: str
|
|
width: int
|
|
height: int
|
|
observations: list[ExpertScore]
|
|
|
|
|
|
def _observation(name: str, supported: bool, score: float | None) -> ExpertScore:
|
|
return {"name": name, "supported": supported, "score": score}
|
|
|
|
|
|
def score_pixels(pixels: NDArray[np.uint8]) -> list[ExpertScore]:
|
|
"""Return explicit fixed, registered, and large observations for RGB PIXELS."""
|
|
if pixels.ndim != 3 or pixels.shape[2] != 3 or pixels.dtype != np.uint8:
|
|
raise ValueError("pixels must be an RGB uint8 array")
|
|
bgr_pixels = np.ascontiguousarray(pixels[:, :, ::-1])
|
|
native = synthid_detector.detect_synthid("decoded-image", image=bgr_pixels, register_scale=False)
|
|
registered = synthid_detector.detect_synthid("decoded-image", image=bgr_pixels, register_scale=True)
|
|
fixed = _observation(FIXED_EXPERT_NAME, False, None)
|
|
large = _observation(LARGE_EXPERT_NAME, False, None)
|
|
native_observation = _observation(native.detector, native.status != "unsupported", native.score)
|
|
if native.detector == FIXED_EXPERT_NAME:
|
|
fixed = native_observation
|
|
elif native.detector == LARGE_EXPERT_NAME:
|
|
large = native_observation
|
|
else:
|
|
raise RuntimeError(f"unexpected default SynthID expert: {native.detector}")
|
|
return [
|
|
fixed,
|
|
_observation(REGISTERED_EXPERT_NAME, registered.status != "unsupported", registered.score),
|
|
large,
|
|
]
|
|
|
|
|
|
def score_path(path: Path) -> ScoredImage:
|
|
"""Decode PATH once and return one hash-pinned observation record."""
|
|
pixels = load_rgb(path)
|
|
height, width = pixels.shape[:2]
|
|
return {
|
|
"id": artifact_sha256(path),
|
|
"path": str(path),
|
|
"width": width,
|
|
"height": height,
|
|
"observations": score_pixels(pixels),
|
|
}
|
|
|
|
|
|
@click.command()
|
|
@click.argument("images", nargs=-1, required=True, type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
|
@click.option("--report-out", type=click.Path(dir_okay=False, path_type=Path), required=True)
|
|
def main(images: tuple[Path, ...], report_out: Path) -> None:
|
|
"""Score IMAGES with every shipped pixel expert."""
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
records = [score_path(path) for path in images]
|
|
report_out.parent.mkdir(parents=True, exist_ok=True)
|
|
report_out.write_text(
|
|
json.dumps(
|
|
{
|
|
"schema_version": 1,
|
|
"experts": [FIXED_EXPERT_NAME, REGISTERED_EXPERT_NAME, LARGE_EXPERT_NAME],
|
|
"records": records,
|
|
},
|
|
indent=2,
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
log.info("Wrote %d three-expert score records: %s", len(records), report_out)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|