From f330dd9c94f38f8f28a0aa4bb5e092287536d0a7 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Wed, 29 Jul 2026 17:28:11 -0700 Subject: [PATCH 01/19] Add video metadata and visible watermark removal --- CLAUDE.md | 6 +- README.md | 54 +- docs/cli.md | 54 ++ docs/index.md | 2 +- docs/known-limitations.md | 34 ++ docs/module-internals.md | 68 +++ docs/python-api.md | 59 ++ docs/supported-signals.md | 14 +- docs/watermarking-landscape.md | 27 +- pyproject.toml | 2 +- src/remove_ai_watermarks/__init__.py | 17 +- src/remove_ai_watermarks/cli.py | 143 ++++- src/remove_ai_watermarks/metadata.py | 45 +- src/remove_ai_watermarks/noai/ebml.py | 167 +++++ src/remove_ai_watermarks/noai/isobmff.py | 212 ++++++- src/remove_ai_watermarks/video.py | 212 +++++++ src/remove_ai_watermarks/video_visible.py | 706 ++++++++++++++++++++++ tests/test_metadata.py | 2 +- tests/test_video.py | 701 +++++++++++++++++++++ 19 files changed, 2487 insertions(+), 38 deletions(-) create mode 100644 src/remove_ai_watermarks/noai/ebml.py create mode 100644 src/remove_ai_watermarks/video.py create mode 100644 src/remove_ai_watermarks/video_visible.py create mode 100644 tests/test_video.py diff --git a/CLAUDE.md b/CLAUDE.md index 4223a78..01b7bea 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,6 +1,6 @@ # Remove-AI-Watermarks -You are a **principal Python engineer** maintaining a CLI tool and library for removing visible and invisible AI watermarks from images. +You are a **principal Python engineer** maintaining a CLI tool and library for removing AI watermarks from images and provenance metadata from video. ## Scope and non-goals @@ -26,6 +26,9 @@ Per-command exit-code semantics (the no-signal / GPU-missing skip branches), tes - `uv run remove-ai-watermarks identify ` — provenance verdict (platform + watermark inventory + confidence); `--json` for machine output, `--no-visible` to skip both registered visible detectors and the optional open invisible-watermark decoder - `uv run remove-ai-watermarks metadata --check` — inspect AI metadata (C2PA, EXIF, PNG chunks) - `uv run remove-ai-watermarks metadata --remove -o ` — strip all AI metadata +- `uv run remove-ai-watermarks video metadata --check` — inspect AI metadata in MP4/MOV/M4V/WebM/MKV +- `uv run remove-ai-watermarks video metadata --remove -o ` — strip verified video metadata without transcoding streams; the experimental video path requires a separate same-container output and defaults to `_clean` +- `uv run remove-ai-watermarks video visible -o ` — remove a temporally recurring Sora mark, or pass `--mark veo` for the current Veo diamond and legacy `Veo` text. It scans the full sequence first, transcodes video through ffmpeg, copies audio, strips AI metadata by default, and writes no output when no stable mark is found. `cv2` is the fast default; `migan`/`lama` improve difficult backgrounds. - `uv run remove-ai-watermarks batch ` — process every supported image in a directory (output defaults to `_clean/`, set with `-o`). `--mode visible|invisible|metadata|all` (default `visible`); the invisible/all path reuses the full `invisible` knob set above, plus `--backend` and `--sensitivity` for the visible localize -> fill pass. Applies the same no-signal skip per image; see the module doc. **Exit code:** non-zero when any image errored OR (mirroring single `all`) a `--mode invisible`/`all` image carried an invisible signal but the GPU extra was absent, so its SynthID scrub was skipped — it emits a loud warning and copies the input through (invisible mode) so the output dir stays complete; a wrapping service can then detect the incomplete run instead of trusting a silent exit 0. ## Test and lint @@ -79,6 +82,7 @@ Compact map. The full per-module detail (design decisions, tuned thresholds, cal - `upscaler.py` — optional Real-ESRGAN pre-diffusion super-resolution for small inputs (extra `esrgan`, spandrel only). Manual opt-in; the default `--upscaler` stays `lanczos` and the engine always falls back to Lanczos on absence/error. ESRGAN can degrade faces and thin text. - `image_io.py` — centralizes Unicode-safe image IO, alpha preservation, content-based format sniffing, and HEIC/AVIF fallbacks. Callers must check `imwrite` success. No-op visible removal preserves original bytes when the output format is unchanged. - `api.py` — the high-level convenience API, re-exported lazily at the package top level via `__init__.__getattr__` (PEP 562, so `import remove_ai_watermarks` stays cheap): `remove_visible(source, output=None, *, sensitivity="auto", backend="auto", strip_metadata=True, write_noop=True) -> (result_bgr, [labels])` (source = path OR BGR ndarray; a PATH auto-reads metadata provenance and preserves alpha, an ARRAY does neither; `write_noop=True` writes a clean passthrough copy when nothing is removed, `False` leaves `output` untouched so a "no mark = produce nothing" caller like the CLI `visible` command does not clobber a pre-existing file there) and `visible_provenance(path) -> frozenset[str]` (the single metadata→vendor-keys mapper; `cli._visible_provenance` is a thin None-guarded wrapper over it). **`remove_visible` is the ONE path the CLI and library share** — `cli.cmd_visible`'s `--mark auto` branch delegates entirely to it (read → provenance → `remove_auto_marks` → write → `strip_metadata`), so there is no CLI-vs-library drift; `strip_metadata` defaults True to match `visible --strip-metadata`. This is where a library caller should start — NOT the engines directly (`GeminiEngine`/`TextMarkEngine` have no `remove_watermark` any more; removal is registry `remove_auto_marks`/`KnownMark.remove`; the old single-strongest `best_auto_mark` is gone — removal takes EVERY mark). `identify` is NOT top-level re-exported (it collides with the `identify` submodule); use `from remove_ai_watermarks.identify import identify`. +- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, and `remove_video_visible(source, output=None, *, mark="sora", backend="cv2", strip_metadata=True) -> VideoVisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM TC260 is read by `noai/ebml.py` and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: fully synthetic Sora and Veo silhouettes propose frame boxes, provider-specific temporal recurrence authorizes them, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying audio. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels and invisible video watermarks are not built yet. For the Doubao alpha-distillation history (why content-image reverse-alpha distillation fails by physics and controlled captures were required), see `docs/research-doubao-distillation.md`. diff --git a/README.md b/README.md index 99ffea4..cde4b17 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,15 @@ # Remove AI Watermarks -Remove AI provenance marks from images you generated yourself: +Remove AI provenance marks from images and video you generated yourself: - known visible labels such as the Gemini sparkle and vendor text marks; - invisible pixel watermarks through diffusion regeneration; - C2PA, EXIF, XMP, IPTC, and related AI metadata. +Video support covers metadata inspection and removal plus experimental visible +Sora-wordmark removal. Invisible video-watermark removal remains a follow-up +stage. + > Try it online at [raiw.cc](https://raiw.cc) if you do not want to install Python > or run diffusion models locally. @@ -28,6 +32,8 @@ Remove AI provenance marks from images you generated yourself: | Remove known visible AI marks | `visible` | No | | Erase a region you select | `erase` | No | | Strip AI metadata | `metadata` | No | +| Strip AI metadata from video | `video metadata` | No | +| Remove a known Sora or Veo mark from video | `video visible` | No | | Regenerate an image to disrupt invisible watermarks | `invisible` | Recommended | | Run visible, invisible, and metadata removal | `all` | Recommended | | Process a directory | `batch` | Depends on mode | @@ -58,6 +64,35 @@ Strip metadata without running visible inpainting or diffusion: remove-ai-watermarks metadata image.png --remove -o clean.png ``` +Inspect or remove AI metadata from an MP4, MOV, M4V, WebM, or MKV file: + +```bash +remove-ai-watermarks video metadata input.mp4 --check +remove-ai-watermarks video metadata input.mp4 --remove -o clean.mp4 +``` + +The metadata command does not transcode video or audio streams. When `-o` is +omitted it writes `_clean` and preserves the original. MP4 and MOV +inspection includes the native TC260 `AIGC` tag in +`moov.udta.meta.keys/ilst`, including a `moov` placed after the media payload. +MKV and WebM inspection reads the normative +`Segment.Tags.Tag.SimpleTag` placement. + +Remove a moving Sora wordmark or a Veo corner mark: + +```bash +remove-ai-watermarks video visible input.mp4 -o clean.mp4 +remove-ai-watermarks video visible veo.mp4 --mark veo -o veo_clean.mp4 +``` + +This path scans the complete sequence before changing pixels. It accepts only a +mark that repeats at a stable position across adjacent frames, then reuses the +same OpenCV, MI-GAN, or LaMa fill backends as image removal. Audio is copied +without re-encoding; the video stream is transcoded because its pixels change. +Sora covers the moving Sora 2 mascot and wordmark. Veo covers both the current +four-point diamond and the legacy `Veo` text in the bottom-right corner. No +output is written when no stable mark is found. + For invisible watermark removal, install the diffusion dependencies: ```bash @@ -184,8 +219,9 @@ Visible removal follows three steps: 3. Fill only the masked region with OpenCV, MI-GAN, or LaMa. Metadata removal uses format aware stripping. JPEG metadata removal preserves -the encoded image scan instead of recompressing it. Other supported containers -use their corresponding metadata path. +the encoded image scan instead of recompressing it. Native MP4/MOV TC260 values +are blanked without changing box sizes or media offsets. Other supported +containers use their corresponding metadata path. Invisible removal is different. It regenerates the image through a diffusion pipeline to disrupt pixel and frequency domain watermarks. This changes the @@ -202,6 +238,11 @@ import remove_ai_watermarks as raiw result, removed = raiw.remove_visible("watermarked.png", "clean.png") print(removed) + +report = raiw.inspect_video_metadata("input.mp4") +cleaned = raiw.remove_video_metadata("input.mp4") +visible = raiw.remove_video_visible("sora.mp4", "sora_clean.mp4") +veo = raiw.remove_video_visible("veo.mp4", "veo_clean.mp4", mark="veo") ``` The high level API accepts a file path or a BGR NumPy array. For path inputs it @@ -226,6 +267,13 @@ invisible removal. and selected fill backend. - Invisible removal changes the whole image and may alter faces, text, or fine detail. +- Visible video removal recognizes the moving Sora 2 wordmark and the current + Veo diamond plus legacy `Veo` text. It does not recognize the older Sora Turbo + corner swirl. + The classical OpenCV backend can smear structured backgrounds; use MI-GAN or + LaMa when recovery quality matters. Video SynthID is not removed yet. + MP4/MOV/M4V metadata stripping currently reads the full container into + memory, so very large videos are not an intended metadata input yet. - `qwen-zimage` requires CUDA. The other diffusion profiles also support the devices listed by `remove-ai-watermarks invisible --help`. - Provider watermark systems can change. Validate important outputs with the diff --git a/docs/cli.md b/docs/cli.md index da76b8e..c8ba68a 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -125,6 +125,60 @@ The command also supports the audio and video containers listed in [supported signals](supported-signals.md). ffmpeg must be available for the non-ISOBMFF audio and video path. +## Strip AI metadata from video + +The experimental video namespace starts with metadata inspection and removal: + +```bash +remove-ai-watermarks video metadata input.mp4 --check +remove-ai-watermarks video metadata input.mp4 --remove -o clean.mp4 +``` + +Supported containers are MP4, MOV, M4V, WebM, and MKV. The operation delegates +to the same verified metadata scanner and stripper as the generic `metadata` +command, so detection and removal stay in parity. Video and audio streams are +not transcoded. For MP4 and MOV, this includes the native TC260 `AIGC` key and +JSON value stored in `moov.udta.meta.keys/ilst`. The inspector seeks past a +large `mdat` to find a tail `moov`; removal blanks the key and value in place so +box sizes and media offsets do not move. + +For MKV and WebM, the inspector reads the native TC260 +`Segment.Tags.Tag.SimpleTag` entry. Removal uses ffmpeg stream copying to +discard container tags and chapters without transcoding the streams. + +When `-o` is omitted, the command writes `_clean` with the same +extension. It never overwrites the source, and it rejects an output with a +different container extension. + +Visible video labels and invisible video watermarks are not handled by this +command. + +## Remove a visible Sora or Veo video mark + +```bash +remove-ai-watermarks video visible input.mp4 -o clean.mp4 +remove-ai-watermarks video visible veo.mp4 --mark veo -o veo_clean.mp4 +``` + +The experimental command supports the moving Sora mascot and wordmark plus two +Veo corner variants: the current four-point diamond and the legacy `Veo` text. +Sora searches the whole frame at multiple scales. Veo searches the +bottom-right corner using separate synthetic silhouettes for the two variants. +Both require a spatially recurring candidate across adjacent frames. Matching +provider provenance may relax the visual score, but metadata alone never +creates a detection. Clean API exports therefore remain untouched. + +The video stream is transcoded and the original audio stream is copied. +Supported input and output containers are MP4, MOV, M4V, WebM, and MKV; the +output extension must match the input. The default `cv2` backend is fast but can +smear structured backgrounds. Select `--backend migan` or `--backend lama` for +a learned fill, or `--backend auto` to choose the best installed backend. + +AI metadata is stripped from the encoded output by default. Use +`--keep-metadata` to retain mapped container metadata. When no temporally stable +mark is found, the command writes no output and exits with the no-visible-mark +status. + ## Remove invisible watermarks Install the diffusion dependencies first: diff --git a/docs/index.md b/docs/index.md index 9128ba7..f096e33 100644 --- a/docs/index.md +++ b/docs/index.md @@ -8,7 +8,7 @@ to run the tool. Use the maintainer references only when changing the code. | Page | Use it when | | --- | --- | | [Installation](installation.md) | You need the CLI, an optional model backend, or a development environment. | -| [CLI guide](cli.md) | You want a command for one image, a directory, or a specific watermark type. | +| [CLI guide](cli.md) | You want a command for an image, video metadata or visible marks, a directory, or a specific watermark type. | | [Python API](python-api.md) | You want to call the package from Python. | | [Supported signals](supported-signals.md) | You need to know which visible marks, metadata formats, and invisible signals are covered. | | [Known limitations](known-limitations.md) | You need the quality, device, format, or verification boundaries. | diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 72c9e7d..35b4a93 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -174,6 +174,40 @@ WebM, Matroska, MP3, WAV, FLAC, OGG, Opus, and AAC container metadata is strippe through ffmpeg with stream copying. The operation fails if ffmpeg is absent or cannot parse the input. +### Visible video removal supports Sora and Veo and is still experimental + +The experimental `video metadata` command and high level video API inspect and +strip supported AI provenance metadata without transcoding streams. + +`video visible` and `remove_video_visible` additionally support the moving +Sora 2 mascot and wordmark, the current Veo four-point diamond, and the legacy +`Veo` text. Detection requires a recurring visual candidate across adjacent +frames. Provider provenance can recover low-contrast runs only after visual +evidence exists, so metadata alone does not erase a clean API export. +Historical Sora Turbo exports use a small OpenAI swirl in the corner rather +than the moving mascot-and-wordmark design; that earlier variant is not +detected by the `sora` video mark. Other provider video labels and proprietary +invisible video watermarks are not supported yet. + +Visible removal transcodes the video stream and copies audio. Its frame-local +fill is not a motion-aware video inpainting model. OpenCV can leave a visible +smear where the mark overlaps a hard edge or structured texture, and the smear +can vary over time. MI-GAN and LaMa improve individual frames but do not +guarantee temporal coherence. The Veo diamond uses a shape mask to limit damage +outside the symbol, but OpenCV may still soften texture inside it. The current +encoder also emits a constant-frame-rate output at the decoded stream rate, so +variable-frame-rate preservation is not yet guaranteed. + +Native TC260 metadata in MP4/MOV is supported at its normative +`moov.udta.meta.keys/ilst` placement, including non-faststart files whose +`moov` follows a large media payload. MKV/WebM is supported at the normative +`Segment.Tags.Tag.SimpleTag` placement and uses ffmpeg for stream-copy removal. +The corresponding FLV and AVI native tags are not parsed yet. + +The current ISOBMFF stripper reads an MP4, MOV, or M4V container into memory +before rewriting its metadata boxes. A streaming box copier is required before +the experimental command is appropriate for very large video files. + ### Metadata transformation is fail safe `remove_ai_metadata` may copy an undecodable file through unchanged instead of diff --git a/docs/module-internals.md b/docs/module-internals.md index 8432062..cb4dba6 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -86,6 +86,70 @@ Regression coverage: - [`test_api.py`](../tests/test_api.py) - [`test_image_io.py`](../tests/test_image_io.py) +[`video.py`](../src/remove_ai_watermarks/video.py) provides the experimental +video entry point: + +- `inspect_video_metadata` +- `remove_video_metadata` +- `remove_video_visible` + +The video API validates both the supported extension and container signature, +then delegates all metadata detection and stripping to `metadata.py`. It +requires a separate same-container output, defaulting to `_clean`, so +the experimental path does not overwrite an original. The package root exposes +both functions lazily. + +Native MP4/MOV TC260 labels follow TC260-PG-20257A: +`moov.udta.meta.keys` maps an `AIGC` key to a raw JSON value in `ilst`. +[`noai/isobmff.py`](../src/remove_ai_watermarks/noai/isobmff.py) walks those +nested boxes by seeking, so detection reaches a tail `moov` without reading the +preceding `mdat`. Removal changes the four-byte key to `free` and blanks only +the validated JSON value with same-length spaces. This preserves every box +size, `stco`/`co64` offset, and encoded stream byte. A generic `AIGC` key whose +value has no TC260 field is ignored. + +[`noai/ebml.py`](../src/remove_ai_watermarks/noai/ebml.py) provides the +corresponding bounded Matroska/WebM reader. It seeks over clusters and accepts +only a `Segment.Tags.Tag.SimpleTag` pairing `TagName=AIGC` with a JSON +`TagString` carrying a TC260 field. The existing ffmpeg stream-copy path removes +those container tags without transcoding the encoded streams. + +[`video_visible.py`](../src/remove_ai_watermarks/video_visible.py) implements +the first pixel stages for Sora and Veo. The Sora detector searches a normalized +frame with a fully synthetic mascot-and-text silhouette at several scales. The +Veo detector uses separate synthetic silhouettes for the current four-point +diamond and legacy `Veo` text, with bounded bottom-right searches calibrated +independently from Sora. A strong relocated-diamond match may bypass the known +layout anchors, but weak free-corner matches never enter the temporal arbiter. +This prevents recurring scene details in clean API exports from being promoted +to a watermark. + +Every per-frame result is untrusted. The provider-specific stabilization +wrappers share one recurrence implementation, while retaining separate visual +floors and minimum-run policy. Provenance can relax a low-contrast run only +after recurring visual evidence exists. Sora transition frames follow the +nearest confirmed moving position only with Sora provenance. A confirmed Veo +run can cover low-contrast frames at its fixed position. This separation keeps +clean API exports from being modified merely because metadata names the same +generator. + +Removal runs in a second decode pass. Sora and legacy Veo text use padded box +masks. The square Veo diamond uses a synthetic shape mask so transparent box +corners do not erase unrelated pixels. Every mask goes through the shared +`watermark_registry.fill` backends. ffmpeg encodes the changed video stream and +copies optional audio. The default OpenCV fill is the speed floor; structured +backgrounds need MI-GAN or LaMa for better reconstruction. Invisible video +stages must continue to reuse the image and metadata implementations rather +than copying their logic. + +The inherited ISOBMFF metadata path currently reads the complete container into +memory; replacing that with a streaming box copier is a prerequisite for large +video inputs. + +Regression coverage: + +- [`test_video.py`](../tests/test_video.py) + ## Metadata and provenance ### C2PA @@ -110,6 +174,10 @@ Key contracts: scan. - ISOBMFF containers use [`noai/isobmff.py`](../src/remove_ai_watermarks/noai/isobmff.py). +- Native MP4/MOV TC260 `AIGC` entries are read from + `moov.udta.meta.keys/ilst` and blanked without changing box sizes. +- Native MKV/WebM TC260 `AIGC` entries are read from + `Segment.Tags.Tag.SimpleTag` and removed through the ffmpeg stream-copy path. - Supported non-ISOBMFF audio and video containers use ffmpeg stream copying. - The low-level remover is fail-safe and can copy an undecodable file through unchanged. diff --git a/docs/python-api.md b/docs/python-api.md index 4bc2af6..24a4fbb 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -128,6 +128,65 @@ path preserves the pixels but drops standard metadata. Treat a nonempty undecodable input through unchanged, so its return alone must not be presented as proof that metadata was removed. +## Inspect and strip video metadata + +The experimental high level video API supports MP4, MOV, M4V, WebM, and MKV: + +```python +import remove_ai_watermarks as raiw + +report = raiw.inspect_video_metadata("input.mp4") +if report.has_ai_metadata: + result = raiw.remove_video_metadata("input.mp4") + if result.remaining: + raise RuntimeError(f"AI metadata remains: {result.remaining}") +``` + +`remove_video_metadata` does not transcode video or audio streams. Its default +output is `input_clean.mp4`, leaving the source untouched. An explicit output +must use the same container extension as the source. + +The returned `VideoMetadataResult` records the source, output, metadata detected +before removal, and any markers remaining after the verified strip. MP4/MOV +inspection recognizes the native TC260 `AIGC` entry in +`moov.udta.meta.keys/ilst`; its removal preserves container size and encoded +stream bytes. MKV/WebM inspection recognizes the corresponding +`Segment.Tags.Tag.SimpleTag` representation; its removal requires ffmpeg for a +stream-copy remux. + +## Remove a visible Sora or Veo video mark + +```python +import remove_ai_watermarks as raiw + +result = raiw.remove_video_visible( + "sora.mp4", + "sora_clean.mp4", + backend="cv2", + strip_metadata=True, +) +if result.output is None: + print("No temporally stable Sora mark was found") + +veo_result = raiw.remove_video_visible( + "veo.mp4", + "veo_clean.mp4", + mark="veo", +) +``` + +`remove_video_visible` scans the complete video before writing output. It +combines synthetic multi-scale visual matching with temporal consistency, so an +isolated lookalike in one frame is not enough to authorize inpainting. The +supported `mark` values are `sora` and `veo`. The Veo detector recognizes the +current four-point diamond and the legacy `Veo` text with separate synthetic +silhouettes. + +The returned `VideoVisibleResult` records the total, detected, and removed frame +counts plus any AI metadata that survived the output encode. The function +returns `output=None` and writes no file when no stable mark is selected. Video +pixels are transcoded through ffmpeg while the source audio stream is copied. + ## Remove invisible watermarks ```python diff --git a/docs/supported-signals.md b/docs/supported-signals.md index d893e61..ba19ef5 100644 --- a/docs/supported-signals.md +++ b/docs/supported-signals.md @@ -29,6 +29,16 @@ the masked area. Marks from other vendors are not detected automatically. Use `erase --region` when you can select the affected area yourself. +### Visible video marks + +| Key | Mark | Motion | Important limit | +| --- | --- | --- | --- | +| `sora` | Sora 2 mascot and wordmark | Moves among frame positions | Requires a temporally recurring visual match; the older Sora Turbo corner swirl is a different unsupported mark. | +| `veo` | Current four-point diamond and legacy `Veo` text | Fixed bottom-right corner | Uses separate silhouettes and requires a recurring match; learned fill is preferable on structured backgrounds. | + +Use `video visible` for this registry. It is separate from the image `visible` +command because selection is made over a sequence rather than one raster. + ## Fill backends | Backend | Install | Behavior | @@ -48,7 +58,9 @@ The inspection and stripping code handles signals in these groups: - EXIF and XMP generator fields; - IPTC AI disclosure fields; - PNG text chunks and embedded generation parameters; -- China TC260 AIGC labels in supported metadata placements; +- China TC260 AIGC labels in supported image placements and the normative + MP4/MOV `moov.udta.meta.keys/ilst` and MKV/WebM + `Segment.Tags.Tag.SimpleTag` placements; - xAI and Grok EXIF signature fields; - Samsung AI editing markers; - Hugging Face job metadata; diff --git a/docs/watermarking-landscape.md b/docs/watermarking-landscape.md index 49f7729..d350f7c 100644 --- a/docs/watermarking-landscape.md +++ b/docs/watermarking-landscape.md @@ -38,7 +38,13 @@ pair. On the ISOBMFF path, `blank_ai_exif_tokens` provides the corresponding in-place scrub for supported EXIF values, TC260 AIGC blocks, and the xAI pair. - **China TC260 AIGC label (caught by `AIGC_MARKERS` / `metadata.aigc_label`, surfaced by `identify` as the `aigc` signal):** China-served generators embed an XMP `{"Label":"1","ContentProducer":...}` block — China's mandatory AI-content labeling (TC260 namespace `tc260.org.cn/ns/AIGC`). -**Doubao** (ByteDance) uses it (verified on a public issue sample; `ContentProducer` `001191110102MACQD9K64010000`, no C2PA/SynthID/imwatermark — the XMP block is the only signal; GitHub attachment upload did NOT strip it). The same standard is mandatory for Jimeng/Kling/Qwen/Ernie etc., so the one marker covers the whole China-AIGC-labeled ecosystem. `aigc_label` reads **four serializations** through a shared `_parse` helper: the HTML-entity-encoded XMP `TC260:AIGC` block in **either RDF form** — the nested element `{...}` (Doubao) or the attribute `TC260:AIGC="{...}"` (**PicWish**, `ContentProducer="picwish"`, verified on compatible samples) — via a container-agnostic raw-byte scan (any JSON object accepted), a raw-JSON PNG `AIGC` tEXt chunk (Doubao also writes the label this way, no namespaced marker at all — confirmed on compatible samples, `ContentProducer="doubao"`), a bare raw-JSON `{"AIGC":{...}}` object embedded in **JPEG EXIF (UserComment)** by some China-served generators, brace-matched from the scan head with `json.JSONDecoder().raw_decode` (no namespaced marker, no PNG chunk — confirmed on compatible samples, `ContentProducer="001191440300708461136T1308L"`), **and** a bare `AIGC{...}` blob (the label glued straight to its JSON, no `"AIGC":` key wrapper) embedded in a **JPEG APP segment near the JFIF header** — confirmed on compatible samples. The two raw-JSON forms are scanned in one loop (`'"AIGC"'` then `AIGC{`) that **falls through on a non-TC260 / undecodable hit instead of returning** — a quoted `"AIGC"` can appear later in an XMP packet while the real label is a bare `AIGC{...}` earlier in the file, so an unconditional early return on the quoted form would shadow the bare form (the exact bug behind the 06-10 misses). All three generic forms (the PNG chunk, the bare `{"AIGC":...}` object, and the bare `AIGC{...}` blob) are gated on at least one TC260 field (`_TC260_FIELDS`) so a generic `AIGC` key cannot false-positive; the namespaced XMP element is unambiguous and needs no gate. `_TC260_FIELDS` covers **two schemas**: the producer-side one (`Label` / `ContentProducer` / `ProduceID` / `ContentPropagator` / `PropagateID`, Doubao and most China gens) and the **service-provider** one (`ServiceProvider` / `ServiceUser`, plus generic `Time` / `ContentId` which are NOT gated on) — **Tencent Cloud's** AIGC variant (`ServiceProvider` = `腾讯云`), embedded in **EXIF `ImageDescription`**, verified on compatible samples. In `identify`, `aigc` fires on the parsed label **or** the `AIGC_MARKERS` byte scan (the latter preserves the laundering-tell case where the JSON payload is truncated). +**Doubao** (ByteDance) uses it (verified on a public issue sample; `ContentProducer` `001191110102MACQD9K64010000`, no C2PA/SynthID/imwatermark — the XMP block is the only signal; GitHub attachment upload did NOT strip it). The same standard is mandatory for Jimeng/Kling/Qwen/Ernie etc., so the one marker covers the whole China-AIGC-labeled ecosystem. `aigc_label` reads **four image serializations** through a shared `_parse` helper: the HTML-entity-encoded XMP `TC260:AIGC` block in **either RDF form** — the nested element `{...}` (Doubao) or the attribute `TC260:AIGC="{...}"` (**PicWish**, `ContentProducer="picwish"`, verified on compatible samples) — via a container-agnostic raw-byte scan (any JSON object accepted), a raw-JSON PNG `AIGC` tEXt chunk (Doubao also writes the label this way, no namespaced marker at all — confirmed on compatible samples, `ContentProducer="doubao"`), a bare raw-JSON `{"AIGC":{...}}` object embedded in **JPEG EXIF (UserComment)** by some China-served generators, brace-matched from the scan head with `json.JSONDecoder().raw_decode` (no namespaced marker, no PNG chunk — confirmed on compatible samples, `ContentProducer="001191440300708461136T1308L"`), **and** a bare `AIGC{...}` blob (the label glued straight to its JSON, no `"AIGC":` key wrapper) embedded in a **JPEG APP segment near the JFIF header** — confirmed on compatible samples. The two raw-JSON forms are scanned in one loop (`'"AIGC"'` then `AIGC{`) that **falls through on a non-TC260 / undecodable hit instead of returning** — a quoted `"AIGC"` can appear later in an XMP packet while the real label is a bare `AIGC{...}` earlier in the file, so an unconditional early return on the quoted form would shadow the bare form (the exact bug behind the 06-10 misses). Native MP4/MOV is a fifth serialization: TC260-PG-20257A stores an `AIGC` key in `moov.udta.meta.keys` and the raw JSON in the matching `ilst` item. The seeking parser reaches a tail `moov` without reading `mdat`; removal replaces the key with `free` and blanks the validated value at the same length so every box size and stream offset stays fixed. All generic forms are gated on at least one TC260 field (`TC260_AIGC_FIELDS`) so a generic `AIGC` key cannot false-positive; the namespaced XMP element is unambiguous and needs no gate. `TC260_AIGC_FIELDS` covers **two schemas**: the producer-side one (`Label` / `ContentProducer` / `ProduceID` / `ContentPropagator` / `PropagateID`, Doubao and most China gens) and the **service-provider** one (`ServiceProvider` / `ServiceUser`, plus generic `Time` / `ContentId` which are NOT gated on) — **Tencent Cloud's** AIGC variant (`ServiceProvider` = `腾讯云`), embedded in **EXIF `ImageDescription`**, verified on compatible samples. In `identify`, `aigc` fires on the parsed label **or** the `AIGC_MARKERS` byte scan (the latter preserves the laundering-tell case where the JSON payload is truncated). + +Native MKV/WebM is a sixth serialization. TC260-PG-20257A stores +`TagName=AIGC` and the raw JSON `TagString` in +`Segment.Tags.Tag.SimpleTag`. The bounded EBML reader skips cluster payloads; +the existing ffmpeg stream-copy path removes the tags without transcoding. + - **HuggingFace-hosted job (caught by `metadata.huggingface_job`, surfaced by `identify` as the `hf_job` signal, MEDIUM confidence):** HuggingFace Jobs / Spaces can stamp generated PNGs with an `hf-job-id` tEXt chunk holding the job UUID. It marks the *hosting job*, not a model, so it lifts an Unknown verdict to a tentative AI via `hf_only` but never overrides a hard metadata signal. `_HF_JOB_CAVEAT` states the limit. Removal drops the chunk through the PNG metadata whitelist. - **No detectable signal on some downloads:** Recraft exports and some hosted FLUX surfaces can arrive without a supported local signal. Midjourney samples @@ -49,13 +55,18 @@ in-place scrub for supported EXIF values, TC260 AIGC blocks, and the xAI pair. - **C2PA 2.4 "Durable Content Credentials" (April 2026; verified against the spec) raise the bar for metadata stripping.** 2.4 defines soft bindings (an invisible watermark or a content fingerprint) plus a server-side manifest repository and a new `c2pa.repository-receipt` assertion. Per the spec: "if a C2PA manifest is removed from an asset, but a copy of that manifest remains in a provenance store elsewhere, the manifest and asset may be matched using available soft bindings." So our local `metadata --remove` deletes the *embedded* manifest, but a fingerprint/watermark soft binding can still re-link the image to its manifest in a repository server-side. Stripping the file is becoming necessary-but-not-sufficient against durable provenance. (Our parsers target the stable embedded-manifest format documented in C2PA 2.1 §11; that format is unchanged in 2.4 -- the new pieces are repository/soft-binding infra, not the on-file box layout, so no parser change is implied.) Spec: https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html We now READ the soft-binding `alg` (`C2PA_SOFT_BINDINGS` / `soft_binding_vendors_in`) to name the forensic-watermark vendor, and locally DECODE the one open scheme, Adobe TrustMark (`trustmark_detector`); the rest (Digimarc/Imatag/Steg.AI/...) stay name-only (proprietary decoders). - **Built in the dated batch:** soft-binding vendor detection, IPTC Photo Metadata AI-disclosure fields, C2PA detection and stripping for supported - ISOBMFF video, and the optional Adobe TrustMark decoder. Visible video-logo - removal and proprietary audio-watermark detection remain outside the package. + ISOBMFF video, the optional Adobe TrustMark decoder, and temporally stabilized + visible Sora and Veo removal. Other visible video logos and proprietary + audio-watermark detection remain outside the package. Metadata stripping for supported audio containers is a separate implemented path. **Box detection window — now handled (v0.6.8):** detection no longer relies on a fixed first-MB read. `metadata.scan_head(path, size)` reads the first `size` bytes and, for ISOBMFF, appends the payloads of late provenance boxes found by `isobmff.scan_c2pa_region` (a file-seeking top-level box walker that skips past `mdat` by size without reading it), so a C2PA/AIGC/IPTC manifest placed AFTER a large `mdat` in a streaming/non-faststart MP4 is now caught. Every C2PA/marker byte scan (`has_ai_metadata`, `aigc_label`, `iptc_ai_system`, `synthid_source`, `exif_generator` XMP, `get_ai_metadata` soft-binding, and `identify`) goes through `scan_head`; it is behavior-neutral for non-ISOBMFF inputs (exactly `f.read(size)`). +Native TC260 MP4/MOV tags do not live in those top-level provenance boxes. +`tc260_aigc_payloads` separately seeks through `moov.udta.meta.keys/ilst`, so +the normative tag is also found when a large `mdat` precedes `moov`. + **Meta-box XMP and EXIF removal are handled in place:** an AI-label XMP packet stored as a meta-box `mime` item is blanked by `isobmff.blank_ai_xmp_packets`. Supported EXIF items are handled by @@ -69,6 +80,16 @@ against primary sources before adding jurisdiction-specific claims. **Google Gemini visible "sparkle" -- tier-dependent, and spec-undocumented by Google.** Google primary sources (the Nano Banana Pro blog and the gemini.google image-generation page, both WebFetch-verified) confirm Gemini images carry BOTH the invisible SynthID (on ALL Google-AI media) AND a visible sparkle, but the visible mark is **tier-gated**: applied for FREE and Google AI **Pro** users, and **REMOVED** for Google AI **Ultra** subscribers, inside **Google AI Studio**, and on **API / dev** output. So a Google-C2PA image with NO visible sparkle is expected (Ultra / API), not evidence it is clean -- this reinforces the `identify` "no visible mark != clean" rule. The ONLY official verifier is the SynthID flow (upload to the Gemini app, ask if it is AI-generated), which reads the INVISIBLE mark; there is **no official visible-sparkle detector**, and Google publishes **no** glyph geometry / size / opacity / color / locale / placement spec. So our capture-based sparkle template is the only source of truth and cannot be validated against a vendor spec -- keep reverse-engineering from real captures (do not expect a published spec). +**Google Veo video marks use two incompatible visible designs.** Public raw +clips verify the current four-point diamond and the legacy bottom-right `Veo` +text. The independent +[VeoWatermarkRemover](https://github.com/allenk/VeoWatermarkRemover) project +reports the same current-versus-legacy split, multiple output layouts, and the +need for cross-frame position agreement. Our implementation copies no logo +pixels or alpha maps from that project: it uses two synthetic silhouettes, +known-layout searches plus a strong relocated-diamond fallback, and a separate +temporal arbiter calibrated against raw watermarked clips and clean API exports. + **The faint-visible-mark precision/recall wall is fundamental, not a heuristic artifact.** The visible-watermark-detection literature has moved to LEARNED segmentation / object-detection (WDNet WACV'21 arXiv:2012.07616; SLBR ACM MM'21, open code+weights; the PRCV'18 large-scale detector; Su et al. survey 2025), but three verified findings bound what a learned detector actually buys: (1) a claim that a confidence threshold "cleanly separates" true from false matches even with a learned CNN front-end was **REFUTED** in verification (arXiv:1705.08593) -- the precision/recall wall persists even with learned features. (2) Learned detectors need a LARGE, pattern-diverse labeled dataset trained on synthetic composites (PRCV'18: 60k images / 80 watermark classes; CLWD: 60k / 160 marks), and off-distribution degradation is a documented real axis (models trained on limited-pattern LVW transfer worse; diversity of training patterns drives generalization). (3) Inference is cheap (WDNet ~8 ms at 256x256) -- the cost is the data pipeline, not runtime. Net: a learned detector shifts the frontier but does NOT remove the wall; for a SINGLE mark the cheapest next step is a small patch classifier (real-sparkle vs false-positive) on top of the existing NCC localizer, not a full segmentation model. SLBR is a ready baseline. The current NCC + false-positive gate (core-ring brightness margin + gradient-NCC crispness + white-core saturation) is a sound operating point, and the residual miss is the information-theoretic wall the literature confirms. **Visible-mark landscape beyond the registry.** Meta stamps a visible "Imagined with AI" mark (bottom-LEFT, a small symbol) on its OWN Meta AI / "Imagine" output; for third-party images it relies on C2PA / IPTC, not a visible mark. Samsung Galaxy AI additionally uses a **four-star icon** variant in a corner alongside the localized text wordmark `samsung_engine` calibrates (only the Italian text variant is covered) -- the icon is a distinct, uncovered variant. Every source agrees visible + metadata marks are trivially removable (crop / screenshot, ~2 s), which is the tool's premise. diff --git a/pyproject.toml b/pyproject.toml index dd846ca..2c96661 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "remove-ai-watermarks" version = "0.20.2" -description = "AI watermark remover: strip visible and invisible AI watermarks (Gemini / Nano Banana sparkle, SynthID) and provenance metadata (C2PA, EXIF) from images" +description = "Remove visible and invisible AI watermarks from images and provenance metadata from media containers" readme = "README.md" requires-python = ">=3.10.1" license = {text = "Apache-2.0"} diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index b7cceda..8ebbbcd 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -6,6 +6,9 @@ High-level API (lazy, so ``import remove_ai_watermarks`` stays cheap):: raiw.remove_visible("in.png", "out.png") # clean a file (provenance auto) result, removed = raiw.remove_visible(bgr_array) # array -> array raiw.visible_provenance("in.png") # -> frozenset of confirmed vendors + raiw.inspect_video_metadata("in.mp4") # -> VideoMetadataReport + raiw.remove_video_metadata("in.mp4", "out.mp4") # verified metadata strip + raiw.remove_video_visible("in.mp4", "out.mp4") # stable Sora or Veo mark removal For a provenance verdict use the ``identify`` submodule:: @@ -27,10 +30,18 @@ _warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*") __version__ = "0.20.2" -__all__ = ["__version__", "remove_visible", "visible_provenance"] +__all__ = [ + "__version__", + "inspect_video_metadata", + "remove_video_metadata", + "remove_video_visible", + "remove_visible", + "visible_provenance", +] if TYPE_CHECKING: from remove_ai_watermarks.api import remove_visible, visible_provenance + from remove_ai_watermarks.video import inspect_video_metadata, remove_video_metadata, remove_video_visible def __getattr__(name: str) -> object: @@ -40,4 +51,8 @@ def __getattr__(name: str) -> object: from remove_ai_watermarks import api return getattr(api, name) + if name in ("inspect_video_metadata", "remove_video_metadata", "remove_video_visible"): + from remove_ai_watermarks import video + + return getattr(video, name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 3f8676c..02f0236 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -4,6 +4,7 @@ Provides commands for: - Visible watermark removal (Gemini sparkle) - works offline, fast - Invisible watermark removal (SynthID etc.) - requires GPU/diffusion models - AI metadata stripping - lightweight, no ML deps needed + - Experimental video metadata and visible-wordmark removal """ from __future__ import annotations @@ -586,7 +587,7 @@ def _should_skip_invisible_scrub(force: bool, image_path: Path) -> bool: @click.option("-v", "--verbose", is_flag=True, help="Enable verbose logging.") @click.pass_context def main(ctx: click.Context, verbose: bool) -> None: - """Remove visible and invisible AI watermarks from images.""" + """Remove visible and invisible AI watermarks from images, plus provenance metadata from video.""" from dotenv import load_dotenv load_dotenv() # Load .env (e.g. HF_TOKEN) @@ -1016,6 +1017,23 @@ def cmd_invisible( # ── Metadata operations ── +def _print_metadata_report(source: Path, has_ai: bool, metadata: dict[str, str]) -> None: + """Render one metadata inspection result for the generic and video commands.""" + if not has_ai: + console.print(f" No AI metadata found in {source.name}") + return + + console.print(f" Warning: AI metadata detected in {source.name}:") + if synthid := metadata.get("synthid_watermark"): + console.print(f" Warning: SynthID watermark (inferred from C2PA metadata) {synthid}") + table = Table(show_header=True, header_style="bold") + table.add_column("Key", style="cyan") + table.add_column("Value") + for key, value in metadata.items(): + table.add_row(key, str(value)[:80]) + console.print(table) + + @main.command("metadata") @click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) @click.option("--check", is_flag=True, help="Check for AI metadata (don't modify).") @@ -1050,19 +1068,8 @@ def cmd_metadata( if check or (not remove): has_ai = has_ai_metadata(source) - if has_ai: - console.print(f" Warning: AI metadata detected in {source.name}:") - meta = get_ai_metadata(source) - if synthid := meta.get("synthid_watermark"): - console.print(f" Warning: SynthID watermark (inferred from C2PA metadata) {synthid}") - table = Table(show_header=True, header_style="bold") - table.add_column("Key", style="cyan") - table.add_column("Value") - for k, v in meta.items(): - table.add_row(k, str(v)[:80]) - console.print(table) - else: - console.print(f" No AI metadata found in {source.name}") + metadata = get_ai_metadata(source) if has_ai else {} + _print_metadata_report(source, has_ai, metadata) if not remove: return @@ -1082,6 +1089,114 @@ def cmd_metadata( console.print(f" AI metadata stripped -> {out}") +# ── Experimental video pipeline ── +@main.group("video") +def cmd_video() -> None: + """Process AI watermarks in video files.""" + + +@cmd_video.command("metadata") +@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) +@click.option("--check", is_flag=True, help="Check for AI metadata (don't modify).") +@click.option("--remove", is_flag=True, help="Remove AI metadata.") +@click.option( + "-o", + "--output", + type=click.Path(path_type=Path), + default=None, + help="Output path (default: _clean with the same container).", +) +@click.option("--keep-standard/--remove-all", default=True, help="Keep standard metadata.") +def cmd_video_metadata( + source: Path, + check: bool, + remove: bool, + output: Path | None, + keep_standard: bool, +) -> None: + """Check or remove AI metadata without transcoding video streams.""" + from remove_ai_watermarks.video import inspect_video_metadata, remove_video_metadata + + _banner() + try: + report = inspect_video_metadata(source) + except (OSError, ValueError) as e: + raise click.ClickException(str(e)) from e + + if check or not remove: + _print_metadata_report(source, report.has_ai_metadata, report.markers) + + if not remove: + return + + try: + result = remove_video_metadata(source, output, keep_standard=keep_standard) + except (OSError, RuntimeError, ValueError) as e: + raise click.ClickException(str(e)) from e + + if result.remaining: + console.print(f" FAILED: {len(result.remaining)} AI metadata marker(s) survived in {result.output}") + console.print(f" still present: {', '.join(sorted(result.remaining))}") + raise SystemExit(1) + console.print(f" AI metadata stripped -> {result.output}") + + +@cmd_video.command("visible") +@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) +@click.option( + "-o", + "--output", + type=click.Path(path_type=Path), + default=None, + help="Output path (default: _clean with the same container).", +) +@click.option( + "--mark", + type=click.Choice(["sora", "veo"]), + default="sora", + help="Visible AI mark to remove.", +) +@click.option( + "--backend", + type=click.Choice(["auto", "cv2", "migan", "lama"]), + default="cv2", + help="Per-frame fill backend. cv2 is the fast default; learned backends improve difficult backgrounds.", +) +@click.option("--strip-metadata/--keep-metadata", default=True, help="Strip AI metadata from the transcoded output.") +def cmd_video_visible( + source: Path, + output: Path | None, + mark: str, + backend: str, + strip_metadata: bool, +) -> None: + """Remove a temporally stable visible AI wordmark from video.""" + from remove_ai_watermarks.video import remove_video_visible + + _banner() + console.print(f" Scanning {source.name} for a temporally stable {mark} mark...") + try: + result = remove_video_visible( + source, + output, + mark=mark, + backend=backend, + strip_metadata=strip_metadata, + ) + except (OSError, RuntimeError, ValueError) as e: + raise click.ClickException(str(e)) from e + + if result.output is None: + console.print(f" No stable {mark} watermark detected; no output written") + raise SystemExit(EXIT_NO_VISIBLE_MARK) + if result.remaining_metadata: + console.print(f" FAILED: {len(result.remaining_metadata)} AI metadata marker(s) survived in {result.output}") + raise SystemExit(1) + console.print( + f" Removed {mark} watermark from {result.removed_frames}/{result.total_frames} frames -> {result.output}" + ) + + # ── Provenance identification ── @main.command("identify") @click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) diff --git a/src/remove_ai_watermarks/metadata.py b/src/remove_ai_watermarks/metadata.py index 9de9459..016cbe7 100644 --- a/src/remove_ai_watermarks/metadata.py +++ b/src/remove_ai_watermarks/metadata.py @@ -143,7 +143,7 @@ AIGC_MARKERS: tuple[bytes, ...] = ( # the same object as a PNG ``tEXt`` chunk keyed ``AIGC`` (raw JSON, not XMP), so # a JSON object carrying at least one of these is accepted as a valid TC260 # label even when the namespaced XMP element is absent. -_TC260_FIELDS: frozenset[str] = frozenset( +TC260_AIGC_FIELDS: frozenset[str] = frozenset( { # Producer-side schema (Doubao and most China-served generators). "Label", @@ -359,10 +359,12 @@ def has_ai_metadata(image_path: Path) -> bool: def aigc_label(image_path: Path) -> dict[str, str] | None: """Parse a China TC260 AI-labeling block, if present. - Three serializations are recognized: + Supported serializations are: - a PNG ``tEXt``/``iTXt`` chunk keyed ``AIGC`` carrying the raw JSON object (as written by Doubao / ByteDance), read via PIL; + - a native MP4/MOV ``AIGC`` key in ``moov.udta.meta.keys`` whose matching + ``ilst`` item carries the raw JSON object; - an XMP ``{...}`` block (HTML-entity encoded text), found by a container-agnostic raw-byte scan (PNG/JPEG/WebP alike); and - a raw-JSON ``{"AIGC":{...}}`` block with no namespace, as embedded in JPEG @@ -375,7 +377,7 @@ def aigc_label(image_path: Path) -> dict[str, str] | None: Returns the decoded JSON (e.g. ``{"Label": "1", "ContentProducer": ...}``) or None. The generic forms (the PNG-chunk key ``AIGC``, the bare ``{"AIGC":...}`` object, and the bare ``AIGC{...}`` blob) are accepted only - if they carry at least one known TC260 field (``_TC260_FIELDS``); the + if they carry at least one known TC260 field (``TC260_AIGC_FIELDS``); the namespaced XMP element is unambiguous, so any JSON object is accepted. """ import html @@ -390,7 +392,7 @@ def aigc_label(image_path: Path) -> dict[str, str] | None: if not isinstance(parsed, dict): return None fields = {str(k): str(v) for k, v in cast("dict[object, object]", parsed).items()} - if require_tc260_field and not (_TC260_FIELDS & fields.keys()): + if require_tc260_field and not (TC260_AIGC_FIELDS & fields.keys()): return None return fields @@ -407,6 +409,25 @@ def aigc_label(image_path: Path) -> dict[str, str] | None: if isinstance(value, str) and (result := _parse(value, require_tc260_field=True)): return result + # Native MP4/MOV TC260 metadata (TC260-PG-20257A): the ``AIGC`` key lives + # in ``moov.udta.meta.keys`` and points to a raw JSON value in ``ilst``. + # Read it through the bounded box walker so a tail ``moov`` after a large + # ``mdat`` is found without loading or scanning the media payload. + from remove_ai_watermarks.noai.isobmff import tc260_aigc_payloads + + for payload in tc260_aigc_payloads(image_path): + if result := _parse(payload.decode("utf-8", "replace"), require_tc260_field=True): + return result + + # Native MKV/WebM TC260 metadata: ``Segment.Tags.Tag.SimpleTag`` carries + # ``TagName=AIGC`` and the raw JSON in ``TagString``. The EBML walker seeks + # over clusters and reads only bounded metadata values. + from remove_ai_watermarks.noai.ebml import tc260_aigc_payloads as ebml_tc260_aigc_payloads + + for payload in ebml_tc260_aigc_payloads(image_path): + if result := _parse(payload.decode("utf-8", "replace"), require_tc260_field=True): + return result + # XMP TC260:AIGC, namespaced (unambiguous) in either serialization RDF allows: # an element {...} or an attribute TC260:AIGC="{...}" # (the attribute form is what PicWish writes). Both are HTML-entity encoded. @@ -750,7 +771,7 @@ def _is_aigc_exif_value(raw: object) -> bool: Mirrors ``aigc_label``'s EXIF path: the ``{"AIGC":{...}}`` wrapper embedded in ``UserComment`` / ``ImageDescription`` by China-served generators (Doubao's producer schema AND Tencent Cloud's service-provider schema, both keyed under - ``_TC260_FIELDS``). Gated on both the ``AIGC`` marker and a TC260 field so a + ``TC260_AIGC_FIELDS``). Gated on both the ``AIGC`` marker and a TC260 field so a coincidental token cannot false-drop a genuine caption/comment. Accepts a ``str`` too (a PNG ``tEXt``/``iTXt`` value), not only EXIF bytes. """ @@ -761,7 +782,7 @@ def _is_aigc_exif_value(raw: object) -> bool: if b"AIGC" not in raw: return False text = bytes(raw).decode("latin-1", "ignore") - return any(field in text for field in _TC260_FIELDS) + return any(field in text for field in TC260_AIGC_FIELDS) def _ai_exif_targets(loaded: dict[str, Any]) -> list[tuple[str, int, bytes, str]]: @@ -1154,6 +1175,7 @@ def remove_ai_metadata( from remove_ai_watermarks.noai.isobmff import ( blank_ai_exif_tokens, blank_ai_xmp_packets, + blank_tc260_aigc_tags, is_isobmff, strip_c2pa_boxes, ) @@ -1164,17 +1186,20 @@ def remove_ai_metadata( data = source_path.read_bytes() # Top-level uuid/jumb boxes (C2PA + AI-label XMP), then the meta-box items # the top-level stripper can't reach (HEIF/AVIF store them in mdat/idat): - # AI-label XMP packets and AI-generator tokens in an Exif item -- both - # blanked in place (same length) so box sizes and iloc offsets stay valid - # and the coded image is untouched. + # Native TC260 tags, AI-label XMP packets, and AI-generator tokens in an + # Exif item are blanked in place (same length) so box sizes and iloc / + # media offsets stay valid and the coded content is untouched. cleaned, stripped = strip_c2pa_boxes(data) + cleaned, tc260_blanked = blank_tc260_aigc_tags(cleaned) cleaned, blanked = blank_ai_xmp_packets(cleaned) cleaned, exif_blanked = blank_ai_exif_tokens(cleaned) output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_bytes(cleaned) logger.info( - "Stripped %d AI-provenance box(es), blanked %d meta-box XMP packet(s) + %d EXIF token(s) → %s", + "Stripped %d AI-provenance box(es), blanked %d native TC260 tag(s) + " + "%d meta-box XMP packet(s) + %d EXIF token(s) → %s", stripped, + tc260_blanked, blanked, exif_blanked, output_path, diff --git a/src/remove_ai_watermarks/noai/ebml.py b/src/remove_ai_watermarks/noai/ebml.py new file mode 100644 index 0000000..6b87412 --- /dev/null +++ b/src/remove_ai_watermarks/noai/ebml.py @@ -0,0 +1,167 @@ +"""Bounded Matroska/WebM metadata reader for native TC260 AIGC labels. + +TC260-PG-20257A stores the label as a Matroska ``SimpleTag`` whose +``TagName`` is ``AIGC`` and whose ``TagString`` is the normative JSON object. +The walker seeks over unrelated elements such as clusters instead of reading +their payloads. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, BinaryIO, cast + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + +from remove_ai_watermarks.metadata import TC260_AIGC_FIELDS + +_EBML_MAGIC = b"\x1aE\xdf\xa3" +_SEGMENT_ID = 0x18538067 +_TAGS_ID = 0x1254C367 +_TAG_ID = 0x7373 +_SIMPLE_TAG_ID = 0x67C8 +_TAG_NAME_ID = 0x45A3 +_TAG_STRING_ID = 0x4487 +_MAX_TAG_NAME_BYTES = 256 +_MAX_TC260_VALUE_BYTES = 1024 * 1024 + + +def _vint_length(first: int, *, maximum: int) -> int | None: + """Return an EBML variable-integer length from its first byte.""" + mask = 0x80 + for length in range(1, maximum + 1): + if first & mask: + return length + mask >>= 1 + return None + + +def _read_element_header( + stream: BinaryIO, + pos: int, + limit: int, +) -> tuple[int, int, int] | None: + """Return ``(element_id, payload_start, element_end)`` inside ``limit``.""" + if pos < 0 or pos >= limit: + return None + stream.seek(pos) + first_raw = stream.read(1) + if not first_raw: + return None + id_length = _vint_length(first_raw[0], maximum=4) + if id_length is None or pos + id_length >= limit: + return None + element_id_raw = first_raw + stream.read(id_length - 1) + if len(element_id_raw) != id_length: + return None + element_id = int.from_bytes(element_id_raw, "big") + + size_first_raw = stream.read(1) + if not size_first_raw: + return None + size_length = _vint_length(size_first_raw[0], maximum=8) + if size_length is None: + return None + size_rest = stream.read(size_length - 1) + if len(size_rest) != size_length - 1: + return None + marker = 1 << (8 - size_length) + size_value = int.from_bytes(bytes([size_first_raw[0] & (marker - 1)]) + size_rest, "big") + payload_start = pos + id_length + size_length + if payload_start > limit: + return None + unknown_size = size_value == (1 << (7 * size_length)) - 1 + element_end = limit if unknown_size else payload_start + size_value + if element_end > limit: + return None + return element_id, payload_start, element_end + + +def _iter_elements( + stream: BinaryIO, + start: int, + end: int, +) -> Iterator[tuple[int, int, int]]: + """Yield valid direct children from one bounded EBML region.""" + pos = start + while pos < end: + header = _read_element_header(stream, pos, end) + if header is None: + return + element_id, payload_start, element_end = header + yield element_id, payload_start, element_end + if element_end <= pos: + return + pos = element_end + + +def _read_bounded( + stream: BinaryIO, + start: int, + end: int, + maximum: int, +) -> bytes | None: + size = end - start + if size < 0 or size > maximum: + return None + stream.seek(start) + value = stream.read(size) + return value if len(value) == size else None + + +def _is_tc260_aigc_json(value: bytes) -> bool: + try: + parsed = json.loads(value.decode("utf-8")) + except (UnicodeDecodeError, ValueError): + return False + if not isinstance(parsed, dict): + return False + fields = cast("dict[object, object]", parsed) + return bool(TC260_AIGC_FIELDS & {str(key) for key in fields}) + + +def _simple_tag_payloads( + stream: BinaryIO, + start: int, + end: int, +) -> tuple[bytes, ...]: + name: bytes | None = None + values: list[bytes] = [] + for element_id, payload_start, element_end in _iter_elements(stream, start, end): + if element_id == _TAG_NAME_ID: + name = _read_bounded(stream, payload_start, element_end, _MAX_TAG_NAME_BYTES) + elif element_id == _TAG_STRING_ID: + value = _read_bounded(stream, payload_start, element_end, _MAX_TC260_VALUE_BYTES) + if value is not None: + values.append(value) + if name != b"AIGC": + return () + return tuple(value for value in values if _is_tc260_aigc_json(value)) + + +def tc260_aigc_payloads(path: str | Path) -> tuple[bytes, ...]: + """Read validated TC260 values from Matroska/WebM ``SimpleTag`` entries.""" + found: list[bytes] = [] + try: + with open(path, "rb") as stream: + if stream.read(4) != _EBML_MAGIC: + return () + stream.seek(0, 2) + file_size = stream.tell() + for element_id, payload_start, element_end in _iter_elements(stream, 0, file_size): + if element_id != _SEGMENT_ID: + continue + for child_id, child_start, child_end in _iter_elements(stream, payload_start, element_end): + if child_id != _TAGS_ID: + continue + for tag_id, tag_start, tag_end in _iter_elements(stream, child_start, child_end): + if tag_id != _TAG_ID: + continue + for simple_id, simple_start, simple_end in _iter_elements(stream, tag_start, tag_end): + if simple_id == _SIMPLE_TAG_ID: + found.extend(_simple_tag_payloads(stream, simple_start, simple_end)) + except OSError: + return () + return tuple(found) diff --git a/src/remove_ai_watermarks/noai/isobmff.py b/src/remove_ai_watermarks/noai/isobmff.py index 2086638..af135f7 100644 --- a/src/remove_ai_watermarks/noai/isobmff.py +++ b/src/remove_ai_watermarks/noai/isobmff.py @@ -1,4 +1,4 @@ -"""Minimal ISOBMFF box walker for stripping C2PA from AVIF / HEIF / MP4 / JPEG-XL. +"""Minimal ISOBMFF box walker for AI provenance in AVIF / HEIF / MP4 / JPEG-XL. The ISO Base Media File Format wraps content in nested ``[size:4][type:4][...]`` boxes. C2PA stores its manifest in a top-level ``uuid`` box keyed by the @@ -8,6 +8,11 @@ carry C2PA, and emit the rest verbatim. The codestream (``mdat`` for ISOBMFF, ``jxlc`` / ``jxlp`` for JPEG-XL) is untouched, so pixel data is preserved bit-for-bit. +TC260-PG-20257A video metadata is nested instead: +``moov.udta.meta.keys/ilst``. Its detector seeks through those boxes without +reading media payloads, and its stripper blanks the validated key/value in +place so fast-start media offsets remain valid. + This file intentionally avoids dependencies on format-specific libraries (pillow-heif, pillow-jxl, pymp4) so it works on systems where they aren't installed. @@ -17,10 +22,12 @@ Reference: ISO/IEC 14496-12 (ISOBMFF) and C2PA 2.1 spec §11. from __future__ import annotations +import io +import json import logging import re import struct -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, BinaryIO, cast if TYPE_CHECKING: from collections.abc import Iterator @@ -31,6 +38,7 @@ from remove_ai_watermarks.metadata import ( C2PA_UUID, IPTC_AI_FIELD_MARKERS, IPTC_AI_MARKERS, + TC260_AIGC_FIELDS, ) logger = logging.getLogger(__name__) @@ -53,6 +61,13 @@ _AI_LABEL_MARKERS: tuple[bytes, ...] = AIGC_MARKERS + IPTC_AI_MARKERS + IPTC_AI_ # blanked in place (see ``blank_ai_xmp_packets``). _XMP_PACKET_RE = re.compile(rb"<\?xpacket begin=.*?<\?xpacket end=[^>]*?\?>", re.DOTALL) +# TC260-PG-20257A stores an MP4/MOV label as an ``AIGC`` key in +# ``moov.udta.meta.keys`` and its JSON value in the corresponding +# ``moov.udta.meta.ilst`` item. The value is intentionally bounded before it is +# read: the normative object is tiny, and a corrupt size must not allocate an +# arbitrary amount of memory during an inspection. +_MAX_TC260_VALUE_BYTES = 1024 * 1024 + def _iter_top_level_boxes(data: bytes) -> Iterator[tuple[int, int, bytes, int]]: """Yield ``(start, end, type, payload_offset)`` for each top-level box. @@ -84,6 +99,199 @@ def _iter_top_level_boxes(data: bytes) -> Iterator[tuple[int, int, bytes, int]]: pos += size +def _read_box_header( + stream: BinaryIO, + pos: int, + limit: int, +) -> tuple[int, bytes, int] | None: + """Return ``(end, type, payload_offset)`` for one box inside ``limit``.""" + if pos < 0 or pos + 8 > limit: + return None + stream.seek(pos) + header = stream.read(8) + if len(header) != 8: + return None + size32 = struct.unpack(">I", header[:4])[0] + box_type = header[4:8] + payload_off = pos + 8 + if size32 == 1: + extended = stream.read(8) + if len(extended) != 8: + return None + size = struct.unpack(">Q", extended)[0] + payload_off = pos + 16 + elif size32 == 0: + size = limit - pos + else: + size = size32 + end = pos + size + if size < payload_off - pos or end > limit: + return None + return end, box_type, payload_off + + +def _iter_file_boxes( + stream: BinaryIO, + start: int, + end: int, +) -> Iterator[tuple[int, int, bytes, int]]: + """Yield valid boxes from one bounded container region.""" + pos = start + while pos + 8 <= end: + header = _read_box_header(stream, pos, end) + if header is None: + return + box_end, box_type, payload_off = header + yield pos, box_end, box_type, payload_off + pos = box_end + + +def _tc260_key_indices( + stream: BinaryIO, + payload_off: int, + box_end: int, +) -> dict[int, tuple[int, int]]: + """Map every exact ``AIGC`` key index to its byte span.""" + if payload_off + 8 > box_end: + return {} + stream.seek(payload_off) + prefix = stream.read(8) + if len(prefix) != 8: + return {} + entry_count = struct.unpack(">I", prefix[4:8])[0] + pos = payload_off + 8 + found: dict[int, tuple[int, int]] = {} + for index in range(1, entry_count + 1): + if pos + 8 > box_end: + return {} + stream.seek(pos) + header = stream.read(8) + if len(header) != 8: + return {} + entry_size = struct.unpack(">I", header[:4])[0] + entry_end = pos + entry_size + if entry_size < 8 or entry_end > box_end: + return {} + name_start = pos + 8 + if entry_end - name_start == 4: + stream.seek(name_start) + if stream.read(4) == b"AIGC": + found[index] = (name_start, entry_end) + pos = entry_end + return found + + +def _is_tc260_aigc_json(value: bytes) -> bool: + """Require a JSON object carrying at least one normative TC260 field.""" + try: + parsed = json.loads(value.decode("utf-8")) + except (UnicodeDecodeError, ValueError): + return False + if not isinstance(parsed, dict): + return False + fields = cast("dict[object, object]", parsed) + return bool(TC260_AIGC_FIELDS & {str(key) for key in fields}) + + +def _tc260_aigc_regions( + stream: BinaryIO, + file_size: int, +) -> list[tuple[int, int, int, int, bytes]]: + """Locate validated native TC260 entries without reading media payloads. + + Each tuple is ``(key_start, key_end, value_start, value_end, value)``. + """ + regions: list[tuple[int, int, int, int, bytes]] = [] + for _moov_start, moov_end, moov_type, moov_payload in _iter_file_boxes(stream, 0, file_size): + if moov_type != b"moov": + continue + for _udta_start, udta_end, udta_type, udta_payload in _iter_file_boxes( + stream, + moov_payload, + moov_end, + ): + if udta_type != b"udta": + continue + for _meta_start, meta_end, meta_type, meta_payload in _iter_file_boxes( + stream, + udta_payload, + udta_end, + ): + if meta_type != b"meta" or meta_payload + 4 > meta_end: + continue + keys: dict[int, tuple[int, int]] = {} + ilst_boxes: list[tuple[int, int]] = [] + for _child_start, child_end, child_type, child_payload in _iter_file_boxes( + stream, + meta_payload + 4, + meta_end, + ): + if child_type == b"keys": + keys.update(_tc260_key_indices(stream, child_payload, child_end)) + elif child_type == b"ilst": + ilst_boxes.append((child_payload, child_end)) + if not keys: + continue + for ilst_payload, ilst_end in ilst_boxes: + for _item_start, item_end, item_type, item_payload in _iter_file_boxes( + stream, + ilst_payload, + ilst_end, + ): + index = int.from_bytes(item_type, "big") + key_span = keys.get(index) + if key_span is None: + continue + for _data_start, data_end, data_type, data_payload in _iter_file_boxes( + stream, + item_payload, + item_end, + ): + value_start = data_payload + 8 + value_size = data_end - value_start + if data_type != b"data" or value_size < 0 or value_size > _MAX_TC260_VALUE_BYTES: + continue + stream.seek(value_start) + value = stream.read(value_size) + if len(value) == value_size and _is_tc260_aigc_json(value): + regions.append((*key_span, value_start, data_end, value)) + return regions + + +def tc260_aigc_payloads(path: str | Path) -> tuple[bytes, ...]: + """Read native TC260 ``AIGC`` JSON values from an MP4/MOV container.""" + try: + with open(path, "rb") as stream: + if not is_isobmff(stream.read(8)): + return () + stream.seek(0, 2) + file_size = stream.tell() + return tuple(region[4] for region in _tc260_aigc_regions(stream, file_size)) + except OSError: + return () + + +def blank_tc260_aigc_tags(data: bytes) -> tuple[bytes, int]: + """Blank native TC260 values in place while preserving every box offset. + + Removing a nested ``ilst`` item would shift ``mdat`` in a fast-start MP4 and + invalidate its chunk offsets. Replacing the four-byte key with ``free`` and + the JSON value with spaces keeps every box size and media offset unchanged. + """ + if not is_isobmff(data): + return data, 0 + regions = _tc260_aigc_regions(io.BytesIO(data), len(data)) + if not regions: + return data, 0 + out = bytearray(data) + key_spans: set[tuple[int, int]] = set() + for key_start, key_end, value_start, value_end, _value in regions: + key_spans.add((key_start, key_end)) + out[key_start:key_end] = b"free" + out[value_start:value_end] = b" " * (value_end - value_start) + return bytes(out), len(key_spans) + + def is_isobmff(data: bytes) -> bool: """Cheap sniff: ISOBMFF files start with an ``ftyp`` box.""" return len(data) >= 8 and data[4:8] == b"ftyp" diff --git a/src/remove_ai_watermarks/video.py b/src/remove_ai_watermarks/video.py new file mode 100644 index 0000000..1fcb364 --- /dev/null +++ b/src/remove_ai_watermarks/video.py @@ -0,0 +1,212 @@ +"""High-level video processing API. + +Supported experimental stages are container-level AI metadata inspection and +removal plus temporally stabilized visible Sora and Veo removal. The pixel path +reuses the image package's shared fill backends. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v", ".webm", ".mkv"}) +_ISOBMFF_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v"}) +_EBML_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".webm", ".mkv"}) +_EBML_MAGIC = b"\x1aE\xdf\xa3" + + +@dataclass(frozen=True) +class VideoMetadataReport: + """AI metadata found in one supported video container.""" + + source: Path + has_ai_metadata: bool + markers: dict[str, str] + + +@dataclass(frozen=True) +class VideoMetadataResult: + """Result of a verified video metadata-removal operation.""" + + source: Path + output: Path + detected: dict[str, str] + remaining: dict[str, str] + + +@dataclass(frozen=True) +class VideoVisibleResult: + """Result of visible AI-watermark removal from a video.""" + + source: Path + output: Path | None + mark: str + total_frames: int + detected_frames: int + removed_frames: int + remaining_metadata: dict[str, str] + + +def _video_source(source: str | Path) -> Path: + path = Path(source) + if not path.exists(): + raise FileNotFoundError(f"Video does not exist: {path}") + if not path.is_file(): + raise ValueError(f"Video source must be a file: {path}") + if path.suffix.lower() not in VIDEO_EXTENSIONS: + supported = ", ".join(sorted(VIDEO_EXTENSIONS)) + raise ValueError(f"Unsupported video format {path.suffix or ''}; expected one of: {supported}") + with path.open("rb") as stream: + head = stream.read(12) + suffix = path.suffix.lower() + matches_container = (suffix in _ISOBMFF_VIDEO_EXTENSIONS and len(head) >= 8 and head[4:8] == b"ftyp") or ( + suffix in _EBML_VIDEO_EXTENSIONS and head.startswith(_EBML_MAGIC) + ) + if not matches_container: + raise ValueError(f"Video content does not match its {suffix} extension: {path}") + return path + + +def _video_output( + source: Path, + output: str | Path | None, + *, + operation: str = "metadata removal", +) -> Path: + path = Path(output) if output is not None else source.with_stem(source.stem + "_clean") + if path.suffix.lower() != source.suffix.lower(): + raise ValueError( + f"Video output container must match the source ({source.suffix}); " + f"{operation} does not change containers to {path.suffix or ''}" + ) + if path.resolve() == source.resolve(): + raise ValueError(f"Video {operation} requires a distinct output path") + return path + + +def inspect_video_metadata(source: str | Path) -> VideoMetadataReport: + """Inspect supported AI-provenance metadata in a video container.""" + from remove_ai_watermarks.metadata import get_ai_metadata, has_ai_metadata + + source_path = _video_source(source) + return VideoMetadataReport( + source=source_path, + has_ai_metadata=has_ai_metadata(source_path), + markers=get_ai_metadata(source_path), + ) + + +def remove_video_metadata( + source: str | Path, + output: str | Path | None = None, + *, + keep_standard: bool = True, +) -> VideoMetadataResult: + """Remove AI metadata without transcoding video or audio streams. + + The default output is ``_clean``. A separate + output is required so an experimental video operation never overwrites the + original file. + """ + from remove_ai_watermarks.metadata import get_ai_metadata, strip_and_verify + + source_path = _video_source(source) + output_path = _video_output(source_path, output) + detected = get_ai_metadata(source_path) + written, remaining = strip_and_verify(source_path, output_path, keep_standard=keep_standard) + return VideoMetadataResult( + source=source_path, + output=written, + detected=detected, + remaining=remaining, + ) + + +def remove_video_visible( + source: str | Path, + output: str | Path | None = None, + *, + mark: str = "sora", + backend: str = "cv2", + strip_metadata: bool = True, +) -> VideoVisibleResult: + """Remove a supported visible AI wordmark from a video. + + Supported marks are ``sora`` and ``veo``. The full sequence is scanned before + pixels change, and only recurring candidates are accepted. Audio is copied + without re-encoding; video is transcoded because the pixels change. When no + stable mark is found, no output is written and ``output`` in the result is + ``None``. + """ + from remove_ai_watermarks.metadata import get_ai_metadata + from remove_ai_watermarks.video_visible import ( + encode_clean_video, + has_sora_provenance, + has_veo_provenance, + scan_sora_video, + scan_veo_video, + stabilize_sora_localizations, + stabilize_veo_localizations, + ) + from remove_ai_watermarks.watermark_registry import resolve_backend + + if mark not in {"sora", "veo"}: + raise ValueError("Unsupported visible video mark; expected 'sora' or 'veo'") + if backend not in {"auto", "cv2", "migan", "lama"}: + raise ValueError("Unsupported fill backend; expected auto, cv2, migan, or lama") + + source_path = _video_source(source) + output_path = _video_output(source_path, output, operation="visible watermark removal") + markers = get_ai_metadata(source_path) + if mark == "sora": + scan = scan_sora_video(source_path) + regions = stabilize_sora_localizations( + scan.detections, + provenance=has_sora_provenance(markers), + ) + padding_fraction = 0.28 + mask_style = "box" + else: + scan = scan_veo_video(source_path) + regions = stabilize_veo_localizations( + scan.detections, + provenance=has_veo_provenance(markers), + ) + padding_fraction = 0.18 + mask_style = "veo" + detected_frames = sum(region is not None for region in regions) + if detected_frames == 0: + return VideoVisibleResult( + source=source_path, + output=None, + mark=mark, + total_frames=len(scan.detections), + detected_frames=0, + removed_frames=0, + remaining_metadata=markers if strip_metadata else {}, + ) + + # Validate optional model availability before ffmpeg creates or overwrites + # the requested output. + resolve_backend(backend) # type: ignore[arg-type] + removed_frames = encode_clean_video( + source_path, + output_path, + scan, + regions, + backend=backend, # type: ignore[arg-type] + strip_metadata=strip_metadata, + padding_fraction=padding_fraction, + mask_style=mask_style, + ) + remaining_metadata = get_ai_metadata(output_path) if strip_metadata else {} + return VideoVisibleResult( + source=source_path, + output=output_path, + mark=mark, + total_frames=len(scan.detections), + detected_frames=detected_frames, + removed_frames=removed_frames, + remaining_metadata=remaining_metadata, + ) diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py new file mode 100644 index 0000000..7b15ca0 --- /dev/null +++ b/src/remove_ai_watermarks/video_visible.py @@ -0,0 +1,706 @@ +"""Visible AI-watermark localization and removal for video. + +Supported marks use fully synthetic silhouettes made from geometric primitives +and Pillow's bundled font. Sora detection searches the full frame because the +wordmark moves. Veo detection covers both the current four-point diamond and the +legacy ``Veo`` text in the bottom-right corner. A single frame is never enough +to authorize removal: the temporal arbiter requires the candidate to recur at +the same location across adjacent frames. This keeps isolated lookalikes in +clean videos from becoming removal masks. + +Video pixels are decoded with OpenCV and encoded with the system ``ffmpeg``. +Audio is stream-copied from the source. The video stream must be transcoded +because visible-mark removal changes pixels. +""" + +# cv2/numpy boundary: these packages do not expose usable types for many array +# operations. Public signatures remain annotated while unknown third-party types +# are relaxed only in this module. +# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false + +from __future__ import annotations + +import logging +import shutil +import subprocess +from dataclasses import dataclass +from functools import lru_cache +from itertools import pairwise +from typing import TYPE_CHECKING, Any, Literal + +import cv2 +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +if TYPE_CHECKING: + from pathlib import Path + + from numpy.typing import NDArray + + from remove_ai_watermarks.watermark_registry import Backend + + +log = logging.getLogger(__name__) + +Region = tuple[int, int, int, int] + +_NORMALIZED_SHORT_SIDE = 480 +_SORA_TEMPLATE_SIZE = (180, 64) +_SORA_RELATIVE_HEIGHTS = (0.065, 0.075, 0.085, 0.095, 0.105) +_SORA_PROVENANCE_WEAK_CONFIDENCE = 0.58 +_SORA_STRICT_WEAK_CONFIDENCE = 0.60 +_SORA_STRONG_CONFIDENCE = 0.65 +_VEO_PROVENANCE_WEAK_CONFIDENCE = 0.45 +_VEO_STRICT_WEAK_CONFIDENCE = 0.50 +_VEO_STRONG_CONFIDENCE = 0.55 +_MIN_STABLE_FRAMES = 5 +_MIN_VEO_STABLE_FRAMES = 12 +_MAX_STABLE_GAP = 2 +_STABLE_IOU = 0.55 +_VEO_REFERENCE_SHORT_SIDE = 720 +_VEO_DIAMOND_PROFILES = ( + (56, 92, 92), + (48, 72, 72), + (44, 29, 40), +) + + +@dataclass(frozen=True) +class FrameLocalization: + """Best untrusted visible-mark candidate found in one decoded frame.""" + + frame_index: int + confidence: float + region: Region | None + + +@dataclass(frozen=True) +class VideoScan: + """Decoded video geometry plus one localization candidate per frame.""" + + width: int + height: int + fps: float + detections: tuple[FrameLocalization, ...] + + +def _scalable_default_font(size: int) -> ImageFont.ImageFont | ImageFont.FreeTypeFont: + """Load Pillow's bundled scalable font, with a Pillow 10.0 fallback.""" + try: + return ImageFont.load_default(size=size) + except TypeError: + # ``size=`` was added after Pillow 10.0, which is still within the + # package's supported dependency range. Resize that bundled bitmap font + # before compositing it into the synthetic template. + return ImageFont.load_default() + + +@lru_cache(maxsize=1) +def _sora_templates() -> tuple[NDArray[Any], NDArray[Any]]: + """Return synthetic full-wordmark and mascot-only silhouettes. + + No source frame or provider logo asset contributes pixels to these templates. + The cloud-like mascot is assembled from primitive shapes and the word is + rendered with Pillow's bundled font. + """ + width, height = _SORA_TEMPLATE_SIZE + canvas = Image.new("L", (width, height), 0) + draw = ImageDraw.Draw(canvas) + draw.rounded_rectangle((4, 8, 58, 57), radius=22, fill=255) + draw.ellipse((0, 18, 26, 50), fill=255) + draw.ellipse((38, 16, 64, 51), fill=255) + draw.ellipse((16, 18, 29, 43), fill=0) + draw.ellipse((35, 18, 48, 43), fill=0) + + font = _scalable_default_font(50) + if isinstance(font, ImageFont.FreeTypeFont): + draw.text((67, 0), "Sora", font=font, fill=255, stroke_width=1, stroke_fill=255) + else: + text_box = font.getbbox("Sora") + text = Image.new("L", (max(1, text_box[2]), max(1, text_box[3])), 0) + ImageDraw.Draw(text).text((0, 0), "Sora", font=font, fill=255) + text = text.resize((102, 50), Image.Resampling.NEAREST) + canvas.paste(text, (67, 0), text) + + full = np.asarray(canvas, dtype=np.uint8) + return full, full[:, :64] + + +@lru_cache(maxsize=1) +def _veo_templates() -> tuple[NDArray[Any], NDArray[Any]]: + """Return synthetic current-diamond and legacy-text Veo silhouettes.""" + size = 256 + diamond_canvas = Image.new("L", (size, size), 0) + diamond_points = ( + (0.50, 0.02), + (0.60, 0.39), + (0.98, 0.50), + (0.60, 0.61), + (0.50, 0.98), + (0.40, 0.61), + (0.02, 0.50), + (0.40, 0.39), + ) + ImageDraw.Draw(diamond_canvas).polygon( + [(round(x * size), round(y * size)) for x, y in diamond_points], + fill=255, + ) + + text_canvas = Image.new("L", (140, 60), 0) + text_draw = ImageDraw.Draw(text_canvas) + text_draw.text((2, 0), "Veo", font=_scalable_default_font(48), fill=255) + text = np.asarray(text_canvas, dtype=np.uint8) + ys, xs = np.where(text > 0) + text = text[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1] + return np.asarray(diamond_canvas, dtype=np.uint8), text + + +def _top_hat(gray: NDArray[Any]) -> NDArray[Any]: + kernel = np.ones((7, 7), dtype=np.uint8) + return cv2.morphologyEx(gray, cv2.MORPH_TOPHAT, kernel) + + +def _normalized_gray(image_bgr: NDArray[Any]) -> tuple[NDArray[Any], float]: + gray = image_bgr if image_bgr.ndim == 2 else cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) + height, width = gray.shape[:2] + short_side = min(height, width) + if short_side <= 0: + return gray, 1.0 + scale = min(1.0, _NORMALIZED_SHORT_SIDE / short_side) + if scale < 1.0: + gray = cv2.resize( + gray, + (max(1, round(width * scale)), max(1, round(height * scale))), + interpolation=cv2.INTER_AREA, + ) + return gray, scale + + +def _expanded_region( + location: tuple[int, int], + template_width: int, + template_height: int, + *, + icon_only: bool, + scale: float, + frame_width: int, + frame_height: int, +) -> Region: + x = round(location[0] / scale) + y = round(location[1] / scale) + height = max(1, round(template_height / scale)) + width = max(1, round(template_width / scale)) + if icon_only: + width = round(height * _SORA_TEMPLATE_SIZE[0] / _SORA_TEMPLATE_SIZE[1]) + width = min(width, frame_width - x) + height = min(height, frame_height - y) + return x, y, max(1, width), max(1, height) + + +def detect_sora_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: + """Locate the strongest synthetic Sora-wordmark match in one frame. + + The returned candidate is intentionally untrusted. Call + :func:`stabilize_sora_localizations` across the full sequence before building + any removal mask. + """ + if image_bgr.size == 0: + return FrameLocalization(frame_index, 0.0, None) + + frame_height, frame_width = image_bgr.shape[:2] + gray, scale = _normalized_gray(image_bgr) + normalized_height, normalized_width = gray.shape[:2] + feature = _top_hat(gray) + best_confidence = 0.0 + best_region: Region | None = None + + for template_index, base_template in enumerate(_sora_templates()): + icon_only = template_index == 1 + for relative_height in _SORA_RELATIVE_HEIGHTS: + template_height = max(16, round(min(normalized_height, normalized_width) * relative_height)) + template_width = max(1, round(base_template.shape[1] * template_height / base_template.shape[0])) + if template_height >= normalized_height or template_width >= normalized_width: + continue + template = cv2.resize( + base_template, + (template_width, template_height), + interpolation=cv2.INTER_AREA, + ) + scores = cv2.matchTemplate(feature, _top_hat(template), cv2.TM_CCOEFF_NORMED) + _, confidence, _, location = cv2.minMaxLoc(scores) + if confidence <= best_confidence: + continue + best_confidence = float(confidence) + best_region = _expanded_region( + location, + template_width, + template_height, + icon_only=icon_only, + scale=scale, + frame_width=frame_width, + frame_height=frame_height, + ) + + return FrameLocalization(frame_index, best_confidence, best_region) + + +def _match_template( + gray: NDArray[Any], + template: NDArray[Any], + *, + region: Region, + kernel_size: int, +) -> tuple[float, Region | None]: + """Match one synthetic silhouette inside a bounded frame region.""" + x, y, width, height = region + roi = gray[y : y + height, x : x + width] + template_height, template_width = template.shape[:2] + if roi.size == 0 or template_height >= roi.shape[0] or template_width >= roi.shape[1]: + return 0.0, None + kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8) + feature = cv2.morphologyEx(roi, cv2.MORPH_TOPHAT, kernel) + template_feature = cv2.morphologyEx(template, cv2.MORPH_TOPHAT, kernel) + scores = cv2.matchTemplate(feature, template_feature, cv2.TM_CCOEFF_NORMED) + _, confidence, _, location = cv2.minMaxLoc(scores) + return float(confidence), ( + x + location[0], + y + location[1], + template_width, + template_height, + ) + + +def detect_veo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: + """Locate the strongest current-diamond or legacy-text Veo candidate.""" + if image_bgr.size == 0: + return FrameLocalization(frame_index, 0.0, None) + + frame_height, frame_width = image_bgr.shape[:2] + gray = image_bgr if image_bgr.ndim == 2 else cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) + short_scale = min(frame_height, frame_width) / _VEO_REFERENCE_SHORT_SIDE + diamond_base, text_base = _veo_templates() + best_confidence = 0.0 + best_region: Region | None = None + + diamond_sizes: set[int] = set() + for base_size, right_margin, bottom_margin in _VEO_DIAMOND_PROFILES: + diamond_size = max(16, round(base_size * short_scale)) + diamond_sizes.add(diamond_size) + template = cv2.resize( + diamond_base, + (diamond_size, diamond_size), + interpolation=cv2.INTER_AREA, + ) + expected_x = round(frame_width - (right_margin + base_size) * short_scale) + expected_y = round(frame_height - (bottom_margin + base_size) * short_scale) + search_padding = max(6, round(diamond_size * 0.25)) + search_x = max(0, expected_x - search_padding) + search_y = max(0, expected_y - search_padding) + search_width = min(frame_width - search_x, diamond_size + search_padding * 2) + search_height = min(frame_height - search_y, diamond_size + search_padding * 2) + confidence, candidate = _match_template( + gray, + template, + region=(search_x, search_y, search_width, search_height), + kernel_size=max(3, round(7 * short_scale) | 1), + ) + if confidence > best_confidence: + best_confidence = confidence + best_region = candidate + + # Provider layouts have moved before. A bounded corner search is a safety + # net for a relocated diamond, but it is admitted only at a much stronger + # per-frame score than the known-profile path. Without this gate, recurring + # bright scene details in clean API exports can become stable false matches. + corner_x = round(frame_width * 0.65) + corner_y = round(frame_height * 0.65) + corner_region = (corner_x, corner_y, frame_width - corner_x, frame_height - corner_y) + for diamond_size in diamond_sizes: + template = cv2.resize( + diamond_base, + (diamond_size, diamond_size), + interpolation=cv2.INTER_AREA, + ) + confidence, candidate = _match_template( + gray, + template, + region=corner_region, + kernel_size=max(3, round(7 * short_scale) | 1), + ) + if confidence >= 0.70 and confidence > best_confidence: + best_confidence = confidence + best_region = candidate + + text_region_width = min(frame_width, max(32, round(180 * short_scale))) + text_region_height = min(frame_height, max(24, round(120 * short_scale))) + text_region = ( + frame_width - text_region_width, + frame_height - text_region_height, + text_region_width, + text_region_height, + ) + text_heights = sorted({max(5, round(height * short_scale)) for height in range(8, 22)}) + for text_height in text_heights: + text_width = max(1, round(text_base.shape[1] * text_height / text_base.shape[0])) + template = cv2.resize( + text_base, + (text_width, text_height), + interpolation=cv2.INTER_AREA, + ) + confidence, candidate = _match_template( + gray, + template, + region=text_region, + kernel_size=max(3, round(3 * short_scale) | 1), + ) + if confidence > best_confidence: + best_confidence = confidence + best_region = candidate + + return FrameLocalization(frame_index, best_confidence, best_region) + + +def _region_iou(left: Region, right: Region) -> float: + lx, ly, lw, lh = left + rx, ry, rw, rh = right + x0 = max(lx, rx) + y0 = max(ly, ry) + x1 = min(lx + lw, rx + rw) + y1 = min(ly + lh, ry + rh) + intersection = max(0, x1 - x0) * max(0, y1 - y0) + union = lw * lh + rw * rh - intersection + return intersection / union if union > 0 else 0.0 + + +def stabilize_sora_localizations( + detections: tuple[FrameLocalization, ...] | list[FrameLocalization], + *, + provenance: bool, +) -> list[Region | None]: + """Accept only spatially recurring Sora candidates and bridge short dropouts. + + Metadata never creates a detection. It only allows a stable visual run whose + scores remain below the strict confidence floor, which covers low-contrast + Sora marks while clean metadata-bearing exports stay untouched. + """ + weak_floor = _SORA_PROVENANCE_WEAK_CONFIDENCE if provenance else _SORA_STRICT_WEAK_CONFIDENCE + return _stabilize_localizations( + detections, + provenance=provenance, + weak_floor=weak_floor, + strong_floor=_SORA_STRONG_CONFIDENCE, + transition_floor=0.45, + min_stable_frames=_MIN_STABLE_FRAMES, + cover_after_confirmation=False, + ) + + +def stabilize_veo_localizations( + detections: tuple[FrameLocalization, ...] | list[FrameLocalization], + *, + provenance: bool, +) -> list[Region | None]: + """Accept temporally recurring current or legacy Veo candidates.""" + weak_floor = _VEO_PROVENANCE_WEAK_CONFIDENCE if provenance else _VEO_STRICT_WEAK_CONFIDENCE + return _stabilize_localizations( + detections, + provenance=provenance, + weak_floor=weak_floor, + strong_floor=_VEO_STRONG_CONFIDENCE, + transition_floor=0.35, + min_stable_frames=_MIN_VEO_STABLE_FRAMES, + cover_after_confirmation=True, + ) + + +def _stabilize_localizations( + detections: tuple[FrameLocalization, ...] | list[FrameLocalization], + *, + provenance: bool, + weak_floor: float, + strong_floor: float, + transition_floor: float, + min_stable_frames: int, + cover_after_confirmation: bool, +) -> list[Region | None]: + """Apply the shared recurrence policy to provider-specific candidates.""" + accepted: list[Region | None] = [None] * len(detections) + runs: list[list[int]] = [] + current: list[int] = [] + + for position, detection in enumerate(detections): + if detection.region is None or detection.confidence < weak_floor: + continue + if current: + previous = detections[current[-1]] + frame_gap = detection.frame_index - previous.frame_index + if ( + previous.region is None + or frame_gap > _MAX_STABLE_GAP + 1 + or _region_iou(previous.region, detection.region) < _STABLE_IOU + ): + runs.append(current) + current = [] + current.append(position) + if current: + runs.append(current) + + for run in runs: + strong = max(detections[position].confidence for position in run) >= strong_floor + if len(run) < min_stable_frames or (not provenance and not strong): + continue + for position in run: + accepted[position] = detections[position].region + for left_position, right_position in pairwise(run): + if right_position - left_position <= 1: + continue + left = detections[left_position] + right = detections[right_position] + if left.region is None or right.region is None or _region_iou(left.region, right.region) < _STABLE_IOU: + continue + for missing_position in range(left_position + 1, right_position): + distance_left = missing_position - left_position + distance_right = right_position - missing_position + accepted[missing_position] = left.region if distance_left <= distance_right else right.region + + # Provider provenance plus a confirmed run establishes a continuously + # watermarked app export rather than a clean API export that merely shares + # the generator name. Veo may also cover the sequence without provenance + # after its longer, strong fixed-position run. Cover low-contrast transition + # frames with the nearest confirmed position. + confirmed_positions = [position for position, region in enumerate(accepted) if region is not None] + if (provenance or cover_after_confirmation) and confirmed_positions: + confirmed_regions = [accepted[position] for position in confirmed_positions] + carry_position = confirmed_positions[0] + for position, region in enumerate(accepted): + if region is not None: + carry_position = position + continue + raw = detections[position] + if ( + raw.region is not None + and raw.confidence >= transition_floor + and any( + confirmed_region is not None and _region_iou(raw.region, confirmed_region) >= _STABLE_IOU + for confirmed_region in confirmed_regions + ) + ): + accepted[position] = raw.region + continue + accepted[position] = accepted[carry_position] + + return accepted + + +def _scan_video( + source: Path, + detector: Any, +) -> VideoScan: + """Decode a video once and collect one untrusted candidate per frame.""" + capture = cv2.VideoCapture(str(source)) + if not capture.isOpened(): + raise RuntimeError(f"OpenCV could not decode video: {source}") + width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = float(capture.get(cv2.CAP_PROP_FPS)) + if width <= 0 or height <= 0 or fps <= 0: + capture.release() + raise RuntimeError(f"Video has invalid stream geometry or frame rate: {source}") + + detections: list[FrameLocalization] = [] + frame_index = 0 + while True: + ok, frame = capture.read() + if not ok: + break + if frame.shape[:2] != (height, width): + capture.release() + raise RuntimeError(f"Video changes frame dimensions at frame {frame_index}: {source}") + detections.append(detector(frame, frame_index=frame_index)) + frame_index += 1 + capture.release() + if not detections: + raise RuntimeError(f"Video contains no decodable frames: {source}") + return VideoScan(width, height, fps, tuple(detections)) + + +def scan_sora_video(source: Path) -> VideoScan: + """Decode a video once and collect one untrusted Sora candidate per frame.""" + return _scan_video(source, detect_sora_frame) + + +def scan_veo_video(source: Path) -> VideoScan: + """Decode a video once and collect one untrusted Veo candidate per frame.""" + return _scan_video(source, detect_veo_frame) + + +def _ffmpeg_video_args(suffix: str) -> list[str]: + if suffix == ".webm": + return ["-c:v", "libvpx-vp9", "-crf", "18", "-b:v", "0"] + return ["-c:v", "libx264", "-preset", "medium", "-crf", "14"] + + +def _mask_for_region( + frame_bgr: NDArray[Any], + region: Region, + *, + padding_fraction: float, + mask_style: Literal["box", "veo"], +) -> NDArray[Any]: + height, width = frame_bgr.shape[:2] + x, y, region_width, region_height = region + mask = np.zeros((height, width), dtype=np.uint8) + if mask_style == "veo" and 0.80 <= region_width / region_height <= 1.25: + diamond_base, _ = _veo_templates() + diamond = cv2.resize( + diamond_base, + (region_width, region_height), + interpolation=cv2.INTER_AREA, + ) + diamond = np.where(diamond >= 24, 255, 0).astype(np.uint8) + dilation = max(2, round(region_height * 0.08)) + kernel_size = dilation * 2 + 1 + diamond = cv2.dilate( + diamond, + cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)), + ) + x1 = min(width, x + region_width) + y1 = min(height, y + region_height) + mask[y:y1, x:x1] = diamond[: y1 - y, : x1 - x] + return mask + + # A glyph-shaped mask leaves a thin translucent rim outside the approximate + # synthetic silhouette. Classical inpainting then pulls that white rim back + # into the hole, recreating the mascot as a bright blob. The measured clean + # floor on real Sora frames is a full box with roughly 0.28 mark-heights of + # context on every side. + padding = max(4, round(region_height * padding_fraction)) + x0 = max(0, x - padding) + y0 = max(0, y - padding) + x1 = min(width, x + region_width + padding) + y1 = min(height, y + region_height + padding) + mask[y0:y1, x0:x1] = 255 + return mask + + +def encode_clean_video( + source: Path, + output: Path, + scan: VideoScan, + regions: list[Region | None], + *, + backend: Backend, + strip_metadata: bool, + padding_fraction: float = 0.28, + mask_style: Literal["box", "veo"] = "box", +) -> int: + """Decode again, fill accepted regions, and encode video while copying audio.""" + from remove_ai_watermarks.watermark_registry import fill, resolve_backend + + ffmpeg = shutil.which("ffmpeg") + if ffmpeg is None: + raise RuntimeError("Visible video removal requires ffmpeg on PATH") + if len(regions) != len(scan.detections): + raise ValueError("Temporal localization count does not match the scanned frame count") + + output.parent.mkdir(parents=True, exist_ok=True) + command = [ + ffmpeg, + "-y", + "-loglevel", + "error", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s:v", + f"{scan.width}x{scan.height}", + "-r", + f"{scan.fps:.12g}", + "-i", + "pipe:0", + "-i", + str(source), + "-map", + "0:v:0", + "-map", + "1:a?", + *_ffmpeg_video_args(output.suffix.lower()), + "-c:a", + "copy", + "-map_metadata", + "-1" if strip_metadata else "1", + "-map_chapters", + "-1" if strip_metadata else "1", + "-shortest", + ] + if output.suffix.lower() in {".mp4", ".mov", ".m4v"}: + command.extend(["-movflags", "+faststart"]) + command.append(str(output)) + log.info("Encoding visible-watermark removal with ffmpeg: command=%s", command) + + process = subprocess.Popen( # noqa: S603 + command, + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + ) + if process.stdin is None or process.stderr is None: + process.kill() + raise RuntimeError("Could not open ffmpeg pipes") + + capture = cv2.VideoCapture(str(source)) + if not capture.isOpened(): + process.kill() + raise RuntimeError(f"OpenCV could not reopen video for removal: {source}") + + removed_frames = 0 + resolved_backend: Literal["cv2", "migan", "lama"] = resolve_backend(backend) + try: + for frame_index, region in enumerate(regions): + ok, frame = capture.read() + if not ok: + raise RuntimeError(f"Video ended while reading frame {frame_index}: {source}") + if region is not None: + frame = fill( + frame, + _mask_for_region( + frame, + region, + padding_fraction=padding_fraction, + mask_style=mask_style, + ), + backend=resolved_backend, + ) + removed_frames += 1 + process.stdin.write(frame.tobytes()) + process.stdin.close() + stderr = process.stderr.read().decode("utf-8", errors="replace") + return_code = process.wait() + except Exception: + process.kill() + process.wait() + raise + finally: + capture.release() + + log.info("ffmpeg visible-watermark encode finished: status=%s stderr=%s", return_code, stderr) + if return_code != 0: + raise RuntimeError(f"ffmpeg failed to encode {output}: {stderr.strip()[:500]}") + return removed_frames + + +def has_sora_provenance(markers: dict[str, str]) -> bool: + """Whether container provenance specifically names the Sora generator.""" + return "sora" in markers.get("claim_generator", "").lower() + + +def has_veo_provenance(markers: dict[str, str]) -> bool: + """Whether container provenance names Google as the AI-video generator.""" + identity = " ".join( + ( + markers.get("claim_generator", ""), + markers.get("issuer", ""), + ) + ).lower() + return "google" in identity and "trainedalgorithmicmedia" in markers.get("source_type", "").lower() diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 85555ba..6cd268f 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -1103,7 +1103,7 @@ class TestAIGCLabel: """Tencent Cloud's AIGC variant uses a service-provider schema (``ServiceProvider`` / ``ServiceUser`` / ``Time`` / ``ContentId``) rather than the producer schema, embedded in EXIF ``ImageDescription`` -- none of - its fields overlap the producer-side ``_TC260_FIELDS``, so the generic + its fields overlap the producer-side ``TC260_AIGC_FIELDS``, so the generic ``"AIGC":{...}`` gate must recognize them too.""" import json diff --git a/tests/test_video.py b/tests/test_video.py new file mode 100644 index 0000000..ebde4be --- /dev/null +++ b/tests/test_video.py @@ -0,0 +1,701 @@ +"""Tests for the video processing API and CLI.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import cv2 +import numpy as np +import pytest +from click.testing import CliRunner +from PIL import Image, ImageDraw, ImageFont + +from remove_ai_watermarks.cli import main +from remove_ai_watermarks.metadata import C2PA_UUID + +if TYPE_CHECKING: + from pathlib import Path + + +_MP4_FTYP = b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00mp42isom" +_VIDEO_PAYLOAD = b"synthetic-video-payload" +_TC260_AIGC = ( + b'{"Label":"1","ContentProducer":"00119144030008867405X210002",' + b'"ProduceID":"sample-001","ReservedCode1":"","ContentPropagator":"",' + b'"PropagateID":"","ReservedCode2":""}' +) + + +def _box(box_type: bytes, payload: bytes) -> bytes: + return (8 + len(payload)).to_bytes(4, "big") + box_type + payload + + +def _video_with_c2pa(path: Path) -> Path: + manifest = C2PA_UUID + b"OpenAI trainedAlgorithmicMedia" + path.write_bytes(_MP4_FTYP + _box(b"uuid", manifest) + _box(b"mdat", _VIDEO_PAYLOAD)) + return path + + +def _metadata_key(name: bytes) -> bytes: + return (8 + len(name)).to_bytes(4, "big") + b"mdta" + name + + +def _metadata_value(index: int, value: bytes) -> bytes: + data = _box(b"data", b"\x00\x00\x00\x01\x00\x00\x00\x00" + value) + return _box(index.to_bytes(4, "big"), data) + + +def _video_with_tc260(path: Path, *, media_payload: bytes = _VIDEO_PAYLOAD) -> Path: + keys = _box( + b"keys", + b"\x00\x00\x00\x00" + (2).to_bytes(4, "big") + _metadata_key(b"AIGC") + _metadata_key(b"title"), + ) + ilst = _box( + b"ilst", + _metadata_value(1, _TC260_AIGC) + _metadata_value(2, b"standard title"), + ) + meta = _box(b"meta", b"\x00\x00\x00\x00" + keys + ilst) + path.write_bytes(_MP4_FTYP + _box(b"mdat", media_payload) + _box(b"moov", _box(b"udta", meta))) + return path + + +def _ebml_size(value: int) -> bytes: + for length in range(1, 9): + if value < (1 << (7 * length)) - 1: + return ((1 << (7 * length)) | value).to_bytes(length, "big") + raise ValueError("EBML test value is too large") + + +def _ebml_element(element_id: bytes, payload: bytes) -> bytes: + return element_id + _ebml_size(len(payload)) + payload + + +def _video_with_tc260_ebml(path: Path, *, value: bytes = _TC260_AIGC) -> Path: + simple_tag = _ebml_element( + b"\x67\xc8", + _ebml_element(b"\x45\xa3", b"AIGC") + _ebml_element(b"\x44\x87", value), + ) + tags = _ebml_element(b"\x12\x54\xc3\x67", _ebml_element(b"\x73\x73", simple_tag)) + segment = _ebml_element(b"\x18\x53\x80\x67", tags) + path.write_bytes(_ebml_element(b"\x1a\x45\xdf\xa3", b"") + segment) + return path + + +class TestVideoMetadataApi: + def test_top_level_api_is_lazy_exported(self): + import remove_ai_watermarks as raiw + + assert raiw.inspect_video_metadata is not None + assert raiw.remove_video_metadata is not None + assert raiw.remove_video_visible is not None + + def test_inspects_video_metadata(self, tmp_path: Path): + from remove_ai_watermarks.video import inspect_video_metadata + + source = _video_with_c2pa(tmp_path / "source.mp4") + + report = inspect_video_metadata(source) + + assert report.source == source + assert report.has_ai_metadata is True + assert report.markers + + def test_removes_metadata_without_touching_video_payload(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_metadata + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + + result = remove_video_metadata(source, output) + + assert result.output == output + assert result.detected + assert result.remaining == {} + assert _VIDEO_PAYLOAD in output.read_bytes() + assert C2PA_UUID not in output.read_bytes() + + def test_default_output_preserves_source(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_metadata + + source = _video_with_c2pa(tmp_path / "source.mp4") + original = source.read_bytes() + + result = remove_video_metadata(source) + + assert result.output == tmp_path / "source_clean.mp4" + assert result.output.exists() + assert source.read_bytes() == original + + @pytest.mark.parametrize("suffix", [".mp4", ".mov"]) + def test_inspects_native_tc260_metadata(self, tmp_path: Path, suffix: str): + from remove_ai_watermarks.video import inspect_video_metadata + + source = _video_with_tc260(tmp_path / f"source{suffix}") + + report = inspect_video_metadata(source) + + assert report.has_ai_metadata is True + assert report.markers["aigc_label"].endswith("producer 00119144030008867405X210002") + + def test_inspects_native_tc260_metadata_after_large_media_payload(self, tmp_path: Path): + from remove_ai_watermarks.video import inspect_video_metadata + + source = _video_with_tc260( + tmp_path / "source.mp4", + media_payload=b"x" * (1024 * 1024), + ) + + report = inspect_video_metadata(source) + + assert report.has_ai_metadata is True + assert "aigc_label" in report.markers + + def test_removes_native_tc260_metadata_without_touching_media_or_standard_tag(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_metadata + + source = _video_with_tc260(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + + result = remove_video_metadata(source, output) + cleaned = output.read_bytes() + + assert result.detected["aigc_label"].startswith("China AIGC label") + assert result.remaining == {} + assert len(cleaned) == source.stat().st_size + assert _VIDEO_PAYLOAD in cleaned + assert b"standard title" in cleaned + assert b"AIGC" not in cleaned + assert _TC260_AIGC not in cleaned + + def test_ignores_generic_mp4_aigc_tag_without_tc260_fields(self, tmp_path: Path): + from remove_ai_watermarks.video import inspect_video_metadata + + source = _video_with_tc260(tmp_path / "source.mp4") + source.write_bytes(source.read_bytes().replace(_TC260_AIGC, b'{"description":"' + b"x" * 146 + b'"}')) + + report = inspect_video_metadata(source) + + assert report.has_ai_metadata is False + assert report.markers == {} + + @pytest.mark.parametrize("suffix", [".mkv", ".webm"]) + def test_inspects_native_tc260_ebml_metadata(self, tmp_path: Path, suffix: str): + from remove_ai_watermarks.video import inspect_video_metadata + + source = _video_with_tc260_ebml(tmp_path / f"source{suffix}") + + report = inspect_video_metadata(source) + + assert report.has_ai_metadata is True + assert report.markers["aigc_label"].endswith("producer 00119144030008867405X210002") + + def test_ignores_generic_ebml_aigc_tag_without_tc260_fields(self, tmp_path: Path): + from remove_ai_watermarks.video import inspect_video_metadata + + source = _video_with_tc260_ebml( + tmp_path / "source.mkv", + value=b'{"description":"ordinary application metadata"}', + ) + + report = inspect_video_metadata(source) + + assert report.has_ai_metadata is False + assert report.markers == {} + + def test_rejects_image_input(self, tmp_clean_png: Path): + from remove_ai_watermarks.video import inspect_video_metadata + + with pytest.raises(ValueError, match="Unsupported video format"): + inspect_video_metadata(tmp_clean_png) + + def test_rejects_image_with_video_extension(self, tmp_clean_png: Path, tmp_path: Path): + from remove_ai_watermarks.video import inspect_video_metadata + + disguised = tmp_path / "image.mp4" + disguised.write_bytes(tmp_clean_png.read_bytes()) + + with pytest.raises(ValueError, match="does not match"): + inspect_video_metadata(disguised) + + def test_rejects_output_container_change(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_metadata + + source = _video_with_c2pa(tmp_path / "source.mp4") + + with pytest.raises(ValueError, match="must match"): + remove_video_metadata(source, tmp_path / "clean.mov") + + +class TestVideoMetadataCli: + def test_help(self): + runner = CliRunner() + + result = runner.invoke(main, ["video", "metadata", "--help"]) + + assert result.exit_code == 0, result.output + assert "AI metadata" in result.output + + def test_check_reports_metadata(self, tmp_path: Path): + runner = CliRunner() + source = _video_with_c2pa(tmp_path / "source.mp4") + + result = runner.invoke(main, ["video", "metadata", str(source), "--check"]) + + assert result.exit_code == 0, result.output + assert "AI metadata detected" in result.output + + def test_remove_reports_output(self, tmp_path: Path): + runner = CliRunner() + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + + result = runner.invoke(main, ["video", "metadata", str(source), "--remove", "-o", str(output)]) + + assert result.exit_code == 0, result.output + assert "AI metadata stripped" in result.output + assert C2PA_UUID not in output.read_bytes() + + def test_rejects_image_input(self, tmp_clean_png: Path): + runner = CliRunner() + + result = runner.invoke(main, ["video", "metadata", str(tmp_clean_png), "--check"]) + + assert result.exit_code != 0 + assert "Unsupported video format" in result.output + + +class TestSoraFrameLocalization: + @staticmethod + def _sora_like_frame() -> tuple[np.ndarray, tuple[int, int, int, int]]: + frame = np.full((480, 840, 3), 36, dtype=np.uint8) + mark = Image.new("L", (180, 64), 0) + draw = ImageDraw.Draw(mark) + draw.ellipse((1, 14, 32, 54), fill=255) + draw.ellipse((25, 8, 62, 58), fill=255) + draw.ellipse((15, 20, 28, 44), fill=0) + draw.ellipse((37, 18, 50, 43), fill=0) + try: + font = ImageFont.load_default(size=49) + except TypeError: + font = ImageFont.load_default() + draw.text((68, 1), "Sora", font=font, fill=255, stroke_width=1) + mark_array = cv2.resize(np.asarray(mark), (124, 44), interpolation=cv2.INTER_AREA) + x, y = 620, 398 + alpha = mark_array.astype(np.float32)[:, :, None] / 255 * 0.78 + crop = frame[y : y + 44, x : x + 124].astype(np.float32) + frame[y : y + 44, x : x + 124] = np.clip(crop * (1 - alpha) + 255 * alpha, 0, 255).astype(np.uint8) + return frame, (x, y, 124, 44) + + def test_localizes_independently_rendered_sora_like_mark(self): + from remove_ai_watermarks.video_visible import _region_iou, detect_sora_frame + + frame, expected = self._sora_like_frame() + + detection = detect_sora_frame(frame) + + assert detection.region is not None + assert detection.confidence >= 0.58 + assert _region_iou(detection.region, expected) >= 0.45 + + def test_empty_frame_is_not_localized(self): + from remove_ai_watermarks.video_visible import detect_sora_frame + + detection = detect_sora_frame(np.empty((0, 0, 3), dtype=np.uint8)) + + assert detection.confidence == 0.0 + assert detection.region is None + + +class TestVeoFrameLocalization: + def test_localizes_independently_rendered_diamond_at_relocated_position(self): + from remove_ai_watermarks.video_visible import _region_iou, detect_veo_frame + + frame = np.full((720, 1280, 3), 28, dtype=np.uint8) + size = 48 + x, y = 1080, 570 + mark = Image.new("L", (size, size), 0) + points = ( + (size // 2, 1), + (round(size * 0.61), round(size * 0.38)), + (size - 2, size // 2), + (round(size * 0.61), round(size * 0.62)), + (size // 2, size - 2), + (round(size * 0.39), round(size * 0.62)), + (1, size // 2), + (round(size * 0.39), round(size * 0.38)), + ) + ImageDraw.Draw(mark).polygon(points, fill=255) + alpha = np.asarray(mark, dtype=np.float32)[:, :, None] / 255 * 0.72 + crop = frame[y : y + size, x : x + size].astype(np.float32) + frame[y : y + size, x : x + size] = np.clip( + crop * (1 - alpha) + 255 * alpha, + 0, + 255, + ).astype(np.uint8) + + detection = detect_veo_frame(frame) + + assert detection.region is not None + assert detection.confidence >= 0.70 + assert _region_iou(detection.region, (x, y, size, size)) >= 0.70 + + def test_localizes_independently_rendered_legacy_text(self): + from remove_ai_watermarks.video_visible import _region_iou, detect_veo_frame + + frame = np.full((720, 1280, 3), 42, dtype=np.uint8) + mark = Image.new("L", (60, 24), 0) + try: + font = ImageFont.load_default(size=19) + except TypeError: + font = ImageFont.load_default() + ImageDraw.Draw(mark).text((1, 0), "Veo", font=font, fill=255) + mark_array = np.asarray(mark) + ys, xs = np.where(mark_array > 0) + mark_array = mark_array[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1] + mark_height, mark_width = mark_array.shape + x = frame.shape[1] - mark_width - 20 + y = frame.shape[0] - mark_height - 18 + alpha = mark_array.astype(np.float32)[:, :, None] / 255 * 0.66 + crop = frame[y : y + mark_height, x : x + mark_width].astype(np.float32) + frame[y : y + mark_height, x : x + mark_width] = np.clip( + crop * (1 - alpha) + 255 * alpha, + 0, + 255, + ).astype(np.uint8) + + detection = detect_veo_frame(frame) + + assert detection.region is not None + assert detection.confidence >= 0.55 + assert _region_iou(detection.region, (x, y, mark_width, mark_height)) >= 0.65 + + def test_empty_frame_is_not_localized(self): + from remove_ai_watermarks.video_visible import detect_veo_frame + + detection = detect_veo_frame(np.empty((0, 0, 3), dtype=np.uint8)) + + assert detection.confidence == 0.0 + assert detection.region is None + + def test_diamond_mask_preserves_transparent_box_corners(self): + from remove_ai_watermarks.video_visible import _mask_for_region + + mask = _mask_for_region( + np.zeros((100, 100, 3), dtype=np.uint8), + (20, 20, 48, 48), + padding_fraction=0.18, + mask_style="veo", + ) + + assert mask[44, 44] == 255 + assert mask[20, 20] == 0 + assert mask[67, 67] == 0 + + +class TestSoraTemporalArbiter: + _BOX = (40, 60, 150, 54) + + def test_four_frame_lookalike_run_is_too_short(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_sora_localizations + + detections = [FrameLocalization(index, 0.70, self._BOX) for index in range(4)] + + assert stabilize_sora_localizations(detections, provenance=False) == [None] * 4 + + def test_provenance_accepts_recurring_low_contrast_visual_match(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_sora_localizations + + detections = [ + FrameLocalization(0, 0.59, self._BOX), + FrameLocalization(1, 0.61, self._BOX), + FrameLocalization(2, 0.62, self._BOX), + FrameLocalization(3, 0.60, self._BOX), + FrameLocalization(4, 0.61, self._BOX), + ] + + assert stabilize_sora_localizations(detections, provenance=True) == [self._BOX] * 5 + + def test_confirmed_provenance_run_covers_transition_frames(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_sora_localizations + + detections = [ + FrameLocalization(0, 0.30, (500, 300, 54, 54)), + FrameLocalization(1, 0.59, self._BOX), + FrameLocalization(2, 0.61, self._BOX), + FrameLocalization(3, 0.62, self._BOX), + FrameLocalization(4, 0.60, self._BOX), + FrameLocalization(5, 0.61, self._BOX), + FrameLocalization(6, 0.30, (300, 100, 54, 54)), + ] + + assert stabilize_sora_localizations(detections, provenance=True) == [self._BOX] * 7 + + def test_transition_prefers_low_score_match_at_a_confirmed_position(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_sora_localizations + + other_box = (500, 300, 150, 54) + detections = [ + FrameLocalization(0, 0.61, self._BOX), + FrameLocalization(1, 0.62, self._BOX), + FrameLocalization(2, 0.63, self._BOX), + FrameLocalization(3, 0.61, self._BOX), + FrameLocalization(4, 0.62, self._BOX), + FrameLocalization(5, 0.20, (250, 180, 54, 54)), + FrameLocalization(6, 0.52, self._BOX), + FrameLocalization(7, 0.61, other_box), + FrameLocalization(8, 0.62, other_box), + FrameLocalization(9, 0.63, other_box), + FrameLocalization(10, 0.61, other_box), + FrameLocalization(11, 0.62, other_box), + ] + + stabilized = stabilize_sora_localizations(detections, provenance=True) + + assert stabilized[6] == self._BOX + assert stabilized[7:] == [other_box] * 5 + + def test_transition_without_a_match_keeps_previous_stable_position(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_sora_localizations + + other_box = (500, 300, 150, 54) + detections = [ + FrameLocalization(0, 0.61, self._BOX), + FrameLocalization(1, 0.62, self._BOX), + FrameLocalization(2, 0.63, self._BOX), + FrameLocalization(3, 0.61, self._BOX), + FrameLocalization(4, 0.62, self._BOX), + FrameLocalization(5, 0.20, (250, 180, 54, 54)), + FrameLocalization(6, 0.20, (300, 200, 54, 54)), + FrameLocalization(7, 0.61, other_box), + FrameLocalization(8, 0.62, other_box), + FrameLocalization(9, 0.63, other_box), + FrameLocalization(10, 0.61, other_box), + FrameLocalization(11, 0.62, other_box), + ] + + stabilized = stabilize_sora_localizations(detections, provenance=True) + + assert stabilized[5:7] == [self._BOX, self._BOX] + + def test_unproven_weak_run_is_rejected(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_sora_localizations + + detections = [ + FrameLocalization(0, 0.61, self._BOX), + FrameLocalization(1, 0.62, self._BOX), + FrameLocalization(2, 0.63, self._BOX), + FrameLocalization(3, 0.62, self._BOX), + FrameLocalization(4, 0.61, self._BOX), + ] + + assert stabilize_sora_localizations(detections, provenance=False) == [None] * 5 + + def test_strong_recurring_visual_run_needs_no_metadata(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_sora_localizations + + detections = [ + FrameLocalization(0, 0.61, self._BOX), + FrameLocalization(1, 0.66, self._BOX), + FrameLocalization(2, 0.62, self._BOX), + FrameLocalization(3, 0.61, self._BOX), + FrameLocalization(4, 0.62, self._BOX), + ] + + assert stabilize_sora_localizations(detections, provenance=False) == [self._BOX] * 5 + + def test_isolated_lookalikes_at_different_positions_are_rejected(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_sora_localizations + + detections = [ + FrameLocalization(0, 0.70, (10, 10, 150, 54)), + FrameLocalization(1, 0.70, (400, 200, 150, 54)), + FrameLocalization(2, 0.70, (650, 400, 150, 54)), + ] + + assert stabilize_sora_localizations(detections, provenance=True) == [None, None, None] + + def test_short_dropout_between_matching_boxes_is_filled(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_sora_localizations + + detections = [ + FrameLocalization(0, 0.66, self._BOX), + FrameLocalization(1, 0.20, (500, 300, 54, 54)), + FrameLocalization(2, 0.67, self._BOX), + FrameLocalization(3, 0.66, self._BOX), + FrameLocalization(4, 0.66, self._BOX), + FrameLocalization(5, 0.66, self._BOX), + ] + + assert stabilize_sora_localizations(detections, provenance=False) == [self._BOX] * 6 + + +class TestVeoTemporalArbiter: + _BOX = (1132, 572, 56, 56) + + def test_eleven_frame_lookalike_run_is_too_short(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_veo_localizations + + detections = [FrameLocalization(index, 0.70, self._BOX) for index in range(11)] + + assert stabilize_veo_localizations(detections, provenance=False) == [None] * 11 + + def test_strong_fixed_run_covers_video_without_metadata(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_veo_localizations + + detections = [FrameLocalization(index, 0.60, self._BOX) for index in range(12)] + detections.extend(FrameLocalization(index, 0.20, (300, 200, 48, 48)) for index in range(12, 15)) + + assert stabilize_veo_localizations(detections, provenance=False) == [self._BOX] * 15 + + def test_google_provenance_accepts_recurring_low_contrast_diamond(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_veo_localizations + + detections = [FrameLocalization(index, 0.47, self._BOX) for index in range(12)] + + assert stabilize_veo_localizations(detections, provenance=True) == [self._BOX] * 12 + + def test_unproven_weak_run_is_rejected(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_veo_localizations + + detections = [FrameLocalization(index, 0.52, self._BOX) for index in range(12)] + + assert stabilize_veo_localizations(detections, provenance=False) == [None] * 12 + + +class TestVideoVisibleApi: + def test_removes_stable_sora_run_and_writes_output(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video import remove_video_visible + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + box = (4, 4, 20, 8) + scan = VideoScan( + width=64, + height=64, + fps=24.0, + detections=tuple(FrameLocalization(index, 0.66, box) for index in range(5)), + ) + monkeypatch.setattr(video_visible, "scan_sora_video", lambda _source: scan) + + def fake_encode( + _source: Path, + target: Path, + _scan: VideoScan, + regions: list[tuple[int, int, int, int] | None], + **_kwargs: object, + ) -> int: + assert regions == [box] * 5 + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return 5 + + monkeypatch.setattr(video_visible, "encode_clean_video", fake_encode) + + result = remove_video_visible(source, output) + + assert result.output == output + assert result.detected_frames == 5 + assert result.removed_frames == 5 + assert result.remaining_metadata == {} + + def test_no_stable_mark_writes_no_output(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video import remove_video_visible + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + scan = VideoScan( + width=64, + height=64, + fps=24.0, + detections=( + FrameLocalization(0, 0.70, (1, 1, 20, 8)), + FrameLocalization(1, 0.70, (30, 30, 20, 8)), + FrameLocalization(2, 0.70, (1, 30, 20, 8)), + ), + ) + monkeypatch.setattr(video_visible, "scan_sora_video", lambda _source: scan) + + result = remove_video_visible(source, output) + + assert result.output is None + assert result.removed_frames == 0 + assert not output.exists() + + def test_dispatches_veo_detector_and_uses_tighter_mask( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video import remove_video_visible + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + box = (4, 4, 20, 20) + scan = VideoScan( + width=64, + height=64, + fps=24.0, + detections=tuple(FrameLocalization(index, 0.60, box) for index in range(12)), + ) + monkeypatch.setattr(video_visible, "scan_veo_video", lambda _source: scan) + + def fake_encode( + _source: Path, + target: Path, + _scan: VideoScan, + regions: list[tuple[int, int, int, int] | None], + **kwargs: object, + ) -> int: + assert regions == [box] * 12 + assert kwargs["padding_fraction"] == 0.18 + assert kwargs["mask_style"] == "veo" + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return 12 + + monkeypatch.setattr(video_visible, "encode_clean_video", fake_encode) + + result = remove_video_visible(source, output, mark="veo") + + assert result.output == output + assert result.mark == "veo" + assert result.detected_frames == 12 + assert result.removed_frames == 12 + + +class TestVideoVisibleCli: + def test_help(self): + result = CliRunner().invoke(main, ["video", "visible", "--help"]) + + assert result.exit_code == 0, result.output + assert "temporally stable" in result.output + assert "sora|veo" in result.output + + def test_reports_removed_frames(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + from remove_ai_watermarks import video + from remove_ai_watermarks.video import VideoVisibleResult + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + monkeypatch.setattr( + video, + "remove_video_visible", + lambda *_args, **_kwargs: VideoVisibleResult( + source=source, + output=output, + mark="sora", + total_frames=12, + detected_frames=10, + removed_frames=10, + remaining_metadata={}, + ), + ) + + result = CliRunner().invoke(main, ["video", "visible", str(source), "-o", str(output)]) + + assert result.exit_code == 0, result.output + assert "10/12 frames" in result.output From 41b6fd06821b68c2cfc1b24c9b909225ac9d052a Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Wed, 29 Jul 2026 17:56:52 -0700 Subject: [PATCH 02/19] Add Seedance and Dola video removal --- CLAUDE.md | 4 +- README.md | 27 ++- docs/cli.md | 15 +- docs/known-limitations.md | 22 ++- docs/module-internals.md | 40 ++-- docs/python-api.md | 20 +- docs/supported-signals.md | 2 + docs/watermarking-landscape.md | 15 +- src/remove_ai_watermarks/__init__.py | 2 +- src/remove_ai_watermarks/cli.py | 2 +- src/remove_ai_watermarks/video.py | 41 ++++- src/remove_ai_watermarks/video_visible.py | 215 +++++++++++++++++++++- tests/test_video.py | 191 ++++++++++++++++++- 13 files changed, 529 insertions(+), 67 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 01b7bea..f8cc4dd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -28,7 +28,7 @@ Per-command exit-code semantics (the no-signal / GPU-missing skip branches), tes - `uv run remove-ai-watermarks metadata --remove -o ` — strip all AI metadata - `uv run remove-ai-watermarks video metadata --check` — inspect AI metadata in MP4/MOV/M4V/WebM/MKV - `uv run remove-ai-watermarks video metadata --remove -o ` — strip verified video metadata without transcoding streams; the experimental video path requires a separate same-container output and defaults to `_clean` -- `uv run remove-ai-watermarks video visible -o ` — remove a temporally recurring Sora mark, or pass `--mark veo` for the current Veo diamond and legacy `Veo` text. It scans the full sequence first, transcodes video through ffmpeg, copies audio, strips AI metadata by default, and writes no output when no stable mark is found. `cv2` is the fast default; `migan`/`lama` improve difficult backgrounds. +- `uv run remove-ai-watermarks video visible -o ` — remove a temporally recurring Sora mark, or pass `--mark veo|seedance|dola` for the current Veo diamond and legacy `Veo` text, the Seedance boxed `AI`, or `Dola AI`. It scans the full sequence first, transcodes video through ffmpeg, copies audio, strips AI metadata by default, and writes no output when no stable mark is found. `cv2` is the fast default; `migan`/`lama` improve difficult backgrounds. - `uv run remove-ai-watermarks batch ` — process every supported image in a directory (output defaults to `_clean/`, set with `-o`). `--mode visible|invisible|metadata|all` (default `visible`); the invisible/all path reuses the full `invisible` knob set above, plus `--backend` and `--sensitivity` for the visible localize -> fill pass. Applies the same no-signal skip per image; see the module doc. **Exit code:** non-zero when any image errored OR (mirroring single `all`) a `--mode invisible`/`all` image carried an invisible signal but the GPU extra was absent, so its SynthID scrub was skipped — it emits a loud warning and copies the input through (invisible mode) so the output dir stays complete; a wrapping service can then detect the incomplete run instead of trusting a silent exit 0. ## Test and lint @@ -82,7 +82,7 @@ Compact map. The full per-module detail (design decisions, tuned thresholds, cal - `upscaler.py` — optional Real-ESRGAN pre-diffusion super-resolution for small inputs (extra `esrgan`, spandrel only). Manual opt-in; the default `--upscaler` stays `lanczos` and the engine always falls back to Lanczos on absence/error. ESRGAN can degrade faces and thin text. - `image_io.py` — centralizes Unicode-safe image IO, alpha preservation, content-based format sniffing, and HEIC/AVIF fallbacks. Callers must check `imwrite` success. No-op visible removal preserves original bytes when the output format is unchanged. - `api.py` — the high-level convenience API, re-exported lazily at the package top level via `__init__.__getattr__` (PEP 562, so `import remove_ai_watermarks` stays cheap): `remove_visible(source, output=None, *, sensitivity="auto", backend="auto", strip_metadata=True, write_noop=True) -> (result_bgr, [labels])` (source = path OR BGR ndarray; a PATH auto-reads metadata provenance and preserves alpha, an ARRAY does neither; `write_noop=True` writes a clean passthrough copy when nothing is removed, `False` leaves `output` untouched so a "no mark = produce nothing" caller like the CLI `visible` command does not clobber a pre-existing file there) and `visible_provenance(path) -> frozenset[str]` (the single metadata→vendor-keys mapper; `cli._visible_provenance` is a thin None-guarded wrapper over it). **`remove_visible` is the ONE path the CLI and library share** — `cli.cmd_visible`'s `--mark auto` branch delegates entirely to it (read → provenance → `remove_auto_marks` → write → `strip_metadata`), so there is no CLI-vs-library drift; `strip_metadata` defaults True to match `visible --strip-metadata`. This is where a library caller should start — NOT the engines directly (`GeminiEngine`/`TextMarkEngine` have no `remove_watermark` any more; removal is registry `remove_auto_marks`/`KnownMark.remove`; the old single-strongest `best_auto_mark` is gone — removal takes EVERY mark). `identify` is NOT top-level re-exported (it collides with the `identify` submodule); use `from remove_ai_watermarks.identify import identify`. -- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, and `remove_video_visible(source, output=None, *, mark="sora", backend="cv2", strip_metadata=True) -> VideoVisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM TC260 is read by `noai/ebml.py` and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: fully synthetic Sora and Veo silhouettes propose frame boxes, provider-specific temporal recurrence authorizes them, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying audio. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels and invisible video watermarks are not built yet. +- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, and `remove_video_visible(source, output=None, *, mark="sora", backend="cv2", strip_metadata=True) -> VideoVisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM TC260 is read by `noai/ebml.py` and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: fully synthetic Sora, Veo, Seedance, and Dola silhouettes propose frame boxes, provider-specific temporal recurrence authorizes them, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying audio. Fixed marks require anchored runs so a slowly drifting scene detail cannot pass on adjacent overlap alone. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Seedance uses the full localized box because an outline mask left part of the real translucent border behind. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels and invisible video watermarks are not built yet. For the Doubao alpha-distillation history (why content-image reverse-alpha distillation fails by physics and controlled captures were required), see `docs/research-doubao-distillation.md`. diff --git a/README.md b/README.md index cde4b17..013830e 100644 --- a/README.md +++ b/README.md @@ -7,8 +7,8 @@ Remove AI provenance marks from images and video you generated yourself: - C2PA, EXIF, XMP, IPTC, and related AI metadata. Video support covers metadata inspection and removal plus experimental visible -Sora-wordmark removal. Invisible video-watermark removal remains a follow-up -stage. +Sora, Veo, Seedance, and Dola mark removal. Invisible video-watermark removal +remains a follow-up stage. > Try it online at [raiw.cc](https://raiw.cc) if you do not want to install Python > or run diffusion models locally. @@ -33,7 +33,7 @@ stage. | Erase a region you select | `erase` | No | | Strip AI metadata | `metadata` | No | | Strip AI metadata from video | `video metadata` | No | -| Remove a known Sora or Veo mark from video | `video visible` | No | +| Remove a known Sora, Veo, Seedance, or Dola mark from video | `video visible` | No | | Regenerate an image to disrupt invisible watermarks | `invisible` | Recommended | | Run visible, invisible, and metadata removal | `all` | Recommended | | Process a directory | `batch` | Depends on mode | @@ -78,11 +78,13 @@ inspection includes the native TC260 `AIGC` tag in MKV and WebM inspection reads the normative `Segment.Tags.Tag.SimpleTag` placement. -Remove a moving Sora wordmark or a Veo corner mark: +Remove a supported visible video mark: ```bash remove-ai-watermarks video visible input.mp4 -o clean.mp4 remove-ai-watermarks video visible veo.mp4 --mark veo -o veo_clean.mp4 +remove-ai-watermarks video visible seedance.mp4 --mark seedance -o seedance_clean.mp4 +remove-ai-watermarks video visible dola.mp4 --mark dola -o dola_clean.mp4 ``` This path scans the complete sequence before changing pixels. It accepts only a @@ -90,8 +92,9 @@ mark that repeats at a stable position across adjacent frames, then reuses the same OpenCV, MI-GAN, or LaMa fill backends as image removal. Audio is copied without re-encoding; the video stream is transcoded because its pixels change. Sora covers the moving Sora 2 mascot and wordmark. Veo covers both the current -four-point diamond and the legacy `Veo` text in the bottom-right corner. No -output is written when no stable mark is found. +four-point diamond and the legacy `Veo` text. Seedance covers the fixed boxed +`AI` label, and Dola covers the fixed `Dola AI` text. No output is written when +no stable mark is found. For invisible watermark removal, install the diffusion dependencies: @@ -243,6 +246,12 @@ report = raiw.inspect_video_metadata("input.mp4") cleaned = raiw.remove_video_metadata("input.mp4") visible = raiw.remove_video_visible("sora.mp4", "sora_clean.mp4") veo = raiw.remove_video_visible("veo.mp4", "veo_clean.mp4", mark="veo") +seedance = raiw.remove_video_visible( + "seedance.mp4", + "seedance_clean.mp4", + mark="seedance", +) +dola = raiw.remove_video_visible("dola.mp4", "dola_clean.mp4", mark="dola") ``` The high level API accepts a file path or a BGR NumPy array. For path inputs it @@ -267,9 +276,9 @@ invisible removal. and selected fill backend. - Invisible removal changes the whole image and may alter faces, text, or fine detail. -- Visible video removal recognizes the moving Sora 2 wordmark and the current - Veo diamond plus legacy `Veo` text. It does not recognize the older Sora Turbo - corner swirl. +- Visible video removal recognizes the moving Sora 2 wordmark, the current Veo + diamond plus legacy `Veo` text, the Seedance boxed `AI` label, and the fixed + `Dola AI` text. It does not recognize the older Sora Turbo corner swirl. The classical OpenCV backend can smear structured backgrounds; use MI-GAN or LaMa when recovery quality matters. Video SynthID is not removed yet. MP4/MOV/M4V metadata stripping currently reads the full container into diff --git a/docs/cli.md b/docs/cli.md index c8ba68a..293082d 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -153,18 +153,21 @@ different container extension. Visible video labels and invisible video watermarks are not handled by this command. -## Remove a visible Sora or Veo video mark +## Remove a supported visible video mark ```bash remove-ai-watermarks video visible input.mp4 -o clean.mp4 remove-ai-watermarks video visible veo.mp4 --mark veo -o veo_clean.mp4 +remove-ai-watermarks video visible seedance.mp4 --mark seedance -o seedance_clean.mp4 +remove-ai-watermarks video visible dola.mp4 --mark dola -o dola_clean.mp4 ``` -The experimental command supports the moving Sora mascot and wordmark plus two -Veo corner variants: the current four-point diamond and the legacy `Veo` text. -Sora searches the whole frame at multiple scales. Veo searches the -bottom-right corner using separate synthetic silhouettes for the two variants. -Both require a spatially recurring candidate across adjacent frames. Matching +The experimental command supports the moving Sora mascot and wordmark, two Veo +corner variants, the Seedance boxed `AI` label, and the `Dola AI` text label. +Sora searches the whole frame at multiple scales. The other detectors search +bounded bottom-right regions with separate synthetic silhouettes. Every mark +requires a spatially recurring candidate across adjacent frames. Fixed marks +must also remain anchored instead of drifting with a scene object. Matching provider provenance may relax the visual score, but metadata alone never creates a detection. Clean API exports therefore remain untouched. diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 35b4a93..627f7dc 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -174,16 +174,19 @@ WebM, Matroska, MP3, WAV, FLAC, OGG, Opus, and AAC container metadata is strippe through ffmpeg with stream copying. The operation fails if ffmpeg is absent or cannot parse the input. -### Visible video removal supports Sora and Veo and is still experimental +### Visible video removal is provider-specific and still experimental The experimental `video metadata` command and high level video API inspect and strip supported AI provenance metadata without transcoding streams. `video visible` and `remove_video_visible` additionally support the moving -Sora 2 mascot and wordmark, the current Veo four-point diamond, and the legacy -`Veo` text. Detection requires a recurring visual candidate across adjacent -frames. Provider provenance can recover low-contrast runs only after visual -evidence exists, so metadata alone does not erase a clean API export. +Sora 2 mascot and wordmark, the current Veo four-point diamond, the legacy +`Veo` text, the Seedance boxed `AI` label, and the fixed `Dola AI` text. +Detection requires a recurring visual candidate across adjacent frames. +Seedance, Dola, and Veo candidates must remain anchored rather than drifting +with a scene object. Provider provenance can recover low-contrast runs only +after visual evidence exists, so metadata alone does not erase a clean API +export. Historical Sora Turbo exports use a small OpenAI swirl in the corner rather than the moving mascot-and-wordmark design; that earlier variant is not detected by the `sora` video mark. Other provider video labels and proprietary @@ -194,9 +197,12 @@ fill is not a motion-aware video inpainting model. OpenCV can leave a visible smear where the mark overlaps a hard edge or structured texture, and the smear can vary over time. MI-GAN and LaMa improve individual frames but do not guarantee temporal coherence. The Veo diamond uses a shape mask to limit damage -outside the symbol, but OpenCV may still soften texture inside it. The current -encoder also emits a constant-frame-rate output at the decoded stream rate, so -variable-frame-rate preservation is not yet guaranteed. +outside the symbol. Seedance fills the full localized box because a synthetic +outline mask left part of the real translucent border visible in an end-to-end +check. OpenCV may therefore soften texture inside that small box; use MI-GAN or +LaMa when reconstruction quality matters. The current encoder also emits a +constant-frame-rate output at the decoded stream rate, so variable-frame-rate +preservation is not yet guaranteed. Native TC260 metadata in MP4/MOV is supported at its normative `moov.udta.meta.keys/ilst` placement, including non-faststart files whose diff --git a/docs/module-internals.md b/docs/module-internals.md index cb4dba6..8713b7f 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -115,32 +115,34 @@ only a `Segment.Tags.Tag.SimpleTag` pairing `TagName=AIGC` with a JSON those container tags without transcoding the encoded streams. [`video_visible.py`](../src/remove_ai_watermarks/video_visible.py) implements -the first pixel stages for Sora and Veo. The Sora detector searches a normalized -frame with a fully synthetic mascot-and-text silhouette at several scales. The -Veo detector uses separate synthetic silhouettes for the current four-point -diamond and legacy `Veo` text, with bounded bottom-right searches calibrated -independently from Sora. A strong relocated-diamond match may bypass the known +the first pixel stages for Sora, Veo, Seedance, and Dola. The Sora detector +searches a normalized frame with a fully synthetic mascot-and-text silhouette +at several scales. The Veo detector uses separate synthetic silhouettes for the +current four-point diamond and legacy `Veo` text. Seedance uses a synthetic +rounded boxed-`AI` silhouette, while Dola uses an OpenCV-font `Dola AI` +silhouette. All fixed-mark searches are bounded to the bottom-right area and +calibrated independently. A strong relocated Veo diamond may bypass the known layout anchors, but weak free-corner matches never enter the temporal arbiter. -This prevents recurring scene details in clean API exports from being promoted -to a watermark. Every per-frame result is untrusted. The provider-specific stabilization wrappers share one recurrence implementation, while retaining separate visual floors and minimum-run policy. Provenance can relax a low-contrast run only after recurring visual evidence exists. Sora transition frames follow the -nearest confirmed moving position only with Sora provenance. A confirmed Veo -run can cover low-contrast frames at its fixed position. This separation keeps -clean API exports from being modified merely because metadata names the same -generator. +nearest confirmed moving position only with Sora provenance. Veo, Seedance, +and Dola additionally require candidates to remain anchored to the start of a +run. This rejects slowly drifting scene details that still have high +frame-to-frame overlap. -Removal runs in a second decode pass. Sora and legacy Veo text use padded box -masks. The square Veo diamond uses a synthetic shape mask so transparent box -corners do not erase unrelated pixels. Every mask goes through the shared -`watermark_registry.fill` backends. ffmpeg encodes the changed video stream and -copies optional audio. The default OpenCV fill is the speed floor; structured -backgrounds need MI-GAN or LaMa for better reconstruction. Invisible video -stages must continue to reuse the image and metadata implementations rather -than copying their logic. +Removal runs in a second decode pass. Sora, legacy Veo text, Dola text, and the +Seedance box use box masks. Seedance deliberately fills the complete localized +box: a synthetic outline mask passed repeat detection but left part of the real +translucent border visible during visual end-to-end review. The square Veo +diamond uses a synthetic shape mask so transparent corners do not erase +unrelated pixels. Every mask goes through the shared `watermark_registry.fill` +backends. ffmpeg encodes the changed video stream and copies optional audio. +The default OpenCV fill is the speed floor; structured backgrounds need MI-GAN +or LaMa for better reconstruction. Invisible video stages must continue to +reuse the image and metadata implementations rather than copying their logic. The inherited ISOBMFF metadata path currently reads the complete container into memory; replacing that with a streaming box copier is a prerequisite for large diff --git a/docs/python-api.md b/docs/python-api.md index 24a4fbb..632fc29 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -154,7 +154,7 @@ stream bytes. MKV/WebM inspection recognizes the corresponding `Segment.Tags.Tag.SimpleTag` representation; its removal requires ffmpeg for a stream-copy remux. -## Remove a visible Sora or Veo video mark +## Remove a supported visible video mark ```python import remove_ai_watermarks as raiw @@ -173,14 +173,26 @@ veo_result = raiw.remove_video_visible( "veo_clean.mp4", mark="veo", ) +seedance_result = raiw.remove_video_visible( + "seedance.mp4", + "seedance_clean.mp4", + mark="seedance", +) +dola_result = raiw.remove_video_visible( + "dola.mp4", + "dola_clean.mp4", + mark="dola", +) ``` `remove_video_visible` scans the complete video before writing output. It combines synthetic multi-scale visual matching with temporal consistency, so an isolated lookalike in one frame is not enough to authorize inpainting. The -supported `mark` values are `sora` and `veo`. The Veo detector recognizes the -current four-point diamond and the legacy `Veo` text with separate synthetic -silhouettes. +supported `mark` values are `sora`, `veo`, `seedance`, and `dola`. The Veo +detector recognizes the current four-point diamond and the legacy `Veo` text. +Seedance recognizes the boxed `AI` label, while Dola recognizes its compact +text label. Each variant has an independent synthetic silhouette and calibrated +temporal policy. The returned `VideoVisibleResult` records the total, detected, and removed frame counts plus any AI metadata that survived the output encode. The function diff --git a/docs/supported-signals.md b/docs/supported-signals.md index ba19ef5..28ed2d7 100644 --- a/docs/supported-signals.md +++ b/docs/supported-signals.md @@ -35,6 +35,8 @@ when you can select the affected area yourself. | --- | --- | --- | --- | | `sora` | Sora 2 mascot and wordmark | Moves among frame positions | Requires a temporally recurring visual match; the older Sora Turbo corner swirl is a different unsupported mark. | | `veo` | Current four-point diamond and legacy `Veo` text | Fixed bottom-right corner | Uses separate silhouettes and requires a recurring match; learned fill is preferable on structured backgrounds. | +| `seedance` | Boxed `AI` label | Fixed bottom-right corner | Requires an anchored recurring match; the full localized box is filled because a thinner synthetic shape mask leaves the real translucent rim behind. | +| `dola` | `Dola AI` text | Fixed bottom-right corner | Requires an anchored recurring match; ByteDance or BytePlus provenance can relax only an existing visual run. | Use `video visible` for this registry. It is separate from the image `visible` command because selection is made over a sequence rather than one raster. diff --git a/docs/watermarking-landscape.md b/docs/watermarking-landscape.md index d350f7c..9e736c2 100644 --- a/docs/watermarking-landscape.md +++ b/docs/watermarking-landscape.md @@ -56,8 +56,8 @@ the existing ffmpeg stream-copy path removes the tags without transcoding. - **Built in the dated batch:** soft-binding vendor detection, IPTC Photo Metadata AI-disclosure fields, C2PA detection and stripping for supported ISOBMFF video, the optional Adobe TrustMark decoder, and temporally stabilized - visible Sora and Veo removal. Other visible video logos and proprietary - audio-watermark detection remain outside the package. + visible Sora, Veo, Seedance, and Dola removal. Other visible video logos and + proprietary audio-watermark detection remain outside the package. Metadata stripping for supported audio containers is a separate implemented path. @@ -90,6 +90,17 @@ pixels or alpha maps from that project: it uses two synthetic silhouettes, known-layout searches plus a strong relocated-diamond fallback, and a separate temporal arbiter calibrated against raw watermarked clips and clean API exports. +**ByteDance video surfaces use distinct visible labels.** Public Seedance +showcase clips contain a fixed rounded box with `AI`, while the Dola sample in +[issue #16](https://github.com/wiltodelta/remove-ai-watermarks/issues/16) uses +fixed `Dola AI` text. The independent +[Seedance remover](https://github.com/SamurAIGPT/seedance-2.0-watermark-remover) +estimates a static corner from a temporal mean frame and edge density. Our +implementation instead matches provider-specific synthetic silhouettes on +every frame, then requires an anchored temporal run. This extra anchor check +was necessary because a moving clean scene detail could retain enough adjacent +overlap to pass a recurrence-only gate. + **The faint-visible-mark precision/recall wall is fundamental, not a heuristic artifact.** The visible-watermark-detection literature has moved to LEARNED segmentation / object-detection (WDNet WACV'21 arXiv:2012.07616; SLBR ACM MM'21, open code+weights; the PRCV'18 large-scale detector; Su et al. survey 2025), but three verified findings bound what a learned detector actually buys: (1) a claim that a confidence threshold "cleanly separates" true from false matches even with a learned CNN front-end was **REFUTED** in verification (arXiv:1705.08593) -- the precision/recall wall persists even with learned features. (2) Learned detectors need a LARGE, pattern-diverse labeled dataset trained on synthetic composites (PRCV'18: 60k images / 80 watermark classes; CLWD: 60k / 160 marks), and off-distribution degradation is a documented real axis (models trained on limited-pattern LVW transfer worse; diversity of training patterns drives generalization). (3) Inference is cheap (WDNet ~8 ms at 256x256) -- the cost is the data pipeline, not runtime. Net: a learned detector shifts the frontier but does NOT remove the wall; for a SINGLE mark the cheapest next step is a small patch classifier (real-sparkle vs false-positive) on top of the existing NCC localizer, not a full segmentation model. SLBR is a ready baseline. The current NCC + false-positive gate (core-ring brightness margin + gradient-NCC crispness + white-core saturation) is a sound operating point, and the residual miss is the information-theoretic wall the literature confirms. **Visible-mark landscape beyond the registry.** Meta stamps a visible "Imagined with AI" mark (bottom-LEFT, a small symbol) on its OWN Meta AI / "Imagine" output; for third-party images it relies on C2PA / IPTC, not a visible mark. Samsung Galaxy AI additionally uses a **four-star icon** variant in a corner alongside the localized text wordmark `samsung_engine` calibrates (only the Italian text variant is covered) -- the icon is a distinct, uncovered variant. Every source agrees visible + metadata marks are trivially removable (crop / screenshot, ~2 s), which is the tool's premise. diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index 8ebbbcd..9e6f545 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -8,7 +8,7 @@ High-level API (lazy, so ``import remove_ai_watermarks`` stays cheap):: raiw.visible_provenance("in.png") # -> frozenset of confirmed vendors raiw.inspect_video_metadata("in.mp4") # -> VideoMetadataReport raiw.remove_video_metadata("in.mp4", "out.mp4") # verified metadata strip - raiw.remove_video_visible("in.mp4", "out.mp4") # stable Sora or Veo mark removal + raiw.remove_video_visible("in.mp4", "out.mp4") # stable visible video-mark removal For a provenance verdict use the ``identify`` submodule:: diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 02f0236..671a80b 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -1152,7 +1152,7 @@ def cmd_video_metadata( ) @click.option( "--mark", - type=click.Choice(["sora", "veo"]), + type=click.Choice(["sora", "veo", "seedance", "dola"]), default="sora", help="Visible AI mark to remove.", ) diff --git a/src/remove_ai_watermarks/video.py b/src/remove_ai_watermarks/video.py index 1fcb364..2928bf5 100644 --- a/src/remove_ai_watermarks/video.py +++ b/src/remove_ai_watermarks/video.py @@ -1,8 +1,8 @@ """High-level video processing API. Supported experimental stages are container-level AI metadata inspection and -removal plus temporally stabilized visible Sora and Veo removal. The pixel path -reuses the image package's shared fill backends. +removal plus temporally stabilized visible Sora, Veo, Seedance, and Dola +removal. The pixel path reuses the image package's shared fill backends. """ from __future__ import annotations @@ -133,26 +133,31 @@ def remove_video_visible( ) -> VideoVisibleResult: """Remove a supported visible AI wordmark from a video. - Supported marks are ``sora`` and ``veo``. The full sequence is scanned before - pixels change, and only recurring candidates are accepted. Audio is copied - without re-encoding; video is transcoded because the pixels change. When no - stable mark is found, no output is written and ``output`` in the result is - ``None``. + Supported marks are ``sora``, ``veo``, ``seedance``, and ``dola``. The full + sequence is scanned before pixels change, and only recurring candidates are + accepted. Audio is copied without re-encoding; video is transcoded because + the pixels change. When no stable mark is found, no output is written and + ``output`` in the result is ``None``. """ from remove_ai_watermarks.metadata import get_ai_metadata from remove_ai_watermarks.video_visible import ( encode_clean_video, + has_bytedance_video_provenance, has_sora_provenance, has_veo_provenance, + scan_dola_video, + scan_seedance_video, scan_sora_video, scan_veo_video, + stabilize_dola_localizations, + stabilize_seedance_localizations, stabilize_sora_localizations, stabilize_veo_localizations, ) from remove_ai_watermarks.watermark_registry import resolve_backend - if mark not in {"sora", "veo"}: - raise ValueError("Unsupported visible video mark; expected 'sora' or 'veo'") + if mark not in {"sora", "veo", "seedance", "dola"}: + raise ValueError("Unsupported visible video mark; expected sora, veo, seedance, or dola") if backend not in {"auto", "cv2", "migan", "lama"}: raise ValueError("Unsupported fill backend; expected auto, cv2, migan, or lama") @@ -167,7 +172,7 @@ def remove_video_visible( ) padding_fraction = 0.28 mask_style = "box" - else: + elif mark == "veo": scan = scan_veo_video(source_path) regions = stabilize_veo_localizations( scan.detections, @@ -175,6 +180,22 @@ def remove_video_visible( ) padding_fraction = 0.18 mask_style = "veo" + elif mark == "seedance": + scan = scan_seedance_video(source_path) + regions = stabilize_seedance_localizations( + scan.detections, + provenance=has_bytedance_video_provenance(markers), + ) + padding_fraction = 0.0 + mask_style = "box" + else: + scan = scan_dola_video(source_path) + regions = stabilize_dola_localizations( + scan.detections, + provenance=has_bytedance_video_provenance(markers), + ) + padding_fraction = 0.20 + mask_style = "box" detected_frames = sum(region is not None for region in regions) if detected_frames == 0: return VideoVisibleResult( diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py index 7b15ca0..28ce876 100644 --- a/src/remove_ai_watermarks/video_visible.py +++ b/src/remove_ai_watermarks/video_visible.py @@ -1,9 +1,10 @@ """Visible AI-watermark localization and removal for video. -Supported marks use fully synthetic silhouettes made from geometric primitives -and Pillow's bundled font. Sora detection searches the full frame because the -wordmark moves. Veo detection covers both the current four-point diamond and the -legacy ``Veo`` text in the bottom-right corner. A single frame is never enough +Supported marks use fully synthetic silhouettes made from geometric primitives, +OpenCV's built-in font, and Pillow's bundled font. Sora detection searches the +full frame because the wordmark moves. Veo detection covers both the current +four-point diamond and legacy ``Veo`` text. Seedance detects the boxed ``AI`` +label, while Dola detects its compact text label. A single frame is never enough to authorize removal: the temporal arbiter requires the candidate to recur at the same location across adjacent frames. This keeps isolated lookalikes in clean videos from becoming removal masks. @@ -53,8 +54,14 @@ _SORA_STRONG_CONFIDENCE = 0.65 _VEO_PROVENANCE_WEAK_CONFIDENCE = 0.45 _VEO_STRICT_WEAK_CONFIDENCE = 0.50 _VEO_STRONG_CONFIDENCE = 0.55 +_SEEDANCE_WEAK_CONFIDENCE = 0.38 +_SEEDANCE_STRONG_CONFIDENCE = 0.43 +_DOLA_PROVENANCE_WEAK_CONFIDENCE = 0.48 +_DOLA_STRICT_WEAK_CONFIDENCE = 0.50 +_DOLA_STRONG_CONFIDENCE = 0.52 _MIN_STABLE_FRAMES = 5 _MIN_VEO_STABLE_FRAMES = 12 +_MIN_FIXED_MARK_STABLE_FRAMES = 12 _MAX_STABLE_GAP = 2 _STABLE_IOU = 0.55 _VEO_REFERENCE_SHORT_SIDE = 720 @@ -63,6 +70,7 @@ _VEO_DIAMOND_PROFILES = ( (48, 72, 72), (44, 29, 40), ) +_DOLA_RELATIVE_HEIGHTS = tuple(value / 1000 for value in range(22, 41)) @dataclass(frozen=True) @@ -155,6 +163,43 @@ def _veo_templates() -> tuple[NDArray[Any], NDArray[Any]]: return np.asarray(diamond_canvas, dtype=np.uint8), text +@lru_cache(maxsize=1) +def _seedance_template() -> NDArray[Any]: + """Return a synthetic boxed-AI silhouette for Seedance exports.""" + canvas = Image.new("L", (160, 120), 0) + draw = ImageDraw.Draw(canvas) + draw.rounded_rectangle((8, 8, 142, 105), radius=28, outline=255, width=7) + draw.text( + (35, 17), + "AI", + font=_scalable_default_font(70), + fill=255, + stroke_width=1, + stroke_fill=255, + ) + draw.rounded_rectangle((142, 94, 157, 109), radius=3, outline=255, width=2) + draw.text((145, 94), "AI", font=_scalable_default_font(8), fill=255) + return np.asarray(canvas, dtype=np.uint8) + + +@lru_cache(maxsize=1) +def _dola_template() -> NDArray[Any]: + """Return a synthetic Dola AI text silhouette using OpenCV's font.""" + canvas = np.zeros((100, 400), dtype=np.uint8) + cv2.putText( + canvas, + "Dola AI", + (2, 72), + cv2.FONT_HERSHEY_DUPLEX, + 2.2, + 255, + 3, + cv2.LINE_AA, + ) + ys, xs = np.where(canvas > 0) + return canvas[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1] + + def _top_hat(gray: NDArray[Any]) -> NDArray[Any]: kernel = np.ones((7, 7), dtype=np.uint8) return cv2.morphologyEx(gray, cv2.MORPH_TOPHAT, kernel) @@ -270,6 +315,105 @@ def _match_template( ) +def _restore_region( + region: Region | None, + *, + scale: float, + frame_width: int, + frame_height: int, +) -> Region | None: + """Map a localization from normalized pixels back to the source frame.""" + if region is None: + return None + x, y, width, height = region + source_x = round(x / scale) + source_y = round(y / scale) + source_width = min(frame_width - source_x, max(1, round(width / scale))) + source_height = min(frame_height - source_y, max(1, round(height / scale))) + return source_x, source_y, source_width, source_height + + +def _detect_fixed_bottom_right_mark( + image_bgr: NDArray[Any], + template: NDArray[Any], + *, + relative_heights: tuple[float, ...], + search_origin: tuple[float, float], + kernel_fraction: float, + frame_index: int, +) -> FrameLocalization: + """Match one fixed bottom-right synthetic mark on a normalized frame.""" + if image_bgr.size == 0: + return FrameLocalization(frame_index, 0.0, None) + + frame_height, frame_width = image_bgr.shape[:2] + gray, scale = _normalized_gray(image_bgr) + normalized_height, normalized_width = gray.shape[:2] + short_side = min(normalized_height, normalized_width) + search_x = round(normalized_width * search_origin[0]) + search_y = round(normalized_height * search_origin[1]) + search_region = ( + search_x, + search_y, + normalized_width - search_x, + normalized_height - search_y, + ) + best_confidence = 0.0 + best_region: Region | None = None + for relative_height in relative_heights: + template_height = max(6, round(short_side * relative_height)) + template_width = max(1, round(template.shape[1] * template_height / template.shape[0])) + resized = cv2.resize( + template, + (template_width, template_height), + interpolation=cv2.INTER_AREA, + ) + confidence, candidate = _match_template( + gray, + resized, + region=search_region, + kernel_size=max(3, round(template_height * kernel_fraction) | 1), + ) + if confidence > best_confidence: + best_confidence = confidence + best_region = candidate + + return FrameLocalization( + frame_index, + best_confidence, + _restore_region( + best_region, + scale=scale, + frame_width=frame_width, + frame_height=frame_height, + ), + ) + + +def detect_seedance_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: + """Locate the strongest fixed Seedance boxed-AI candidate.""" + return _detect_fixed_bottom_right_mark( + image_bgr, + _seedance_template(), + relative_heights=(0.065, 0.075, 0.085, 0.095, 0.105), + search_origin=(0.68, 0.72), + kernel_fraction=0.12, + frame_index=frame_index, + ) + + +def detect_dola_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: + """Locate the strongest fixed Dola AI text candidate.""" + return _detect_fixed_bottom_right_mark( + image_bgr, + _dola_template(), + relative_heights=_DOLA_RELATIVE_HEIGHTS, + search_origin=(0.65, 0.85), + kernel_fraction=0.50, + frame_index=frame_index, + ) + + def detect_veo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: """Locate the strongest current-diamond or legacy-text Veo candidate.""" if image_bgr.size == 0: @@ -413,6 +557,43 @@ def stabilize_veo_localizations( ) +def stabilize_seedance_localizations( + detections: tuple[FrameLocalization, ...] | list[FrameLocalization], + *, + provenance: bool, +) -> list[Region | None]: + """Accept a recurring Seedance boxed-AI mark at a fixed position.""" + return _stabilize_localizations( + detections, + provenance=provenance, + weak_floor=_SEEDANCE_WEAK_CONFIDENCE, + strong_floor=_SEEDANCE_STRONG_CONFIDENCE, + transition_floor=0.30, + min_stable_frames=_MIN_FIXED_MARK_STABLE_FRAMES, + cover_after_confirmation=True, + anchor_iou=0.80, + ) + + +def stabilize_dola_localizations( + detections: tuple[FrameLocalization, ...] | list[FrameLocalization], + *, + provenance: bool, +) -> list[Region | None]: + """Accept a recurring Dola AI text mark at a fixed position.""" + weak_floor = _DOLA_PROVENANCE_WEAK_CONFIDENCE if provenance else _DOLA_STRICT_WEAK_CONFIDENCE + return _stabilize_localizations( + detections, + provenance=provenance, + weak_floor=weak_floor, + strong_floor=_DOLA_STRONG_CONFIDENCE, + transition_floor=0.40, + min_stable_frames=_MIN_FIXED_MARK_STABLE_FRAMES, + cover_after_confirmation=True, + anchor_iou=0.80, + ) + + def _stabilize_localizations( detections: tuple[FrameLocalization, ...] | list[FrameLocalization], *, @@ -422,6 +603,7 @@ def _stabilize_localizations( transition_floor: float, min_stable_frames: int, cover_after_confirmation: bool, + anchor_iou: float | None = None, ) -> list[Region | None]: """Apply the shared recurrence policy to provider-specific candidates.""" accepted: list[Region | None] = [None] * len(detections) @@ -433,11 +615,14 @@ def _stabilize_localizations( continue if current: previous = detections[current[-1]] + anchor = detections[current[0]] frame_gap = detection.frame_index - previous.frame_index if ( previous.region is None + or anchor.region is None or frame_gap > _MAX_STABLE_GAP + 1 or _region_iou(previous.region, detection.region) < _STABLE_IOU + or (anchor_iou is not None and _region_iou(anchor.region, detection.region) < anchor_iou) ): runs.append(current) current = [] @@ -534,6 +719,16 @@ def scan_veo_video(source: Path) -> VideoScan: return _scan_video(source, detect_veo_frame) +def scan_seedance_video(source: Path) -> VideoScan: + """Decode a video once and collect one untrusted Seedance candidate per frame.""" + return _scan_video(source, detect_seedance_frame) + + +def scan_dola_video(source: Path) -> VideoScan: + """Decode a video once and collect one untrusted Dola candidate per frame.""" + return _scan_video(source, detect_dola_frame) + + def _ffmpeg_video_args(suffix: str) -> list[str]: if suffix == ".webm": return ["-c:v", "libvpx-vp9", "-crf", "18", "-b:v", "0"] @@ -704,3 +899,15 @@ def has_veo_provenance(markers: dict[str, str]) -> bool: ) ).lower() return "google" in identity and "trainedalgorithmicmedia" in markers.get("source_type", "").lower() + + +def has_bytedance_video_provenance(markers: dict[str, str]) -> bool: + """Whether container provenance names ByteDance or BytePlus AI video.""" + identity = " ".join( + ( + markers.get("claim_generator", ""), + markers.get("issuer", ""), + ) + ).lower() + source_type = markers.get("source_type", "").lower() + return ("bytedance" in identity or "byteplus" in identity) and "trainedalgorithmicmedia" in source_type diff --git a/tests/test_video.py b/tests/test_video.py index ebde4be..7b0d73e 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -392,6 +392,85 @@ class TestVeoFrameLocalization: assert mask[67, 67] == 0 +class TestByteDanceFrameLocalization: + def test_localizes_independently_rendered_seedance_box(self): + from remove_ai_watermarks.video_visible import _region_iou, detect_seedance_frame + + frame = np.full((720, 1280, 3), 30, dtype=np.uint8) + mark = Image.new("L", (80, 60), 0) + draw = ImageDraw.Draw(mark) + draw.rounded_rectangle((2, 2, 70, 53), radius=14, outline=255, width=4) + try: + font = ImageFont.load_default(size=35) + except TypeError: + font = ImageFont.load_default() + draw.text((18, 8), "AI", font=font, fill=255) + mark_array = np.asarray(mark, dtype=np.float32) + x, y = 1130, 620 + alpha = mark_array[:, :, None] / 255 * 0.65 + crop = frame[y : y + 60, x : x + 80].astype(np.float32) + frame[y : y + 60, x : x + 80] = np.clip( + crop * (1 - alpha) + 255 * alpha, + 0, + 255, + ).astype(np.uint8) + + detection = detect_seedance_frame(frame) + + assert detection.region is not None + assert detection.confidence >= 0.43 + assert _region_iou(detection.region, (x, y, 80, 60)) >= 0.75 + + def test_localizes_independently_rendered_dola_text(self): + from remove_ai_watermarks.video_visible import _region_iou, detect_dola_frame + + frame = np.full((720, 1280, 3), 35, dtype=np.uint8) + mark = np.zeros((40, 150), dtype=np.uint8) + cv2.putText( + mark, + "Dola AI", + (2, 28), + cv2.FONT_HERSHEY_SIMPLEX, + 0.9, + 255, + 2, + cv2.LINE_AA, + ) + ys, xs = np.where(mark > 0) + mark = mark[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1] + mark_height, mark_width = mark.shape + x = frame.shape[1] - mark_width - 18 + y = frame.shape[0] - mark_height - 14 + alpha = mark.astype(np.float32)[:, :, None] / 255 * 0.75 + crop = frame[y : y + mark_height, x : x + mark_width].astype(np.float32) + frame[y : y + mark_height, x : x + mark_width] = np.clip( + crop * (1 - alpha) + 255 * alpha, + 0, + 255, + ).astype(np.uint8) + + detection = detect_dola_frame(frame) + + assert detection.region is not None + assert detection.confidence >= 0.52 + assert _region_iou(detection.region, (x, y, mark_width, mark_height)) >= 0.75 + + def test_seedance_box_mask_covers_the_full_localized_mark(self): + from remove_ai_watermarks.video_visible import _mask_for_region + + mask = _mask_for_region( + np.zeros((120, 160, 3), dtype=np.uint8), + (20, 20, 80, 60), + padding_fraction=0.0, + mask_style="box", + ) + + assert mask[15, 15] == 0 + assert mask[16, 16] == 255 + assert mask[83, 103] == 255 + assert mask[84, 104] == 0 + + class TestSoraTemporalArbiter: _BOX = (40, 60, 150, 54) @@ -562,6 +641,65 @@ class TestVeoTemporalArbiter: assert stabilize_veo_localizations(detections, provenance=False) == [None] * 12 +class TestByteDanceTemporalArbiter: + _SEEDANCE_BOX = (1110, 610, 90, 66) + _DOLA_BOX = (1160, 680, 96, 22) + + def test_seedance_requires_twelve_recurring_frames(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_seedance_localizations + + detections = [FrameLocalization(index, 0.50, self._SEEDANCE_BOX) for index in range(11)] + + assert stabilize_seedance_localizations(detections, provenance=False) == [None] * 11 + + def test_seedance_strong_run_covers_low_contrast_frames(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_seedance_localizations + + detections = [FrameLocalization(index, 0.45, self._SEEDANCE_BOX) for index in range(12)] + detections.extend(FrameLocalization(index, 0.20, (200, 100, 80, 60)) for index in range(12, 15)) + + assert stabilize_seedance_localizations(detections, provenance=False) == [self._SEEDANCE_BOX] * 15 + + def test_seedance_rejects_a_slowly_drifting_scene_detail(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_seedance_localizations + + detections = [FrameLocalization(index, 0.46, (1110 - index * 3, 610, 90, 66)) for index in range(14)] + + assert stabilize_seedance_localizations(detections, provenance=False) == [None] * 14 + + def test_dola_requires_twelve_recurring_frames(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_dola_localizations + + detections = [FrameLocalization(index, 0.60, self._DOLA_BOX) for index in range(11)] + + assert stabilize_dola_localizations(detections, provenance=True) == [None] * 11 + + def test_dola_provenance_accepts_recurring_low_contrast_text(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_dola_localizations + + detections = [FrameLocalization(index, 0.49, self._DOLA_BOX) for index in range(12)] + + assert stabilize_dola_localizations(detections, provenance=True) == [self._DOLA_BOX] * 12 + + def test_dola_without_provenance_needs_a_strong_frame(self): + from remove_ai_watermarks.video_visible import FrameLocalization, stabilize_dola_localizations + + detections = [FrameLocalization(index, 0.51, self._DOLA_BOX) for index in range(12)] + + assert stabilize_dola_localizations(detections, provenance=False) == [None] * 12 + + def test_bytedance_provenance_requires_ai_source_type(self): + from remove_ai_watermarks.video_visible import has_bytedance_video_provenance + + assert has_bytedance_video_provenance( + { + "issuer": "BytePlus (ByteDance)", + "source_type": "trainedAlgorithmicMedia (AI-generated)", + } + ) + assert not has_bytedance_video_provenance({"issuer": "BytePlus (ByteDance)"}) + + class TestVideoVisibleApi: def test_removes_stable_sora_run_and_writes_output(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): from remove_ai_watermarks import video_visible @@ -666,6 +804,57 @@ class TestVideoVisibleApi: assert result.detected_frames == 12 assert result.removed_frames == 12 + @pytest.mark.parametrize( + ("mark", "scan_name", "mask_style"), + [ + ("seedance", "scan_seedance_video", "box"), + ("dola", "scan_dola_video", "box"), + ], + ) + def test_dispatches_bytedance_detectors( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + mark: str, + scan_name: str, + mask_style: str, + ): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video import remove_video_visible + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + box = (40, 40, 20, 12) + scan = VideoScan( + width=64, + height=64, + fps=24.0, + detections=tuple(FrameLocalization(index, 0.60, box) for index in range(12)), + ) + monkeypatch.setattr(video_visible, scan_name, lambda _source: scan) + + def fake_encode( + _source: Path, + target: Path, + _scan: VideoScan, + regions: list[tuple[int, int, int, int] | None], + **kwargs: object, + ) -> int: + assert regions == [box] * 12 + assert kwargs["mask_style"] == mask_style + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return 12 + + monkeypatch.setattr(video_visible, "encode_clean_video", fake_encode) + + result = remove_video_visible(source, output, mark=mark) + + assert result.output == output + assert result.mark == mark + assert result.detected_frames == 12 + assert result.removed_frames == 12 + class TestVideoVisibleCli: def test_help(self): @@ -673,7 +862,7 @@ class TestVideoVisibleCli: assert result.exit_code == 0, result.output assert "temporally stable" in result.output - assert "sora|veo" in result.output + assert "sora|veo|seedance|dola" in result.output def test_reports_removed_frames(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): from remove_ai_watermarks import video From c9b85851056aefdfb7d0be37a71682c58df2a4c9 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Wed, 29 Jul 2026 21:03:46 -0700 Subject: [PATCH 03/19] Add video SynthID experiment harness --- README.md | 11 + docs/known-limitations.md | 12 + docs/synthid.md | 51 +++- docs/verification-plan.md | 16 + scripts/video_synthid_sweep.py | 470 ++++++++++++++++++++++++++++++ tests/test_video_synthid_sweep.py | 67 +++++ 6 files changed, 622 insertions(+), 5 deletions(-) create mode 100644 scripts/video_synthid_sweep.py create mode 100644 tests/test_video_synthid_sweep.py diff --git a/README.md b/README.md index 013830e..009d00b 100644 --- a/README.md +++ b/README.md @@ -288,6 +288,17 @@ invisible removal. - Provider watermark systems can change. Validate important outputs with the provider's own verifier when one is available. +Video SynthID work is currently oracle-gated research, not a removal command. +`scripts/video_synthid_sweep.py` builds a matched re-encode control plus +VAE-regenerated candidates and leaves the verifier verdict blank: + +```bash +uv run --extra gpu python scripts/video_synthid_sweep.py input.mp4 -o sweep/ +``` + +The control must still be SynthID-positive before a negative candidate can +count as removal evidence. + ## Documentation Start with the [documentation index](docs/index.md). diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 627f7dc..7f5a766 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -79,6 +79,18 @@ For important outputs: Provider systems can change, so a result verified on one file, seed, or version is not a permanent certification. +### Video regeneration is research only + +The package does not expose a `video invisible` command. The separate +`scripts/video_synthid_sweep.py` harness generates a matched transcode control +and VAE-regenerated candidates for external verification. It deliberately +leaves verdicts empty because neither the metadata proxy nor a visual quality +metric can prove that the pixel watermark is gone. + +The control must use the same clip, frame rate, dimensions, and final codec as +the candidates. If the control is not detected by the matching provider oracle, +the experiment cannot attribute a quiet candidate to regeneration. + ### Strength is content and seed dependent For SDXL and ControlNet, the CLI resolves an unset strength from the detected diff --git a/docs/synthid.md b/docs/synthid.md index 86aaa13..d93bf36 100644 --- a/docs/synthid.md +++ b/docs/synthid.md @@ -1,4 +1,4 @@ -# SynthID-Image: technical reference +# SynthID: technical reference > Technical research reference. Current package behavior is defined by the > [supported signals](supported-signals.md), [known limitations](known-limitations.md), @@ -6,10 +6,10 @@ > historical evidence and should not be read as current CLI defaults. This document covers how Google SynthID for images works mechanically, what it -survives, what removes it, and the current deployment landscape. It is written -for engineers working on watermark detection and removal -- specifically to -inform decisions about strength settings, test methodology, and what oracle -results mean. +survives, what removes it, the external video-verification workflow, and the +current deployment landscape. It is written for engineers working on watermark +detection and removal -- specifically to inform decisions about strength +settings, test methodology, and what oracle results mean. Primary sources are cited inline. Marketing-only claims are flagged separately from independently-verified results. @@ -326,6 +326,36 @@ A Google-SynthID image reads clean on openai.com/verify. An OpenAI image reads clean in the Gemini oracle. They are different payloads within the same framework. +### 3.4 Video verification and attack harness + +Gemini's verification flow can report the portions of a video where it detects +Google SynthID. This is still a proprietary oracle: a normal Gemini answer that +describes visual clues, metadata, or an unavailable decoder is not a pixel +verdict. Use the dedicated verification flow offered to an eligible signed-in +account; some versions expose an explicit `@synthid` trigger. + +The research harness `scripts/video_synthid_sweep.py` tests a VAE regeneration +attack without pretending to detect success locally. It emits: + +1. a re-encode control using the same sampled frames, dimensions, frame rate, + and codec as the candidates; +2. VAE round-trip candidates with one spatial latent-noise field shared across + time; +3. paired PSNR and motion-compensated temporal-residual metrics; +4. an empty oracle column for the external verdict. + +The control is the first oracle submission. If it is not SynthID-positive, stop: +the surrounding transcode already changed the verifier result. Only a +control-positive, candidate-negative pair is evidence about the regeneration +attack. PSNR and temporal residual measure fidelity and flicker, never watermark +presence. + +The VAE perturbation follows the general regeneration-attack construction from +Zhao et al. The video-specific control and temporal metric are local additions. +VideoMarkBench motivates testing frame aggregation and matched perturbations, +but it does not evaluate Google's proprietary SynthID, so its findings cannot +stand in for the Gemini oracle. + --- ## 4. Adoption and current state (as of June 2026) @@ -628,3 +658,14 @@ seed dependent, so reproducible verification requires a fixed seed. 5. OpenAI. **Verify tool for AI-generated images.** openai.com/research/verify. Accessed 2026-05-31. + +6. Google. **Verify AI-generated images, videos, and audio.** + https://support.google.com/gemini/answer/16722517 + +7. Zhao et al. (2024). **Invisible Image Watermarks Are Provably Removable + Using Generative AI.** NeurIPS 2024, arXiv:2306.01953. + https://arxiv.org/abs/2306.01953 + +8. Jiang et al. (2025). **VideoMarkBench: Benchmarking Robustness of Video + Watermarking.** arXiv:2505.21620. + https://arxiv.org/abs/2505.21620 diff --git a/docs/verification-plan.md b/docs/verification-plan.md index d5f98f5..9163c51 100644 --- a/docs/verification-plan.md +++ b/docs/verification-plan.md @@ -216,6 +216,22 @@ for the wrong reason reads exactly like success. the control passes -- but that is Google's claim about their own decoder, not our measurement, so it is a hypothesis to test, not a reason to skip the control. +### D4. Video candidates require a matched transcode control + +Video experiments add frame sampling, resizing, frame-rate conversion, and a +final video codec around the actual attack. `scripts/video_synthid_sweep.py` +therefore emits `control.mp4` from the same selected frames and encoder settings +as every VAE candidate. + +Verify the control first. Continue only when the provider oracle still detects +SynthID in it. A generic Gemini response that discusses visual clues, metadata, +or says the chat model lacks a decoder is not an oracle result. Record only the +explicit SynthID verification verdict in the generated CSV. + +The harness shares one latent-noise field across the sequence to avoid adding +independent frame noise. Its temporal-residual metric is a fidelity check, not a +watermark detector. + ## Tier E -- robustness and adversarial inputs Malformed and hostile inputs, including truncated files: diff --git a/scripts/video_synthid_sweep.py b/scripts/video_synthid_sweep.py new file mode 100644 index 0000000..9dca286 --- /dev/null +++ b/scripts/video_synthid_sweep.py @@ -0,0 +1,470 @@ +"""Build oracle-gated video regeneration candidates for SynthID research. + +This is a research harness, not a shipped removal command. Google does not +publish a local video SynthID decoder, so the script cannot label a candidate as +clean. It produces: + +* a re-encode control with the same duration, frame rate, dimensions, and codec; +* one VAE-regenerated video per requested latent-noise level; +* paired fidelity and temporal-residual measurements; +* a CSV column for the external Gemini SynthID verdict. + +The control is load-bearing. If it reads clean, the experiment is invalid: +resize, frame-rate conversion, or H.264 compression already silenced the oracle, +so a VAE candidate cannot be credited with removal. + +The regeneration attack follows the general encode, perturb, reconstruct family +from WatermarkAttacker (NeurIPS 2024). A single spatial latent-noise sample is +shared by every frame. Independent per-frame noise creates avoidable flicker and +does not test the video-specific question. + +Run with the project's GPU extra: + + uv run --extra gpu python scripts/video_synthid_sweep.py input.mp4 -o out/ + +Then upload ``control.mp4`` and each candidate through Gemini's SynthID +verification flow. Some eligible versions expose an explicit ``@synthid`` +trigger. A generic chat answer that says it lacks a decoder is not an oracle +verdict. Only a control-positive, candidate-negative pair is removal evidence. +""" + +from __future__ import annotations + +# torch/diffusers/cv2 expose incomplete types at this boundary. Pure helpers +# remain annotated while third-party tensor and image calls are relaxed here. +# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false +import csv +import hashlib +import logging +import math +import shutil +import subprocess +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import click +import cv2 +import numpy as np + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + +log = logging.getLogger(__name__) + +DEFAULT_VAE = "stabilityai/sd-vae-ft-mse" +_LATENT_MULTIPLE = 8 + + +def _fit_size(width: int, height: int, long_side: int) -> tuple[int, int]: + """Fit dimensions to ``long_side`` while preserving aspect and VAE alignment.""" + if width <= 0 or height <= 0: + raise ValueError("Video dimensions must be positive") + if long_side < _LATENT_MULTIPLE: + raise ValueError(f"Long side must be at least {_LATENT_MULTIPLE}") + scale = long_side / max(width, height) + fitted_width = max(_LATENT_MULTIPLE, round(width * scale) // _LATENT_MULTIPLE * _LATENT_MULTIPLE) + fitted_height = max(_LATENT_MULTIPLE, round(height * scale) // _LATENT_MULTIPLE * _LATENT_MULTIPLE) + return fitted_width, fitted_height + + +def _parse_noise_levels(values: str) -> tuple[float, ...]: + levels = tuple(float(value.strip()) for value in values.split(",") if value.strip()) + if not levels: + raise click.BadParameter("At least one noise level is required") + if any(not 0.0 <= value <= 1.0 for value in levels): + raise click.BadParameter("Noise levels must be between 0 and 1") + return levels + + +def _pick_device(requested: str) -> str: + import torch + + if requested != "auto": + return requested + if torch.cuda.is_available(): + return "cuda" + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return "mps" + return "cpu" + + +def _shared_latent_noise( + spatial_shape: Sequence[int], + *, + seed: int, + device: str, + dtype: Any, +) -> Any: + """Return one deterministic spatial noise field for reuse across time.""" + import torch + + if len(spatial_shape) != 3 or any(size <= 0 for size in spatial_shape): + raise ValueError("Expected a positive CHW latent shape") + generator = torch.Generator(device="cpu").manual_seed(seed) + noise = torch.randn((1, *spatial_shape), generator=generator, dtype=torch.float32) + return noise.to(device=device, dtype=dtype) + + +def _psnr(reference: np.ndarray, candidate: np.ndarray) -> float: + """Return paired PSNR over uint8 frame stacks.""" + if reference.shape != candidate.shape: + raise ValueError("PSNR inputs must have matching shapes") + mse = float(np.mean((reference.astype(np.float32) - candidate.astype(np.float32)) ** 2)) + if mse == 0.0: + return math.inf + return 20.0 * math.log10(255.0 / math.sqrt(mse)) + + +def _backward_map(current_gray: np.ndarray, previous_gray: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Build a remap from a previous frame into current coordinates.""" + flow = cv2.calcOpticalFlowFarneback( + current_gray, + previous_gray, + None, + 0.5, + 3, + 15, + 3, + 5, + 1.2, + 0, + ) + height, width = current_gray.shape + grid_x, grid_y = np.meshgrid(np.arange(width, dtype=np.float32), np.arange(height, dtype=np.float32)) + return grid_x + flow[..., 0], grid_y + flow[..., 1] + + +def _backward_warp(image: np.ndarray, maps: tuple[np.ndarray, np.ndarray]) -> np.ndarray: + """Apply a precomputed backward optical-flow map.""" + return cv2.remap( + image, + maps[0], + maps[1], + interpolation=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_REFLECT, + ) + + +def _temporal_reference( + reference: Sequence[np.ndarray], +) -> tuple[tuple[tuple[np.ndarray, np.ndarray], ...], float]: + """Precompute source motion maps and its mean residual.""" + if len(reference) < 2: + raise ValueError("Temporal metric needs at least two frames") + maps: list[tuple[np.ndarray, np.ndarray]] = [] + reference_residuals: list[float] = [] + for index in range(1, len(reference)): + current_gray = cv2.cvtColor(reference[index], cv2.COLOR_BGR2GRAY) + previous_gray = cv2.cvtColor(reference[index - 1], cv2.COLOR_BGR2GRAY) + frame_maps = _backward_map(current_gray, previous_gray) + maps.append(frame_maps) + warped_reference = _backward_warp(reference[index - 1], frame_maps) + reference_residuals.append( + float(np.mean(np.abs(reference[index].astype(np.float32) - warped_reference.astype(np.float32)))) + ) + return tuple(maps), float(np.mean(reference_residuals)) + + +def _temporal_residual_ratio( + candidate: Sequence[np.ndarray], + maps: Sequence[tuple[np.ndarray, np.ndarray]], + baseline: float, +) -> float: + """Measure candidate flicker against a precomputed source residual.""" + if len(candidate) != len(maps) + 1: + raise ValueError("Temporal metric needs one map per adjacent frame pair") + candidate_residuals: list[float] = [] + for index, frame_maps in enumerate(maps, start=1): + warped_candidate = _backward_warp(candidate[index - 1], frame_maps) + candidate_residuals.append( + float(np.mean(np.abs(candidate[index].astype(np.float32) - warped_candidate.astype(np.float32)))) + ) + measured = float(np.mean(candidate_residuals)) + return measured / max(baseline, 1e-6) + + +def _read_frames( + source: Path, + *, + duration: float, + output_fps: float, + size: tuple[int, int], +) -> tuple[list[np.ndarray], float]: + """Read a uniformly sampled prefix and resize it to the experiment geometry.""" + capture = cv2.VideoCapture(str(source)) + if not capture.isOpened(): + raise ValueError(f"Could not open video: {source}") + source_fps = float(capture.get(cv2.CAP_PROP_FPS)) + if source_fps <= 0.0: + capture.release() + raise ValueError(f"Video has no usable frame rate: {source}") + effective_fps = min(output_fps, source_fps) + sample_period = 1.0 / effective_fps + next_sample_time = 0.0 + frames: list[np.ndarray] = [] + frame_index = 0 + try: + while True: + ok, frame = capture.read() + if not ok: + break + timestamp = frame_index / source_fps + if timestamp + 1e-9 >= duration: + break + if timestamp + 1e-9 >= next_sample_time: + frames.append(cv2.resize(frame, size, interpolation=cv2.INTER_LANCZOS4)) + next_sample_time += sample_period + frame_index += 1 + finally: + capture.release() + if len(frames) < 2: + raise ValueError("The selected clip produced fewer than two frames") + return frames, effective_fps + + +def _frame_batches(frames: Sequence[np.ndarray], batch_size: int) -> Iterable[Sequence[np.ndarray]]: + for start in range(0, len(frames), batch_size): + yield frames[start : start + batch_size] + + +def _encode_frame_latents( + frames: Sequence[np.ndarray], + *, + vae: Any, + device: str, + batch_size: int, +) -> list[Any]: + """Encode source frames once so every noise level reuses identical latents.""" + import torch + + latent_batches: list[Any] = [] + scaling_factor = float(vae.config.scaling_factor) + with torch.inference_mode(): + for batch in _frame_batches(frames, batch_size): + rgb = np.stack([frame[:, :, ::-1] for frame in batch]) + tensor = torch.from_numpy(np.ascontiguousarray(rgb)).permute(0, 3, 1, 2) + tensor = tensor.to(device=device, dtype=vae.dtype) / 127.5 - 1.0 + latents = vae.encode(tensor).latent_dist.mode() * scaling_factor + latent_batches.append(latents) + return latent_batches + + +def _decode_frame_latents( + latent_batches: Sequence[Any], + *, + vae: Any, + noise_std: float, + shared_noise: Any, +) -> list[np.ndarray]: + """Decode cached latents with one perturbation shared across time.""" + import torch + + output: list[np.ndarray] = [] + scaling_factor = float(vae.config.scaling_factor) + with torch.inference_mode(): + for latents in latent_batches: + perturbed = latents + noise_std * shared_noise.expand(latents.shape[0], -1, -1, -1) + decoded = vae.decode(perturbed / scaling_factor).sample + decoded = ((decoded / 2.0 + 0.5).clamp(0.0, 1.0) * 255.0).round().to(torch.uint8) + decoded = decoded.permute(0, 2, 3, 1).cpu().numpy() + output.extend(np.ascontiguousarray(frame[:, :, ::-1]) for frame in decoded) + return output + + +def _write_png_frames(frames: Sequence[np.ndarray], directory: Path) -> None: + directory.mkdir(parents=True, exist_ok=True) + for index, frame in enumerate(frames, start=1): + path = directory / f"{index:06d}.png" + if not cv2.imwrite(str(path), frame): + raise OSError(f"Failed to write frame: {path}") + + +def _encode_video( + frames: Sequence[np.ndarray], + source: Path, + output: Path, + *, + fps: float, + duration: float, +) -> None: + ffmpeg = shutil.which("ffmpeg") + if ffmpeg is None: + raise RuntimeError("ffmpeg is required on PATH") + with tempfile.TemporaryDirectory(prefix="video-synthid-") as temp_dir: + frame_dir = Path(temp_dir) + _write_png_frames(frames, frame_dir) + command = [ + ffmpeg, + "-hide_banner", + "-loglevel", + "error", + "-y", + "-framerate", + f"{fps:.8g}", + "-i", + str(frame_dir / "%06d.png"), + "-i", + str(source), + "-map", + "0:v:0", + "-map", + "1:a:0?", + "-t", + f"{duration:.8g}", + "-c:v", + "libx264", + "-crf", + "18", + "-pix_fmt", + "yuv420p", + "-c:a", + "aac", + "-movflags", + "+faststart", + str(output), + ] + log.info("Encoding %s", output.name) + subprocess.run(command, check=True) # noqa: S603 + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _write_manifest(output_dir: Path, rows: Sequence[dict[str, str]]) -> Path: + path = output_dir / "sweep.csv" + fieldnames = [ + "variant", + "noise_std", + "psnr_db", + "temporal_residual_ratio", + "file", + "sha256", + "synthid_oracle", + ] + with path.open("w", newline="", encoding="utf-8") as stream: + writer = csv.DictWriter(stream, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + return path + + +@click.command() +@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) +@click.option("-o", "--output-dir", required=True, type=click.Path(file_okay=False, path_type=Path)) +@click.option("--noise-levels", default="0,0.025,0.05,0.1", show_default=True) +@click.option("--duration", type=click.FloatRange(min=0.1), default=2.0, show_default=True) +@click.option("--fps", type=click.FloatRange(min=1.0), default=12.0, show_default=True) +@click.option("--long-side", type=click.IntRange(min=_LATENT_MULTIPLE), default=512, show_default=True) +@click.option("--batch-size", type=click.IntRange(min=1), default=4, show_default=True) +@click.option("--seed", type=int, default=0, show_default=True) +@click.option("--model", default=DEFAULT_VAE, show_default=True) +@click.option("--device", type=click.Choice(["auto", "cuda", "mps", "cpu"]), default="auto", show_default=True) +def main( + source: Path, + output_dir: Path, + noise_levels: str, + duration: float, + fps: float, + long_side: int, + batch_size: int, + seed: int, + model: str, + device: str, +) -> None: + """Generate VAE video candidates from the prefix of SOURCE.""" + logging.basicConfig(level=logging.INFO, format="%(message)s") + import torch + from diffusers import AutoencoderKL + + levels = _parse_noise_levels(noise_levels) + capture = cv2.VideoCapture(str(source)) + if not capture.isOpened(): + raise click.ClickException(f"Could not open video: {source}") + width = round(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = round(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) + capture.release() + size = _fit_size(width, height, long_side) + frames, effective_fps = _read_frames(source, duration=duration, output_fps=fps, size=size) + effective_duration = len(frames) / effective_fps + + output_dir.mkdir(parents=True, exist_ok=True) + control_path = output_dir / "control.mp4" + _encode_video(frames, source, control_path, fps=effective_fps, duration=effective_duration) + rows: list[dict[str, str]] = [ + { + "variant": "control", + "noise_std": "", + "psnr_db": "inf", + "temporal_residual_ratio": "1", + "file": control_path.name, + "sha256": _sha256(control_path), + "synthid_oracle": "", + } + ] + + resolved_device = _pick_device(device) + dtype = torch.float16 if resolved_device == "cuda" else torch.float32 + log.info("Loading %s on %s", model, resolved_device) + vae = AutoencoderKL.from_pretrained(model, torch_dtype=dtype).to(resolved_device) + vae.eval() + vae.enable_slicing() + + log.info("Encoding source frames") + latent_batches = _encode_frame_latents( + frames, + vae=vae, + device=resolved_device, + batch_size=batch_size, + ) + first_latents = latent_batches[0] + shared_noise = _shared_latent_noise( + first_latents.shape[1:], + seed=seed, + device=resolved_device, + dtype=first_latents.dtype, + ) + reference_stack = np.stack(frames) + temporal_maps, temporal_baseline = _temporal_reference(frames) + for level in levels: + log.info("Decoding latent noise %.4f", level) + regenerated = _decode_frame_latents( + latent_batches, + vae=vae, + noise_std=level, + shared_noise=shared_noise, + ) + output_path = output_dir / f"vae-noise-{level:.4f}.mp4" + _encode_video( + regenerated, + source, + output_path, + fps=effective_fps, + duration=effective_duration, + ) + psnr = _psnr(reference_stack, np.stack(regenerated)) + temporal_ratio = _temporal_residual_ratio(regenerated, temporal_maps, temporal_baseline) + rows.append( + { + "variant": "vae", + "noise_std": f"{level:.4f}", + "psnr_db": f"{psnr:.4f}", + "temporal_residual_ratio": f"{temporal_ratio:.4f}", + "file": output_path.name, + "sha256": _sha256(output_path), + "synthid_oracle": "", + } + ) + + manifest = _write_manifest(output_dir, rows) + log.info("Wrote %s", manifest) + log.info("Verify control.mp4 first in Gemini's SynthID flow; stop if the control is not detected.") + + +if __name__ == "__main__": + main() diff --git a/tests/test_video_synthid_sweep.py b/tests/test_video_synthid_sweep.py new file mode 100644 index 0000000..329ce58 --- /dev/null +++ b/tests/test_video_synthid_sweep.py @@ -0,0 +1,67 @@ +"""Pure regression tests for the oracle-gated video SynthID experiment.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import pytest +import torch + +if TYPE_CHECKING: + from types import ModuleType + +_SCRIPT = Path(__file__).parent.parent / "scripts" / "video_synthid_sweep.py" + + +@pytest.fixture(scope="module") +def sweep() -> ModuleType: + spec = importlib.util.spec_from_file_location("video_synthid_sweep", _SCRIPT) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_fit_size_preserves_landscape_aspect_and_vae_alignment(sweep: ModuleType) -> None: + assert sweep._fit_size(1280, 720, 512) == (512, 288) + + +def test_fit_size_rejects_invalid_dimensions(sweep: ModuleType) -> None: + with pytest.raises(ValueError, match="positive"): + sweep._fit_size(0, 720, 512) + + +def test_shared_latent_noise_is_one_spatial_field(sweep: ModuleType) -> None: + noise = sweep._shared_latent_noise( + (4, 8, 8), + seed=7, + device="cpu", + dtype=torch.float32, + ) + assert noise.shape == (1, 4, 8, 8) + + +def test_shared_latent_noise_is_seeded(sweep: ModuleType) -> None: + first = sweep._shared_latent_noise((4, 8, 8), seed=7, device="cpu", dtype=torch.float32) + repeated = sweep._shared_latent_noise((4, 8, 8), seed=7, device="cpu", dtype=torch.float32) + other = sweep._shared_latent_noise((4, 8, 8), seed=8, device="cpu", dtype=torch.float32) + assert torch.equal(first, repeated) + assert not torch.equal(first, other) + + +def test_psnr_is_infinite_for_identical_frames(sweep: ModuleType) -> None: + frame = np.full((2, 8, 8, 3), 120, dtype=np.uint8) + assert sweep._psnr(frame, frame.copy()) == pytest.approx(float("inf")) + + +def test_temporal_residual_ratio_is_one_for_identical_sequences(sweep: ModuleType) -> None: + first = np.zeros((32, 32, 3), dtype=np.uint8) + second = first.copy() + second[:, 8:16] = 80 + sequence = [first, second] + maps, baseline = sweep._temporal_reference(sequence) + assert sweep._temporal_residual_ratio([frame.copy() for frame in sequence], maps, baseline) == pytest.approx(1.0) From 63cde7f32aef8cc720616d8ed99dc32abc75727d Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Wed, 29 Jul 2026 22:03:23 -0700 Subject: [PATCH 04/19] Add streaming video SynthID regeneration --- CLAUDE.md | 4 +- README.md | 41 +- docs/cli.md | 32 ++ docs/known-limitations.md | 37 +- docs/module-internals.md | 27 +- docs/python-api.md | 28 ++ docs/supported-signals.md | 5 + docs/synthid.md | 15 +- docs/verification-plan.md | 8 + scripts/video_synthid_sweep.py | 337 ++----------- src/remove_ai_watermarks/__init__.py | 11 +- src/remove_ai_watermarks/cli.py | 88 ++++ src/remove_ai_watermarks/video.py | 107 +++- src/remove_ai_watermarks/video_encoding.py | 108 +++++ src/remove_ai_watermarks/video_invisible.py | 513 ++++++++++++++++++++ src/remove_ai_watermarks/video_synthid.py | 10 + src/remove_ai_watermarks/video_visible.py | 94 ++-- tests/test_video.py | 123 +++++ tests/test_video_invisible.py | 71 +++ tests/test_video_synthid_sweep.py | 6 +- 20 files changed, 1278 insertions(+), 387 deletions(-) create mode 100644 src/remove_ai_watermarks/video_encoding.py create mode 100644 src/remove_ai_watermarks/video_invisible.py create mode 100644 src/remove_ai_watermarks/video_synthid.py create mode 100644 tests/test_video_invisible.py diff --git a/CLAUDE.md b/CLAUDE.md index f8cc4dd..e7b75b0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -82,7 +82,9 @@ Compact map. The full per-module detail (design decisions, tuned thresholds, cal - `upscaler.py` — optional Real-ESRGAN pre-diffusion super-resolution for small inputs (extra `esrgan`, spandrel only). Manual opt-in; the default `--upscaler` stays `lanczos` and the engine always falls back to Lanczos on absence/error. ESRGAN can degrade faces and thin text. - `image_io.py` — centralizes Unicode-safe image IO, alpha preservation, content-based format sniffing, and HEIC/AVIF fallbacks. Callers must check `imwrite` success. No-op visible removal preserves original bytes when the output format is unchanged. - `api.py` — the high-level convenience API, re-exported lazily at the package top level via `__init__.__getattr__` (PEP 562, so `import remove_ai_watermarks` stays cheap): `remove_visible(source, output=None, *, sensitivity="auto", backend="auto", strip_metadata=True, write_noop=True) -> (result_bgr, [labels])` (source = path OR BGR ndarray; a PATH auto-reads metadata provenance and preserves alpha, an ARRAY does neither; `write_noop=True` writes a clean passthrough copy when nothing is removed, `False` leaves `output` untouched so a "no mark = produce nothing" caller like the CLI `visible` command does not clobber a pre-existing file there) and `visible_provenance(path) -> frozenset[str]` (the single metadata→vendor-keys mapper; `cli._visible_provenance` is a thin None-guarded wrapper over it). **`remove_visible` is the ONE path the CLI and library share** — `cli.cmd_visible`'s `--mark auto` branch delegates entirely to it (read → provenance → `remove_auto_marks` → write → `strip_metadata`), so there is no CLI-vs-library drift; `strip_metadata` defaults True to match `visible --strip-metadata`. This is where a library caller should start — NOT the engines directly (`GeminiEngine`/`TextMarkEngine` have no `remove_watermark` any more; removal is registry `remove_auto_marks`/`KnownMark.remove`; the old single-strongest `best_auto_mark` is gone — removal takes EVERY mark). `identify` is NOT top-level re-exported (it collides with the `identify` submodule); use `from remove_ai_watermarks.identify import identify`. -- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, and `remove_video_visible(source, output=None, *, mark="sora", backend="cv2", strip_metadata=True) -> VideoVisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM TC260 is read by `noai/ebml.py` and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: fully synthetic Sora, Veo, Seedance, and Dola silhouettes propose frame boxes, provider-specific temporal recurrence authorizes them, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying audio. Fixed marks require anchored runs so a slowly drifting scene detail cannot pass on adjacent overlap alone. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Seedance uses the full localized box because an outline mask left part of the real translucent border behind. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels and invisible video watermarks are not built yet. +- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, `remove_video_visible(source, output=None, *, mark="sora", backend="cv2", strip_metadata=True) -> VideoVisibleResult`, and `remove_video_invisible(source, output=None, ...) -> VideoInvisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM TC260 is read by `noai/ebml.py` and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: fully synthetic Sora, Veo, Seedance, and Dola silhouettes propose frame boxes, provider-specific temporal recurrence authorizes them, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying audio. Fixed marks require anchored runs so a slowly drifting scene detail cannot pass on adjacent overlap alone. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Seedance uses the full localized box because an outline mask left part of the real translucent border behind. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels are not built yet. +- `video_encoding.py` — the shared raw-BGR ffmpeg command and pipe lifecycle for both visible removal and invisible regeneration. It copies optional audio, controls metadata/chapter retention explicitly, and keeps container-specific codec arguments in one place. +- `video_invisible.py` — the oracle-gated video SynthID candidate engine for MP4/MOV/M4V. It regenerates frames through `stabilityai/sd-vae-ft-mse` with one seeded latent-noise field shared across time, retains only one configured batch, updates PSNR and motion-compensated temporal residuals incrementally, streams pixels directly to ffmpeg, copies audio, strips metadata, and atomically publishes the completed output. Every result remains explicitly unverified because Google exposes no local decoder. The 2026-07-29 Gemini Flash calibration found the default `0.10` candidate negative on both control-positive public Veo clips, while `0.05` remained positive on one, so the operating point is content-dependent and requires the matching external oracle. For the Doubao alpha-distillation history (why content-image reverse-alpha distillation fails by physics and controlled captures were required), see `docs/research-doubao-distillation.md`. diff --git a/README.md b/README.md index 009d00b..c3f658f 100644 --- a/README.md +++ b/README.md @@ -6,9 +6,9 @@ Remove AI provenance marks from images and video you generated yourself: - invisible pixel watermarks through diffusion regeneration; - C2PA, EXIF, XMP, IPTC, and related AI metadata. -Video support covers metadata inspection and removal plus experimental visible -Sora, Veo, Seedance, and Dola mark removal. Invisible video-watermark removal -remains a follow-up stage. +Video support covers metadata inspection and removal, visible Sora, Veo, +Seedance, and Dola mark removal, and experimental VAE regeneration that +produces a video SynthID candidate for external verification. > Try it online at [raiw.cc](https://raiw.cc) if you do not want to install Python > or run diffusion models locally. @@ -34,6 +34,7 @@ remains a follow-up stage. | Strip AI metadata | `metadata` | No | | Strip AI metadata from video | `video metadata` | No | | Remove a known Sora, Veo, Seedance, or Dola mark from video | `video visible` | No | +| Generate an externally verifiable video SynthID candidate | `video invisible` | Recommended | | Regenerate an image to disrupt invisible watermarks | `invisible` | Recommended | | Run visible, invisible, and metadata removal | `all` | Recommended | | Process a directory | `batch` | Depends on mode | @@ -96,6 +97,19 @@ four-point diamond and the legacy `Veo` text. Seedance covers the fixed boxed `AI` label, and Dola covers the fixed `Dola AI` text. No output is written when no stable mark is found. +Generate a video SynthID candidate: + +```bash +uv tool install --force "remove-ai-watermarks[gpu]" +remove-ai-watermarks video invisible input.mp4 -o candidate.mp4 +``` + +This path regenerates the complete sequence with one latent-noise field shared +across time, copies audio, and strips source metadata. It cannot verify +Google's proprietary pixel watermark locally. The command therefore labels +every output `UNVERIFIED` and prints the exact Gemini Flash verification +prompt. + For invisible watermark removal, install the diffusion dependencies: ```bash @@ -244,6 +258,7 @@ print(removed) report = raiw.inspect_video_metadata("input.mp4") cleaned = raiw.remove_video_metadata("input.mp4") +candidate = raiw.remove_video_invisible("input.mp4", "candidate.mp4") visible = raiw.remove_video_visible("sora.mp4", "sora_clean.mp4") veo = raiw.remove_video_visible("veo.mp4", "veo_clean.mp4", mark="veo") seedance = raiw.remove_video_visible( @@ -280,24 +295,32 @@ invisible removal. diamond plus legacy `Veo` text, the Seedance boxed `AI` label, and the fixed `Dola AI` text. It does not recognize the older Sora Turbo corner swirl. The classical OpenCV backend can smear structured backgrounds; use MI-GAN or - LaMa when recovery quality matters. Video SynthID is not removed yet. - MP4/MOV/M4V metadata stripping currently reads the full container into + LaMa when recovery quality matters. +- Video SynthID regeneration changes resolution, frame rate, and image detail. + It produces a candidate, not a locally verified clean video. The matching + Google verifier is still required for every important output. +- MP4/MOV/M4V metadata stripping currently reads the full container into memory, so very large videos are not an intended metadata input yet. - `qwen-zimage` requires CUDA. The other diffusion profiles also support the devices listed by `remove-ai-watermarks invisible --help`. - Provider watermark systems can change. Validate important outputs with the provider's own verifier when one is available. -Video SynthID work is currently oracle-gated research, not a removal command. -`scripts/video_synthid_sweep.py` builds a matched re-encode control plus -VAE-regenerated candidates and leaves the verifier verdict blank: +The shipped `video invisible` command uses the candidate-producing side of the +oracle-gated workflow. The companion `scripts/video_synthid_sweep.py` harness +builds a matched re-encode control plus VAE-regenerated candidates and leaves +the verifier verdict blank: ```bash uv run --extra gpu python scripts/video_synthid_sweep.py input.mp4 -o sweep/ ``` The control must still be SynthID-positive before a negative candidate can -count as removal evidence. +count as removal evidence. In the 2026-07-29 calibration, matched controls from +two public Veo videos remained positive and the default stronger candidate was +negative on both. A weaker candidate was negative on only one, demonstrating +that the removal threshold is content dependent. This is calibration evidence, +not a universal guarantee for new videos or future verifier versions. ## Documentation diff --git a/docs/cli.md b/docs/cli.md index 293082d..76f680e 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -153,6 +153,38 @@ different container extension. Visible video labels and invisible video watermarks are not handled by this command. +## Generate a video SynthID candidate + +```bash +uv tool install --force "remove-ai-watermarks[gpu]" +remove-ai-watermarks video invisible input.mp4 -o candidate.mp4 +``` + +The experimental command supports MP4, MOV, and M4V. It samples the complete +sequence at the configured frame rate, resizes frames to the configured long +side, regenerates them through a VAE, and applies one deterministic latent-noise +field to every frame. Reusing one spatial field avoids the unnecessary flicker +caused by independent per-frame noise. Frames are regenerated in bounded +batches and streamed directly to ffmpeg, which encodes the result, copies +audio, and drops source metadata. + +The output is always an unverified candidate. The project has no local video +SynthID decoder, and PSNR or temporal-residual metrics cannot prove watermark +absence. After generation, upload the candidate to Gemini Flash and ask: + +> Was this uploaded video created or edited by Google AI? Use the built-in +> content verification result. + +Only an explicit built-in verification result is an oracle verdict. A response +based on the visible logo, content appearance, or metadata is not. The command +prints `UNVERIFIED` even when generation succeeds. + +The default output is `_synthid_candidate` in the same container. The +source is never overwritten. Use `--noise-std`, `--long-side`, `--fps`, +`--batch-size`, `--seed`, and `--device` to control the regeneration. The +defaults are calibrated operating points, not a guarantee for every carrier or +future verifier version. + ## Remove a supported visible video mark ```bash diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 7f5a766..0902441 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -79,17 +79,31 @@ For important outputs: Provider systems can change, so a result verified on one file, seed, or version is not a permanent certification. -### Video regeneration is research only +### Video regeneration produces an unverified candidate -The package does not expose a `video invisible` command. The separate -`scripts/video_synthid_sweep.py` harness generates a matched transcode control -and VAE-regenerated candidates for external verification. It deliberately -leaves verdicts empty because neither the metadata proxy nor a visual quality -metric can prove that the pixel watermark is gone. +The `video invisible` command and `remove_video_invisible` API regenerate video +pixels through a VAE, but cannot verify the proprietary SynthID payload locally. +Every successful result remains explicitly unverified and requires Google's +matching content-verification flow. A quiet metadata scan, paired PSNR, and the +temporal-residual metric are not substitutes for that oracle. The control must use the same clip, frame rate, dimensions, and final codec as -the candidates. If the control is not detected by the matching provider oracle, -the experiment cannot attribute a quiet candidate to regeneration. +the candidates. The separate `scripts/video_synthid_sweep.py` harness produces +that matched control. If the control is not detected by the matching provider +oracle, the experiment cannot attribute a quiet candidate to regeneration. + +The 2026-07-29 calibration used two public Veo clips. Both matched controls were +SynthID-positive. The default stronger VAE candidate was negative on both, while +a weaker candidate remained positive on one. This small calibration validates +the mechanism and exposes its content dependence; it does not certify all +videos, full-length sequences, alternate settings, or future verifier versions. + +The shipped engine streams sampled frames in bounded batches, computes its +fidelity metrics incrementally, and pipes regenerated pixels directly to +ffmpeg. Its frame and latent memory is therefore bounded by `--batch-size` +rather than clip duration. Runtime still grows linearly with duration, and the +separate multi-candidate research sweep deliberately retains its short sampled +prefix so it can reuse identical latents across candidate strengths. ### Strength is content and seed dependent @@ -186,7 +200,7 @@ WebM, Matroska, MP3, WAV, FLAC, OGG, Opus, and AAC container metadata is strippe through ffmpeg with stream copying. The operation fails if ffmpeg is absent or cannot parse the input. -### Visible video removal is provider-specific and still experimental +### Video pixel removal is provider-specific and still experimental The experimental `video metadata` command and high level video API inspect and strip supported AI provenance metadata without transcoding streams. @@ -201,8 +215,9 @@ after visual evidence exists, so metadata alone does not erase a clean API export. Historical Sora Turbo exports use a small OpenAI swirl in the corner rather than the moving mascot-and-wordmark design; that earlier variant is not -detected by the `sora` video mark. Other provider video labels and proprietary -invisible video watermarks are not supported yet. +detected by the `sora` video mark. Other provider video labels are not +supported yet. Google video SynthID has a candidate-producing VAE path, while +other proprietary invisible video watermarks have no registered attack. Visible removal transcodes the video stream and copies audio. Its frame-local fill is not a motion-aware video inpainting model. OpenCV can leave a visible diff --git a/docs/module-internals.md b/docs/module-internals.md index 8713b7f..c5ccb29 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -90,6 +90,7 @@ Regression coverage: video entry point: - `inspect_video_metadata` +- `remove_video_invisible` - `remove_video_metadata` - `remove_video_visible` @@ -97,7 +98,7 @@ The video API validates both the supported extension and container signature, then delegates all metadata detection and stripping to `metadata.py`. It requires a separate same-container output, defaulting to `_clean`, so the experimental path does not overwrite an original. The package root exposes -both functions lazily. +all four functions lazily. Native MP4/MOV TC260 labels follow TC260-PG-20257A: `moov.udta.meta.keys` maps an `AIGC` key to a raw JSON value in `ilst`. @@ -114,6 +115,30 @@ only a `Segment.Tags.Tag.SimpleTag` pairing `TagName=AIGC` with a JSON `TagString` carrying a TC260 field. The existing ffmpeg stream-copy path removes those container tags without transcoding the encoded streams. +[`video_encoding.py`](../src/remove_ai_watermarks/video_encoding.py) owns the +raw-BGR ffmpeg command and pipe lifecycle shared by visible removal and +invisible regeneration. It centralizes container codecs, optional audio stream +copying, metadata/chapter policy, and encode-failure reporting. + +[`video_invisible.py`](../src/remove_ai_watermarks/video_invisible.py) +implements the oracle-gated video SynthID candidate engine. It samples frames +uniformly, resizes to a VAE-aligned geometry, encodes each frame to latent +space, applies one seeded spatial-noise field across the entire sequence, and +decodes fresh pixels. Reusing a single noise field avoids independent +frame-to-frame noise. The shipped path retains only one configured frame batch, +updates PSNR and temporal residuals incrementally, and streams BGR frames +directly to ffmpeg. ffmpeg encodes H.264 video, maps optional source audio, and +drops all source metadata. The result is written through a same-directory +temporary file and atomically replaced only after a successful encode. + +The engine returns PSNR and a motion-compensated temporal-residual ratio as +quality measurements. Neither is a watermark detector. The high-level +`VideoInvisibleResult.requires_external_verification` flag is always true, and +the CLI prints an `UNVERIFIED` warning plus the Gemini Flash verification +prompt. The companion `scripts/video_synthid_sweep.py` imports the same engine +helpers to build a matched control and candidate grid, preventing research and +shipped regeneration paths from drifting. + [`video_visible.py`](../src/remove_ai_watermarks/video_visible.py) implements the first pixel stages for Sora, Veo, Seedance, and Dola. The Sora detector searches a normalized frame with a fully synthetic mascot-and-text silhouette diff --git a/docs/python-api.md b/docs/python-api.md index 632fc29..9af954c 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -154,6 +154,34 @@ stream bytes. MKV/WebM inspection recognizes the corresponding `Segment.Tags.Tag.SimpleTag` representation; its removal requires ffmpeg for a stream-copy remux. +## Generate a video SynthID candidate + +```python +import remove_ai_watermarks as raiw + +result = raiw.remove_video_invisible( + "input.mp4", + "candidate.mp4", + device="auto", +) +assert result.requires_external_verification +if result.remaining_metadata: + raise RuntimeError(f"AI metadata remains: {result.remaining_metadata}") +``` + +`remove_video_invisible` supports MP4, MOV, and M4V. It regenerates the complete +video through a VAE in bounded batches, shares one seeded latent-noise field +across all frames, streams pixels to ffmpeg, copies audio, and strips source +metadata. The default output is +`input_synthid_candidate.mp4`; a distinct same-container output is required. + +The returned `VideoInvisibleResult` includes output geometry, frame rate, frame +count, paired PSNR, and the motion-compensated temporal-residual ratio. Those +fields measure fidelity and flicker only. They are not a SynthID detector. +`requires_external_verification` is always true because Google does not publish +a local decoder for this video payload. Verify the candidate with Gemini +Flash's built-in content verification before treating it as watermark-negative. + ## Remove a supported visible video mark ```python diff --git a/docs/supported-signals.md b/docs/supported-signals.md index 28ed2d7..15b9bd3 100644 --- a/docs/supported-signals.md +++ b/docs/supported-signals.md @@ -113,6 +113,10 @@ SynthID does not have a public local pixel decoder in this project. The tool can infer likely presence from supported provenance metadata, but after that metadata is removed a local negative result is inconclusive. +For MP4, MOV, and M4V, `video invisible` can regenerate the video through a VAE +and strip source metadata. This is a candidate-producing attack, not a local +decoder. Every result still requires Gemini Flash's built-in verification. + The optional `detect` extra is different: it provides a local decoder for the open DWT-DCT watermark used by some Stable Diffusion, SDXL, and FLUX workflows. That signal is carrier and transformation sensitive, so a negative is still @@ -123,6 +127,7 @@ not a universal clean verdict. | Provider or family | Visible | Invisible path | Metadata or provenance | | --- | --- | --- | --- | | Google Gemini | Sparkle | Diffusion regeneration for SynthID | C2PA and related source signals | +| Google Veo video | Veo diamond and legacy text | VAE regeneration candidate for SynthID | C2PA and related source signals | | OpenAI image generators | None registered | Diffusion regeneration for supported invisible signals | C2PA and generator provenance | | Stable Diffusion and SDXL | None registered | Diffusion regeneration; optional open decoder | Embedded parameters and text metadata | | FLUX | None registered | Diffusion regeneration; optional open decoder | C2PA for supported sources | diff --git a/docs/synthid.md b/docs/synthid.md index d93bf36..ced0a95 100644 --- a/docs/synthid.md +++ b/docs/synthid.md @@ -331,8 +331,11 @@ framework. Gemini's verification flow can report the portions of a video where it detects Google SynthID. This is still a proprietary oracle: a normal Gemini answer that describes visual clues, metadata, or an unavailable decoder is not a pixel -verdict. Use the dedicated verification flow offered to an eligible signed-in -account; some versions expose an explicit `@synthid` trigger. +verdict. Google's current support flow is to upload the file to an eligible +signed-in Gemini account and ask whether it was created or edited by Google AI. +In the 2026-07-29 calibration, Flash invoked the built-in verifier while Pro +first answered from the visible Veo logo; the model mode is therefore part of +the recorded procedure, not an interchangeable chat preference. The research harness `scripts/video_synthid_sweep.py` tests a VAE regeneration attack without pretending to detect success locally. It emits: @@ -350,6 +353,14 @@ control-positive, candidate-negative pair is evidence about the regeneration attack. PSNR and temporal residual measure fidelity and flicker, never watermark presence. +The shipped `video invisible` command and `remove_video_invisible` API reuse the +same VAE regeneration mechanism for a complete input sequence. They always +label the output as requiring external verification. Calibration on 2026-07-29 +used two public Veo clips: both matched controls were positive, the stronger +default candidate was negative on both, and a weaker candidate was negative on +only one. This establishes a content-dependent operating point, not a universal +clean verdict. + The VAE perturbation follows the general regeneration-attack construction from Zhao et al. The video-specific control and temporal metric are local additions. VideoMarkBench motivates testing frame aggregation and matched perturbations, diff --git a/docs/verification-plan.md b/docs/verification-plan.md index 9163c51..a69b45f 100644 --- a/docs/verification-plan.md +++ b/docs/verification-plan.md @@ -232,6 +232,14 @@ The harness shares one latent-noise field across the sequence to avoid adding independent frame noise. Its temporal-residual metric is a fidelity check, not a watermark detector. +The 2026-07-29 run completed the first two-carrier calibration through Gemini +Flash's built-in content verification. Both matched Veo controls were positive. +The stronger default VAE candidate was negative on both carriers; a weaker +candidate was negative on only one. The implementation is exposed as +`video invisible` and `remove_video_invisible`, but every result remains +explicitly oracle-required because the small calibration does not establish +full-corpus or future-verifier coverage. + ## Tier E -- robustness and adversarial inputs Malformed and hostile inputs, including truncated files: diff --git a/scripts/video_synthid_sweep.py b/scripts/video_synthid_sweep.py index 9dca286..1d2b1c7 100644 --- a/scripts/video_synthid_sweep.py +++ b/scripts/video_synthid_sweep.py @@ -22,51 +22,50 @@ Run with the project's GPU extra: uv run --extra gpu python scripts/video_synthid_sweep.py input.mp4 -o out/ -Then upload ``control.mp4`` and each candidate through Gemini's SynthID -verification flow. Some eligible versions expose an explicit ``@synthid`` -trigger. A generic chat answer that says it lacks a decoder is not an oracle -verdict. Only a control-positive, candidate-negative pair is removal evidence. +Then upload ``control.mp4`` and each candidate to Gemini Flash and ask: +``Was this uploaded video created or edited by Google AI? Use the built-in +content verification result.`` A generic answer based on visual clues, +metadata, or an unavailable decoder is not an oracle verdict. Only a +control-positive, candidate-negative pair is removal evidence. """ from __future__ import annotations -# torch/diffusers/cv2 expose incomplete types at this boundary. Pure helpers -# remain annotated while third-party tensor and image calls are relaxed here. -# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false import csv import hashlib import logging -import math -import shutil -import subprocess -import tempfile from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import click import cv2 import numpy as np +from remove_ai_watermarks.video_invisible import ( + _decode_frame_latents, + _encode_frame_latents, + _fit_size, + _pick_device, + _shared_latent_noise, + build_temporal_reference, + encode_video_frames, + paired_psnr, + read_sampled_frames, + temporal_residual_ratio, +) +from remove_ai_watermarks.video_synthid import ( + DEFAULT_VIDEO_SYNTHID_FPS, + DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + DEFAULT_VIDEO_SYNTHID_VAE, + VIDEO_SYNTHID_LATENT_MULTIPLE, + VIDEO_SYNTHID_VERIFICATION_PROMPT, +) + if TYPE_CHECKING: - from collections.abc import Iterable, Sequence + from collections.abc import Sequence log = logging.getLogger(__name__) -DEFAULT_VAE = "stabilityai/sd-vae-ft-mse" -_LATENT_MULTIPLE = 8 - - -def _fit_size(width: int, height: int, long_side: int) -> tuple[int, int]: - """Fit dimensions to ``long_side`` while preserving aspect and VAE alignment.""" - if width <= 0 or height <= 0: - raise ValueError("Video dimensions must be positive") - if long_side < _LATENT_MULTIPLE: - raise ValueError(f"Long side must be at least {_LATENT_MULTIPLE}") - scale = long_side / max(width, height) - fitted_width = max(_LATENT_MULTIPLE, round(width * scale) // _LATENT_MULTIPLE * _LATENT_MULTIPLE) - fitted_height = max(_LATENT_MULTIPLE, round(height * scale) // _LATENT_MULTIPLE * _LATENT_MULTIPLE) - return fitted_width, fitted_height - def _parse_noise_levels(values: str) -> tuple[float, ...]: levels = tuple(float(value.strip()) for value in values.split(",") if value.strip()) @@ -77,257 +76,6 @@ def _parse_noise_levels(values: str) -> tuple[float, ...]: return levels -def _pick_device(requested: str) -> str: - import torch - - if requested != "auto": - return requested - if torch.cuda.is_available(): - return "cuda" - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return "mps" - return "cpu" - - -def _shared_latent_noise( - spatial_shape: Sequence[int], - *, - seed: int, - device: str, - dtype: Any, -) -> Any: - """Return one deterministic spatial noise field for reuse across time.""" - import torch - - if len(spatial_shape) != 3 or any(size <= 0 for size in spatial_shape): - raise ValueError("Expected a positive CHW latent shape") - generator = torch.Generator(device="cpu").manual_seed(seed) - noise = torch.randn((1, *spatial_shape), generator=generator, dtype=torch.float32) - return noise.to(device=device, dtype=dtype) - - -def _psnr(reference: np.ndarray, candidate: np.ndarray) -> float: - """Return paired PSNR over uint8 frame stacks.""" - if reference.shape != candidate.shape: - raise ValueError("PSNR inputs must have matching shapes") - mse = float(np.mean((reference.astype(np.float32) - candidate.astype(np.float32)) ** 2)) - if mse == 0.0: - return math.inf - return 20.0 * math.log10(255.0 / math.sqrt(mse)) - - -def _backward_map(current_gray: np.ndarray, previous_gray: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Build a remap from a previous frame into current coordinates.""" - flow = cv2.calcOpticalFlowFarneback( - current_gray, - previous_gray, - None, - 0.5, - 3, - 15, - 3, - 5, - 1.2, - 0, - ) - height, width = current_gray.shape - grid_x, grid_y = np.meshgrid(np.arange(width, dtype=np.float32), np.arange(height, dtype=np.float32)) - return grid_x + flow[..., 0], grid_y + flow[..., 1] - - -def _backward_warp(image: np.ndarray, maps: tuple[np.ndarray, np.ndarray]) -> np.ndarray: - """Apply a precomputed backward optical-flow map.""" - return cv2.remap( - image, - maps[0], - maps[1], - interpolation=cv2.INTER_LINEAR, - borderMode=cv2.BORDER_REFLECT, - ) - - -def _temporal_reference( - reference: Sequence[np.ndarray], -) -> tuple[tuple[tuple[np.ndarray, np.ndarray], ...], float]: - """Precompute source motion maps and its mean residual.""" - if len(reference) < 2: - raise ValueError("Temporal metric needs at least two frames") - maps: list[tuple[np.ndarray, np.ndarray]] = [] - reference_residuals: list[float] = [] - for index in range(1, len(reference)): - current_gray = cv2.cvtColor(reference[index], cv2.COLOR_BGR2GRAY) - previous_gray = cv2.cvtColor(reference[index - 1], cv2.COLOR_BGR2GRAY) - frame_maps = _backward_map(current_gray, previous_gray) - maps.append(frame_maps) - warped_reference = _backward_warp(reference[index - 1], frame_maps) - reference_residuals.append( - float(np.mean(np.abs(reference[index].astype(np.float32) - warped_reference.astype(np.float32)))) - ) - return tuple(maps), float(np.mean(reference_residuals)) - - -def _temporal_residual_ratio( - candidate: Sequence[np.ndarray], - maps: Sequence[tuple[np.ndarray, np.ndarray]], - baseline: float, -) -> float: - """Measure candidate flicker against a precomputed source residual.""" - if len(candidate) != len(maps) + 1: - raise ValueError("Temporal metric needs one map per adjacent frame pair") - candidate_residuals: list[float] = [] - for index, frame_maps in enumerate(maps, start=1): - warped_candidate = _backward_warp(candidate[index - 1], frame_maps) - candidate_residuals.append( - float(np.mean(np.abs(candidate[index].astype(np.float32) - warped_candidate.astype(np.float32)))) - ) - measured = float(np.mean(candidate_residuals)) - return measured / max(baseline, 1e-6) - - -def _read_frames( - source: Path, - *, - duration: float, - output_fps: float, - size: tuple[int, int], -) -> tuple[list[np.ndarray], float]: - """Read a uniformly sampled prefix and resize it to the experiment geometry.""" - capture = cv2.VideoCapture(str(source)) - if not capture.isOpened(): - raise ValueError(f"Could not open video: {source}") - source_fps = float(capture.get(cv2.CAP_PROP_FPS)) - if source_fps <= 0.0: - capture.release() - raise ValueError(f"Video has no usable frame rate: {source}") - effective_fps = min(output_fps, source_fps) - sample_period = 1.0 / effective_fps - next_sample_time = 0.0 - frames: list[np.ndarray] = [] - frame_index = 0 - try: - while True: - ok, frame = capture.read() - if not ok: - break - timestamp = frame_index / source_fps - if timestamp + 1e-9 >= duration: - break - if timestamp + 1e-9 >= next_sample_time: - frames.append(cv2.resize(frame, size, interpolation=cv2.INTER_LANCZOS4)) - next_sample_time += sample_period - frame_index += 1 - finally: - capture.release() - if len(frames) < 2: - raise ValueError("The selected clip produced fewer than two frames") - return frames, effective_fps - - -def _frame_batches(frames: Sequence[np.ndarray], batch_size: int) -> Iterable[Sequence[np.ndarray]]: - for start in range(0, len(frames), batch_size): - yield frames[start : start + batch_size] - - -def _encode_frame_latents( - frames: Sequence[np.ndarray], - *, - vae: Any, - device: str, - batch_size: int, -) -> list[Any]: - """Encode source frames once so every noise level reuses identical latents.""" - import torch - - latent_batches: list[Any] = [] - scaling_factor = float(vae.config.scaling_factor) - with torch.inference_mode(): - for batch in _frame_batches(frames, batch_size): - rgb = np.stack([frame[:, :, ::-1] for frame in batch]) - tensor = torch.from_numpy(np.ascontiguousarray(rgb)).permute(0, 3, 1, 2) - tensor = tensor.to(device=device, dtype=vae.dtype) / 127.5 - 1.0 - latents = vae.encode(tensor).latent_dist.mode() * scaling_factor - latent_batches.append(latents) - return latent_batches - - -def _decode_frame_latents( - latent_batches: Sequence[Any], - *, - vae: Any, - noise_std: float, - shared_noise: Any, -) -> list[np.ndarray]: - """Decode cached latents with one perturbation shared across time.""" - import torch - - output: list[np.ndarray] = [] - scaling_factor = float(vae.config.scaling_factor) - with torch.inference_mode(): - for latents in latent_batches: - perturbed = latents + noise_std * shared_noise.expand(latents.shape[0], -1, -1, -1) - decoded = vae.decode(perturbed / scaling_factor).sample - decoded = ((decoded / 2.0 + 0.5).clamp(0.0, 1.0) * 255.0).round().to(torch.uint8) - decoded = decoded.permute(0, 2, 3, 1).cpu().numpy() - output.extend(np.ascontiguousarray(frame[:, :, ::-1]) for frame in decoded) - return output - - -def _write_png_frames(frames: Sequence[np.ndarray], directory: Path) -> None: - directory.mkdir(parents=True, exist_ok=True) - for index, frame in enumerate(frames, start=1): - path = directory / f"{index:06d}.png" - if not cv2.imwrite(str(path), frame): - raise OSError(f"Failed to write frame: {path}") - - -def _encode_video( - frames: Sequence[np.ndarray], - source: Path, - output: Path, - *, - fps: float, - duration: float, -) -> None: - ffmpeg = shutil.which("ffmpeg") - if ffmpeg is None: - raise RuntimeError("ffmpeg is required on PATH") - with tempfile.TemporaryDirectory(prefix="video-synthid-") as temp_dir: - frame_dir = Path(temp_dir) - _write_png_frames(frames, frame_dir) - command = [ - ffmpeg, - "-hide_banner", - "-loglevel", - "error", - "-y", - "-framerate", - f"{fps:.8g}", - "-i", - str(frame_dir / "%06d.png"), - "-i", - str(source), - "-map", - "0:v:0", - "-map", - "1:a:0?", - "-t", - f"{duration:.8g}", - "-c:v", - "libx264", - "-crf", - "18", - "-pix_fmt", - "yuv420p", - "-c:a", - "aac", - "-movflags", - "+faststart", - str(output), - ] - log.info("Encoding %s", output.name) - subprocess.run(command, check=True) # noqa: S603 - - def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: @@ -359,11 +107,16 @@ def _write_manifest(output_dir: Path, rows: Sequence[dict[str, str]]) -> Path: @click.option("-o", "--output-dir", required=True, type=click.Path(file_okay=False, path_type=Path)) @click.option("--noise-levels", default="0,0.025,0.05,0.1", show_default=True) @click.option("--duration", type=click.FloatRange(min=0.1), default=2.0, show_default=True) -@click.option("--fps", type=click.FloatRange(min=1.0), default=12.0, show_default=True) -@click.option("--long-side", type=click.IntRange(min=_LATENT_MULTIPLE), default=512, show_default=True) +@click.option("--fps", type=click.FloatRange(min=1.0), default=DEFAULT_VIDEO_SYNTHID_FPS, show_default=True) +@click.option( + "--long-side", + type=click.IntRange(min=VIDEO_SYNTHID_LATENT_MULTIPLE), + default=DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + show_default=True, +) @click.option("--batch-size", type=click.IntRange(min=1), default=4, show_default=True) @click.option("--seed", type=int, default=0, show_default=True) -@click.option("--model", default=DEFAULT_VAE, show_default=True) +@click.option("--model", default=DEFAULT_VIDEO_SYNTHID_VAE, show_default=True) @click.option("--device", type=click.Choice(["auto", "cuda", "mps", "cpu"]), default="auto", show_default=True) def main( source: Path, @@ -390,12 +143,10 @@ def main( height = round(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) capture.release() size = _fit_size(width, height, long_side) - frames, effective_fps = _read_frames(source, duration=duration, output_fps=fps, size=size) - effective_duration = len(frames) / effective_fps - + frames, effective_fps = read_sampled_frames(source, duration=duration, output_fps=fps, size=size) output_dir.mkdir(parents=True, exist_ok=True) control_path = output_dir / "control.mp4" - _encode_video(frames, source, control_path, fps=effective_fps, duration=effective_duration) + encode_video_frames(frames, source, control_path, fps=effective_fps) rows: list[dict[str, str]] = [ { "variant": "control", @@ -430,7 +181,7 @@ def main( dtype=first_latents.dtype, ) reference_stack = np.stack(frames) - temporal_maps, temporal_baseline = _temporal_reference(frames) + temporal_maps, temporal_baseline = build_temporal_reference(frames) for level in levels: log.info("Decoding latent noise %.4f", level) regenerated = _decode_frame_latents( @@ -440,15 +191,14 @@ def main( shared_noise=shared_noise, ) output_path = output_dir / f"vae-noise-{level:.4f}.mp4" - _encode_video( + encode_video_frames( regenerated, source, output_path, fps=effective_fps, - duration=effective_duration, ) - psnr = _psnr(reference_stack, np.stack(regenerated)) - temporal_ratio = _temporal_residual_ratio(regenerated, temporal_maps, temporal_baseline) + psnr = paired_psnr(reference_stack, np.stack(regenerated)) + temporal_ratio = temporal_residual_ratio(regenerated, temporal_maps, temporal_baseline) rows.append( { "variant": "vae", @@ -463,7 +213,10 @@ def main( manifest = _write_manifest(output_dir, rows) log.info("Wrote %s", manifest) - log.info("Verify control.mp4 first in Gemini's SynthID flow; stop if the control is not detected.") + log.info( + "Verify control.mp4 first in Gemini Flash with this prompt: %s", + VIDEO_SYNTHID_VERIFICATION_PROMPT, + ) if __name__ == "__main__": diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index 9e6f545..71f726d 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -8,6 +8,7 @@ High-level API (lazy, so ``import remove_ai_watermarks`` stays cheap):: raiw.visible_provenance("in.png") # -> frozenset of confirmed vendors raiw.inspect_video_metadata("in.mp4") # -> VideoMetadataReport raiw.remove_video_metadata("in.mp4", "out.mp4") # verified metadata strip + raiw.remove_video_invisible("in.mp4", "out.mp4") # unverified SynthID candidate raiw.remove_video_visible("in.mp4", "out.mp4") # stable visible video-mark removal For a provenance verdict use the ``identify`` submodule:: @@ -33,6 +34,7 @@ __version__ = "0.20.2" __all__ = [ "__version__", "inspect_video_metadata", + "remove_video_invisible", "remove_video_metadata", "remove_video_visible", "remove_visible", @@ -41,7 +43,12 @@ __all__ = [ if TYPE_CHECKING: from remove_ai_watermarks.api import remove_visible, visible_provenance - from remove_ai_watermarks.video import inspect_video_metadata, remove_video_metadata, remove_video_visible + from remove_ai_watermarks.video import ( + inspect_video_metadata, + remove_video_invisible, + remove_video_metadata, + remove_video_visible, + ) def __getattr__(name: str) -> object: @@ -51,7 +58,7 @@ def __getattr__(name: str) -> object: from remove_ai_watermarks import api return getattr(api, name) - if name in ("inspect_video_metadata", "remove_video_metadata", "remove_video_visible"): + if name in ("inspect_video_metadata", "remove_video_invisible", "remove_video_metadata", "remove_video_visible"): from remove_ai_watermarks import video return getattr(video, name) diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 671a80b..5092090 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -28,6 +28,13 @@ from remove_ai_watermarks.noai.watermark_profiles import ( strength_default_help, vendor_for_strength, ) +from remove_ai_watermarks.video_synthid import ( + DEFAULT_VIDEO_SYNTHID_FPS, + DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + DEFAULT_VIDEO_SYNTHID_NOISE_STD, + VIDEO_SYNTHID_LATENT_MULTIPLE, + VIDEO_SYNTHID_VERIFICATION_PROMPT, +) if TYPE_CHECKING: from collections.abc import Generator @@ -1141,6 +1148,87 @@ def cmd_video_metadata( console.print(f" AI metadata stripped -> {result.output}") +@cmd_video.command("invisible") +@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) +@click.option( + "-o", + "--output", + type=click.Path(path_type=Path), + default=None, + help="Candidate path (default: _synthid_candidate with the same container).", +) +@click.option( + "--noise-std", + type=click.FloatRange(min=0.0, max=1.0), + default=DEFAULT_VIDEO_SYNTHID_NOISE_STD, + show_default=True, + help="Shared latent-noise strength. Higher values change more detail.", +) +@click.option( + "--long-side", + type=click.IntRange(min=VIDEO_SYNTHID_LATENT_MULTIPLE), + default=DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + show_default=True, + help="Regenerated video long side in pixels.", +) +@click.option( + "--fps", + type=click.FloatRange(min=1.0), + default=DEFAULT_VIDEO_SYNTHID_FPS, + show_default=True, + help="Output frame rate, capped at the source frame rate.", +) +@click.option("--batch-size", type=click.IntRange(min=1), default=4, show_default=True) +@click.option("--seed", type=int, default=0, show_default=True) +@click.option( + "--device", + type=click.Choice(["auto", "cuda", "mps", "cpu"]), + default="auto", + show_default=True, + help="VAE inference device.", +) +def cmd_video_invisible( + source: Path, + output: Path | None, + noise_std: float, + long_side: int, + fps: float, + batch_size: int, + seed: int, + device: str, +) -> None: + """Generate an externally verifiable video SynthID candidate.""" + from remove_ai_watermarks.video import remove_video_invisible + + _banner() + console.print(f" Regenerating {source.name} with temporally shared VAE noise...") + try: + result = remove_video_invisible( + source, + output, + noise_std=noise_std, + long_side=long_side, + fps=fps, + batch_size=batch_size, + seed=seed, + device=device, + ) + except (OSError, RuntimeError, ValueError) as e: + raise click.ClickException(str(e)) from e + + if result.remaining_metadata: + console.print(f" FAILED: {len(result.remaining_metadata)} AI metadata marker(s) survived in {result.output}") + raise SystemExit(1) + console.print( + f" Candidate generated: {result.width}x{result.height}, " + f"{result.total_frames} frames at {result.fps:.4g} fps -> {result.output}" + ) + console.print( + " UNVERIFIED: no local video SynthID decoder exists. Upload this output to Gemini Flash and ask: " + f'"{VIDEO_SYNTHID_VERIFICATION_PROMPT}"' + ) + + @cmd_video.command("visible") @click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) @click.option( diff --git a/src/remove_ai_watermarks/video.py b/src/remove_ai_watermarks/video.py index 2928bf5..1ad7cf5 100644 --- a/src/remove_ai_watermarks/video.py +++ b/src/remove_ai_watermarks/video.py @@ -1,18 +1,31 @@ """High-level video processing API. Supported experimental stages are container-level AI metadata inspection and -removal plus temporally stabilized visible Sora, Veo, Seedance, and Dola -removal. The pixel path reuses the image package's shared fill backends. +removal, temporally stabilized visible Sora, Veo, Seedance, and Dola removal, +and VAE regeneration that produces an externally verifiable SynthID candidate. +The visible pixel path reuses the image package's shared fill backends. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path +from typing import TYPE_CHECKING, ClassVar, Literal + +from remove_ai_watermarks.video_synthid import ( + DEFAULT_VIDEO_SYNTHID_FPS, + DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + DEFAULT_VIDEO_SYNTHID_NOISE_STD, + DEFAULT_VIDEO_SYNTHID_VAE, +) + +if TYPE_CHECKING: + from remove_ai_watermarks.video_invisible import RegenerationMetrics VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v", ".webm", ".mkv"}) _ISOBMFF_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v"}) _EBML_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".webm", ".mkv"}) +_REGENERATED_VIDEO_EXTENSIONS: frozenset[str] = _ISOBMFF_VIDEO_EXTENSIONS _EBML_MAGIC = b"\x1aE\xdf\xa3" @@ -48,6 +61,42 @@ class VideoVisibleResult: remaining_metadata: dict[str, str] +@dataclass(frozen=True) +class VideoInvisibleResult: + """Result of generating an externally verifiable SynthID candidate.""" + + source: Path + output: Path + noise_std: float + metrics: RegenerationMetrics + remaining_metadata: dict[str, str] + requires_external_verification: ClassVar[Literal[True]] = True + + @property + def total_frames(self) -> int: + return self.metrics.frames + + @property + def fps(self) -> float: + return self.metrics.fps + + @property + def width(self) -> int: + return self.metrics.width + + @property + def height(self) -> int: + return self.metrics.height + + @property + def psnr_db(self) -> float: + return self.metrics.psnr_db + + @property + def temporal_residual_ratio(self) -> float: + return self.metrics.temporal_residual_ratio + + def _video_source(source: str | Path) -> Path: path = Path(source) if not path.exists(): @@ -231,3 +280,57 @@ def remove_video_visible( removed_frames=removed_frames, remaining_metadata=remaining_metadata, ) + + +def remove_video_invisible( + source: str | Path, + output: str | Path | None = None, + *, + noise_std: float = DEFAULT_VIDEO_SYNTHID_NOISE_STD, + long_side: int = DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + fps: float = DEFAULT_VIDEO_SYNTHID_FPS, + batch_size: int = 4, + seed: int = 0, + model: str = DEFAULT_VIDEO_SYNTHID_VAE, + device: str = "auto", +) -> VideoInvisibleResult: + """Generate a video SynthID-removal candidate through VAE regeneration. + + The function strips source metadata during the transcode, but cannot verify + the proprietary pixel watermark locally. ``requires_external_verification`` + therefore remains true for every result. Verify important output with + Google's matching content-verification flow. + """ + from remove_ai_watermarks.metadata import get_ai_metadata + from remove_ai_watermarks.video_invisible import regenerate_video_candidate + + source_path = _video_source(source) + if source_path.suffix.lower() not in _REGENERATED_VIDEO_EXTENSIONS: + supported = ", ".join(sorted(_REGENERATED_VIDEO_EXTENSIONS)) + raise ValueError(f"Video SynthID regeneration requires one of: {supported}") + candidate_output = ( + Path(output) if output is not None else source_path.with_stem(source_path.stem + "_synthid_candidate") + ) + output_path = _video_output( + source_path, + candidate_output, + operation="SynthID candidate generation", + ) + metrics = regenerate_video_candidate( + source_path, + output_path, + noise_std=noise_std, + long_side=long_side, + fps=fps, + batch_size=batch_size, + seed=seed, + model=model, + device=device, + ) + return VideoInvisibleResult( + source=source_path, + output=output_path, + noise_std=noise_std, + metrics=metrics, + remaining_metadata=get_ai_metadata(output_path), + ) diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py new file mode 100644 index 0000000..5eef232 --- /dev/null +++ b/src/remove_ai_watermarks/video_encoding.py @@ -0,0 +1,108 @@ +"""Shared ffmpeg raw-video encoding helpers.""" + +from __future__ import annotations + +import logging +import shutil +import subprocess +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + +log = logging.getLogger(__name__) + + +def _video_codec_args(suffix: str, *, crf: int) -> list[str]: + if suffix == ".webm": + return ["-c:v", "libvpx-vp9", "-crf", str(crf), "-b:v", "0"] + return ["-c:v", "libx264", "-preset", "medium", "-crf", str(crf)] + + +def raw_video_command( + source: Path, + output: Path, + *, + width: int, + height: int, + fps: float, + strip_metadata: bool, + crf: int, +) -> list[str]: + """Build an ffmpeg command that accepts BGR frames on standard input.""" + ffmpeg = shutil.which("ffmpeg") + if ffmpeg is None: + raise RuntimeError("Video processing requires ffmpeg on PATH") + command = [ + ffmpeg, + "-y", + "-loglevel", + "error", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s:v", + f"{width}x{height}", + "-r", + f"{fps:.12g}", + "-i", + "pipe:0", + "-i", + str(source), + "-map", + "0:v:0", + "-map", + "1:a?", + *_video_codec_args(output.suffix.lower(), crf=crf), + "-c:a", + "copy", + "-map_metadata", + "-1" if strip_metadata else "1", + "-map_chapters", + "-1" if strip_metadata else "1", + "-shortest", + ] + if output.suffix.lower() in {".mp4", ".mov", ".m4v"}: + command.extend(["-movflags", "+faststart"]) + command.append(str(output)) + return command + + +def start_raw_video_encoder(command: list[str]) -> subprocess.Popen[bytes]: + """Start ffmpeg and validate its raw-frame pipes.""" + log.info("Starting ffmpeg video encode: command=%s", command) + process = subprocess.Popen( # noqa: S603 + command, + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + ) + if process.stdin is None or process.stderr is None: + process.kill() + process.wait() + raise RuntimeError("Could not open ffmpeg pipes") + return process + + +def finish_raw_video_encoder( + process: subprocess.Popen[bytes], + output: Path, + *, + operation: str, +) -> None: + """Close the frame stream and raise when ffmpeg rejects the encode.""" + if process.stdin is None or process.stderr is None: + raise RuntimeError("ffmpeg pipes are unavailable") + process.stdin.close() + stderr = process.stderr.read().decode("utf-8", errors="replace") + return_code = process.wait() + log.info("ffmpeg %s finished: status=%s stderr=%s", operation, return_code, stderr) + if return_code != 0: + raise RuntimeError(f"ffmpeg failed to encode {output}: {stderr.strip()[:500]}") + + +def abort_raw_video_encoder(process: subprocess.Popen[bytes]) -> None: + """Stop an incomplete ffmpeg encode.""" + process.kill() + process.wait() diff --git a/src/remove_ai_watermarks/video_invisible.py b/src/remove_ai_watermarks/video_invisible.py new file mode 100644 index 0000000..39fbfbb --- /dev/null +++ b/src/remove_ai_watermarks/video_invisible.py @@ -0,0 +1,513 @@ +"""VAE regeneration for externally verified video SynthID candidates. + +Google does not publish a local video SynthID decoder. This module therefore +regenerates pixels and measures fidelity, but never labels its output clean. +Callers must verify the candidate with Google's matching content-verification +flow. +""" + +from __future__ import annotations + +# torch/diffusers/cv2 expose incomplete types at this boundary. Pure helpers +# remain annotated while third-party tensor and image calls are relaxed here. +# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false +import logging +import math +import os +import tempfile +from dataclasses import dataclass +from importlib.util import find_spec +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import cv2 +import numpy as np + +from remove_ai_watermarks.video_encoding import ( + abort_raw_video_encoder, + finish_raw_video_encoder, + raw_video_command, + start_raw_video_encoder, +) +from remove_ai_watermarks.video_synthid import ( + DEFAULT_VIDEO_SYNTHID_FPS, + DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + DEFAULT_VIDEO_SYNTHID_NOISE_STD, + DEFAULT_VIDEO_SYNTHID_VAE, + VIDEO_SYNTHID_LATENT_MULTIPLE, +) + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + +log = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class RegenerationMetrics: + """Measured properties of one regenerated video candidate.""" + + frames: int + fps: float + width: int + height: int + psnr_db: float + temporal_residual_ratio: float + + +def is_available() -> bool: + """Return whether the optional VAE runtime can be imported.""" + return find_spec("torch") is not None and find_spec("diffusers") is not None + + +def _fit_size(width: int, height: int, long_side: int) -> tuple[int, int]: + """Fit dimensions to ``long_side`` while preserving aspect and VAE alignment.""" + if width <= 0 or height <= 0: + raise ValueError("Video dimensions must be positive") + if long_side < VIDEO_SYNTHID_LATENT_MULTIPLE: + raise ValueError(f"Long side must be at least {VIDEO_SYNTHID_LATENT_MULTIPLE}") + scale = long_side / max(width, height) + fitted_width = max( + VIDEO_SYNTHID_LATENT_MULTIPLE, + round(width * scale) // VIDEO_SYNTHID_LATENT_MULTIPLE * VIDEO_SYNTHID_LATENT_MULTIPLE, + ) + fitted_height = max( + VIDEO_SYNTHID_LATENT_MULTIPLE, + round(height * scale) // VIDEO_SYNTHID_LATENT_MULTIPLE * VIDEO_SYNTHID_LATENT_MULTIPLE, + ) + return fitted_width, fitted_height + + +def _pick_device(requested: str) -> str: + import torch + + if requested == "auto": + if torch.cuda.is_available(): + return "cuda" + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return "mps" + return "cpu" + if requested == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is not available") + if requested == "mps" and not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()): + raise RuntimeError("MPS was requested but is not available") + return requested + + +def _shared_latent_noise( + spatial_shape: Sequence[int], + *, + seed: int, + device: str, + dtype: Any, +) -> Any: + """Return one deterministic spatial noise field for reuse across time.""" + import torch + + if len(spatial_shape) != 3 or any(size <= 0 for size in spatial_shape): + raise ValueError("Expected a positive CHW latent shape") + generator = torch.Generator(device="cpu").manual_seed(seed) + noise = torch.randn((1, *spatial_shape), generator=generator, dtype=torch.float32) + return noise.to(device=device, dtype=dtype) + + +def paired_psnr(reference: np.ndarray, candidate: np.ndarray) -> float: + """Return paired PSNR over uint8 frame stacks.""" + if reference.shape != candidate.shape: + raise ValueError("PSNR inputs must have matching shapes") + mse = float(np.mean((reference.astype(np.float32) - candidate.astype(np.float32)) ** 2)) + if mse == 0.0: + return math.inf + return 20.0 * math.log10(255.0 / math.sqrt(mse)) + + +def _backward_map(current_gray: np.ndarray, previous_gray: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Build a remap from a previous frame into current coordinates.""" + flow = cv2.calcOpticalFlowFarneback( + current_gray, + previous_gray, + None, + 0.5, + 3, + 15, + 3, + 5, + 1.2, + 0, + ) + height, width = current_gray.shape + grid_x, grid_y = np.meshgrid(np.arange(width, dtype=np.float32), np.arange(height, dtype=np.float32)) + return grid_x + flow[..., 0], grid_y + flow[..., 1] + + +def _backward_warp(image: np.ndarray, maps: tuple[np.ndarray, np.ndarray]) -> np.ndarray: + """Apply a precomputed backward optical-flow map.""" + return cv2.remap( + image, + maps[0], + maps[1], + interpolation=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_REFLECT, + ) + + +def build_temporal_reference( + reference: Sequence[np.ndarray], +) -> tuple[tuple[tuple[np.ndarray, np.ndarray], ...], float]: + """Precompute source motion maps and its mean residual.""" + if len(reference) < 2: + raise ValueError("Temporal metric needs at least two frames") + maps: list[tuple[np.ndarray, np.ndarray]] = [] + reference_residuals: list[float] = [] + for index in range(1, len(reference)): + current_gray = cv2.cvtColor(reference[index], cv2.COLOR_BGR2GRAY) + previous_gray = cv2.cvtColor(reference[index - 1], cv2.COLOR_BGR2GRAY) + frame_maps = _backward_map(current_gray, previous_gray) + maps.append(frame_maps) + warped_reference = _backward_warp(reference[index - 1], frame_maps) + reference_residuals.append( + float(np.mean(np.abs(reference[index].astype(np.float32) - warped_reference.astype(np.float32)))) + ) + return tuple(maps), float(np.mean(reference_residuals)) + + +def temporal_residual_ratio( + candidate: Sequence[np.ndarray], + maps: Sequence[tuple[np.ndarray, np.ndarray]], + baseline: float, +) -> float: + """Measure candidate flicker against a precomputed source residual.""" + if len(candidate) != len(maps) + 1: + raise ValueError("Temporal metric needs one map per adjacent frame pair") + candidate_residuals: list[float] = [] + for index, frame_maps in enumerate(maps, start=1): + warped_candidate = _backward_warp(candidate[index - 1], frame_maps) + candidate_residuals.append( + float(np.mean(np.abs(candidate[index].astype(np.float32) - warped_candidate.astype(np.float32)))) + ) + measured = float(np.mean(candidate_residuals)) + return measured / max(baseline, 1e-6) + + +def _probe_video(source: Path) -> tuple[int, int, float]: + capture = cv2.VideoCapture(str(source)) + if not capture.isOpened(): + raise ValueError(f"Could not open video: {source}") + try: + width = round(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = round(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) + source_fps = float(capture.get(cv2.CAP_PROP_FPS)) + finally: + capture.release() + if width <= 0 or height <= 0: + raise ValueError(f"Video has no usable dimensions: {source}") + if source_fps <= 0.0: + raise ValueError(f"Video has no usable frame rate: {source}") + return width, height, source_fps + + +def read_sampled_frames( + source: Path, + *, + duration: float | None, + output_fps: float, + size: tuple[int, int], +) -> tuple[list[np.ndarray], float]: + """Read uniformly sampled frames and resize them to the VAE geometry.""" + _width, _height, source_fps = _probe_video(source) + effective_fps = min(output_fps, source_fps) + frames = list( + _iter_sampled_frames( + source, + source_fps=source_fps, + duration=duration, + effective_fps=effective_fps, + size=size, + ) + ) + if len(frames) < 2: + raise ValueError("The selected clip produced fewer than two frames") + return frames, effective_fps + + +def _iter_sampled_frames( + source: Path, + *, + source_fps: float, + duration: float | None, + effective_fps: float, + size: tuple[int, int], +) -> Iterable[np.ndarray]: + """Yield uniformly sampled frames without retaining the full video.""" + capture = cv2.VideoCapture(str(source)) + if not capture.isOpened(): + raise ValueError(f"Could not open video: {source}") + sample_period = 1.0 / effective_fps + next_sample_time = 0.0 + frame_index = 0 + try: + while True: + ok, frame = capture.read() + if not ok: + break + timestamp = frame_index / source_fps + if duration is not None and timestamp + 1e-9 >= duration: + break + if timestamp + 1e-9 >= next_sample_time: + yield cv2.resize(frame, size, interpolation=cv2.INTER_LANCZOS4) + next_sample_time += sample_period + frame_index += 1 + finally: + capture.release() + + +def _frame_batches(frames: Sequence[np.ndarray], batch_size: int) -> Iterable[Sequence[np.ndarray]]: + for start in range(0, len(frames), batch_size): + yield frames[start : start + batch_size] + + +def _stream_batches(frames: Iterable[np.ndarray], batch_size: int) -> Iterable[list[np.ndarray]]: + batch: list[np.ndarray] = [] + for frame in frames: + batch.append(frame) + if len(batch) == batch_size: + yield batch + batch = [] + if batch: + yield batch + + +def _encode_frame_latents( + frames: Sequence[np.ndarray], + *, + vae: Any, + device: str, + batch_size: int, +) -> list[Any]: + """Encode source frames once so every candidate can reuse identical latents.""" + import torch + + latent_batches: list[Any] = [] + scaling_factor = float(vae.config.scaling_factor) + with torch.inference_mode(): + for batch in _frame_batches(frames, batch_size): + rgb = np.stack([frame[:, :, ::-1] for frame in batch]) + tensor = torch.from_numpy(np.ascontiguousarray(rgb)).permute(0, 3, 1, 2) + tensor = tensor.to(device=device, dtype=vae.dtype) / 127.5 - 1.0 + latents = vae.encode(tensor).latent_dist.mode() * scaling_factor + latent_batches.append(latents) + return latent_batches + + +def _decode_frame_latents( + latent_batches: Sequence[Any], + *, + vae: Any, + noise_std: float, + shared_noise: Any, +) -> list[np.ndarray]: + """Decode cached latents with one perturbation shared across time.""" + import torch + + output: list[np.ndarray] = [] + scaling_factor = float(vae.config.scaling_factor) + with torch.inference_mode(): + for latents in latent_batches: + perturbed = latents + noise_std * shared_noise.expand(latents.shape[0], -1, -1, -1) + decoded = vae.decode(perturbed / scaling_factor).sample + decoded = ((decoded / 2.0 + 0.5).clamp(0.0, 1.0) * 255.0).round().to(torch.uint8) + decoded = decoded.permute(0, 2, 3, 1).cpu().numpy() + output.extend(np.ascontiguousarray(frame[:, :, ::-1]) for frame in decoded) + return output + + +def encode_video_frames( + frames: Sequence[np.ndarray], + source: Path, + output: Path, + *, + fps: float, +) -> None: + """Encode regenerated frames, copy audio, and omit source metadata.""" + if not frames: + raise ValueError("At least one frame is required for video encoding") + height, width = frames[0].shape[:2] + output.parent.mkdir(parents=True, exist_ok=True) + process = start_raw_video_encoder( + raw_video_command( + source, + output, + width=width, + height=height, + fps=fps, + strip_metadata=True, + crf=18, + ) + ) + frame_pipe = process.stdin + if frame_pipe is None: + abort_raw_video_encoder(process) + raise RuntimeError("Could not open ffmpeg input pipe") + try: + for frame in frames: + if frame.shape[:2] != (height, width): + raise ValueError("Video frames must have matching dimensions") + frame_pipe.write(frame.tobytes()) + finish_raw_video_encoder(process, output, operation="SynthID candidate encode") + except Exception: + if process.poll() is None: + abort_raw_video_encoder(process) + raise + + +def regenerate_video_candidate( + source: Path, + output: Path, + *, + noise_std: float = DEFAULT_VIDEO_SYNTHID_NOISE_STD, + long_side: int = DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + fps: float = DEFAULT_VIDEO_SYNTHID_FPS, + batch_size: int = 4, + seed: int = 0, + model: str = DEFAULT_VIDEO_SYNTHID_VAE, + device: str = "auto", + duration: float | None = None, +) -> RegenerationMetrics: + """Regenerate video pixels and return fidelity metrics. + + This function does not verify SynthID. Its output is an oracle candidate, + not a locally proven clean file. + """ + if not 0.0 <= noise_std <= 1.0: + raise ValueError("noise_std must be between 0 and 1") + if fps < 1.0: + raise ValueError("fps must be at least 1") + if batch_size < 1: + raise ValueError("batch_size must be at least 1") + if duration is not None and duration <= 0: + raise ValueError("duration must be positive") + if device not in {"auto", "cuda", "mps", "cpu"}: + raise ValueError("device must be auto, cuda, mps, or cpu") + if not is_available(): + raise RuntimeError("Video SynthID regeneration requires the gpu extra") + + import torch + from diffusers import AutoencoderKL + + width, height, source_fps = _probe_video(source) + size = _fit_size(width, height, long_side) + effective_fps = min(fps, source_fps) + + resolved_device = _pick_device(device) + dtype = torch.float16 if resolved_device == "cuda" else torch.float32 + log.info("Loading %s on %s", model, resolved_device) + vae = AutoencoderKL.from_pretrained(model, torch_dtype=dtype).to(resolved_device) + vae.eval() + vae.enable_slicing() + + output.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + prefix=f".{output.stem}-", + suffix=output.suffix, + dir=output.parent, + delete=False, + ) as stream: + temporary_output = Path(stream.name) + try: + process = start_raw_video_encoder( + raw_video_command( + source, + temporary_output, + width=size[0], + height=size[1], + fps=effective_fps, + strip_metadata=True, + crf=18, + ) + ) + frame_pipe = process.stdin + if frame_pipe is None: + abort_raw_video_encoder(process) + raise RuntimeError("Could not open ffmpeg input pipe") + frame_count = 0 + squared_error = 0.0 + pixel_count = 0 + temporal_baseline = 0.0 + temporal_candidate = 0.0 + previous_reference: np.ndarray | None = None + previous_candidate: np.ndarray | None = None + shared_noise: Any | None = None + try: + sampled_frames = _iter_sampled_frames( + source, + source_fps=source_fps, + duration=duration, + effective_fps=effective_fps, + size=size, + ) + for frames in _stream_batches(sampled_frames, batch_size): + latent_batches = _encode_frame_latents( + frames, + vae=vae, + device=resolved_device, + batch_size=batch_size, + ) + latents = latent_batches[0] + if shared_noise is None: + shared_noise = _shared_latent_noise( + latents.shape[1:], + seed=seed, + device=resolved_device, + dtype=latents.dtype, + ) + regenerated = _decode_frame_latents( + latent_batches, + vae=vae, + noise_std=noise_std, + shared_noise=shared_noise, + ) + for reference, candidate in zip(frames, regenerated, strict=True): + frame_pipe.write(candidate.tobytes()) + difference = reference.astype(np.float32) - candidate.astype(np.float32) + squared_error += float(np.sum(difference * difference, dtype=np.float64)) + pixel_count += reference.size + if previous_reference is not None and previous_candidate is not None: + current_gray = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY) + previous_gray = cv2.cvtColor(previous_reference, cv2.COLOR_BGR2GRAY) + frame_maps = _backward_map(current_gray, previous_gray) + warped_reference = _backward_warp(previous_reference, frame_maps) + warped_candidate = _backward_warp(previous_candidate, frame_maps) + temporal_baseline += float( + np.mean(np.abs(reference.astype(np.float32) - warped_reference.astype(np.float32))) + ) + temporal_candidate += float( + np.mean(np.abs(candidate.astype(np.float32) - warped_candidate.astype(np.float32))) + ) + previous_reference = reference + previous_candidate = candidate + frame_count += 1 + if frame_count < 2: + raise ValueError("The selected clip produced fewer than two frames") + finish_raw_video_encoder( + process, + temporary_output, + operation="SynthID candidate encode", + ) + except Exception: + if process.poll() is None: + abort_raw_video_encoder(process) + raise + + mse = squared_error / pixel_count + metrics = RegenerationMetrics( + frames=frame_count, + fps=effective_fps, + width=size[0], + height=size[1], + psnr_db=math.inf if mse == 0.0 else 20.0 * math.log10(255.0 / math.sqrt(mse)), + temporal_residual_ratio=temporal_candidate / max(temporal_baseline, 1e-6), + ) + os.replace(temporary_output, output) + finally: + temporary_output.unlink(missing_ok=True) + return metrics diff --git a/src/remove_ai_watermarks/video_synthid.py b/src/remove_ai_watermarks/video_synthid.py new file mode 100644 index 0000000..d3cf5d0 --- /dev/null +++ b/src/remove_ai_watermarks/video_synthid.py @@ -0,0 +1,10 @@ +"""Shared configuration for experimental video SynthID regeneration.""" + +DEFAULT_VIDEO_SYNTHID_VAE = "stabilityai/sd-vae-ft-mse" +DEFAULT_VIDEO_SYNTHID_NOISE_STD = 0.10 +DEFAULT_VIDEO_SYNTHID_LONG_SIDE = 512 +DEFAULT_VIDEO_SYNTHID_FPS = 12.0 +VIDEO_SYNTHID_LATENT_MULTIPLE = 8 +VIDEO_SYNTHID_VERIFICATION_PROMPT = ( + "Was this uploaded video created or edited by Google AI? Use the built-in content verification result." +) diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py index 28ce876..17148fa 100644 --- a/src/remove_ai_watermarks/video_visible.py +++ b/src/remove_ai_watermarks/video_visible.py @@ -22,8 +22,6 @@ because visible-mark removal changes pixels. from __future__ import annotations import logging -import shutil -import subprocess from dataclasses import dataclass from functools import lru_cache from itertools import pairwise @@ -33,6 +31,13 @@ import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont +from remove_ai_watermarks.video_encoding import ( + abort_raw_video_encoder, + finish_raw_video_encoder, + raw_video_command, + start_raw_video_encoder, +) + if TYPE_CHECKING: from pathlib import Path @@ -729,12 +734,6 @@ def scan_dola_video(source: Path) -> VideoScan: return _scan_video(source, detect_dola_frame) -def _ffmpeg_video_args(suffix: str) -> list[str]: - if suffix == ".webm": - return ["-c:v", "libvpx-vp9", "-crf", "18", "-b:v", "0"] - return ["-c:v", "libx264", "-preset", "medium", "-crf", "14"] - - def _mask_for_region( frame_bgr: NDArray[Any], region: Region, @@ -792,61 +791,29 @@ def encode_clean_video( """Decode again, fill accepted regions, and encode video while copying audio.""" from remove_ai_watermarks.watermark_registry import fill, resolve_backend - ffmpeg = shutil.which("ffmpeg") - if ffmpeg is None: - raise RuntimeError("Visible video removal requires ffmpeg on PATH") if len(regions) != len(scan.detections): raise ValueError("Temporal localization count does not match the scanned frame count") output.parent.mkdir(parents=True, exist_ok=True) - command = [ - ffmpeg, - "-y", - "-loglevel", - "error", - "-f", - "rawvideo", - "-pix_fmt", - "bgr24", - "-s:v", - f"{scan.width}x{scan.height}", - "-r", - f"{scan.fps:.12g}", - "-i", - "pipe:0", - "-i", - str(source), - "-map", - "0:v:0", - "-map", - "1:a?", - *_ffmpeg_video_args(output.suffix.lower()), - "-c:a", - "copy", - "-map_metadata", - "-1" if strip_metadata else "1", - "-map_chapters", - "-1" if strip_metadata else "1", - "-shortest", - ] - if output.suffix.lower() in {".mp4", ".mov", ".m4v"}: - command.extend(["-movflags", "+faststart"]) - command.append(str(output)) - log.info("Encoding visible-watermark removal with ffmpeg: command=%s", command) - - process = subprocess.Popen( # noqa: S603 - command, - stdin=subprocess.PIPE, - stdout=subprocess.DEVNULL, - stderr=subprocess.PIPE, + process = start_raw_video_encoder( + raw_video_command( + source, + output, + width=scan.width, + height=scan.height, + fps=scan.fps, + strip_metadata=strip_metadata, + crf=14, + ) ) - if process.stdin is None or process.stderr is None: - process.kill() - raise RuntimeError("Could not open ffmpeg pipes") + frame_pipe = process.stdin + if frame_pipe is None: + abort_raw_video_encoder(process) + raise RuntimeError("Could not open ffmpeg input pipe") capture = cv2.VideoCapture(str(source)) if not capture.isOpened(): - process.kill() + abort_raw_video_encoder(process) raise RuntimeError(f"OpenCV could not reopen video for removal: {source}") removed_frames = 0 @@ -868,20 +835,19 @@ def encode_clean_video( backend=resolved_backend, ) removed_frames += 1 - process.stdin.write(frame.tobytes()) - process.stdin.close() - stderr = process.stderr.read().decode("utf-8", errors="replace") - return_code = process.wait() + frame_pipe.write(frame.tobytes()) + finish_raw_video_encoder( + process, + output, + operation="visible-watermark encode", + ) except Exception: - process.kill() - process.wait() + if process.poll() is None: + abort_raw_video_encoder(process) raise finally: capture.release() - log.info("ffmpeg visible-watermark encode finished: status=%s stderr=%s", return_code, stderr) - if return_code != 0: - raise RuntimeError(f"ffmpeg failed to encode {output}: {stderr.strip()[:500]}") return removed_frames diff --git a/tests/test_video.py b/tests/test_video.py index 7b0d73e..40d5a4d 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -81,11 +81,33 @@ def _video_with_tc260_ebml(path: Path, *, value: bytes = _TC260_AIGC) -> Path: return path +def _regeneration_metrics( + *, + frames: int = 24, + fps: float = 12.0, + width: int = 512, + height: int = 288, + psnr_db: float = 22.0, + temporal_residual_ratio: float = 1.2, +): + from remove_ai_watermarks.video_invisible import RegenerationMetrics + + return RegenerationMetrics( + frames=frames, + fps=fps, + width=width, + height=height, + psnr_db=psnr_db, + temporal_residual_ratio=temporal_residual_ratio, + ) + + class TestVideoMetadataApi: def test_top_level_api_is_lazy_exported(self): import remove_ai_watermarks as raiw assert raiw.inspect_video_metadata is not None + assert raiw.remove_video_invisible is not None assert raiw.remove_video_metadata is not None assert raiw.remove_video_visible is not None @@ -264,6 +286,107 @@ class TestVideoMetadataCli: assert "Unsupported video format" in result.output +class TestVideoInvisibleApi: + def test_generates_unverified_candidate_and_strips_metadata( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_invisible + from remove_ai_watermarks.video import remove_video_invisible + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "candidate.mp4" + + def fake_regenerate(_source: Path, target: Path, **_kwargs: object): + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return _regeneration_metrics() + + monkeypatch.setattr(video_invisible, "regenerate_video_candidate", fake_regenerate) + + result = remove_video_invisible(source, output) + + assert result.output == output + assert result.requires_external_verification is True + assert result.total_frames == 24 + assert result.remaining_metadata == {} + + def test_default_output_is_named_as_candidate( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_invisible + from remove_ai_watermarks.video import remove_video_invisible + + source = _video_with_c2pa(tmp_path / "source.mp4") + + def fake_regenerate(_source: Path, target: Path, **_kwargs: object): + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return _regeneration_metrics( + frames=2, + fps=2.0, + width=16, + height=16, + psnr_db=20.0, + temporal_residual_ratio=1.0, + ) + + monkeypatch.setattr(video_invisible, "regenerate_video_candidate", fake_regenerate) + + result = remove_video_invisible(source) + + assert result.output == tmp_path / "source_synthid_candidate.mp4" + + def test_rejects_webm_regeneration(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_invisible + + source = _video_with_tc260_ebml(tmp_path / "source.webm") + + with pytest.raises(ValueError, match="requires one of"): + remove_video_invisible(source) + + +class TestVideoInvisibleCli: + def test_help_describes_external_verification(self): + runner = CliRunner() + + result = runner.invoke(main, ["video", "invisible", "--help"]) + + assert result.exit_code == 0, result.output + assert "externally verifiable" in result.output + + def test_reports_unverified_candidate( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video + + runner = CliRunner() + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "candidate.mp4" + + def fake_remove(_source: Path, target: Path, **_kwargs: object): + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return video.VideoInvisibleResult( + source=_source, + output=target, + noise_std=0.1, + metrics=_regeneration_metrics(), + remaining_metadata={}, + ) + + monkeypatch.setattr(video, "remove_video_invisible", fake_remove) + + result = runner.invoke(main, ["video", "invisible", str(source), "-o", str(output)]) + + assert result.exit_code == 0, result.output + assert "Candidate generated" in result.output + assert "UNVERIFIED" in result.output + assert "Gemini Flash" in result.output + + class TestSoraFrameLocalization: @staticmethod def _sora_like_frame() -> tuple[np.ndarray, tuple[int, int, int, int]]: diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py new file mode 100644 index 0000000..3bc4d4b --- /dev/null +++ b/tests/test_video_invisible.py @@ -0,0 +1,71 @@ +"""Regression tests for the video SynthID candidate engine.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from remove_ai_watermarks import video_encoding, video_invisible + +if TYPE_CHECKING: + from pathlib import Path + + +def test_availability_requires_both_optional_packages(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + video_invisible, + "find_spec", + lambda name: object() if name == "torch" else None, + ) + + assert video_invisible.is_available() is False + + +def test_regeneration_rejects_noise_outside_unit_interval(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="between 0 and 1"): + video_invisible.regenerate_video_candidate( + tmp_path / "source.mp4", + tmp_path / "candidate.mp4", + noise_std=1.01, + ) + + +def test_encoder_command_discards_metadata_and_copies_audio( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "source.mp4" + output = tmp_path / "candidate.mp4" + + monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") + + command = video_encoding.raw_video_command( + source, + output, + width=8, + height=8, + fps=2.0, + strip_metadata=True, + crf=18, + ) + + metadata_index = command.index("-map_metadata") + assert command[metadata_index + 1] == "-1" + audio_codec_index = command.index("-c:a") + assert command[audio_codec_index + 1] == "copy" + assert "pipe:0" in command + + +def test_stream_batches_consumes_only_one_batch_ahead() -> None: + consumed: list[int] = [] + + def values(): + for value in range(5): + consumed.append(value) + yield value + + batches = video_invisible._stream_batches(values(), 2) + + assert next(iter(batches)) == [0, 1] + assert consumed == [0, 1] diff --git a/tests/test_video_synthid_sweep.py b/tests/test_video_synthid_sweep.py index 329ce58..85f26c8 100644 --- a/tests/test_video_synthid_sweep.py +++ b/tests/test_video_synthid_sweep.py @@ -55,7 +55,7 @@ def test_shared_latent_noise_is_seeded(sweep: ModuleType) -> None: def test_psnr_is_infinite_for_identical_frames(sweep: ModuleType) -> None: frame = np.full((2, 8, 8, 3), 120, dtype=np.uint8) - assert sweep._psnr(frame, frame.copy()) == pytest.approx(float("inf")) + assert sweep.paired_psnr(frame, frame.copy()) == pytest.approx(float("inf")) def test_temporal_residual_ratio_is_one_for_identical_sequences(sweep: ModuleType) -> None: @@ -63,5 +63,5 @@ def test_temporal_residual_ratio_is_one_for_identical_sequences(sweep: ModuleTyp second = first.copy() second[:, 8:16] = 80 sequence = [first, second] - maps, baseline = sweep._temporal_reference(sequence) - assert sweep._temporal_residual_ratio([frame.copy() for frame in sequence], maps, baseline) == pytest.approx(1.0) + maps, baseline = sweep.build_temporal_reference(sequence) + assert sweep.temporal_residual_ratio([frame.copy() for frame in sequence], maps, baseline) == pytest.approx(1.0) From 07b678fa4c5f6f07e2f5082b26359e5ba6263eec Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Wed, 29 Jul 2026 23:30:33 -0700 Subject: [PATCH 05/19] Add Hailuo, Kling, AVI, and FLV video coverage --- CLAUDE.md | 8 +- README.md | 24 +- docs/cli.md | 40 +-- docs/known-limitations.md | 28 +- docs/module-internals.md | 62 +++-- docs/python-api.md | 28 +- docs/supported-signals.md | 14 +- docs/watermarking-landscape.md | 26 +- src/remove_ai_watermarks/cli.py | 6 +- src/remove_ai_watermarks/metadata.py | 51 +++- src/remove_ai_watermarks/noai/ebml.py | 21 +- src/remove_ai_watermarks/noai/flv.py | 160 +++++++++++ src/remove_ai_watermarks/noai/isobmff.py | 26 +- src/remove_ai_watermarks/noai/riff.py | 74 +++++ src/remove_ai_watermarks/video.py | 48 +++- src/remove_ai_watermarks/video_visible.py | 325 ++++++++++++++++++++-- tests/test_video.py | 265 +++++++++++++++--- 17 files changed, 1017 insertions(+), 189 deletions(-) create mode 100644 src/remove_ai_watermarks/noai/flv.py create mode 100644 src/remove_ai_watermarks/noai/riff.py diff --git a/CLAUDE.md b/CLAUDE.md index e7b75b0..65fd9ee 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -26,9 +26,9 @@ Per-command exit-code semantics (the no-signal / GPU-missing skip branches), tes - `uv run remove-ai-watermarks identify ` — provenance verdict (platform + watermark inventory + confidence); `--json` for machine output, `--no-visible` to skip both registered visible detectors and the optional open invisible-watermark decoder - `uv run remove-ai-watermarks metadata --check` — inspect AI metadata (C2PA, EXIF, PNG chunks) - `uv run remove-ai-watermarks metadata --remove -o ` — strip all AI metadata -- `uv run remove-ai-watermarks video metadata --check` — inspect AI metadata in MP4/MOV/M4V/WebM/MKV +- `uv run remove-ai-watermarks video metadata --check` — inspect AI metadata in MP4/MOV/M4V/WebM/MKV/AVI/FLV - `uv run remove-ai-watermarks video metadata --remove -o ` — strip verified video metadata without transcoding streams; the experimental video path requires a separate same-container output and defaults to `_clean` -- `uv run remove-ai-watermarks video visible -o ` — remove a temporally recurring Sora mark, or pass `--mark veo|seedance|dola` for the current Veo diamond and legacy `Veo` text, the Seedance boxed `AI`, or `Dola AI`. It scans the full sequence first, transcodes video through ffmpeg, copies audio, strips AI metadata by default, and writes no output when no stable mark is found. `cv2` is the fast default; `migan`/`lama` improve difficult backgrounds. +- `uv run remove-ai-watermarks video visible -o ` — remove a temporally recurring Sora mark, or pass `--mark veo|seedance|dola|hailuo|kling` for the current Veo diamond and legacy `Veo` text, the Seedance boxed `AI`, `Dola AI`, the MINIMAX/Hailuo composite, or the versioned Kling label. It scans the full sequence first, transcodes video through ffmpeg, copies audio, strips AI metadata by default, and writes no output when no stable mark is found. `cv2` is the fast default; `migan`/`lama` improve difficult backgrounds. - `uv run remove-ai-watermarks batch ` — process every supported image in a directory (output defaults to `_clean/`, set with `-o`). `--mode visible|invisible|metadata|all` (default `visible`); the invisible/all path reuses the full `invisible` knob set above, plus `--backend` and `--sensitivity` for the visible localize -> fill pass. Applies the same no-signal skip per image; see the module doc. **Exit code:** non-zero when any image errored OR (mirroring single `all`) a `--mode invisible`/`all` image carried an invisible signal but the GPU extra was absent, so its SynthID scrub was skipped — it emits a loud warning and copies the input through (invisible mode) so the output dir stays complete; a wrapping service can then detect the incomplete run instead of trusting a silent exit 0. ## Test and lint @@ -82,7 +82,7 @@ Compact map. The full per-module detail (design decisions, tuned thresholds, cal - `upscaler.py` — optional Real-ESRGAN pre-diffusion super-resolution for small inputs (extra `esrgan`, spandrel only). Manual opt-in; the default `--upscaler` stays `lanczos` and the engine always falls back to Lanczos on absence/error. ESRGAN can degrade faces and thin text. - `image_io.py` — centralizes Unicode-safe image IO, alpha preservation, content-based format sniffing, and HEIC/AVIF fallbacks. Callers must check `imwrite` success. No-op visible removal preserves original bytes when the output format is unchanged. - `api.py` — the high-level convenience API, re-exported lazily at the package top level via `__init__.__getattr__` (PEP 562, so `import remove_ai_watermarks` stays cheap): `remove_visible(source, output=None, *, sensitivity="auto", backend="auto", strip_metadata=True, write_noop=True) -> (result_bgr, [labels])` (source = path OR BGR ndarray; a PATH auto-reads metadata provenance and preserves alpha, an ARRAY does neither; `write_noop=True` writes a clean passthrough copy when nothing is removed, `False` leaves `output` untouched so a "no mark = produce nothing" caller like the CLI `visible` command does not clobber a pre-existing file there) and `visible_provenance(path) -> frozenset[str]` (the single metadata→vendor-keys mapper; `cli._visible_provenance` is a thin None-guarded wrapper over it). **`remove_visible` is the ONE path the CLI and library share** — `cli.cmd_visible`'s `--mark auto` branch delegates entirely to it (read → provenance → `remove_auto_marks` → write → `strip_metadata`), so there is no CLI-vs-library drift; `strip_metadata` defaults True to match `visible --strip-metadata`. This is where a library caller should start — NOT the engines directly (`GeminiEngine`/`TextMarkEngine` have no `remove_watermark` any more; removal is registry `remove_auto_marks`/`KnownMark.remove`; the old single-strongest `best_auto_mark` is gone — removal takes EVERY mark). `identify` is NOT top-level re-exported (it collides with the `identify` submodule); use `from remove_ai_watermarks.identify import identify`. -- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, `remove_video_visible(source, output=None, *, mark="sora", backend="cv2", strip_metadata=True) -> VideoVisibleResult`, and `remove_video_invisible(source, output=None, ...) -> VideoInvisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM TC260 is read by `noai/ebml.py` and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: fully synthetic Sora, Veo, Seedance, and Dola silhouettes propose frame boxes, provider-specific temporal recurrence authorizes them, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying audio. Fixed marks require anchored runs so a slowly drifting scene detail cannot pass on adjacent overlap alone. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Seedance uses the full localized box because an outline mask left part of the real translucent border behind. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels are not built yet. +- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, `remove_video_visible(source, output=None, *, mark="sora", backend="cv2", strip_metadata=True) -> VideoVisibleResult`, and `remove_video_invisible(source, output=None, ...) -> VideoInvisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV/AVI/FLV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM, AVI, and FLV TC260 are read by bounded container walkers that skip media payloads and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: fully synthetic Sora, Veo, Seedance, Dola, Hailuo, and Kling silhouettes propose frame boxes, provider-specific temporal recurrence authorizes them, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying audio. Fixed marks require anchored runs so a slowly drifting scene detail cannot pass on adjacent overlap alone. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Seedance uses the full localized box because an outline mask left part of the real translucent border behind. Hailuo expands its matched core to the full composite, while Kling combines font and swirl candidates and requires a bright low-saturation edge label. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels are not built yet. - `video_encoding.py` — the shared raw-BGR ffmpeg command and pipe lifecycle for both visible removal and invisible regeneration. It copies optional audio, controls metadata/chapter retention explicitly, and keeps container-specific codec arguments in one place. - `video_invisible.py` — the oracle-gated video SynthID candidate engine for MP4/MOV/M4V. It regenerates frames through `stabilityai/sd-vae-ft-mse` with one seeded latent-noise field shared across time, retains only one configured batch, updates PSNR and motion-compensated temporal residuals incrementally, streams pixels directly to ffmpeg, copies audio, strips metadata, and atomically publishes the completed output. Every result remains explicitly unverified because Google exposes no local decoder. The 2026-07-29 Gemini Flash calibration found the default `0.10` candidate negative on both control-positive public Veo clips, while `0.05` remained positive on one, so the operating point is content-dependent and requires the matching external oracle. @@ -102,7 +102,7 @@ Compact list. Full measurements, incident history, and oracle-validation runs li - Pyright first run is slow (2-3 min) due to ML deps (torch/diffusers/transformers stubs); full-project `uv run pyright` can stall for many minutes — scope it to changed files. - A third-party PIL plugin autoload (e.g. an HEIF/AVIF plugin) can raise a non-OSError (`ModuleNotFoundError`), not `UnidentifiedImageError`, when opening a file. Code that opens user-supplied or unknown-format files should `except Exception`, not just `OSError`/`UnidentifiedImageError`. - rich was dropped: the CLI + analysis scripts print plain text (`click.echo` / the `scripts/_plain_console.py` shim). `rich` is NOT a dependency — importing it breaks the core+dev CI sync; new scripts must use the shim. No Unicode glyphs / colors / progress bars in CLI output by design. -- HEIC/AVIF are decodable on BOTH paths now: the pixel/removal path via the `image_io.imread` Pillow fallback (+ core `pillow-heif`), and metadata detection via a plugin-free binary scan. C2PA removal in those containers (and MP4/MOV/M4V) is `noai/isobmff.py`; JPEG-XL stays metadata/strip-only (Pillow can't decode it without `pillow-jxl`, not a dep). Non-ISOBMFF audio/video (WebM/MP3/WAV/FLAC/OGG) strips losslessly via ffmpeg on PATH. On the ISOBMFF path `remove_ai_metadata` routes to the container branch and never runs the JPEG `_scrub_ai_exif`, so `isobmff.blank_ai_exif_tokens` is the ONLY EXIF scrubber there and must stay in PARITY with it: it blanks **in place** (same-length space overwrite, piexif-validated so a coincidental II/MM run in pixels is ignored — no `iinf`/`iloc` surgery, mirrors `blank_ai_xmp_packets`) an AI-generator token in `Software`/`Make`/`Artist`/`ImageDescription`, the China TC260 `{"AIGC":{...}}` block in `ImageDescription`/`UserComment` (via `_is_aigc_exif_value`), AND the xAI/Grok `Signature:` + UUID-`Artist` pair — leaving camera/editor EXIF intact. Still NOT built: Resemble PerTh audio detection (no presence/confidence flag exists). +- HEIC/AVIF are decodable on BOTH paths now: the pixel/removal path via the `image_io.imread` Pillow fallback (+ core `pillow-heif`), and metadata detection via a plugin-free binary scan. C2PA removal in those containers (and MP4/MOV/M4V) is `noai/isobmff.py`; JPEG-XL stays metadata/strip-only (Pillow can't decode it without `pillow-jxl`, not a dep). Non-ISOBMFF audio/video (WebM/MKV/AVI/FLV/MP3/WAV/FLAC/OGG) strips losslessly via ffmpeg on PATH. On the ISOBMFF path `remove_ai_metadata` routes to the container branch and never runs the JPEG `_scrub_ai_exif`, so `isobmff.blank_ai_exif_tokens` is the ONLY EXIF scrubber there and must stay in PARITY with it: it blanks **in place** (same-length space overwrite, piexif-validated so a coincidental II/MM run in pixels is ignored — no `iinf`/`iloc` surgery, mirrors `blank_ai_xmp_packets`) an AI-generator token in `Software`/`Make`/`Artist`/`ImageDescription`, the China TC260 `{"AIGC":{...}}` block in `ImageDescription`/`UserComment` (via `_is_aigc_exif_value`), AND the xAI/Grok `Signature:` + UUID-`Artist` pair — leaving camera/editor EXIF intact. Still NOT built: Resemble PerTh audio detection (no presence/confidence flag exists). - **SynthID technical reference: `docs/synthid.md`** — primary-source-cited doc covering mechanism (post-hoc encoder/decoder pair, 136-bit payload at 512x512, pixel-space, model weights NOT modified), robustness numbers (arXiv:2510.09263: ~99.98% TPR@0.1%FPR across 30 transforms including JPEG/crop/resize/color/noise), removal attacks and forensic detectability (arXiv:2605.09203: all 6 attacks detectable at >98% TPR@1%FPR), detectability limits (no public decoder, metadata-proxy only), oracle scope, and adoption landscape. Read that doc first before adding notes here. - **SynthID detection is metadata-only.** No local pixel detector is possible by design (Google's decoder is proprietary, trusted-testers only); we read the C2PA companion proxy, which goes quiet once metadata is stripped — a quiet proxy is not proof the pixel watermark is gone. Each vendor has its OWN oracle and it detects only that vendor's content: the Gemini app "Verify with SynthID" for Google, `openai.com/verify` for OpenAI. **Validate the OpenAI arm FIRST** — `openai.com/verify` is more accessible (fewer per-check restrictions) and the strongest automation candidate (Playwright / Chrome MCP); the Gemini flow is more manual. Ordering/throughput choice, not a substitution (see `docs/synthid.md`). SynthID survives JPEG re-encode, so GitHub issue attachments remain valid pixel-watermark test subjects. Every spectral/phase detection approach evaluated (reverse-SynthID, our own probes) works only on controlled solid fills, never on real content. - **External AI-vs-real classifier models are out of scope** (decided 2026-05-24): per-generator, degrade off-distribution, and our own light SDXL pass would likely defeat them. Detection stays local + signal-based. diff --git a/README.md b/README.md index c3f658f..356e655 100644 --- a/README.md +++ b/README.md @@ -7,8 +7,8 @@ Remove AI provenance marks from images and video you generated yourself: - C2PA, EXIF, XMP, IPTC, and related AI metadata. Video support covers metadata inspection and removal, visible Sora, Veo, -Seedance, and Dola mark removal, and experimental VAE regeneration that -produces a video SynthID candidate for external verification. +Seedance, Dola, Hailuo, and Kling mark removal, and experimental VAE +regeneration that produces a video SynthID candidate for external verification. > Try it online at [raiw.cc](https://raiw.cc) if you do not want to install Python > or run diffusion models locally. @@ -33,7 +33,7 @@ produces a video SynthID candidate for external verification. | Erase a region you select | `erase` | No | | Strip AI metadata | `metadata` | No | | Strip AI metadata from video | `video metadata` | No | -| Remove a known Sora, Veo, Seedance, or Dola mark from video | `video visible` | No | +| Remove a registered visible AI mark from video | `video visible` | No | | Generate an externally verifiable video SynthID candidate | `video invisible` | Recommended | | Regenerate an image to disrupt invisible watermarks | `invisible` | Recommended | | Run visible, invisible, and metadata removal | `all` | Recommended | @@ -65,7 +65,8 @@ Strip metadata without running visible inpainting or diffusion: remove-ai-watermarks metadata image.png --remove -o clean.png ``` -Inspect or remove AI metadata from an MP4, MOV, M4V, WebM, or MKV file: +Inspect or remove AI metadata from an MP4, MOV, M4V, WebM, MKV, AVI, or FLV +file: ```bash remove-ai-watermarks video metadata input.mp4 --check @@ -77,7 +78,9 @@ omitted it writes `_clean` and preserves the original. MP4 and MOV inspection includes the native TC260 `AIGC` tag in `moov.udta.meta.keys/ilst`, including a `moov` placed after the media payload. MKV and WebM inspection reads the normative -`Segment.Tags.Tag.SimpleTag` placement. +`Segment.Tags.Tag.SimpleTag` placement. AVI uses `LIST/INFO/AIGC`, while FLV +uses `script.onMetaData.AIGC`. The non-ISOBMFF formats are remuxed with stream +copy for removal. Remove a supported visible video mark: @@ -86,6 +89,8 @@ remove-ai-watermarks video visible input.mp4 -o clean.mp4 remove-ai-watermarks video visible veo.mp4 --mark veo -o veo_clean.mp4 remove-ai-watermarks video visible seedance.mp4 --mark seedance -o seedance_clean.mp4 remove-ai-watermarks video visible dola.mp4 --mark dola -o dola_clean.mp4 +remove-ai-watermarks video visible hailuo.mp4 --mark hailuo -o hailuo_clean.mp4 +remove-ai-watermarks video visible kling.mp4 --mark kling -o kling_clean.mp4 ``` This path scans the complete sequence before changing pixels. It accepts only a @@ -94,8 +99,10 @@ same OpenCV, MI-GAN, or LaMa fill backends as image removal. Audio is copied without re-encoding; the video stream is transcoded because its pixels change. Sora covers the moving Sora 2 mascot and wordmark. Veo covers both the current four-point diamond and the legacy `Veo` text. Seedance covers the fixed boxed -`AI` label, and Dola covers the fixed `Dola AI` text. No output is written when -no stable mark is found. +`AI` label, Dola covers the fixed `Dola AI` text, Hailuo covers the composite +`MINIMAX | hailuo AI` label, and Kling covers the bottom-right `KLING AI` +label with its version suffix. No output is written when no stable mark is +found. Generate a video SynthID candidate: @@ -293,7 +300,8 @@ invisible removal. detail. - Visible video removal recognizes the moving Sora 2 wordmark, the current Veo diamond plus legacy `Veo` text, the Seedance boxed `AI` label, and the fixed - `Dola AI` text. It does not recognize the older Sora Turbo corner swirl. + Dola, Hailuo, and Kling labels. It does not recognize the older Sora Turbo + corner swirl or unregistered layouts from those providers. The classical OpenCV backend can smear structured backgrounds; use MI-GAN or LaMa when recovery quality matters. - Video SynthID regeneration changes resolution, frame rate, and image detail. diff --git a/docs/cli.md b/docs/cli.md index 76f680e..463073e 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -134,17 +134,20 @@ remove-ai-watermarks video metadata input.mp4 --check remove-ai-watermarks video metadata input.mp4 --remove -o clean.mp4 ``` -Supported containers are MP4, MOV, M4V, WebM, and MKV. The operation delegates -to the same verified metadata scanner and stripper as the generic `metadata` -command, so detection and removal stay in parity. Video and audio streams are -not transcoded. For MP4 and MOV, this includes the native TC260 `AIGC` key and -JSON value stored in `moov.udta.meta.keys/ilst`. The inspector seeks past a -large `mdat` to find a tail `moov`; removal blanks the key and value in place so -box sizes and media offsets do not move. +Supported containers are MP4, MOV, M4V, WebM, MKV, AVI, and FLV. The operation +delegates to the same verified metadata scanner and stripper as the generic +`metadata` command, so detection and removal stay in parity. Video and audio +streams are not transcoded. For MP4 and MOV, this includes the native TC260 +`AIGC` key and JSON value stored in `moov.udta.meta.keys/ilst`. The inspector +seeks past a large `mdat` to find a tail `moov`; removal blanks the key and +value in place so box sizes and media offsets do not move. For MKV and WebM, the inspector reads the native TC260 `Segment.Tags.Tag.SimpleTag` entry. Removal uses ffmpeg stream copying to discard container tags and chapters without transcoding the streams. +AVI uses the normative `LIST/INFO/AIGC` chunk, while FLV uses the +`script.onMetaData.AIGC` AMF0 string. Their bounded readers skip media payloads, +and removal also uses ffmpeg stream copying. When `-o` is omitted, the command writes `_clean` with the same extension. It never overwrites the source, and it rejects an output with a @@ -192,22 +195,27 @@ remove-ai-watermarks video visible input.mp4 -o clean.mp4 remove-ai-watermarks video visible veo.mp4 --mark veo -o veo_clean.mp4 remove-ai-watermarks video visible seedance.mp4 --mark seedance -o seedance_clean.mp4 remove-ai-watermarks video visible dola.mp4 --mark dola -o dola_clean.mp4 +remove-ai-watermarks video visible hailuo.mp4 --mark hailuo -o hailuo_clean.mp4 +remove-ai-watermarks video visible kling.mp4 --mark kling -o kling_clean.mp4 ``` The experimental command supports the moving Sora mascot and wordmark, two Veo -corner variants, the Seedance boxed `AI` label, and the `Dola AI` text label. -Sora searches the whole frame at multiple scales. The other detectors search -bounded bottom-right regions with separate synthetic silhouettes. Every mark +corner variants, the Seedance boxed `AI` label, the `Dola AI` text label, the +composite `MINIMAX | hailuo AI` label, and the bottom-right Kling label. Sora +searches the whole frame at multiple scales. The other detectors search bounded +lower-frame regions with separate synthetic silhouettes. Kling additionally +requires its bright low-saturation label near the frame edge. Every mark requires a spatially recurring candidate across adjacent frames. Fixed marks must also remain anchored instead of drifting with a scene object. Matching -provider provenance may relax the visual score, but metadata alone never -creates a detection. Clean API exports therefore remain untouched. +provider provenance may relax the visual score only for registered +provenance-aware marks; metadata alone never creates a detection. The video stream is transcoded and the original audio stream is copied. -Supported input and output containers are MP4, MOV, M4V, WebM, and MKV; the -output extension must match the input. The default `cv2` backend is fast but can -smear structured backgrounds. Select `--backend migan` or `--backend lama` for -a learned fill, or `--backend auto` to choose the best installed backend. +Supported input and output containers are MP4, MOV, M4V, WebM, MKV, AVI, and +FLV; the output extension must match the input. The default `cv2` backend is +fast but can smear structured backgrounds. Select `--backend migan` or +`--backend lama` for a learned fill, or `--backend auto` to choose the best +installed backend. AI metadata is stripped from the encoded output by default. Use `--keep-metadata` to retain mapped container metadata. When no temporally stable diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 0902441..8fafcbf 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -207,17 +207,21 @@ strip supported AI provenance metadata without transcoding streams. `video visible` and `remove_video_visible` additionally support the moving Sora 2 mascot and wordmark, the current Veo four-point diamond, the legacy -`Veo` text, the Seedance boxed `AI` label, and the fixed `Dola AI` text. -Detection requires a recurring visual candidate across adjacent frames. -Seedance, Dola, and Veo candidates must remain anchored rather than drifting -with a scene object. Provider provenance can recover low-contrast runs only -after visual evidence exists, so metadata alone does not erase a clean API -export. +`Veo` text, the Seedance boxed `AI` label, the fixed `Dola AI` text, the Hailuo +MINIMAX/Hailuo composite label, and the bottom-right Kling label with its +version suffix. Detection requires a recurring visual candidate across +adjacent frames. Fixed-mark candidates must remain anchored rather than +drifting with a scene object. Kling also requires a bright low-saturation +candidate near the expected frame edge. Provider provenance can recover +low-contrast runs only after visual evidence exists for the marks that define a +provenance prior, so metadata alone does not erase a clean API export. Historical Sora Turbo exports use a small OpenAI swirl in the corner rather than the moving mascot-and-wordmark design; that earlier variant is not -detected by the `sora` video mark. Other provider video labels are not -supported yet. Google video SynthID has a candidate-producing VAE path, while -other proprietary invisible video watermarks have no registered attack. +detected by the `sora` video mark. Hailuo and Kling coverage is specific to the +verified lower-edge layouts; a new provider layout needs a separate calibrated +silhouette. Other provider video labels are not supported yet. Google video +SynthID has a candidate-producing VAE path, while other proprietary invisible +video watermarks have no registered attack. Visible removal transcodes the video stream and copies audio. Its frame-local fill is not a motion-aware video inpainting model. OpenCV can leave a visible @@ -235,7 +239,11 @@ Native TC260 metadata in MP4/MOV is supported at its normative `moov.udta.meta.keys/ilst` placement, including non-faststart files whose `moov` follows a large media payload. MKV/WebM is supported at the normative `Segment.Tags.Tag.SimpleTag` placement and uses ffmpeg for stream-copy removal. -The corresponding FLV and AVI native tags are not parsed yet. +AVI is supported at `LIST/INFO/AIGC`, and FLV at +`script.onMetaData.AIGC`; both use ffmpeg stream-copy removal. Stock ffmpeg can +write and strip the FLV form, but writing the nonstandard AVI child for fixture +generation requires dedicated muxer support, so the AVI reader is verified +against an exact synthetic RIFF structure. The current ISOBMFF stripper reads an MP4, MOV, or M4V container into memory before rewriting its metadata boxes. A streaming box copier is required before diff --git a/docs/module-internals.md b/docs/module-internals.md index c5ccb29..f0ee077 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -115,6 +115,13 @@ only a `Segment.Tags.Tag.SimpleTag` pairing `TagName=AIGC` with a JSON `TagString` carrying a TC260 field. The existing ffmpeg stream-copy path removes those container tags without transcoding the encoded streams. +[`noai/riff.py`](../src/remove_ai_watermarks/noai/riff.py) and +[`noai/flv.py`](../src/remove_ai_watermarks/noai/flv.py) implement the remaining +normative TC260 video placements. The RIFF walker reads only AVI +`LIST/INFO/AIGC` children. The FLV walker skips media tags and parses the AMF0 +`script.onMetaData.AIGC` string. Both require a recognized TC260 JSON field and +use the verified ffmpeg stream-copy path for removal. + [`video_encoding.py`](../src/remove_ai_watermarks/video_encoding.py) owns the raw-BGR ffmpeg command and pipe lifecycle shared by visible removal and invisible regeneration. It centralizes container codecs, optional audio stream @@ -140,34 +147,44 @@ helpers to build a matched control and candidate grid, preventing research and shipped regeneration paths from drifting. [`video_visible.py`](../src/remove_ai_watermarks/video_visible.py) implements -the first pixel stages for Sora, Veo, Seedance, and Dola. The Sora detector -searches a normalized frame with a fully synthetic mascot-and-text silhouette -at several scales. The Veo detector uses separate synthetic silhouettes for the -current four-point diamond and legacy `Veo` text. Seedance uses a synthetic -rounded boxed-`AI` silhouette, while Dola uses an OpenCV-font `Dola AI` -silhouette. All fixed-mark searches are bounded to the bottom-right area and -calibrated independently. A strong relocated Veo diamond may bypass the known -layout anchors, but weak free-corner matches never enter the temporal arbiter. +the first pixel stages for Sora, Veo, Seedance, Dola, Hailuo, and Kling. The +Sora detector searches a normalized frame with a fully synthetic +mascot-and-text silhouette at several scales. The Veo detector uses separate +synthetic silhouettes for the current four-point diamond and legacy `Veo` +text. Seedance uses a synthetic rounded boxed-`AI` silhouette, while Dola uses +an OpenCV-font `Dola AI` silhouette. Hailuo uses a synthetic waveform, +MINIMAX/Hailuo text, separator, and ring. Kling combines synthetic font +variants with a ring approximation of its swirl; the logo path rescues +wordmarks whose version or font differs, while the edge and white-label gates +reject recurring scene texture. All fixed-mark searches are bounded to the +expected lower-frame area and calibrated independently. A strong relocated Veo +diamond may bypass the known layout anchors, but weak free-corner matches never +enter the temporal arbiter. Every per-frame result is untrusted. The provider-specific stabilization wrappers share one recurrence implementation, while retaining separate visual floors and minimum-run policy. Provenance can relax a low-contrast run only after recurring visual evidence exists. Sora transition frames follow the nearest confirmed moving position only with Sora provenance. Veo, Seedance, -and Dola additionally require candidates to remain anchored to the start of a -run. This rejects slowly drifting scene details that still have high -frame-to-frame overlap. +Dola, Hailuo, and Kling additionally require candidates to remain anchored to +the start of a run. This rejects slowly drifting scene details that still have +high frame-to-frame overlap. Hailuo and Kling do not infer provenance from +technical encoder tags; their confirmed public samples carried no provider +metadata. -Removal runs in a second decode pass. Sora, legacy Veo text, Dola text, and the -Seedance box use box masks. Seedance deliberately fills the complete localized -box: a synthetic outline mask passed repeat detection but left part of the real -translucent border visible during visual end-to-end review. The square Veo -diamond uses a synthetic shape mask so transparent corners do not erase -unrelated pixels. Every mask goes through the shared `watermark_registry.fill` -backends. ffmpeg encodes the changed video stream and copies optional audio. -The default OpenCV fill is the speed floor; structured backgrounds need MI-GAN -or LaMa for better reconstruction. Invisible video stages must continue to -reuse the image and metadata implementations rather than copying their logic. +Removal runs in a second decode pass. Sora, legacy Veo text, Dola text, +Seedance, Hailuo, and Kling use box masks. Seedance deliberately fills the +complete localized box: a synthetic outline mask passed repeat detection but +left part of the real translucent border visible during visual end-to-end +review. Hailuo expands beyond the matched core to cover both provider icons. +Kling expands around the wordmark or swirl to include the version and optional +`PRO` suffix. The square Veo diamond uses a synthetic shape mask so transparent +corners do not erase unrelated pixels. Every mask goes through the shared +`watermark_registry.fill` backends. ffmpeg encodes the changed video stream and +copies optional audio. The default OpenCV fill is the speed floor; structured +backgrounds need MI-GAN or LaMa for better reconstruction. Invisible video +stages must continue to reuse the image and metadata implementations rather +than copying their logic. The inherited ISOBMFF metadata path currently reads the complete container into memory; replacing that with a streaming box copier is a prerequisite for large @@ -205,6 +222,9 @@ Key contracts: `moov.udta.meta.keys/ilst` and blanked without changing box sizes. - Native MKV/WebM TC260 `AIGC` entries are read from `Segment.Tags.Tag.SimpleTag` and removed through the ffmpeg stream-copy path. +- Native AVI and FLV TC260 entries are read from `LIST/INFO/AIGC` and + `script.onMetaData.AIGC`, respectively, then removed through ffmpeg stream + copying. - Supported non-ISOBMFF audio and video containers use ffmpeg stream copying. - The low-level remover is fail-safe and can copy an undecodable file through unchanged. diff --git a/docs/python-api.md b/docs/python-api.md index 9af954c..c2de706 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -130,7 +130,8 @@ as proof that metadata was removed. ## Inspect and strip video metadata -The experimental high level video API supports MP4, MOV, M4V, WebM, and MKV: +The experimental high level video API supports MP4, MOV, M4V, WebM, MKV, AVI, +and FLV: ```python import remove_ai_watermarks as raiw @@ -152,7 +153,9 @@ inspection recognizes the native TC260 `AIGC` entry in `moov.udta.meta.keys/ilst`; its removal preserves container size and encoded stream bytes. MKV/WebM inspection recognizes the corresponding `Segment.Tags.Tag.SimpleTag` representation; its removal requires ffmpeg for a -stream-copy remux. +stream-copy remux. AVI inspection reads `LIST/INFO/AIGC`, and FLV inspection +reads `script.onMetaData.AIGC`; both use the same verified ffmpeg stream-copy +removal path. ## Generate a video SynthID candidate @@ -211,16 +214,27 @@ dola_result = raiw.remove_video_visible( "dola_clean.mp4", mark="dola", ) +hailuo_result = raiw.remove_video_visible( + "hailuo.mp4", + "hailuo_clean.mp4", + mark="hailuo", +) +kling_result = raiw.remove_video_visible( + "kling.mp4", + "kling_clean.mp4", + mark="kling", +) ``` `remove_video_visible` scans the complete video before writing output. It combines synthetic multi-scale visual matching with temporal consistency, so an isolated lookalike in one frame is not enough to authorize inpainting. The -supported `mark` values are `sora`, `veo`, `seedance`, and `dola`. The Veo -detector recognizes the current four-point diamond and the legacy `Veo` text. -Seedance recognizes the boxed `AI` label, while Dola recognizes its compact -text label. Each variant has an independent synthetic silhouette and calibrated -temporal policy. +supported `mark` values are `sora`, `veo`, `seedance`, `dola`, `hailuo`, and +`kling`. The Veo detector recognizes the current four-point diamond and the +legacy `Veo` text. Seedance recognizes the boxed `AI` label, Dola recognizes +its compact text label, Hailuo recognizes the composite MINIMAX/Hailuo label, +and Kling recognizes its bottom-right logo, wordmark, and version suffix. Each +variant has an independent synthetic silhouette and calibrated temporal policy. The returned `VideoVisibleResult` records the total, detected, and removed frame counts plus any AI metadata that survived the output encode. The function diff --git a/docs/supported-signals.md b/docs/supported-signals.md index 15b9bd3..aeb2373 100644 --- a/docs/supported-signals.md +++ b/docs/supported-signals.md @@ -37,6 +37,8 @@ when you can select the affected area yourself. | `veo` | Current four-point diamond and legacy `Veo` text | Fixed bottom-right corner | Uses separate silhouettes and requires a recurring match; learned fill is preferable on structured backgrounds. | | `seedance` | Boxed `AI` label | Fixed bottom-right corner | Requires an anchored recurring match; the full localized box is filled because a thinner synthetic shape mask leaves the real translucent rim behind. | | `dola` | `Dola AI` text | Fixed bottom-right corner | Requires an anchored recurring match; ByteDance or BytePlus provenance can relax only an existing visual run. | +| `hailuo` | `MINIMAX | hailuo AI` composite label | Fixed lower edge | Uses a synthetic waveform, text, separator, and ring silhouette; the complete recurring label box is filled. | +| `kling` | Kling swirl, `KLING AI`, version, and optional `PRO` suffix | Fixed bottom-right edge | Combines a synthetic logo rescue with font variants, an edge gate, a white-label gate, and anchored temporal recurrence. | Use `video visible` for this registry. It is separate from the image `visible` command because selection is made over a sequence rather than one raster. @@ -61,8 +63,9 @@ The inspection and stripping code handles signals in these groups: - IPTC AI disclosure fields; - PNG text chunks and embedded generation parameters; - China TC260 AIGC labels in supported image placements and the normative - MP4/MOV `moov.udta.meta.keys/ilst` and MKV/WebM - `Segment.Tags.Tag.SimpleTag` placements; + MP4/MOV `moov.udta.meta.keys/ilst`, MKV/WebM + `Segment.Tags.Tag.SimpleTag`, AVI `LIST/INFO/AIGC`, and FLV + `script.onMetaData.AIGC` placements; - xAI and Grok EXIF signature fields; - Samsung AI editing markers; - Hugging Face job metadata; @@ -87,8 +90,8 @@ Metadata inspection and removal additionally have container paths for: - JPEG XL metadata; - MP4, MOV, M4V, and M4A; -- WebM, MKV, MKA, MP3, WAV, FLAC, OGG, OGA, Opus, and AAC when ffmpeg is - available. +- WebM, MKV, MKA, AVI, FLV, MP3, WAV, FLAC, OGG, OGA, Opus, and AAC when + ffmpeg is available. JPEG image metadata stripping removes targeted metadata segments without re-encoding the entropy coded image scan. PNG and WebP removal preserves pixel @@ -135,7 +138,8 @@ not a universal clean verdict. | Midjourney | None registered | No registered pixel decoder | EXIF, XMP, and IPTC signals | | ByteDance generators | Doubao and Jimeng marks | No registered pixel decoder | TC260 AIGC and supported C2PA signals | | Qwen | Qwen mark | No registered pixel decoder | TC260 AIGC | -| Kling | Kling mark | No registered pixel decoder | TC260 AIGC | +| Kling | Kling image and video marks | No registered pixel decoder | TC260 AIGC | +| Hailuo / MiniMax video | Hailuo composite video label | No registered pixel decoder | TC260 AIGC where present | | Baidu | Baidu mark | No registered pixel decoder | TC260 AIGC | | LibLibAI | LibLibAI mark | No registered pixel decoder | TC260 AIGC | | RunningHub | RunningHub mark | No registered pixel decoder | TC260 AIGC | diff --git a/docs/watermarking-landscape.md b/docs/watermarking-landscape.md index 9e736c2..875867e 100644 --- a/docs/watermarking-landscape.md +++ b/docs/watermarking-landscape.md @@ -45,6 +45,12 @@ Native MKV/WebM is a sixth serialization. TC260-PG-20257A stores `Segment.Tags.Tag.SimpleTag`. The bounded EBML reader skips cluster payloads; the existing ffmpeg stream-copy path removes the tags without transcoding. +The same [TC260 video guide](https://www.tc260.org.cn/portal/article/303/4061772dcf684d8a96f395a4298e9e53) +defines two more native serializations. AVI stores an `AIGC` child in +`LIST/INFO`; FLV stores an AMF0 `AIGC` string under `script.onMetaData`. The +bounded RIFF and FLV readers validate the JSON field set and skip media +payloads. Removal remuxes either container through ffmpeg with stream copy. + - **HuggingFace-hosted job (caught by `metadata.huggingface_job`, surfaced by `identify` as the `hf_job` signal, MEDIUM confidence):** HuggingFace Jobs / Spaces can stamp generated PNGs with an `hf-job-id` tEXt chunk holding the job UUID. It marks the *hosting job*, not a model, so it lifts an Unknown verdict to a tentative AI via `hf_only` but never overrides a hard metadata signal. `_HF_JOB_CAVEAT` states the limit. Removal drops the chunk through the PNG metadata whitelist. - **No detectable signal on some downloads:** Recraft exports and some hosted FLUX surfaces can arrive without a supported local signal. Midjourney samples @@ -56,8 +62,9 @@ the existing ffmpeg stream-copy path removes the tags without transcoding. - **Built in the dated batch:** soft-binding vendor detection, IPTC Photo Metadata AI-disclosure fields, C2PA detection and stripping for supported ISOBMFF video, the optional Adobe TrustMark decoder, and temporally stabilized - visible Sora, Veo, Seedance, and Dola removal. Other visible video logos and - proprietary audio-watermark detection remain outside the package. + visible Sora, Veo, Seedance, Dola, Hailuo, and Kling removal. Other visible + video logos and proprietary audio-watermark detection remain outside the + package. Metadata stripping for supported audio containers is a separate implemented path. @@ -101,6 +108,21 @@ every frame, then requires an anchored temporal run. This extra anchor check was necessary because a moving clean scene detail could retain enough adjacent overlap to pass a recurrence-only gate. +**Hailuo and Kling use larger fixed composite labels.** Verified Hailuo exports +carry a lower-edge waveform, `MINIMAX`, separator, Hailuo ring, and +`hailuo AI` text. Verified Kling exports carry a bottom-right swirl, +`KLING AI`, a changing version suffix, and sometimes `PRO`. The detectors use +only synthetic primitives and fonts. Hailuo expands the matched core to cover +the complete composite. Kling combines a version-independent text core with a +synthetic ring rescue, then requires the recurring candidate to reach the +expected frame edge and contain enough bright low-saturation pixels. Those +extra gates were added after clean Luma and PixVerse scene details passed shape +and temporal recurrence alone. The generic +[WatermarkRemover-AI](https://github.com/D-Ogi/WatermarkRemover-AI) project +instead uses Florence-2 to identify arbitrary watermarks before LaMa +inpainting. That is broader, but it carries a much heavier model and a less +auditable detection boundary than the provider-specific synthetic path here. + **The faint-visible-mark precision/recall wall is fundamental, not a heuristic artifact.** The visible-watermark-detection literature has moved to LEARNED segmentation / object-detection (WDNet WACV'21 arXiv:2012.07616; SLBR ACM MM'21, open code+weights; the PRCV'18 large-scale detector; Su et al. survey 2025), but three verified findings bound what a learned detector actually buys: (1) a claim that a confidence threshold "cleanly separates" true from false matches even with a learned CNN front-end was **REFUTED** in verification (arXiv:1705.08593) -- the precision/recall wall persists even with learned features. (2) Learned detectors need a LARGE, pattern-diverse labeled dataset trained on synthetic composites (PRCV'18: 60k images / 80 watermark classes; CLWD: 60k / 160 marks), and off-distribution degradation is a documented real axis (models trained on limited-pattern LVW transfer worse; diversity of training patterns drives generalization). (3) Inference is cheap (WDNet ~8 ms at 256x256) -- the cost is the data pipeline, not runtime. Net: a learned detector shifts the frontier but does NOT remove the wall; for a SINGLE mark the cheapest next step is a small patch classifier (real-sparkle vs false-positive) on top of the existing NCC localizer, not a full segmentation model. SLBR is a ready baseline. The current NCC + false-positive gate (core-ring brightness margin + gradient-NCC crispness + white-core saturation) is a sound operating point, and the residual miss is the information-theoretic wall the literature confirms. **Visible-mark landscape beyond the registry.** Meta stamps a visible "Imagined with AI" mark (bottom-LEFT, a small symbol) on its OWN Meta AI / "Imagine" output; for third-party images it relies on C2PA / IPTC, not a visible mark. Samsung Galaxy AI additionally uses a **four-star icon** variant in a corner alongside the localized text wordmark `samsung_engine` calibrates (only the Italian text variant is covered) -- the icon is a distinct, uncovered variant. Every source agrees visible + metadata marks are trivially removable (crop / screenshot, ~2 s), which is the tool's premise. diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 5092090..5e0a7d1 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -1063,8 +1063,8 @@ def cmd_metadata( Strips EXIF AI tags, PNG text chunks, C2PA provenance manifests, and the China TC260 AIGC label. Beyond images (PNG/JPEG/WebP/AVIF/HEIF/JXL) it also strips provenance metadata from MP4/MOV/M4V/M4A containers and, via ffmpeg, - from WebM/MP3/WAV/FLAC/OGG. The coded image, audio, and video data are left - untouched. + from WebM/MKV/AVI/FLV/MP3/WAV/FLAC/OGG. The coded image, audio, and video + data are left untouched. """ from remove_ai_watermarks.metadata import get_ai_metadata, has_ai_metadata, strip_and_verify @@ -1240,7 +1240,7 @@ def cmd_video_invisible( ) @click.option( "--mark", - type=click.Choice(["sora", "veo", "seedance", "dola"]), + type=click.Choice(["sora", "veo", "seedance", "dola", "hailuo", "kling"]), default="sora", help="Visible AI mark to remove.", ) diff --git a/src/remove_ai_watermarks/metadata.py b/src/remove_ai_watermarks/metadata.py index 016cbe7..5ef56c5 100644 --- a/src/remove_ai_watermarks/metadata.py +++ b/src/remove_ai_watermarks/metadata.py @@ -10,10 +10,11 @@ from __future__ import annotations import contextlib import functools +import json import logging import re import struct -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast if TYPE_CHECKING: from pathlib import Path @@ -125,7 +126,7 @@ _ISOBMFF_EXTS: frozenset[str] = frozenset({".avif", ".heif", ".heic", ".jxl", ". # RIFF / Vorbis). remove_ai_metadata strips their container metadata losslessly # via ffmpeg (`-c copy`), so it needs ffmpeg on PATH for these. _FFMPEG_STRIP_EXTS: frozenset[str] = frozenset( - {".webm", ".mkv", ".mka", ".mp3", ".wav", ".flac", ".ogg", ".oga", ".opus", ".aac"} + {".webm", ".mkv", ".mka", ".avi", ".flv", ".mp3", ".wav", ".flac", ".ogg", ".oga", ".opus", ".aac"} ) # China's mandatory AI-content labeling (TC260, the national cybersecurity @@ -161,6 +162,22 @@ TC260_AIGC_FIELDS: frozenset[str] = frozenset( "ServiceUser", } ) +MAX_TC260_VALUE_BYTES = 1024 * 1024 + + +def parse_tc260_aigc_json(value: bytes) -> dict[str, str] | None: + """Parse a bounded JSON object carrying at least one normative TC260 field.""" + if len(value) > MAX_TC260_VALUE_BYTES: + return None + try: + parsed = json.loads(value.rstrip(b"\x00 ").decode("utf-8")) + except (UnicodeDecodeError, ValueError): + return None + if not isinstance(parsed, dict): + return None + fields = {str(key): str(item) for key, item in cast("dict[object, object]", parsed).items()} + return fields if TC260_AIGC_FIELDS & fields.keys() else None + # HuggingFace-hosted GPU jobs (Jobs / Spaces) stamp generated PNGs with this # ``tEXt`` chunk key holding the job UUID. It marks the hosting job, not a @@ -365,6 +382,8 @@ def aigc_label(image_path: Path) -> dict[str, str] | None: (as written by Doubao / ByteDance), read via PIL; - a native MP4/MOV ``AIGC`` key in ``moov.udta.meta.keys`` whose matching ``ilst`` item carries the raw JSON object; + - a native AVI ``LIST/INFO/AIGC`` chunk or FLV + ``script.onMetaData.AIGC`` string carrying the raw JSON object; - an XMP ``{...}`` block (HTML-entity encoded text), found by a container-agnostic raw-byte scan (PNG/JPEG/WebP alike); and - a raw-JSON ``{"AIGC":{...}}`` block with no namespace, as embedded in JPEG @@ -381,20 +400,17 @@ def aigc_label(image_path: Path) -> dict[str, str] | None: namespaced XMP element is unambiguous, so any JSON object is accepted. """ import html - import json - from typing import cast def _parse(text: str, *, require_tc260_field: bool) -> dict[str, str] | None: + if require_tc260_field: + return parse_tc260_aigc_json(text.encode("utf-8")) try: parsed = json.loads(text) except ValueError: return None if not isinstance(parsed, dict): return None - fields = {str(k): str(v) for k, v in cast("dict[object, object]", parsed).items()} - if require_tc260_field and not (TC260_AIGC_FIELDS & fields.keys()): - return None - return fields + return {str(k): str(v) for k, v in cast("dict[object, object]", parsed).items()} # PNG tEXt chunk keyed "AIGC" with raw JSON (Doubao and other China gens). # The key is generic, so require a TC260 field to avoid a false positive. @@ -428,6 +444,22 @@ def aigc_label(image_path: Path) -> dict[str, str] | None: if result := _parse(payload.decode("utf-8", "replace"), require_tc260_field=True): return result + # Native AVI and FLV TC260 metadata. Both readers walk their container + # structures and skip media payloads instead of relying on a raw substring + # that could collide inside compressed video. + legacy_payloads: tuple[bytes, ...] = () + if image_path.suffix.lower() == ".avi": + from remove_ai_watermarks.noai.riff import tc260_aigc_payloads as riff_tc260_aigc_payloads + + legacy_payloads = riff_tc260_aigc_payloads(image_path) + elif image_path.suffix.lower() == ".flv": + from remove_ai_watermarks.noai.flv import tc260_aigc_payloads as flv_tc260_aigc_payloads + + legacy_payloads = flv_tc260_aigc_payloads(image_path) + for payload in legacy_payloads: + if result := _parse(payload.decode("utf-8", "replace"), require_tc260_field=True): + return result + # XMP TC260:AIGC, namespaced (unambiguous) in either serialization RDF allows: # an element {...} or an attribute TC260:AIGC="{...}" # (the attribute form is what PicWish writes). Both are HTML-entity encoded. @@ -1206,7 +1238,8 @@ def remove_ai_metadata( ) return output_path - # Non-ISOBMFF audio/video (WebM/Matroska EBML, MP3 ID3, WAV/FLAC/OGG): the + # Non-ISOBMFF audio/video (WebM/Matroska EBML, AVI/FLV, MP3 ID3, + # WAV/FLAC/OGG): the # box walker can't reach these, so strip container metadata losslessly via # ffmpeg (-c copy -- codec data untouched, only tags/chapters dropped). if source_path.suffix.lower() in _FFMPEG_STRIP_EXTS: diff --git a/src/remove_ai_watermarks/noai/ebml.py b/src/remove_ai_watermarks/noai/ebml.py index 6b87412..d5a2c21 100644 --- a/src/remove_ai_watermarks/noai/ebml.py +++ b/src/remove_ai_watermarks/noai/ebml.py @@ -8,14 +8,13 @@ their payloads. from __future__ import annotations -import json -from typing import TYPE_CHECKING, BinaryIO, cast +from typing import TYPE_CHECKING, BinaryIO if TYPE_CHECKING: from collections.abc import Iterator from pathlib import Path -from remove_ai_watermarks.metadata import TC260_AIGC_FIELDS +from remove_ai_watermarks.metadata import MAX_TC260_VALUE_BYTES, parse_tc260_aigc_json _EBML_MAGIC = b"\x1aE\xdf\xa3" _SEGMENT_ID = 0x18538067 @@ -25,7 +24,6 @@ _SIMPLE_TAG_ID = 0x67C8 _TAG_NAME_ID = 0x45A3 _TAG_STRING_ID = 0x4487 _MAX_TAG_NAME_BYTES = 256 -_MAX_TC260_VALUE_BYTES = 1024 * 1024 def _vint_length(first: int, *, maximum: int) -> int | None: @@ -111,17 +109,6 @@ def _read_bounded( return value if len(value) == size else None -def _is_tc260_aigc_json(value: bytes) -> bool: - try: - parsed = json.loads(value.decode("utf-8")) - except (UnicodeDecodeError, ValueError): - return False - if not isinstance(parsed, dict): - return False - fields = cast("dict[object, object]", parsed) - return bool(TC260_AIGC_FIELDS & {str(key) for key in fields}) - - def _simple_tag_payloads( stream: BinaryIO, start: int, @@ -133,12 +120,12 @@ def _simple_tag_payloads( if element_id == _TAG_NAME_ID: name = _read_bounded(stream, payload_start, element_end, _MAX_TAG_NAME_BYTES) elif element_id == _TAG_STRING_ID: - value = _read_bounded(stream, payload_start, element_end, _MAX_TC260_VALUE_BYTES) + value = _read_bounded(stream, payload_start, element_end, MAX_TC260_VALUE_BYTES) if value is not None: values.append(value) if name != b"AIGC": return () - return tuple(value for value in values if _is_tc260_aigc_json(value)) + return tuple(value for value in values if parse_tc260_aigc_json(value) is not None) def tc260_aigc_payloads(path: str | Path) -> tuple[bytes, ...]: diff --git a/src/remove_ai_watermarks/noai/flv.py b/src/remove_ai_watermarks/noai/flv.py new file mode 100644 index 0000000..78a6dc1 --- /dev/null +++ b/src/remove_ai_watermarks/noai/flv.py @@ -0,0 +1,160 @@ +"""Bounded FLV metadata reader for native TC260 AIGC labels. + +TC260-PG-20257A stores the label in the ``onMetaData`` script tag as an AMF0 +property named ``AIGC`` whose string value is the normative JSON object. Media +tag payloads are skipped without being loaded. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + +from remove_ai_watermarks.metadata import MAX_TC260_VALUE_BYTES, parse_tc260_aigc_json + +_SCRIPT_TAG = 18 +_MAX_SCRIPT_BYTES = 4 * 1024 * 1024 + + +def _u24(value: bytes) -> int: + return int.from_bytes(value, "big") + + +def _amf0_string(data: bytes, position: int, *, long: bool = False) -> tuple[bytes, int] | None: + length_size = 4 if long else 2 + if position + length_size > len(data): + return None + length = int.from_bytes(data[position : position + length_size], "big") + start = position + length_size + end = start + length + if end > len(data): + return None + return data[start:end], end + + +def _skip_amf0(data: bytes, position: int, depth: int = 0) -> int | None: + if position >= len(data) or depth > 8: + return None + value_type = data[position] + position += 1 + if value_type == 0: + return position + 8 if position + 8 <= len(data) else None + if value_type == 1: + return position + 1 if position + 1 <= len(data) else None + if value_type == 2: + parsed = _amf0_string(data, position) + return parsed[1] if parsed is not None else None + if value_type in {5, 6}: + return position + if value_type == 7: + return position + 2 if position + 2 <= len(data) else None + if value_type == 11: + return position + 10 if position + 10 <= len(data) else None + if value_type == 12: + parsed = _amf0_string(data, position, long=True) + return parsed[1] if parsed is not None else None + if value_type == 10: + if position + 4 > len(data): + return None + count = int.from_bytes(data[position : position + 4], "big") + position += 4 + for _ in range(count): + next_position = _skip_amf0(data, position, depth + 1) + if next_position is None: + return None + position = next_position + return position + if value_type in {3, 8}: + if value_type == 8: + if position + 4 > len(data): + return None + position += 4 + while position + 3 <= len(data): + name_length = int.from_bytes(data[position : position + 2], "big") + position += 2 + if name_length == 0 and data[position] == 9: + return position + 1 + position += name_length + if position > len(data): + return None + next_position = _skip_amf0(data, position, depth + 1) + if next_position is None: + return None + position = next_position + return None + return None + + +def _script_payloads(data: bytes) -> tuple[bytes, ...]: + first = _amf0_string(data, 1) if data[:1] == b"\x02" else None + if first is None or first[0] != b"onMetaData": + return () + position = first[1] + if position >= len(data) or data[position] not in {3, 8}: + return () + if data[position] == 8: + position += 5 + else: + position += 1 + found: list[bytes] = [] + while position + 3 <= len(data): + name_length = int.from_bytes(data[position : position + 2], "big") + position += 2 + if name_length == 0 and data[position] == 9: + break + name_end = position + name_length + if name_end > len(data): + break + name = data[position:name_end] + position = name_end + if name == b"AIGC" and position < len(data) and data[position] in {2, 12}: + long = data[position] == 12 + parsed = _amf0_string(data, position + 1, long=long) + if parsed is None: + break + value, position = parsed + if len(value) <= MAX_TC260_VALUE_BYTES and parse_tc260_aigc_json(value) is not None: + found.append(value) + continue + next_position = _skip_amf0(data, position) + if next_position is None: + break + position = next_position + return tuple(found) + + +def tc260_aigc_payloads(path: str | Path) -> tuple[bytes, ...]: + """Read validated TC260 values from FLV ``script.onMetaData.AIGC``.""" + found: list[bytes] = [] + try: + with open(path, "rb") as stream: + header = stream.read(9) + if len(header) != 9 or header[:3] != b"FLV": + return () + data_offset = int.from_bytes(header[5:9], "big") + stream.seek(0, 2) + file_size = stream.tell() + position = data_offset + 4 + while position + 11 <= file_size: + stream.seek(position) + tag_header = stream.read(11) + if len(tag_header) != 11: + break + tag_type = tag_header[0] & 0x1F + data_size = _u24(tag_header[1:4]) + payload_start = position + 11 + payload_end = payload_start + data_size + if payload_end + 4 > file_size: + break + if tag_type == _SCRIPT_TAG and data_size <= _MAX_SCRIPT_BYTES: + payload = stream.read(data_size) + if len(payload) == data_size: + found.extend(_script_payloads(payload)) + if found: + return tuple(found) + position = payload_end + 4 + except OSError: + return () + return tuple(found) diff --git a/src/remove_ai_watermarks/noai/isobmff.py b/src/remove_ai_watermarks/noai/isobmff.py index af135f7..ddfd8fe 100644 --- a/src/remove_ai_watermarks/noai/isobmff.py +++ b/src/remove_ai_watermarks/noai/isobmff.py @@ -23,11 +23,10 @@ Reference: ISO/IEC 14496-12 (ISOBMFF) and C2PA 2.1 spec §11. from __future__ import annotations import io -import json import logging import re import struct -from typing import TYPE_CHECKING, Any, BinaryIO, cast +from typing import TYPE_CHECKING, Any, BinaryIO if TYPE_CHECKING: from collections.abc import Iterator @@ -38,7 +37,8 @@ from remove_ai_watermarks.metadata import ( C2PA_UUID, IPTC_AI_FIELD_MARKERS, IPTC_AI_MARKERS, - TC260_AIGC_FIELDS, + MAX_TC260_VALUE_BYTES, + parse_tc260_aigc_json, ) logger = logging.getLogger(__name__) @@ -61,14 +61,12 @@ _AI_LABEL_MARKERS: tuple[bytes, ...] = AIGC_MARKERS + IPTC_AI_MARKERS + IPTC_AI_ # blanked in place (see ``blank_ai_xmp_packets``). _XMP_PACKET_RE = re.compile(rb"<\?xpacket begin=.*?<\?xpacket end=[^>]*?\?>", re.DOTALL) + # TC260-PG-20257A stores an MP4/MOV label as an ``AIGC`` key in # ``moov.udta.meta.keys`` and its JSON value in the corresponding # ``moov.udta.meta.ilst`` item. The value is intentionally bounded before it is # read: the normative object is tiny, and a corrupt size must not allocate an # arbitrary amount of memory during an inspection. -_MAX_TC260_VALUE_BYTES = 1024 * 1024 - - def _iter_top_level_boxes(data: bytes) -> Iterator[tuple[int, int, bytes, int]]: """Yield ``(start, end, type, payload_offset)`` for each top-level box. @@ -181,18 +179,6 @@ def _tc260_key_indices( return found -def _is_tc260_aigc_json(value: bytes) -> bool: - """Require a JSON object carrying at least one normative TC260 field.""" - try: - parsed = json.loads(value.decode("utf-8")) - except (UnicodeDecodeError, ValueError): - return False - if not isinstance(parsed, dict): - return False - fields = cast("dict[object, object]", parsed) - return bool(TC260_AIGC_FIELDS & {str(key) for key in fields}) - - def _tc260_aigc_regions( stream: BinaryIO, file_size: int, @@ -249,11 +235,11 @@ def _tc260_aigc_regions( ): value_start = data_payload + 8 value_size = data_end - value_start - if data_type != b"data" or value_size < 0 or value_size > _MAX_TC260_VALUE_BYTES: + if data_type != b"data" or value_size < 0 or value_size > MAX_TC260_VALUE_BYTES: continue stream.seek(value_start) value = stream.read(value_size) - if len(value) == value_size and _is_tc260_aigc_json(value): + if len(value) == value_size and parse_tc260_aigc_json(value) is not None: regions.append((*key_span, value_start, data_end, value)) return regions diff --git a/src/remove_ai_watermarks/noai/riff.py b/src/remove_ai_watermarks/noai/riff.py new file mode 100644 index 0000000..1034664 --- /dev/null +++ b/src/remove_ai_watermarks/noai/riff.py @@ -0,0 +1,74 @@ +"""Bounded AVI metadata reader for native TC260 AIGC labels. + +TC260-PG-20257A stores the label in an AVI ``LIST/INFO`` chunk whose child +chunk ID is ``AIGC`` and whose value is the normative JSON object. The walker +seeks over media chunks and reads only bounded metadata values. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, BinaryIO + +if TYPE_CHECKING: + from pathlib import Path + +from remove_ai_watermarks.metadata import MAX_TC260_VALUE_BYTES, parse_tc260_aigc_json + + +def _info_payloads( + stream: BinaryIO, + start: int, + end: int, +) -> tuple[bytes, ...]: + found: list[bytes] = [] + position = start + while position + 8 <= end: + stream.seek(position) + chunk_id = stream.read(4) + size_raw = stream.read(4) + if len(chunk_id) != 4 or len(size_raw) != 4: + break + size = int.from_bytes(size_raw, "little") + payload_start = position + 8 + payload_end = payload_start + size + if payload_end > end: + break + if chunk_id == b"AIGC" and size <= MAX_TC260_VALUE_BYTES: + value = stream.read(size) + if len(value) == size and parse_tc260_aigc_json(value) is not None: + found.append(value.rstrip(b"\x00 ")) + position = payload_end + (size & 1) + return tuple(found) + + +def tc260_aigc_payloads(path: str | Path) -> tuple[bytes, ...]: + """Read validated TC260 values from an AVI ``LIST/INFO/AIGC`` chunk.""" + found: list[bytes] = [] + try: + with open(path, "rb") as stream: + header = stream.read(12) + if len(header) != 12 or header[:4] != b"RIFF" or header[8:12] != b"AVI ": + return () + stream.seek(0, 2) + file_size = stream.tell() + declared_end = min(8 + int.from_bytes(header[4:8], "little"), file_size) + position = 12 + while position + 8 <= declared_end: + stream.seek(position) + chunk_id = stream.read(4) + size_raw = stream.read(4) + if len(chunk_id) != 4 or len(size_raw) != 4: + break + size = int.from_bytes(size_raw, "little") + payload_start = position + 8 + payload_end = payload_start + size + if payload_end > declared_end: + break + if chunk_id == b"LIST" and size >= 4: + list_type = stream.read(4) + if list_type == b"INFO": + found.extend(_info_payloads(stream, payload_start + 4, payload_end)) + position = payload_end + (size & 1) + except OSError: + return () + return tuple(found) diff --git a/src/remove_ai_watermarks/video.py b/src/remove_ai_watermarks/video.py index 1ad7cf5..7b34bbf 100644 --- a/src/remove_ai_watermarks/video.py +++ b/src/remove_ai_watermarks/video.py @@ -1,9 +1,10 @@ """High-level video processing API. Supported experimental stages are container-level AI metadata inspection and -removal, temporally stabilized visible Sora, Veo, Seedance, and Dola removal, -and VAE regeneration that produces an externally verifiable SynthID candidate. -The visible pixel path reuses the image package's shared fill backends. +removal, temporally stabilized visible Sora, Veo, Seedance, Dola, Hailuo, and +Kling removal, and VAE regeneration that produces an externally verifiable +SynthID candidate. The visible pixel path reuses the image package's shared +fill backends. """ from __future__ import annotations @@ -22,9 +23,11 @@ from remove_ai_watermarks.video_synthid import ( if TYPE_CHECKING: from remove_ai_watermarks.video_invisible import RegenerationMetrics -VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v", ".webm", ".mkv"}) +VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v", ".webm", ".mkv", ".avi", ".flv"}) _ISOBMFF_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v"}) _EBML_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".webm", ".mkv"}) +_RIFF_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".avi"}) +_FLV_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".flv"}) _REGENERATED_VIDEO_EXTENSIONS: frozenset[str] = _ISOBMFF_VIDEO_EXTENSIONS _EBML_MAGIC = b"\x1aE\xdf\xa3" @@ -109,8 +112,11 @@ def _video_source(source: str | Path) -> Path: with path.open("rb") as stream: head = stream.read(12) suffix = path.suffix.lower() - matches_container = (suffix in _ISOBMFF_VIDEO_EXTENSIONS and len(head) >= 8 and head[4:8] == b"ftyp") or ( - suffix in _EBML_VIDEO_EXTENSIONS and head.startswith(_EBML_MAGIC) + matches_container = ( + (suffix in _ISOBMFF_VIDEO_EXTENSIONS and len(head) >= 8 and head[4:8] == b"ftyp") + or (suffix in _EBML_VIDEO_EXTENSIONS and head.startswith(_EBML_MAGIC)) + or (suffix in _RIFF_VIDEO_EXTENSIONS and len(head) >= 12 and head[:4] == b"RIFF" and head[8:12] == b"AVI ") + or (suffix in _FLV_VIDEO_EXTENSIONS and head.startswith(b"FLV")) ) if not matches_container: raise ValueError(f"Video content does not match its {suffix} extension: {path}") @@ -182,11 +188,11 @@ def remove_video_visible( ) -> VideoVisibleResult: """Remove a supported visible AI wordmark from a video. - Supported marks are ``sora``, ``veo``, ``seedance``, and ``dola``. The full - sequence is scanned before pixels change, and only recurring candidates are - accepted. Audio is copied without re-encoding; video is transcoded because - the pixels change. When no stable mark is found, no output is written and - ``output`` in the result is ``None``. + Supported marks are ``sora``, ``veo``, ``seedance``, ``dola``, ``hailuo``, + and ``kling``. The full sequence is scanned before pixels change, and only + recurring candidates are accepted. Audio is copied without re-encoding; + video is transcoded because the pixels change. When no stable mark is found, + no output is written and ``output`` in the result is ``None``. """ from remove_ai_watermarks.metadata import get_ai_metadata from remove_ai_watermarks.video_visible import ( @@ -195,18 +201,22 @@ def remove_video_visible( has_sora_provenance, has_veo_provenance, scan_dola_video, + scan_hailuo_video, + scan_kling_video, scan_seedance_video, scan_sora_video, scan_veo_video, stabilize_dola_localizations, + stabilize_hailuo_localizations, + stabilize_kling_localizations, stabilize_seedance_localizations, stabilize_sora_localizations, stabilize_veo_localizations, ) from remove_ai_watermarks.watermark_registry import resolve_backend - if mark not in {"sora", "veo", "seedance", "dola"}: - raise ValueError("Unsupported visible video mark; expected sora, veo, seedance, or dola") + if mark not in {"sora", "veo", "seedance", "dola", "hailuo", "kling"}: + raise ValueError("Unsupported visible video mark; expected sora, veo, seedance, dola, hailuo, or kling") if backend not in {"auto", "cv2", "migan", "lama"}: raise ValueError("Unsupported fill backend; expected auto, cv2, migan, or lama") @@ -237,7 +247,7 @@ def remove_video_visible( ) padding_fraction = 0.0 mask_style = "box" - else: + elif mark == "dola": scan = scan_dola_video(source_path) regions = stabilize_dola_localizations( scan.detections, @@ -245,6 +255,16 @@ def remove_video_visible( ) padding_fraction = 0.20 mask_style = "box" + elif mark == "hailuo": + scan = scan_hailuo_video(source_path) + regions = stabilize_hailuo_localizations(scan.detections) + padding_fraction = 0.12 + mask_style = "box" + else: + scan = scan_kling_video(source_path) + regions = stabilize_kling_localizations(scan.detections) + padding_fraction = 0.12 + mask_style = "box" detected_frames = sum(region is not None for region in regions) if detected_frames == 0: return VideoVisibleResult( diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py index 17148fa..e7b6949 100644 --- a/src/remove_ai_watermarks/video_visible.py +++ b/src/remove_ai_watermarks/video_visible.py @@ -4,10 +4,11 @@ Supported marks use fully synthetic silhouettes made from geometric primitives, OpenCV's built-in font, and Pillow's bundled font. Sora detection searches the full frame because the wordmark moves. Veo detection covers both the current four-point diamond and legacy ``Veo`` text. Seedance detects the boxed ``AI`` -label, while Dola detects its compact text label. A single frame is never enough -to authorize removal: the temporal arbiter requires the candidate to recur at -the same location across adjacent frames. This keeps isolated lookalikes in -clean videos from becoming removal masks. +label, Dola detects its compact text label, Hailuo detects the composite +MINIMAX/Hailuo label, and Kling detects its version-independent wordmark core. +A single frame is never enough to authorize removal: the temporal arbiter +requires the candidate to recur at the same location across adjacent frames. +This keeps isolated lookalikes in clean videos from becoming removal masks. Video pixels are decoded with OpenCV and encoded with the system ``ffmpeg``. Audio is stream-copied from the source. The video stream must be transcoded @@ -64,6 +65,11 @@ _SEEDANCE_STRONG_CONFIDENCE = 0.43 _DOLA_PROVENANCE_WEAK_CONFIDENCE = 0.48 _DOLA_STRICT_WEAK_CONFIDENCE = 0.50 _DOLA_STRONG_CONFIDENCE = 0.52 +_HAILUO_WEAK_CONFIDENCE = 0.30 +_HAILUO_STRONG_CONFIDENCE = 0.34 +_KLING_WEAK_CONFIDENCE = 0.20 +_KLING_STRONG_CONFIDENCE = 0.24 +_KLING_MIN_WHITE_FRACTION = 0.02 _MIN_STABLE_FRAMES = 5 _MIN_VEO_STABLE_FRAMES = 12 _MIN_FIXED_MARK_STABLE_FRAMES = 12 @@ -76,6 +82,8 @@ _VEO_DIAMOND_PROFILES = ( (44, 29, 40), ) _DOLA_RELATIVE_HEIGHTS = tuple(value / 1000 for value in range(22, 41)) +_HAILUO_RELATIVE_HEIGHTS = tuple(value / 1000 for value in range(28, 56, 3)) +_KLING_RELATIVE_HEIGHTS = tuple(value / 1000 for value in range(24, 49, 3)) @dataclass(frozen=True) @@ -108,6 +116,14 @@ def _scalable_default_font(size: int) -> ImageFont.ImageFont | ImageFont.FreeTyp return ImageFont.load_default() +def _crop_nonzero(image: NDArray[Any]) -> NDArray[Any]: + """Crop a synthetic template to its nonzero footprint.""" + ys, xs = np.where(image > 0) + if len(xs) == 0: + return image + return image[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1] + + @lru_cache(maxsize=1) def _sora_templates() -> tuple[NDArray[Any], NDArray[Any]]: """Return synthetic full-wordmark and mascot-only silhouettes. @@ -162,9 +178,7 @@ def _veo_templates() -> tuple[NDArray[Any], NDArray[Any]]: text_canvas = Image.new("L", (140, 60), 0) text_draw = ImageDraw.Draw(text_canvas) text_draw.text((2, 0), "Veo", font=_scalable_default_font(48), fill=255) - text = np.asarray(text_canvas, dtype=np.uint8) - ys, xs = np.where(text > 0) - text = text[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1] + text = _crop_nonzero(np.asarray(text_canvas, dtype=np.uint8)) return np.asarray(diamond_canvas, dtype=np.uint8), text @@ -201,8 +215,75 @@ def _dola_template() -> NDArray[Any]: 3, cv2.LINE_AA, ) - ys, xs = np.where(canvas > 0) - return canvas[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1] + return _crop_nonzero(canvas) + + +@lru_cache(maxsize=1) +def _hailuo_template() -> NDArray[Any]: + """Return a synthetic MINIMAX/Hailuo composite-label silhouette.""" + canvas = Image.new("L", (680, 112), 0) + draw = ImageDraw.Draw(canvas) + + # Five symmetric waveform strokes approximate the provider-independent + # geometry at the left edge without copying pixels from an export. + waveform_heights = (42, 70, 96, 70, 42) + center_y = 56 + for index, height in enumerate(waveform_heights): + x = 10 + index * 13 + draw.rounded_rectangle( + (x, center_y - height // 2, x + 5, center_y + height // 2), + radius=2, + fill=255, + ) + + font = _scalable_default_font(58) + draw.text((84, 18), "MINIMAX", font=font, fill=255, stroke_width=1, stroke_fill=255) + draw.rectangle((342, 18, 347, 92), fill=255) + + # The Hailuo symbol is a ring with a small offset highlight. + draw.ellipse((370, 20, 446, 96), outline=255, width=12) + draw.ellipse((397, 38, 434, 76), fill=255) + draw.ellipse((397, 29, 420, 52), fill=0) + draw.text((456, 18), "hailuo AI", font=font, fill=255, stroke_width=1, stroke_fill=255) + return np.asarray(canvas, dtype=np.uint8) + + +@lru_cache(maxsize=1) +def _kling_templates() -> tuple[NDArray[Any], ...]: + """Return synthetic font variants for the Kling wordmark core.""" + canvas = Image.new("L", (430, 104), 0) + draw = ImageDraw.Draw(canvas) + draw.text( + (2, 12), + "KLING AI", + font=_scalable_default_font(64), + fill=255, + stroke_width=1, + stroke_fill=255, + ) + templates = [_crop_nonzero(np.asarray(canvas, dtype=np.uint8))] + for font in (cv2.FONT_HERSHEY_SIMPLEX, cv2.FONT_HERSHEY_DUPLEX): + cv_template = np.zeros((100, 500), dtype=np.uint8) + cv2.putText( + cv_template, + "KLING AI", + (2, 72), + font, + 2.2, + 255, + 3, + cv2.LINE_AA, + ) + templates.append(_crop_nonzero(cv_template)) + return tuple(templates) + + +@lru_cache(maxsize=1) +def _kling_logo_template() -> NDArray[Any]: + """Return a synthetic ring approximation of the Kling swirl.""" + template = np.zeros((100, 100), dtype=np.uint8) + cv2.circle(template, (50, 50), 34, 255, 14, cv2.LINE_AA) + return _crop_nonzero(template) def _top_hat(gray: NDArray[Any]) -> NDArray[Any]: @@ -338,21 +419,43 @@ def _restore_region( return source_x, source_y, source_width, source_height -def _detect_fixed_bottom_right_mark( +def _bounded_region( + x: int, + y: int, + width: int, + height: int, + *, + frame_width: int, + frame_height: int, +) -> Region: + """Clip an expanded region to the frame without changing its anchor.""" + bounded_x = max(0, x) + bounded_y = max(0, y) + return ( + bounded_x, + bounded_y, + min(frame_width - bounded_x, width + x - bounded_x), + min(frame_height - bounded_y, height + y - bounded_y), + ) + + +def _detect_fixed_mark( image_bgr: NDArray[Any], template: NDArray[Any], *, relative_heights: tuple[float, ...], search_origin: tuple[float, float], kernel_fraction: float, + prefer_larger_within: float = 0.0, + normalized: tuple[NDArray[Any], float] | None = None, frame_index: int, ) -> FrameLocalization: - """Match one fixed bottom-right synthetic mark on a normalized frame.""" + """Match one fixed synthetic mark inside a normalized-frame search region.""" if image_bgr.size == 0: return FrameLocalization(frame_index, 0.0, None) frame_height, frame_width = image_bgr.shape[:2] - gray, scale = _normalized_gray(image_bgr) + gray, scale = normalized if normalized is not None else _normalized_gray(image_bgr) normalized_height, normalized_width = gray.shape[:2] short_side = min(normalized_height, normalized_width) search_x = round(normalized_width * search_origin[0]) @@ -363,8 +466,7 @@ def _detect_fixed_bottom_right_mark( normalized_width - search_x, normalized_height - search_y, ) - best_confidence = 0.0 - best_region: Region | None = None + matches: list[tuple[float, Region]] = [] for relative_height in relative_heights: template_height = max(6, round(short_side * relative_height)) template_width = max(1, round(template.shape[1] * template_height / template.shape[0])) @@ -379,9 +481,21 @@ def _detect_fixed_bottom_right_mark( region=search_region, kernel_size=max(3, round(template_height * kernel_fraction) | 1), ) - if confidence > best_confidence: - best_confidence = confidence - best_region = candidate + if candidate is not None and confidence > 0: + matches.append((confidence, candidate)) + if not matches: + return FrameLocalization(frame_index, 0.0, None) + best_confidence, best_region = max(matches, key=lambda match: match[0]) + if prefer_larger_within > 0: + eligible = [ + (confidence, candidate) + for confidence, candidate in matches + if confidence >= best_confidence - prefer_larger_within + ] + best_confidence, best_region = max( + eligible, + key=lambda match: (match[1][2] * match[1][3], match[0]), + ) return FrameLocalization( frame_index, @@ -397,7 +511,7 @@ def _detect_fixed_bottom_right_mark( def detect_seedance_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: """Locate the strongest fixed Seedance boxed-AI candidate.""" - return _detect_fixed_bottom_right_mark( + return _detect_fixed_mark( image_bgr, _seedance_template(), relative_heights=(0.065, 0.075, 0.085, 0.095, 0.105), @@ -409,7 +523,7 @@ def detect_seedance_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> F def detect_dola_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: """Locate the strongest fixed Dola AI text candidate.""" - return _detect_fixed_bottom_right_mark( + return _detect_fixed_mark( image_bgr, _dola_template(), relative_heights=_DOLA_RELATIVE_HEIGHTS, @@ -419,6 +533,137 @@ def detect_dola_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> Frame ) +def detect_hailuo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: + """Locate the strongest fixed MINIMAX/Hailuo composite-label candidate.""" + detection = _detect_fixed_mark( + image_bgr, + _hailuo_template(), + relative_heights=_HAILUO_RELATIVE_HEIGHTS, + search_origin=(0.28, 0.76), + kernel_fraction=0.18, + frame_index=frame_index, + ) + if detection.region is None: + return detection + frame_width = image_bgr.shape[1] + x, y, width, height = detection.region + horizontal_padding = round(height * 1.25) + region = _bounded_region( + x - horizontal_padding, + y, + width + horizontal_padding * 2, + height, + frame_width=frame_width, + frame_height=image_bgr.shape[0], + ) + return FrameLocalization( + frame_index, + detection.confidence, + region, + ) + + +def detect_kling_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: + """Locate the fixed Kling wordmark core and include its version suffix.""" + if image_bgr.size == 0: + return FrameLocalization(frame_index, 0.0, None) + frame_height, frame_width = image_bgr.shape[:2] + normalized = _normalized_gray(image_bgr) + expanded: list[FrameLocalization] = [] + for template in _kling_templates(): + detection = _detect_fixed_mark( + image_bgr, + template, + relative_heights=_KLING_RELATIVE_HEIGHTS, + search_origin=(0.64, 0.84), + kernel_fraction=0.18, + prefer_larger_within=0.06, + normalized=normalized, + frame_index=frame_index, + ) + if detection.region is None: + continue + x, y, width, height = detection.region + left_padding = round(height * 1.8) + right_padding = round(height * 4.0) + vertical_padding = round(height * 0.4) + region = _bounded_region( + x - left_padding, + y - vertical_padding, + width + left_padding + right_padding, + height + vertical_padding * 2, + frame_width=frame_width, + frame_height=frame_height, + ) + expanded.append( + FrameLocalization( + frame_index, + detection.confidence, + region, + ) + ) + if not expanded: + return FrameLocalization(frame_index, 0.0, None) + edge_candidates = [ + candidate + for candidate in expanded + if candidate.region is not None + and candidate.region[0] + candidate.region[2] >= frame_width * 0.96 + and candidate.region[1] + candidate.region[3] >= frame_height * 0.94 + ] + font_candidate = None if not edge_candidates else max(edge_candidates, key=lambda candidate: candidate.confidence) + + logo = _detect_fixed_mark( + image_bgr, + _kling_logo_template(), + relative_heights=tuple(value / 1000 for value in range(20, 61, 3)), + search_origin=(0.62, 0.90), + kernel_fraction=0.18, + prefer_larger_within=0.03, + normalized=normalized, + frame_index=frame_index, + ) + logo_candidate: FrameLocalization | None = None + logo_x: int | None = None + if logo.region is not None and logo.confidence >= 0.44: + logo_x, logo_y, _, logo_height = logo.region + logo_candidate = FrameLocalization( + frame_index, + logo.confidence, + _bounded_region( + logo_x - round(logo_height * 0.2), + logo_y - round(logo_height * 0.25), + round(logo_height * 7.8), + round(logo_height * 1.5), + frame_width=frame_width, + frame_height=frame_height, + ), + ) + + best = font_candidate or logo_candidate + if ( + font_candidate is not None + and font_candidate.region is not None + and logo_candidate is not None + and logo_x is not None + ): + font_x, _, font_width, _ = font_candidate.region + if logo_x <= font_x + round(font_width * 0.50): + best = logo_candidate + if best is None or best.region is None: + return FrameLocalization(frame_index, 0.0, None) + x, y, width, height = best.region + roi = image_bgr[y : y + height, x : x + width] + if roi.ndim == 2: + white_fraction = float(np.mean(roi >= 180)) + else: + hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV) + white_fraction = float(np.mean((hsv[:, :, 1] <= 55) & (hsv[:, :, 2] >= 180))) + if white_fraction < _KLING_MIN_WHITE_FRACTION: + return FrameLocalization(frame_index, 0.0, None) + return best + + def detect_veo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: """Locate the strongest current-diamond or legacy-text Veo candidate.""" if image_bgr.size == 0: @@ -599,6 +844,38 @@ def stabilize_dola_localizations( ) +def stabilize_hailuo_localizations( + detections: tuple[FrameLocalization, ...] | list[FrameLocalization], +) -> list[Region | None]: + """Accept a recurring MINIMAX/Hailuo label at a fixed position.""" + return _stabilize_localizations( + detections, + provenance=False, + weak_floor=_HAILUO_WEAK_CONFIDENCE, + strong_floor=_HAILUO_STRONG_CONFIDENCE, + transition_floor=0.28, + min_stable_frames=_MIN_FIXED_MARK_STABLE_FRAMES, + cover_after_confirmation=True, + anchor_iou=0.80, + ) + + +def stabilize_kling_localizations( + detections: tuple[FrameLocalization, ...] | list[FrameLocalization], +) -> list[Region | None]: + """Accept a recurring versioned Kling label at a fixed position.""" + return _stabilize_localizations( + detections, + provenance=False, + weak_floor=_KLING_WEAK_CONFIDENCE, + strong_floor=_KLING_STRONG_CONFIDENCE, + transition_floor=0.30, + min_stable_frames=_MIN_FIXED_MARK_STABLE_FRAMES, + cover_after_confirmation=True, + anchor_iou=0.80, + ) + + def _stabilize_localizations( detections: tuple[FrameLocalization, ...] | list[FrameLocalization], *, @@ -734,6 +1011,16 @@ def scan_dola_video(source: Path) -> VideoScan: return _scan_video(source, detect_dola_frame) +def scan_hailuo_video(source: Path) -> VideoScan: + """Decode a video once and collect one untrusted Hailuo candidate per frame.""" + return _scan_video(source, detect_hailuo_frame) + + +def scan_kling_video(source: Path) -> VideoScan: + """Decode a video once and collect one untrusted Kling candidate per frame.""" + return _scan_video(source, detect_kling_frame) + + def _mask_for_region( frame_bgr: NDArray[Any], region: Region, diff --git a/tests/test_video.py b/tests/test_video.py index 40d5a4d..2450891 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -14,6 +14,7 @@ from remove_ai_watermarks.cli import main from remove_ai_watermarks.metadata import C2PA_UUID if TYPE_CHECKING: + from collections.abc import Callable from pathlib import Path @@ -81,6 +82,71 @@ def _video_with_tc260_ebml(path: Path, *, value: bytes = _TC260_AIGC) -> Path: return path +def _riff_chunk(chunk_id: bytes, payload: bytes) -> bytes: + return chunk_id + len(payload).to_bytes(4, "little") + payload + (b"\x00" if len(payload) & 1 else b"") + + +def _video_with_tc260_avi(path: Path, *, value: bytes = _TC260_AIGC) -> Path: + info = _riff_chunk(b"AIGC", value) + _riff_chunk(b"INAM", b"standard title\x00") + body = b"AVI " + _riff_chunk(b"LIST", b"INFO" + info) + _riff_chunk(b"JUNK", _VIDEO_PAYLOAD) + path.write_bytes(b"RIFF" + len(body).to_bytes(4, "little") + body) + return path + + +def _amf0_string(value: bytes) -> bytes: + return b"\x02" + len(value).to_bytes(2, "big") + value + + +def _video_with_tc260_flv(path: Path, *, value: bytes = _TC260_AIGC) -> Path: + payload = ( + _amf0_string(b"onMetaData") + + b"\x08\x00\x00\x00\x02" + + len(b"AIGC").to_bytes(2, "big") + + b"AIGC" + + _amf0_string(value) + + len(b"duration").to_bytes(2, "big") + + b"duration" + + b"\x00" + + b"\x00\x00\x00\x00\x00\x00\x00\x00" + + b"\x00\x00\x09" + ) + tag_header = b"\x12" + len(payload).to_bytes(3, "big") + b"\x00" * 7 + path.write_bytes( + b"FLV\x01\x05\x00\x00\x00\x09" + + b"\x00\x00\x00\x00" + + tag_header + + payload + + (11 + len(payload)).to_bytes(4, "big") + ) + return path + + +_LEGACY_VIDEO_CASES = ( + (".avi", _video_with_tc260_avi), + (".flv", _video_with_tc260_flv), +) + + +def _stamp_gray_mark( + frame: np.ndarray, + mark: Image.Image | np.ndarray, + *, + x: int, + y: int, + opacity: float, +) -> None: + """Alpha-composite a grayscale synthetic mark onto a BGR test frame.""" + mark_array = np.asarray(mark, dtype=np.float32) + height, width = mark_array.shape + alpha = mark_array[:, :, None] / 255 * opacity + crop = frame[y : y + height, x : x + width].astype(np.float32) + frame[y : y + height, x : x + width] = np.clip( + crop * (1 - alpha) + 255 * alpha, + 0, + 255, + ).astype(np.uint8) + + def _regeneration_metrics( *, frames: int = 24, @@ -224,6 +290,47 @@ class TestVideoMetadataApi: assert report.has_ai_metadata is False assert report.markers == {} + @pytest.mark.parametrize( + ("suffix", "factory"), + _LEGACY_VIDEO_CASES, + ) + def test_inspects_native_tc260_legacy_video_metadata( + self, + tmp_path: Path, + suffix: str, + factory: Callable[..., Path], + ): + from remove_ai_watermarks.video import inspect_video_metadata + + source = factory(tmp_path / f"source{suffix}") + + report = inspect_video_metadata(source) + + assert report.has_ai_metadata is True + assert report.markers["aigc_label"].endswith("producer 00119144030008867405X210002") + + @pytest.mark.parametrize( + ("suffix", "factory"), + _LEGACY_VIDEO_CASES, + ) + def test_ignores_generic_legacy_video_aigc_tag( + self, + tmp_path: Path, + suffix: str, + factory: Callable[..., Path], + ): + from remove_ai_watermarks.video import inspect_video_metadata + + source = factory( + tmp_path / f"source{suffix}", + value=b'{"description":"ordinary application metadata"}', + ) + + report = inspect_video_metadata(source) + + assert report.has_ai_metadata is False + assert report.markers == {} + def test_rejects_image_input(self, tmp_clean_png: Path): from remove_ai_watermarks.video import inspect_video_metadata @@ -404,9 +511,7 @@ class TestSoraFrameLocalization: draw.text((68, 1), "Sora", font=font, fill=255, stroke_width=1) mark_array = cv2.resize(np.asarray(mark), (124, 44), interpolation=cv2.INTER_AREA) x, y = 620, 398 - alpha = mark_array.astype(np.float32)[:, :, None] / 255 * 0.78 - crop = frame[y : y + 44, x : x + 124].astype(np.float32) - frame[y : y + 44, x : x + 124] = np.clip(crop * (1 - alpha) + 255 * alpha, 0, 255).astype(np.uint8) + _stamp_gray_mark(frame, mark_array, x=x, y=y, opacity=0.78) return frame, (x, y, 124, 44) def test_localizes_independently_rendered_sora_like_mark(self): @@ -448,13 +553,7 @@ class TestVeoFrameLocalization: (round(size * 0.39), round(size * 0.38)), ) ImageDraw.Draw(mark).polygon(points, fill=255) - alpha = np.asarray(mark, dtype=np.float32)[:, :, None] / 255 * 0.72 - crop = frame[y : y + size, x : x + size].astype(np.float32) - frame[y : y + size, x : x + size] = np.clip( - crop * (1 - alpha) + 255 * alpha, - 0, - 255, - ).astype(np.uint8) + _stamp_gray_mark(frame, mark, x=x, y=y, opacity=0.72) detection = detect_veo_frame(frame) @@ -478,13 +577,7 @@ class TestVeoFrameLocalization: mark_height, mark_width = mark_array.shape x = frame.shape[1] - mark_width - 20 y = frame.shape[0] - mark_height - 18 - alpha = mark_array.astype(np.float32)[:, :, None] / 255 * 0.66 - crop = frame[y : y + mark_height, x : x + mark_width].astype(np.float32) - frame[y : y + mark_height, x : x + mark_width] = np.clip( - crop * (1 - alpha) + 255 * alpha, - 0, - 255, - ).astype(np.uint8) + _stamp_gray_mark(frame, mark_array, x=x, y=y, opacity=0.66) detection = detect_veo_frame(frame) @@ -528,15 +621,8 @@ class TestByteDanceFrameLocalization: except TypeError: font = ImageFont.load_default() draw.text((18, 8), "AI", font=font, fill=255) - mark_array = np.asarray(mark, dtype=np.float32) x, y = 1130, 620 - alpha = mark_array[:, :, None] / 255 * 0.65 - crop = frame[y : y + 60, x : x + 80].astype(np.float32) - frame[y : y + 60, x : x + 80] = np.clip( - crop * (1 - alpha) + 255 * alpha, - 0, - 255, - ).astype(np.uint8) + _stamp_gray_mark(frame, mark, x=x, y=y, opacity=0.65) detection = detect_seedance_frame(frame) @@ -564,13 +650,7 @@ class TestByteDanceFrameLocalization: mark_height, mark_width = mark.shape x = frame.shape[1] - mark_width - 18 y = frame.shape[0] - mark_height - 14 - alpha = mark.astype(np.float32)[:, :, None] / 255 * 0.75 - crop = frame[y : y + mark_height, x : x + mark_width].astype(np.float32) - frame[y : y + mark_height, x : x + mark_width] = np.clip( - crop * (1 - alpha) + 255 * alpha, - 0, - 255, - ).astype(np.uint8) + _stamp_gray_mark(frame, mark, x=x, y=y, opacity=0.75) detection = detect_dola_frame(frame) @@ -594,6 +674,92 @@ class TestByteDanceFrameLocalization: assert mask[84, 104] == 0 +class TestAdditionalProviderFrameLocalization: + def test_localizes_independently_rendered_hailuo_label(self): + from remove_ai_watermarks.video_visible import _region_iou, detect_hailuo_frame + + frame = np.full((720, 1280, 3), 32, dtype=np.uint8) + mark = Image.new("L", (330, 54), 0) + draw = ImageDraw.Draw(mark) + try: + font = ImageFont.load_default(size=28) + except TypeError: + font = ImageFont.load_default() + for index, height in enumerate((20, 34, 46, 34, 20)): + x = 4 + index * 6 + draw.rounded_rectangle((x, 27 - height // 2, x + 2, 27 + height // 2), radius=1, fill=255) + draw.text((39, 8), "MINIMAX", font=font, fill=255) + draw.rectangle((164, 7, 166, 47), fill=255) + draw.ellipse((178, 8, 224, 50), outline=255, width=5) + draw.text((228, 8), "hailuo AI", font=font, fill=255) + x, y = 930, 650 + _stamp_gray_mark(frame, mark, x=x, y=y, opacity=0.75) + + detection = detect_hailuo_frame(frame) + + assert detection.region is not None + assert detection.confidence >= 0.24 + assert _region_iou(detection.region, (x, y, 330, 54)) >= 0.45 + + def test_localizes_kling_core_and_covers_version_suffix(self): + from remove_ai_watermarks.video_visible import detect_kling_frame + + frame = np.full((720, 1280, 3), 28, dtype=np.uint8) + mark = np.zeros((42, 245), dtype=np.uint8) + cv2.ellipse(mark, (20, 21), (15, 15), 0, 20, 330, 255, 4, cv2.LINE_AA) + cv2.putText( + mark, + "KLING AI 1.6", + (43, 31), + cv2.FONT_HERSHEY_SIMPLEX, + 0.9, + 255, + 2, + cv2.LINE_AA, + ) + x, y = 1018, 664 + _stamp_gray_mark(frame, mark, x=x, y=y, opacity=0.72) + + detection = detect_kling_frame(frame) + glyph_ys, glyph_xs = np.where(mark > 0) + glyph_box = ( + x + int(glyph_xs.min()), + y + int(glyph_ys.min()), + int(glyph_xs.max() - glyph_xs.min() + 1), + int(glyph_ys.max() - glyph_ys.min() + 1), + ) + + assert detection.region is not None + assert detection.confidence >= 0.24 + detected_x, detected_y, detected_width, detected_height = detection.region + glyph_x, glyph_y, glyph_width, glyph_height = glyph_box + assert detected_x <= glyph_x + assert detected_y <= glyph_y + assert detected_x + detected_width >= glyph_x + glyph_width + assert detected_y + detected_height >= glyph_y + glyph_height + + def test_rejects_a_saturated_fixed_kling_shape(self): + from remove_ai_watermarks.video_visible import detect_kling_frame + + frame = np.full((720, 1280, 3), 24, dtype=np.uint8) + cv2.circle(frame, (1040, 670), 15, (0, 220, 0), 5, cv2.LINE_AA) + cv2.putText( + frame, + "KLING AI 1.6", + (1065, 681), + cv2.FONT_HERSHEY_SIMPLEX, + 0.9, + (0, 220, 0), + 2, + cv2.LINE_AA, + ) + + detection = detect_kling_frame(frame) + + assert detection.confidence == 0.0 + assert detection.region is None + + class TestSoraTemporalArbiter: _BOX = (40, 60, 150, 54) @@ -823,6 +989,35 @@ class TestByteDanceTemporalArbiter: assert not has_bytedance_video_provenance({"issuer": "BytePlus (ByteDance)"}) +class TestAdditionalProviderTemporalArbiter: + _HAILUO_BOX = (930, 650, 330, 54) + _KLING_BOX = (1018, 664, 245, 42) + + @pytest.mark.parametrize( + ("stabilizer_name", "box", "weak_score", "strong_score"), + [ + ("stabilize_hailuo_localizations", _HAILUO_BOX, 0.31, 0.35), + ("stabilize_kling_localizations", _KLING_BOX, 0.21, 0.25), + ], + ) + def test_requires_a_strong_anchored_twelve_frame_run( + self, + stabilizer_name: str, + box: tuple[int, int, int, int], + weak_score: float, + strong_score: float, + ): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video_visible import FrameLocalization + + stabilize = getattr(video_visible, stabilizer_name) + weak = [FrameLocalization(index, weak_score, box) for index in range(12)] + strong = [FrameLocalization(index, strong_score, box) for index in range(12)] + + assert stabilize(weak) == [None] * 12 + assert stabilize(strong) == [box] * 12 + + class TestVideoVisibleApi: def test_removes_stable_sora_run_and_writes_output(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): from remove_ai_watermarks import video_visible @@ -932,9 +1127,11 @@ class TestVideoVisibleApi: [ ("seedance", "scan_seedance_video", "box"), ("dola", "scan_dola_video", "box"), + ("hailuo", "scan_hailuo_video", "box"), + ("kling", "scan_kling_video", "box"), ], ) - def test_dispatches_bytedance_detectors( + def test_dispatches_fixed_mark_detectors( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -985,7 +1182,7 @@ class TestVideoVisibleCli: assert result.exit_code == 0, result.output assert "temporally stable" in result.output - assert "sora|veo|seedance|dola" in result.output + assert "sora|veo|seedance|dola|hailuo|kling" in result.output def test_reports_removed_frames(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): from remove_ai_watermarks import video From 8975916b1143b2b8e9ded5a8378b92697e775d2e Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Thu, 30 Jul 2026 18:36:04 -0700 Subject: [PATCH 06/19] Harden automatic video watermark removal for product use --- CLAUDE.md | 8 +- README.md | 16 +- docs/cli.md | 13 +- docs/known-limitations.md | 29 +- docs/module-internals.md | 12 +- docs/python-api.md | 32 +- docs/supported-signals.md | 5 +- src/remove_ai_watermarks/cli.py | 10 +- src/remove_ai_watermarks/video.py | 150 +++++--- src/remove_ai_watermarks/video_encoding.py | 25 +- src/remove_ai_watermarks/video_invisible.py | 21 +- src/remove_ai_watermarks/video_visible.py | 363 ++++++++++++++------ tests/test_video.py | 284 ++++++++++++++- tests/test_video_invisible.py | 1 + 14 files changed, 737 insertions(+), 232 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 65fd9ee..62a97eb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -28,7 +28,7 @@ Per-command exit-code semantics (the no-signal / GPU-missing skip branches), tes - `uv run remove-ai-watermarks metadata --remove -o ` — strip all AI metadata - `uv run remove-ai-watermarks video metadata --check` — inspect AI metadata in MP4/MOV/M4V/WebM/MKV/AVI/FLV - `uv run remove-ai-watermarks video metadata --remove -o ` — strip verified video metadata without transcoding streams; the experimental video path requires a separate same-container output and defaults to `_clean` -- `uv run remove-ai-watermarks video visible -o ` — remove a temporally recurring Sora mark, or pass `--mark veo|seedance|dola|hailuo|kling` for the current Veo diamond and legacy `Veo` text, the Seedance boxed `AI`, `Dola AI`, the MINIMAX/Hailuo composite, or the versioned Kling label. It scans the full sequence first, transcodes video through ffmpeg, copies audio, strips AI metadata by default, and writes no output when no stable mark is found. `cv2` is the fast default; `migan`/`lama` improve difficult backgrounds. +- `uv run remove-ai-watermarks video visible -o ` — remove a temporally recurring Sora, Veo, Seedance, Dola, Hailuo, or Kling mark. `--mark auto` is the default: all providers share one decode pass, then the first stable result wins in specificity order; an explicit mark restricts the scan. It transcodes video through ffmpeg, copies the complete audio stream, strips AI metadata by default, publishes atomically, and writes no output when no stable mark is found. `cv2` is the fast default; `migan`/`lama` improve difficult backgrounds. - `uv run remove-ai-watermarks batch ` — process every supported image in a directory (output defaults to `_clean/`, set with `-o`). `--mode visible|invisible|metadata|all` (default `visible`); the invisible/all path reuses the full `invisible` knob set above, plus `--backend` and `--sensitivity` for the visible localize -> fill pass. Applies the same no-signal skip per image; see the module doc. **Exit code:** non-zero when any image errored OR (mirroring single `all`) a `--mode invisible`/`all` image carried an invisible signal but the GPU extra was absent, so its SynthID scrub was skipped — it emits a loud warning and copies the input through (invisible mode) so the output dir stays complete; a wrapping service can then detect the incomplete run instead of trusting a silent exit 0. ## Test and lint @@ -82,9 +82,9 @@ Compact map. The full per-module detail (design decisions, tuned thresholds, cal - `upscaler.py` — optional Real-ESRGAN pre-diffusion super-resolution for small inputs (extra `esrgan`, spandrel only). Manual opt-in; the default `--upscaler` stays `lanczos` and the engine always falls back to Lanczos on absence/error. ESRGAN can degrade faces and thin text. - `image_io.py` — centralizes Unicode-safe image IO, alpha preservation, content-based format sniffing, and HEIC/AVIF fallbacks. Callers must check `imwrite` success. No-op visible removal preserves original bytes when the output format is unchanged. - `api.py` — the high-level convenience API, re-exported lazily at the package top level via `__init__.__getattr__` (PEP 562, so `import remove_ai_watermarks` stays cheap): `remove_visible(source, output=None, *, sensitivity="auto", backend="auto", strip_metadata=True, write_noop=True) -> (result_bgr, [labels])` (source = path OR BGR ndarray; a PATH auto-reads metadata provenance and preserves alpha, an ARRAY does neither; `write_noop=True` writes a clean passthrough copy when nothing is removed, `False` leaves `output` untouched so a "no mark = produce nothing" caller like the CLI `visible` command does not clobber a pre-existing file there) and `visible_provenance(path) -> frozenset[str]` (the single metadata→vendor-keys mapper; `cli._visible_provenance` is a thin None-guarded wrapper over it). **`remove_visible` is the ONE path the CLI and library share** — `cli.cmd_visible`'s `--mark auto` branch delegates entirely to it (read → provenance → `remove_auto_marks` → write → `strip_metadata`), so there is no CLI-vs-library drift; `strip_metadata` defaults True to match `visible --strip-metadata`. This is where a library caller should start — NOT the engines directly (`GeminiEngine`/`TextMarkEngine` have no `remove_watermark` any more; removal is registry `remove_auto_marks`/`KnownMark.remove`; the old single-strongest `best_auto_mark` is gone — removal takes EVERY mark). `identify` is NOT top-level re-exported (it collides with the `identify` submodule); use `from remove_ai_watermarks.identify import identify`. -- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, `remove_video_visible(source, output=None, *, mark="sora", backend="cv2", strip_metadata=True) -> VideoVisibleResult`, and `remove_video_invisible(source, output=None, ...) -> VideoInvisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV/AVI/FLV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM, AVI, and FLV TC260 are read by bounded container walkers that skip media payloads and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: fully synthetic Sora, Veo, Seedance, Dola, Hailuo, and Kling silhouettes propose frame boxes, provider-specific temporal recurrence authorizes them, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying audio. Fixed marks require anchored runs so a slowly drifting scene detail cannot pass on adjacent overlap alone. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Seedance uses the full localized box because an outline mask left part of the real translucent border behind. Hailuo expands its matched core to the full composite, while Kling combines font and swirl candidates and requires a bright low-saturation edge label. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels are not built yet. -- `video_encoding.py` — the shared raw-BGR ffmpeg command and pipe lifecycle for both visible removal and invisible regeneration. It copies optional audio, controls metadata/chapter retention explicitly, and keeps container-specific codec arguments in one place. -- `video_invisible.py` — the oracle-gated video SynthID candidate engine for MP4/MOV/M4V. It regenerates frames through `stabilityai/sd-vae-ft-mse` with one seeded latent-noise field shared across time, retains only one configured batch, updates PSNR and motion-compensated temporal residuals incrementally, streams pixels directly to ffmpeg, copies audio, strips metadata, and atomically publishes the completed output. Every result remains explicitly unverified because Google exposes no local decoder. The 2026-07-29 Gemini Flash calibration found the default `0.10` candidate negative on both control-positive public Veo clips, while `0.05` remained positive on one, so the operating point is content-dependent and requires the matching external oracle. +- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, `remove_video_visible(source, output=None, *, mark="auto", backend="cv2", strip_metadata=True) -> VideoVisibleResult`, and `remove_video_invisible(source, output=None, ...) -> VideoInvisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV/AVI/FLV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM, AVI, and FLV TC260 are read by bounded container walkers that skip media payloads and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: auto scans all six synthetic provider silhouettes in one decode pass and selects the first stable result in specificity order because provider confidence scales are not comparable. Provider-specific temporal recurrence authorizes frame boxes, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying complete audio. Fixed marks require anchored runs so a slowly drifting scene detail cannot pass on adjacent overlap alone. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Seedance uses the full localized box because an outline mask left part of the real translucent border behind. Hailuo expands its matched core to the full composite, while Kling combines font and swirl candidates and requires a bright low-saturation edge label. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels are not built yet. +- `video_encoding.py` — the shared raw-BGR ffmpeg command, atomic sibling-temporary publication, and pipe lifecycle for both visible removal and invisible regeneration. It copies complete optional audio without shortening a tail, controls metadata/chapter retention explicitly, and keeps container-specific codec arguments in one place. +- `video_invisible.py` — the oracle-gated video SynthID candidate engine for MP4/MOV/M4V. It regenerates frames through `stabilityai/sd-vae-ft-mse` with one seeded latent-noise field shared across time, retains only one configured batch, updates PSNR and motion-compensated temporal residuals incrementally, streams pixels directly to ffmpeg, copies complete audio, strips metadata, and atomically publishes the completed output. Every result remains explicitly unverified because Google exposes no local decoder. The 2026-07-29 Gemini Flash calibration found the default `0.10` candidate negative on both control-positive public Veo clips, while `0.05` remained positive on one, so the operating point is content-dependent and requires the matching external oracle. For the Doubao alpha-distillation history (why content-image reverse-alpha distillation fails by physics and controlled captures were required), see `docs/research-doubao-distillation.md`. diff --git a/README.md b/README.md index 356e655..431e3a0 100644 --- a/README.md +++ b/README.md @@ -96,13 +96,17 @@ remove-ai-watermarks video visible kling.mp4 --mark kling -o kling_clean.mp4 This path scans the complete sequence before changing pixels. It accepts only a mark that repeats at a stable position across adjacent frames, then reuses the same OpenCV, MI-GAN, or LaMa fill backends as image removal. Audio is copied -without re-encoding; the video stream is transcoded because its pixels change. +without re-encoding and is allowed to reach its natural end; the video stream +is transcoded because its pixels change. The default `--mark auto` scans all +providers in one decode pass and selects the first stable match in the +specificity order shown below. Pass an explicit mark to restrict detection to +one provider. Sora covers the moving Sora 2 mascot and wordmark. Veo covers both the current four-point diamond and the legacy `Veo` text. Seedance covers the fixed boxed `AI` label, Dola covers the fixed `Dola AI` text, Hailuo covers the composite `MINIMAX | hailuo AI` label, and Kling covers the bottom-right `KLING AI` -label with its version suffix. No output is written when no stable mark is -found. +label with its version suffix. A completed encode is published atomically. No +output is written when no stable mark is found. Generate a video SynthID candidate: @@ -112,7 +116,8 @@ remove-ai-watermarks video invisible input.mp4 -o candidate.mp4 ``` This path regenerates the complete sequence with one latent-noise field shared -across time, copies audio, and strips source metadata. It cannot verify +across time, copies complete audio, strips source metadata, and publishes the +completed encode atomically. It cannot verify Google's proprietary pixel watermark locally. The command therefore labels every output `UNVERIFIED` and prints the exact Gemini Flash verification prompt. @@ -266,7 +271,8 @@ print(removed) report = raiw.inspect_video_metadata("input.mp4") cleaned = raiw.remove_video_metadata("input.mp4") candidate = raiw.remove_video_invisible("input.mp4", "candidate.mp4") -visible = raiw.remove_video_visible("sora.mp4", "sora_clean.mp4") +visible = raiw.remove_video_visible("input.mp4", "clean.mp4") +print(visible.mark) veo = raiw.remove_video_visible("veo.mp4", "veo_clean.mp4", mark="veo") seedance = raiw.remove_video_visible( "seedance.mp4", diff --git a/docs/cli.md b/docs/cli.md index 463073e..5c23400 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -210,7 +210,15 @@ must also remain anchored instead of drifting with a scene object. Matching provider provenance may relax the visual score only for registered provenance-aware marks; metadata alone never creates a detection. -The video stream is transcoded and the original audio stream is copied. +`--mark auto` is the default. It evaluates all providers in one decode pass and +selects the first stable match in specificity order: Sora, Veo, Seedance, Dola, +Hailuo, then Kling. Their confidence scores are independently calibrated and +are not compared across providers. Pass an explicit `--mark` to scan only that +provider. + +The video stream is transcoded and the complete original audio stream is +copied without truncating an audio tail that extends beyond the final video +frame. Supported input and output containers are MP4, MOV, M4V, WebM, MKV, AVI, and FLV; the output extension must match the input. The default `cv2` backend is fast but can smear structured backgrounds. Select `--backend migan` or @@ -220,7 +228,8 @@ installed backend. AI metadata is stripped from the encoded output by default. Use `--keep-metadata` to retain mapped container metadata. When no temporally stable mark is found, the command writes no output and exits with the no-visible-mark -status. +status. The final path is replaced atomically only after ffmpeg completes, so a +failed encode does not overwrite an existing result. ## Remove invisible watermarks diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 8fafcbf..8d279a8 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -215,6 +215,11 @@ drifting with a scene object. Kling also requires a bright low-saturation candidate near the expected frame edge. Provider provenance can recover low-contrast runs only after visual evidence exists for the marks that define a provenance prior, so metadata alone does not erase a clean API export. +The default auto-router evaluates all detectors in one decode pass but does not +rank their raw confidence values. Those scores are provider-specific and known +to cross-match in some layouts, so the router applies the independent temporal +policies and selects the first stable result in specificity order. Use an +explicit mark when the provider is already known. Historical Sora Turbo exports use a small OpenAI swirl in the corner rather than the moving mascot-and-wordmark design; that earlier variant is not detected by the `sora` video mark. Hailuo and Kling coverage is specific to the @@ -223,17 +228,19 @@ silhouette. Other provider video labels are not supported yet. Google video SynthID has a candidate-producing VAE path, while other proprietary invisible video watermarks have no registered attack. -Visible removal transcodes the video stream and copies audio. Its frame-local -fill is not a motion-aware video inpainting model. OpenCV can leave a visible -smear where the mark overlaps a hard edge or structured texture, and the smear -can vary over time. MI-GAN and LaMa improve individual frames but do not -guarantee temporal coherence. The Veo diamond uses a shape mask to limit damage -outside the symbol. Seedance fills the full localized box because a synthetic -outline mask left part of the real translucent border visible in an end-to-end -check. OpenCV may therefore soften texture inside that small box; use MI-GAN or -LaMa when reconstruction quality matters. The current encoder also emits a -constant-frame-rate output at the decoded stream rate, so variable-frame-rate -preservation is not yet guaranteed. +Visible removal transcodes the video stream and copies the complete audio +stream without shortening an audio tail. Completed visible and invisible +encodes are published atomically, so an encode failure preserves an existing +output. Its frame-local fill is not a motion-aware video inpainting model. +OpenCV can leave a visible smear where the mark overlaps a hard edge or +structured texture, and the smear can vary over time. MI-GAN and LaMa improve +individual frames but do not guarantee temporal coherence. The Veo diamond +uses a shape mask to limit damage outside the symbol. Seedance fills the full +localized box because a synthetic outline mask left part of the real +translucent border visible in an end-to-end check. OpenCV may therefore soften +texture inside that small box; use MI-GAN or LaMa when reconstruction quality +matters. The current encoder also emits a constant-frame-rate output at the +decoded stream rate, so variable-frame-rate preservation is not yet guaranteed. Native TC260 metadata in MP4/MOV is supported at its normative `moov.udta.meta.keys/ilst` placement, including non-faststart files whose diff --git a/docs/module-internals.md b/docs/module-internals.md index f0ee077..1679e67 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -125,7 +125,9 @@ use the verified ffmpeg stream-copy path for removal. [`video_encoding.py`](../src/remove_ai_watermarks/video_encoding.py) owns the raw-BGR ffmpeg command and pipe lifecycle shared by visible removal and invisible regeneration. It centralizes container codecs, optional audio stream -copying, metadata/chapter policy, and encode-failure reporting. +copying, metadata/chapter policy, encode-failure reporting, and atomic +same-directory publication. Each mapped stream is allowed to reach its own end, +so a copied audio tail is not shortened to the raw-video input duration. [`video_invisible.py`](../src/remove_ai_watermarks/video_invisible.py) implements the oracle-gated video SynthID candidate engine. It samples frames @@ -161,6 +163,14 @@ expected lower-frame area and calibrated independently. A strong relocated Veo diamond may bypass the known layout anchors, but weak free-corner matches never enter the temporal arbiter. +The default `auto` route decodes each frame once, shares its grayscale and +normalized representations across all detectors, and caches resized synthetic +template features for the fixed stream geometry. Provider confidence scales +are not comparable: selection applies each provider's temporal arbiter and +takes the first stable result in specificity order (`sora`, `veo`, `seedance`, +`dola`, `hailuo`, `kling`). An explicit mark uses the same scan path with one +candidate. + Every per-frame result is untrusted. The provider-specific stabilization wrappers share one recurrence implementation, while retaining separate visual floors and minimum-run policy. Provenance can relax a low-contrast run only diff --git a/docs/python-api.md b/docs/python-api.md index c2de706..444bf09 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -174,8 +174,8 @@ if result.remaining_metadata: `remove_video_invisible` supports MP4, MOV, and M4V. It regenerates the complete video through a VAE in bounded batches, shares one seeded latent-noise field -across all frames, streams pixels to ffmpeg, copies audio, and strips source -metadata. The default output is +across all frames, streams pixels to ffmpeg, copies complete audio, strips +source metadata, and publishes atomically. The default output is `input_synthid_candidate.mp4`; a distinct same-container output is required. The returned `VideoInvisibleResult` includes output geometry, frame rate, frame @@ -191,13 +191,15 @@ Flash's built-in content verification before treating it as watermark-negative. import remove_ai_watermarks as raiw result = raiw.remove_video_visible( - "sora.mp4", - "sora_clean.mp4", + "input.mp4", + "clean.mp4", backend="cv2", strip_metadata=True, ) if result.output is None: - print("No temporally stable Sora mark was found") + print("No temporally stable supported mark was found") +else: + print(result.mark) veo_result = raiw.remove_video_visible( "veo.mp4", @@ -228,18 +230,24 @@ kling_result = raiw.remove_video_visible( `remove_video_visible` scans the complete video before writing output. It combines synthetic multi-scale visual matching with temporal consistency, so an -isolated lookalike in one frame is not enough to authorize inpainting. The -supported `mark` values are `sora`, `veo`, `seedance`, `dola`, `hailuo`, and -`kling`. The Veo detector recognizes the current four-point diamond and the +isolated lookalike in one frame is not enough to authorize inpainting. +`mark="auto"` is the default: it evaluates all providers in one decode pass and +selects the first stable match in specificity order (`sora`, `veo`, `seedance`, +`dola`, `hailuo`, `kling`). Provider confidence values are calibrated +independently and are not compared across detectors. Pass one of those explicit +values to restrict the scan to a single provider. The Veo detector recognizes +the current four-point diamond and the legacy `Veo` text. Seedance recognizes the boxed `AI` label, Dola recognizes its compact text label, Hailuo recognizes the composite MINIMAX/Hailuo label, and Kling recognizes its bottom-right logo, wordmark, and version suffix. Each variant has an independent synthetic silhouette and calibrated temporal policy. -The returned `VideoVisibleResult` records the total, detected, and removed frame -counts plus any AI metadata that survived the output encode. The function -returns `output=None` and writes no file when no stable mark is selected. Video -pixels are transcoded through ffmpeg while the source audio stream is copied. +The returned `VideoVisibleResult` records the selected `mark`, the total, +detected, and removed frame counts, plus any AI metadata that survived the +output encode. The function returns `output=None` and writes no file when no +stable mark is selected. Video pixels are transcoded through ffmpeg while the +complete source audio stream is copied. A failed encode preserves any existing +output; only a completed result is published atomically. ## Remove invisible watermarks diff --git a/docs/supported-signals.md b/docs/supported-signals.md index aeb2373..73063c4 100644 --- a/docs/supported-signals.md +++ b/docs/supported-signals.md @@ -41,7 +41,10 @@ when you can select the affected area yourself. | `kling` | Kling swirl, `KLING AI`, version, and optional `PRO` suffix | Fixed bottom-right edge | Combines a synthetic logo rescue with font variants, an edge gate, a white-label gate, and anchored temporal recurrence. | Use `video visible` for this registry. It is separate from the image `visible` -command because selection is made over a sequence rather than one raster. +command because selection is made over a sequence rather than one raster. Its +default `auto` mode scans all six entries in one decode pass and selects the +first temporally stable match in table order; an explicit mark restricts the +scan to that row. ## Fill backends diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 5e0a7d1..bd07d4a 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -28,6 +28,7 @@ from remove_ai_watermarks.noai.watermark_profiles import ( strength_default_help, vendor_for_strength, ) +from remove_ai_watermarks.video import VIDEO_VISIBLE_MARKS from remove_ai_watermarks.video_synthid import ( DEFAULT_VIDEO_SYNTHID_FPS, DEFAULT_VIDEO_SYNTHID_LONG_SIDE, @@ -1240,9 +1241,9 @@ def cmd_video_invisible( ) @click.option( "--mark", - type=click.Choice(["sora", "veo", "seedance", "dola", "hailuo", "kling"]), - default="sora", - help="Visible AI mark to remove.", + type=click.Choice(["auto", *VIDEO_VISIBLE_MARKS]), + default="auto", + help="Visible AI mark to remove. Auto scans every supported provider in one decode pass.", ) @click.option( "--backend", @@ -1281,7 +1282,8 @@ def cmd_video_visible( console.print(f" FAILED: {len(result.remaining_metadata)} AI metadata marker(s) survived in {result.output}") raise SystemExit(1) console.print( - f" Removed {mark} watermark from {result.removed_frames}/{result.total_frames} frames -> {result.output}" + f" Removed {result.mark} watermark from " + f"{result.removed_frames}/{result.total_frames} frames -> {result.output}" ) diff --git a/src/remove_ai_watermarks/video.py b/src/remove_ai_watermarks/video.py index 7b34bbf..5dba122 100644 --- a/src/remove_ai_watermarks/video.py +++ b/src/remove_ai_watermarks/video.py @@ -22,8 +22,10 @@ from remove_ai_watermarks.video_synthid import ( if TYPE_CHECKING: from remove_ai_watermarks.video_invisible import RegenerationMetrics + from remove_ai_watermarks.video_visible import VideoScan VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v", ".webm", ".mkv", ".avi", ".flv"}) +VIDEO_VISIBLE_MARKS = ("sora", "veo", "seedance", "dola", "hailuo", "kling") _ISOBMFF_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v"}) _EBML_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".webm", ".mkv"}) _RIFF_VIDEO_EXTENSIONS: frozenset[str] = frozenset({".avi"}) @@ -182,17 +184,20 @@ def remove_video_visible( source: str | Path, output: str | Path | None = None, *, - mark: str = "sora", + mark: str = "auto", backend: str = "cv2", strip_metadata: bool = True, ) -> VideoVisibleResult: """Remove a supported visible AI wordmark from a video. - Supported marks are ``sora``, ``veo``, ``seedance``, ``dola``, ``hailuo``, - and ``kling``. The full sequence is scanned before pixels change, and only - recurring candidates are accepted. Audio is copied without re-encoding; - video is transcoded because the pixels change. When no stable mark is found, - no output is written and ``output`` in the result is ``None``. + ``mark="auto"`` scans every supported provider in one decode pass and selects + the first stable match in specificity order. Explicit marks are ``sora``, + ``veo``, ``seedance``, ``dola``, ``hailuo``, and ``kling``. The full + sequence is scanned before pixels change, and only recurring candidates are + accepted. Complete audio is copied without re-encoding; video is transcoded + because the pixels change. Completed output is published atomically. When + no stable mark is found, no output is written and ``output`` in the result + is ``None``. """ from remove_ai_watermarks.metadata import get_ai_metadata from remove_ai_watermarks.video_visible import ( @@ -200,12 +205,7 @@ def remove_video_visible( has_bytedance_video_provenance, has_sora_provenance, has_veo_provenance, - scan_dola_video, - scan_hailuo_video, - scan_kling_video, - scan_seedance_video, - scan_sora_video, - scan_veo_video, + scan_video_marks, stabilize_dola_localizations, stabilize_hailuo_localizations, stabilize_kling_localizations, @@ -215,56 +215,98 @@ def remove_video_visible( ) from remove_ai_watermarks.watermark_registry import resolve_backend - if mark not in {"sora", "veo", "seedance", "dola", "hailuo", "kling"}: - raise ValueError("Unsupported visible video mark; expected sora, veo, seedance, dola, hailuo, or kling") + if mark not in {"auto", *VIDEO_VISIBLE_MARKS}: + raise ValueError("Unsupported visible video mark; expected auto, sora, veo, seedance, dola, hailuo, or kling") if backend not in {"auto", "cv2", "migan", "lama"}: raise ValueError("Unsupported fill backend; expected auto, cv2, migan, or lama") source_path = _video_source(source) output_path = _video_output(source_path, output, operation="visible watermark removal") markers = get_ai_metadata(source_path) - if mark == "sora": - scan = scan_sora_video(source_path) - regions = stabilize_sora_localizations( - scan.detections, - provenance=has_sora_provenance(markers), + + def removal_plan( + selected_mark: str, + selected_scan: VideoScan, + ) -> tuple[list[tuple[int, int, int, int] | None], float, Literal["box", "veo"]]: + if selected_mark == "sora": + return ( + stabilize_sora_localizations( + selected_scan.detections, + provenance=has_sora_provenance(markers), + ), + 0.28, + "box", + ) + if selected_mark == "veo": + return ( + stabilize_veo_localizations( + selected_scan.detections, + provenance=has_veo_provenance(markers), + ), + 0.18, + "veo", + ) + if selected_mark == "seedance": + return ( + stabilize_seedance_localizations( + selected_scan.detections, + provenance=has_bytedance_video_provenance(markers), + ), + 0.0, + "box", + ) + if selected_mark == "dola": + return ( + stabilize_dola_localizations( + selected_scan.detections, + provenance=has_bytedance_video_provenance(markers), + ), + 0.20, + "box", + ) + if selected_mark == "hailuo": + return stabilize_hailuo_localizations(selected_scan.detections), 0.12, "box" + return stabilize_kling_localizations(selected_scan.detections), 0.12, "box" + + candidate_marks = VIDEO_VISIBLE_MARKS if mark == "auto" else (mark,) + scans = scan_video_marks(source_path, candidate_marks) + selected: ( + tuple[ + str, + VideoScan, + list[tuple[int, int, int, int] | None], + float, + Literal["box", "veo"], + ] + | None + ) = None + for candidate_mark in candidate_marks: + candidate_scan = scans[candidate_mark] + candidate_regions, candidate_padding, candidate_mask_style = removal_plan( + candidate_mark, + candidate_scan, ) - padding_fraction = 0.28 - mask_style = "box" - elif mark == "veo": - scan = scan_veo_video(source_path) - regions = stabilize_veo_localizations( - scan.detections, - provenance=has_veo_provenance(markers), + if any(region is not None for region in candidate_regions): + selected = ( + candidate_mark, + candidate_scan, + candidate_regions, + candidate_padding, + candidate_mask_style, + ) + break + if selected is None: + scan = scans[candidate_marks[0]] + return VideoVisibleResult( + source=source_path, + output=None, + mark=mark, + total_frames=len(scan.detections), + detected_frames=0, + removed_frames=0, + remaining_metadata=markers if strip_metadata else {}, ) - padding_fraction = 0.18 - mask_style = "veo" - elif mark == "seedance": - scan = scan_seedance_video(source_path) - regions = stabilize_seedance_localizations( - scan.detections, - provenance=has_bytedance_video_provenance(markers), - ) - padding_fraction = 0.0 - mask_style = "box" - elif mark == "dola": - scan = scan_dola_video(source_path) - regions = stabilize_dola_localizations( - scan.detections, - provenance=has_bytedance_video_provenance(markers), - ) - padding_fraction = 0.20 - mask_style = "box" - elif mark == "hailuo": - scan = scan_hailuo_video(source_path) - regions = stabilize_hailuo_localizations(scan.detections) - padding_fraction = 0.12 - mask_style = "box" - else: - scan = scan_kling_video(source_path) - regions = stabilize_kling_localizations(scan.detections) - padding_fraction = 0.12 - mask_style = "box" + mark, scan, regions, padding_fraction, mask_style = selected detected_frames = sum(region is not None for region in regions) if detected_frames == 0: return VideoVisibleResult( diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index 5eef232..84e903c 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -3,16 +3,38 @@ from __future__ import annotations import logging +import os import shutil import subprocess +import tempfile +from contextlib import contextmanager +from pathlib import Path from typing import TYPE_CHECKING if TYPE_CHECKING: - from pathlib import Path + from collections.abc import Generator log = logging.getLogger(__name__) +@contextmanager +def atomic_video_output(output: Path) -> Generator[Path]: + """Yield a sibling temporary path and publish it only after success.""" + output.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + prefix=f".{output.stem}-", + suffix=output.suffix, + dir=output.parent, + delete=False, + ) as stream: + temporary_output = Path(stream.name) + try: + yield temporary_output + os.replace(temporary_output, output) + finally: + temporary_output.unlink(missing_ok=True) + + def _video_codec_args(suffix: str, *, crf: int) -> list[str]: if suffix == ".webm": return ["-c:v", "libvpx-vp9", "-crf", str(crf), "-b:v", "0"] @@ -61,7 +83,6 @@ def raw_video_command( "-1" if strip_metadata else "1", "-map_chapters", "-1" if strip_metadata else "1", - "-shortest", ] if output.suffix.lower() in {".mp4", ".mov", ".m4v"}: command.extend(["-movflags", "+faststart"]) diff --git a/src/remove_ai_watermarks/video_invisible.py b/src/remove_ai_watermarks/video_invisible.py index 39fbfbb..c9bd66a 100644 --- a/src/remove_ai_watermarks/video_invisible.py +++ b/src/remove_ai_watermarks/video_invisible.py @@ -13,11 +13,8 @@ from __future__ import annotations # pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false import logging import math -import os -import tempfile from dataclasses import dataclass from importlib.util import find_spec -from pathlib import Path from typing import TYPE_CHECKING, Any import cv2 @@ -25,6 +22,7 @@ import numpy as np from remove_ai_watermarks.video_encoding import ( abort_raw_video_encoder, + atomic_video_output, finish_raw_video_encoder, raw_video_command, start_raw_video_encoder, @@ -39,6 +37,7 @@ from remove_ai_watermarks.video_synthid import ( if TYPE_CHECKING: from collections.abc import Iterable, Sequence + from pathlib import Path log = logging.getLogger(__name__) @@ -405,15 +404,7 @@ def regenerate_video_candidate( vae.eval() vae.enable_slicing() - output.parent.mkdir(parents=True, exist_ok=True) - with tempfile.NamedTemporaryFile( - prefix=f".{output.stem}-", - suffix=output.suffix, - dir=output.parent, - delete=False, - ) as stream: - temporary_output = Path(stream.name) - try: + with atomic_video_output(output) as temporary_output: process = start_raw_video_encoder( raw_video_command( source, @@ -499,7 +490,7 @@ def regenerate_video_candidate( raise mse = squared_error / pixel_count - metrics = RegenerationMetrics( + return RegenerationMetrics( frames=frame_count, fps=effective_fps, width=size[0], @@ -507,7 +498,3 @@ def regenerate_video_candidate( psnr_db=math.inf if mse == 0.0 else 20.0 * math.log10(255.0 / math.sqrt(mse)), temporal_residual_ratio=temporal_candidate / max(temporal_baseline, 1e-6), ) - os.replace(temporary_output, output) - finally: - temporary_output.unlink(missing_ok=True) - return metrics diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py index e7b6949..e7bd2ed 100644 --- a/src/remove_ai_watermarks/video_visible.py +++ b/src/remove_ai_watermarks/video_visible.py @@ -11,8 +11,8 @@ requires the candidate to recur at the same location across adjacent frames. This keeps isolated lookalikes in clean videos from becoming removal masks. Video pixels are decoded with OpenCV and encoded with the system ``ffmpeg``. -Audio is stream-copied from the source. The video stream must be transcoded -because visible-mark removal changes pixels. +Complete audio is stream-copied from the source. The video stream must be +transcoded because visible-mark removal changes pixels. """ # cv2/numpy boundary: these packages do not expose usable types for many array @@ -32,8 +32,10 @@ import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont +from remove_ai_watermarks.video import VIDEO_VISIBLE_MARKS from remove_ai_watermarks.video_encoding import ( abort_raw_video_encoder, + atomic_video_output, finish_raw_video_encoder, raw_video_command, start_raw_video_encoder, @@ -105,6 +107,15 @@ class VideoScan: detections: tuple[FrameLocalization, ...] +@dataclass(frozen=True) +class _PreparedFrame: + """Frame representations shared by every detector in one scan pass.""" + + gray: NDArray[Any] + normalized_gray: NDArray[Any] + normalized_scale: float + + def _scalable_default_font(size: int) -> ImageFont.ImageFont | ImageFont.FreeTypeFont: """Load Pillow's bundled scalable font, with a Pillow 10.0 fallback.""" try: @@ -291,6 +302,43 @@ def _top_hat(gray: NDArray[Any]) -> NDArray[Any]: return cv2.morphologyEx(gray, cv2.MORPH_TOPHAT, kernel) +@lru_cache(maxsize=1) +def _template_sources() -> dict[str, NDArray[Any]]: + """Return every immutable synthetic template by detector-local key.""" + sora_word, sora_icon = _sora_templates() + veo_diamond, veo_text = _veo_templates() + sources = { + "sora-word": sora_word, + "sora-icon": sora_icon, + "veo-diamond": veo_diamond, + "veo-text": veo_text, + "seedance": _seedance_template(), + "dola": _dola_template(), + "hailuo": _hailuo_template(), + "kling-logo": _kling_logo_template(), + } + sources.update({f"kling-{index}": template for index, template in enumerate(_kling_templates())}) + return sources + + +@lru_cache(maxsize=512) +def _resized_template_feature( + template_key: str, + width: int, + height: int, + kernel_size: int, +) -> tuple[NDArray[Any], NDArray[Any]]: + """Resize one template and cache its invariant top-hat representation.""" + template = cv2.resize( + _template_sources()[template_key], + (width, height), + interpolation=cv2.INTER_AREA, + ) + kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8) + feature = cv2.morphologyEx(template, cv2.MORPH_TOPHAT, kernel) + return template, feature + + def _normalized_gray(image_bgr: NDArray[Any]) -> tuple[NDArray[Any], float]: gray = image_bgr if image_bgr.ndim == 2 else cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) height, width = gray.shape[:2] @@ -307,6 +355,13 @@ def _normalized_gray(image_bgr: NDArray[Any]) -> tuple[NDArray[Any], float]: return gray, scale +def _prepare_frame(image_bgr: NDArray[Any]) -> _PreparedFrame: + """Compute the shared grayscale representations for one decoded frame.""" + gray = image_bgr if image_bgr.ndim == 2 else cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) + normalized_gray, normalized_scale = _normalized_gray(gray) + return _PreparedFrame(gray, normalized_gray, normalized_scale) + + def _expanded_region( location: tuple[int, int], template_width: int, @@ -328,7 +383,12 @@ def _expanded_region( return x, y, max(1, width), max(1, height) -def detect_sora_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: +def detect_sora_frame( + image_bgr: NDArray[Any], + *, + frame_index: int = 0, + prepared: _PreparedFrame | None = None, +) -> FrameLocalization: """Locate the strongest synthetic Sora-wordmark match in one frame. The returned candidate is intentionally untrusted. Call @@ -339,7 +399,8 @@ def detect_sora_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> Frame return FrameLocalization(frame_index, 0.0, None) frame_height, frame_width = image_bgr.shape[:2] - gray, scale = _normalized_gray(image_bgr) + prepared = prepared or _prepare_frame(image_bgr) + gray, scale = prepared.normalized_gray, prepared.normalized_scale normalized_height, normalized_width = gray.shape[:2] feature = _top_hat(gray) best_confidence = 0.0 @@ -347,17 +408,19 @@ def detect_sora_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> Frame for template_index, base_template in enumerate(_sora_templates()): icon_only = template_index == 1 + template_key = "sora-icon" if icon_only else "sora-word" for relative_height in _SORA_RELATIVE_HEIGHTS: template_height = max(16, round(min(normalized_height, normalized_width) * relative_height)) template_width = max(1, round(base_template.shape[1] * template_height / base_template.shape[0])) if template_height >= normalized_height or template_width >= normalized_width: continue - template = cv2.resize( - base_template, - (template_width, template_height), - interpolation=cv2.INTER_AREA, + _, template_feature = _resized_template_feature( + template_key, + template_width, + template_height, + 7, ) - scores = cv2.matchTemplate(feature, _top_hat(template), cv2.TM_CCOEFF_NORMED) + scores = cv2.matchTemplate(feature, template_feature, cv2.TM_CCOEFF_NORMED) _, confidence, _, location = cv2.minMaxLoc(scores) if confidence <= best_confidence: continue @@ -381,6 +444,7 @@ def _match_template( *, region: Region, kernel_size: int, + template_feature: NDArray[Any] | None = None, ) -> tuple[float, Region | None]: """Match one synthetic silhouette inside a bounded frame region.""" x, y, width, height = region @@ -390,7 +454,8 @@ def _match_template( return 0.0, None kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8) feature = cv2.morphologyEx(roi, cv2.MORPH_TOPHAT, kernel) - template_feature = cv2.morphologyEx(template, cv2.MORPH_TOPHAT, kernel) + if template_feature is None: + template_feature = cv2.morphologyEx(template, cv2.MORPH_TOPHAT, kernel) scores = cv2.matchTemplate(feature, template_feature, cv2.TM_CCOEFF_NORMED) _, confidence, _, location = cv2.minMaxLoc(scores) return float(confidence), ( @@ -441,7 +506,7 @@ def _bounded_region( def _detect_fixed_mark( image_bgr: NDArray[Any], - template: NDArray[Any], + template_key: str, *, relative_heights: tuple[float, ...], search_origin: tuple[float, float], @@ -467,19 +532,23 @@ def _detect_fixed_mark( normalized_height - search_y, ) matches: list[tuple[float, Region]] = [] + base_template = _template_sources()[template_key] for relative_height in relative_heights: template_height = max(6, round(short_side * relative_height)) - template_width = max(1, round(template.shape[1] * template_height / template.shape[0])) - resized = cv2.resize( - template, - (template_width, template_height), - interpolation=cv2.INTER_AREA, + template_width = max(1, round(base_template.shape[1] * template_height / base_template.shape[0])) + kernel_size = max(3, round(template_height * kernel_fraction) | 1) + resized, template_feature = _resized_template_feature( + template_key, + template_width, + template_height, + kernel_size, ) confidence, candidate = _match_template( gray, resized, region=search_region, - kernel_size=max(3, round(template_height * kernel_fraction) | 1), + kernel_size=kernel_size, + template_feature=template_feature, ) if candidate is not None and confidence > 0: matches.append((confidence, candidate)) @@ -509,38 +578,56 @@ def _detect_fixed_mark( ) -def detect_seedance_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: +def detect_seedance_frame( + image_bgr: NDArray[Any], + *, + frame_index: int = 0, + prepared: _PreparedFrame | None = None, +) -> FrameLocalization: """Locate the strongest fixed Seedance boxed-AI candidate.""" return _detect_fixed_mark( image_bgr, - _seedance_template(), + "seedance", relative_heights=(0.065, 0.075, 0.085, 0.095, 0.105), search_origin=(0.68, 0.72), kernel_fraction=0.12, + normalized=None if prepared is None else (prepared.normalized_gray, prepared.normalized_scale), frame_index=frame_index, ) -def detect_dola_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: +def detect_dola_frame( + image_bgr: NDArray[Any], + *, + frame_index: int = 0, + prepared: _PreparedFrame | None = None, +) -> FrameLocalization: """Locate the strongest fixed Dola AI text candidate.""" return _detect_fixed_mark( image_bgr, - _dola_template(), + "dola", relative_heights=_DOLA_RELATIVE_HEIGHTS, search_origin=(0.65, 0.85), kernel_fraction=0.50, + normalized=None if prepared is None else (prepared.normalized_gray, prepared.normalized_scale), frame_index=frame_index, ) -def detect_hailuo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: +def detect_hailuo_frame( + image_bgr: NDArray[Any], + *, + frame_index: int = 0, + prepared: _PreparedFrame | None = None, +) -> FrameLocalization: """Locate the strongest fixed MINIMAX/Hailuo composite-label candidate.""" detection = _detect_fixed_mark( image_bgr, - _hailuo_template(), + "hailuo", relative_heights=_HAILUO_RELATIVE_HEIGHTS, search_origin=(0.28, 0.76), kernel_fraction=0.18, + normalized=None if prepared is None else (prepared.normalized_gray, prepared.normalized_scale), frame_index=frame_index, ) if detection.region is None: @@ -563,17 +650,23 @@ def detect_hailuo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> Fra ) -def detect_kling_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: +def detect_kling_frame( + image_bgr: NDArray[Any], + *, + frame_index: int = 0, + prepared: _PreparedFrame | None = None, +) -> FrameLocalization: """Locate the fixed Kling wordmark core and include its version suffix.""" if image_bgr.size == 0: return FrameLocalization(frame_index, 0.0, None) frame_height, frame_width = image_bgr.shape[:2] - normalized = _normalized_gray(image_bgr) + prepared = prepared or _prepare_frame(image_bgr) + normalized = prepared.normalized_gray, prepared.normalized_scale expanded: list[FrameLocalization] = [] - for template in _kling_templates(): + for template_index, _template in enumerate(_kling_templates()): detection = _detect_fixed_mark( image_bgr, - template, + f"kling-{template_index}", relative_heights=_KLING_RELATIVE_HEIGHTS, search_origin=(0.64, 0.84), kernel_fraction=0.18, @@ -615,7 +708,7 @@ def detect_kling_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> Fram logo = _detect_fixed_mark( image_bgr, - _kling_logo_template(), + "kling-logo", relative_heights=tuple(value / 1000 for value in range(20, 61, 3)), search_origin=(0.62, 0.90), kernel_fraction=0.18, @@ -664,15 +757,20 @@ def detect_kling_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> Fram return best -def detect_veo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameLocalization: +def detect_veo_frame( + image_bgr: NDArray[Any], + *, + frame_index: int = 0, + prepared: _PreparedFrame | None = None, +) -> FrameLocalization: """Locate the strongest current-diamond or legacy-text Veo candidate.""" if image_bgr.size == 0: return FrameLocalization(frame_index, 0.0, None) frame_height, frame_width = image_bgr.shape[:2] - gray = image_bgr if image_bgr.ndim == 2 else cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) + gray = (prepared or _prepare_frame(image_bgr)).gray short_scale = min(frame_height, frame_width) / _VEO_REFERENCE_SHORT_SIDE - diamond_base, text_base = _veo_templates() + _, text_base = _veo_templates() best_confidence = 0.0 best_region: Region | None = None @@ -680,10 +778,12 @@ def detect_veo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameL for base_size, right_margin, bottom_margin in _VEO_DIAMOND_PROFILES: diamond_size = max(16, round(base_size * short_scale)) diamond_sizes.add(diamond_size) - template = cv2.resize( - diamond_base, - (diamond_size, diamond_size), - interpolation=cv2.INTER_AREA, + kernel_size = max(3, round(7 * short_scale) | 1) + template, template_feature = _resized_template_feature( + "veo-diamond", + diamond_size, + diamond_size, + kernel_size, ) expected_x = round(frame_width - (right_margin + base_size) * short_scale) expected_y = round(frame_height - (bottom_margin + base_size) * short_scale) @@ -696,7 +796,8 @@ def detect_veo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameL gray, template, region=(search_x, search_y, search_width, search_height), - kernel_size=max(3, round(7 * short_scale) | 1), + kernel_size=kernel_size, + template_feature=template_feature, ) if confidence > best_confidence: best_confidence = confidence @@ -710,16 +811,19 @@ def detect_veo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameL corner_y = round(frame_height * 0.65) corner_region = (corner_x, corner_y, frame_width - corner_x, frame_height - corner_y) for diamond_size in diamond_sizes: - template = cv2.resize( - diamond_base, - (diamond_size, diamond_size), - interpolation=cv2.INTER_AREA, + kernel_size = max(3, round(7 * short_scale) | 1) + template, template_feature = _resized_template_feature( + "veo-diamond", + diamond_size, + diamond_size, + kernel_size, ) confidence, candidate = _match_template( gray, template, region=corner_region, - kernel_size=max(3, round(7 * short_scale) | 1), + kernel_size=kernel_size, + template_feature=template_feature, ) if confidence >= 0.70 and confidence > best_confidence: best_confidence = confidence @@ -736,16 +840,19 @@ def detect_veo_frame(image_bgr: NDArray[Any], *, frame_index: int = 0) -> FrameL text_heights = sorted({max(5, round(height * short_scale)) for height in range(8, 22)}) for text_height in text_heights: text_width = max(1, round(text_base.shape[1] * text_height / text_base.shape[0])) - template = cv2.resize( - text_base, - (text_width, text_height), - interpolation=cv2.INTER_AREA, + kernel_size = max(3, round(3 * short_scale) | 1) + template, template_feature = _resized_template_feature( + "veo-text", + text_width, + text_height, + kernel_size, ) confidence, candidate = _match_template( gray, template, region=text_region, - kernel_size=max(3, round(3 * short_scale) | 1), + kernel_size=kernel_size, + template_feature=template_feature, ) if confidence > best_confidence: best_confidence = confidence @@ -959,11 +1066,11 @@ def _stabilize_localizations( return accepted -def _scan_video( +def _scan_video_detectors( source: Path, - detector: Any, -) -> VideoScan: - """Decode a video once and collect one untrusted candidate per frame.""" + detectors: dict[str, Any], +) -> dict[str, VideoScan]: + """Decode once and collect one untrusted candidate per detector and frame.""" capture = cv2.VideoCapture(str(source)) if not capture.isOpened(): raise RuntimeError(f"OpenCV could not decode video: {source}") @@ -974,7 +1081,7 @@ def _scan_video( capture.release() raise RuntimeError(f"Video has invalid stream geometry or frame rate: {source}") - detections: list[FrameLocalization] = [] + detections: dict[str, list[FrameLocalization]] = {mark: [] for mark in detectors} frame_index = 0 while True: ok, frame = capture.read() @@ -983,12 +1090,56 @@ def _scan_video( if frame.shape[:2] != (height, width): capture.release() raise RuntimeError(f"Video changes frame dimensions at frame {frame_index}: {source}") - detections.append(detector(frame, frame_index=frame_index)) + prepared = _prepare_frame(frame) + for mark, detector in detectors.items(): + detections[mark].append( + detector( + frame, + frame_index=frame_index, + prepared=prepared, + ) + ) frame_index += 1 capture.release() - if not detections: + if frame_index == 0: raise RuntimeError(f"Video contains no decodable frames: {source}") - return VideoScan(width, height, fps, tuple(detections)) + return {mark: VideoScan(width, height, fps, tuple(mark_detections)) for mark, mark_detections in detections.items()} + + +def _scan_video( + source: Path, + detector: Any, +) -> VideoScan: + """Decode a video once and collect one untrusted candidate per frame.""" + return _scan_video_detectors(source, {"selected": detector})["selected"] + + +def scan_video_marks( + source: Path, + marks: tuple[str, ...] = VIDEO_VISIBLE_MARKS, +) -> dict[str, VideoScan]: + """Decode once and collect candidates for every requested provider mark.""" + detectors = dict( + zip( + VIDEO_VISIBLE_MARKS, + ( + detect_sora_frame, + detect_veo_frame, + detect_seedance_frame, + detect_dola_frame, + detect_hailuo_frame, + detect_kling_frame, + ), + strict=True, + ) + ) + unsupported = sorted(set(marks) - detectors.keys()) + if unsupported: + raise ValueError(f"Unsupported visible video mark: {', '.join(unsupported)}") + return _scan_video_detectors( + source, + {mark: detectors[mark] for mark in marks}, + ) def scan_sora_video(source: Path) -> VideoScan: @@ -1075,65 +1226,65 @@ def encode_clean_video( padding_fraction: float = 0.28, mask_style: Literal["box", "veo"] = "box", ) -> int: - """Decode again, fill accepted regions, and encode video while copying audio.""" + """Decode again, fill accepted regions, and atomically encode with complete audio.""" from remove_ai_watermarks.watermark_registry import fill, resolve_backend if len(regions) != len(scan.detections): raise ValueError("Temporal localization count does not match the scanned frame count") - output.parent.mkdir(parents=True, exist_ok=True) - process = start_raw_video_encoder( - raw_video_command( - source, - output, - width=scan.width, - height=scan.height, - fps=scan.fps, - strip_metadata=strip_metadata, - crf=14, + with atomic_video_output(output) as temporary_output: + process = start_raw_video_encoder( + raw_video_command( + source, + temporary_output, + width=scan.width, + height=scan.height, + fps=scan.fps, + strip_metadata=strip_metadata, + crf=14, + ) ) - ) - frame_pipe = process.stdin - if frame_pipe is None: - abort_raw_video_encoder(process) - raise RuntimeError("Could not open ffmpeg input pipe") - - capture = cv2.VideoCapture(str(source)) - if not capture.isOpened(): - abort_raw_video_encoder(process) - raise RuntimeError(f"OpenCV could not reopen video for removal: {source}") - - removed_frames = 0 - resolved_backend: Literal["cv2", "migan", "lama"] = resolve_backend(backend) - try: - for frame_index, region in enumerate(regions): - ok, frame = capture.read() - if not ok: - raise RuntimeError(f"Video ended while reading frame {frame_index}: {source}") - if region is not None: - frame = fill( - frame, - _mask_for_region( - frame, - region, - padding_fraction=padding_fraction, - mask_style=mask_style, - ), - backend=resolved_backend, - ) - removed_frames += 1 - frame_pipe.write(frame.tobytes()) - finish_raw_video_encoder( - process, - output, - operation="visible-watermark encode", - ) - except Exception: - if process.poll() is None: + frame_pipe = process.stdin + if frame_pipe is None: abort_raw_video_encoder(process) - raise - finally: - capture.release() + raise RuntimeError("Could not open ffmpeg input pipe") + + capture = cv2.VideoCapture(str(source)) + if not capture.isOpened(): + abort_raw_video_encoder(process) + raise RuntimeError(f"OpenCV could not reopen video for removal: {source}") + + removed_frames = 0 + resolved_backend: Literal["cv2", "migan", "lama"] = resolve_backend(backend) + try: + for frame_index, region in enumerate(regions): + ok, frame = capture.read() + if not ok: + raise RuntimeError(f"Video ended while reading frame {frame_index}: {source}") + if region is not None: + frame = fill( + frame, + _mask_for_region( + frame, + region, + padding_fraction=padding_fraction, + mask_style=mask_style, + ), + backend=resolved_backend, + ) + removed_frames += 1 + frame_pipe.write(frame.tobytes()) + finish_raw_video_encoder( + process, + temporary_output, + operation="visible-watermark encode", + ) + except Exception: + if process.poll() is None: + abort_raw_video_encoder(process) + raise + finally: + capture.release() return removed_frames diff --git a/tests/test_video.py b/tests/test_video.py index 2450891..31ec243 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -2,6 +2,7 @@ from __future__ import annotations +import io from typing import TYPE_CHECKING import cv2 @@ -1018,7 +1019,249 @@ class TestAdditionalProviderTemporalArbiter: assert stabilize(strong) == [box] * 12 +class TestVideoVisibleScan: + def test_auto_prepares_each_frame_once_for_every_detector( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video_visible import FrameLocalization, scan_video_marks + + frame = np.zeros((8, 12, 3), dtype=np.uint8) + + class FakeCapture: + def __init__(self) -> None: + self._read = False + + def isOpened(self) -> bool: + return True + + def get(self, property_id: int) -> float: + return { + video_visible.cv2.CAP_PROP_FRAME_WIDTH: 12.0, + video_visible.cv2.CAP_PROP_FRAME_HEIGHT: 8.0, + video_visible.cv2.CAP_PROP_FPS: 24.0, + }[property_id] + + def read(self) -> tuple[bool, np.ndarray | None]: + if self._read: + return False, None + self._read = True + return True, frame + + def release(self) -> None: + pass + + prepared_ids: list[int] = [] + + def fake_detector( + _frame: np.ndarray, + *, + frame_index: int, + prepared: object, + ) -> FrameLocalization: + assert prepared is not None + prepared_ids.append(id(prepared)) + return FrameLocalization(frame_index, 0.0, None) + + monkeypatch.setattr(video_visible.cv2, "VideoCapture", lambda _path: FakeCapture()) + for detector_name in ( + "detect_sora_frame", + "detect_veo_frame", + "detect_seedance_frame", + "detect_dola_frame", + "detect_hailuo_frame", + "detect_kling_frame", + ): + monkeypatch.setattr(video_visible, detector_name, fake_detector) + + scans = scan_video_marks( + tmp_path / "synthetic.mp4", + ("sora", "veo", "seedance", "dola", "hailuo", "kling"), + ) + + assert set(scans) == {"sora", "veo", "seedance", "dola", "hailuo", "kling"} + assert len(prepared_ids) == 6 + assert len(set(prepared_ids)) == 1 + + +class TestVideoVisibleEncoding: + @staticmethod + def _patch_single_frame_encode( + monkeypatch: pytest.MonkeyPatch, + *, + encoded_bytes: bytes, + fail: bool, + ) -> tuple[object, list[Path]]: + from remove_ai_watermarks import video_visible, watermark_registry + + frame = np.full((8, 8, 3), 32, dtype=np.uint8) + + class FakeCapture: + def __init__(self) -> None: + self._read = False + + def isOpened(self) -> bool: + return True + + def read(self) -> tuple[bool, np.ndarray | None]: + if self._read: + return False, None + self._read = True + return True, frame.copy() + + def release(self) -> None: + pass + + class FakeProcess: + def __init__(self) -> None: + self.stdin = io.BytesIO() + + def poll(self) -> int: + return 1 + + targets: list[Path] = [] + + def fake_command(_source: Path, target: Path, **_kwargs: object) -> list[str]: + targets.append(target) + return ["ffmpeg", str(target)] + + def fake_finish(_process: object, target: Path, **_kwargs: object) -> None: + target.write_bytes(encoded_bytes) + if fail: + raise RuntimeError("synthetic encode failure") + + process = FakeProcess() + monkeypatch.setattr(video_visible.cv2, "VideoCapture", lambda _path: FakeCapture()) + monkeypatch.setattr(video_visible, "raw_video_command", fake_command) + monkeypatch.setattr(video_visible, "start_raw_video_encoder", lambda _command: process) + monkeypatch.setattr(video_visible, "finish_raw_video_encoder", fake_finish) + monkeypatch.setattr(watermark_registry, "resolve_backend", lambda _backend: "cv2") + return process, targets + + def test_publishes_completed_encode_atomically( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan, encode_clean_video + + source = tmp_path / "source.mp4" + source.write_bytes(b"source") + output = tmp_path / "clean.mp4" + scan = VideoScan(8, 8, 24.0, (FrameLocalization(0, 0.0, None),)) + _process, targets = self._patch_single_frame_encode( + monkeypatch, + encoded_bytes=b"complete", + fail=False, + ) + + encode_clean_video( + source, + output, + scan, + [None], + backend="cv2", + strip_metadata=True, + ) + + assert output.read_bytes() == b"complete" + assert targets[0] != output + assert targets[0].suffix == output.suffix + assert not targets[0].exists() + + def test_failed_encode_preserves_existing_output( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan, encode_clean_video + + source = tmp_path / "source.mp4" + source.write_bytes(b"source") + output = tmp_path / "clean.mp4" + output.write_bytes(b"previous") + scan = VideoScan(8, 8, 24.0, (FrameLocalization(0, 0.0, None),)) + _process, targets = self._patch_single_frame_encode( + monkeypatch, + encoded_bytes=b"partial", + fail=True, + ) + + with pytest.raises(RuntimeError, match="synthetic encode failure"): + encode_clean_video( + source, + output, + scan, + [None], + backend="cv2", + strip_metadata=True, + ) + + assert output.read_bytes() == b"previous" + assert not targets[0].exists() + + class TestVideoVisibleApi: + def test_auto_prefers_specific_sora_run_over_hailuo_cross_match( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video import remove_video_visible + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + sora_box = (4, 4, 20, 8) + hailuo_box = (30, 40, 28, 12) + sora_scan = VideoScan( + width=64, + height=64, + fps=24.0, + detections=tuple(FrameLocalization(index, 0.66, sora_box) for index in range(12)), + ) + hailuo_scan = VideoScan( + width=64, + height=64, + fps=24.0, + detections=tuple(FrameLocalization(index, 0.35, hailuo_box) for index in range(12)), + ) + + def fake_scan(_source: Path, marks: tuple[str, ...]): + assert marks == ("sora", "veo", "seedance", "dola", "hailuo", "kling") + return { + "sora": sora_scan, + "veo": sora_scan, + "seedance": sora_scan, + "dola": sora_scan, + "hailuo": hailuo_scan, + "kling": sora_scan, + } + + monkeypatch.setattr(video_visible, "scan_video_marks", fake_scan) + + def fake_encode( + _source: Path, + target: Path, + _scan: VideoScan, + regions: list[tuple[int, int, int, int] | None], + **kwargs: object, + ) -> int: + assert regions == [sora_box] * 12 + assert kwargs["padding_fraction"] == 0.28 + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return 12 + + monkeypatch.setattr(video_visible, "encode_clean_video", fake_encode) + + result = remove_video_visible(source, output) + + assert result.output == output + assert result.mark == "sora" + def test_removes_stable_sora_run_and_writes_output(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): from remove_ai_watermarks import video_visible from remove_ai_watermarks.video import remove_video_visible @@ -1033,7 +1276,11 @@ class TestVideoVisibleApi: fps=24.0, detections=tuple(FrameLocalization(index, 0.66, box) for index in range(5)), ) - monkeypatch.setattr(video_visible, "scan_sora_video", lambda _source: scan) + monkeypatch.setattr( + video_visible, + "scan_video_marks", + lambda _source, marks: {"sora": scan} if marks == ("sora",) else {}, + ) def fake_encode( _source: Path, @@ -1048,7 +1295,7 @@ class TestVideoVisibleApi: monkeypatch.setattr(video_visible, "encode_clean_video", fake_encode) - result = remove_video_visible(source, output) + result = remove_video_visible(source, output, mark="sora") assert result.output == output assert result.detected_frames == 5 @@ -1072,9 +1319,13 @@ class TestVideoVisibleApi: FrameLocalization(2, 0.70, (1, 30, 20, 8)), ), ) - monkeypatch.setattr(video_visible, "scan_sora_video", lambda _source: scan) + monkeypatch.setattr( + video_visible, + "scan_video_marks", + lambda _source, marks: {"sora": scan} if marks == ("sora",) else {}, + ) - result = remove_video_visible(source, output) + result = remove_video_visible(source, output, mark="sora") assert result.output is None assert result.removed_frames == 0 @@ -1098,7 +1349,11 @@ class TestVideoVisibleApi: fps=24.0, detections=tuple(FrameLocalization(index, 0.60, box) for index in range(12)), ) - monkeypatch.setattr(video_visible, "scan_veo_video", lambda _source: scan) + monkeypatch.setattr( + video_visible, + "scan_video_marks", + lambda _source, marks: {"veo": scan} if marks == ("veo",) else {}, + ) def fake_encode( _source: Path, @@ -1123,12 +1378,12 @@ class TestVideoVisibleApi: assert result.removed_frames == 12 @pytest.mark.parametrize( - ("mark", "scan_name", "mask_style"), + ("mark", "mask_style"), [ - ("seedance", "scan_seedance_video", "box"), - ("dola", "scan_dola_video", "box"), - ("hailuo", "scan_hailuo_video", "box"), - ("kling", "scan_kling_video", "box"), + ("seedance", "box"), + ("dola", "box"), + ("hailuo", "box"), + ("kling", "box"), ], ) def test_dispatches_fixed_mark_detectors( @@ -1136,7 +1391,6 @@ class TestVideoVisibleApi: tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mark: str, - scan_name: str, mask_style: str, ): from remove_ai_watermarks import video_visible @@ -1152,7 +1406,11 @@ class TestVideoVisibleApi: fps=24.0, detections=tuple(FrameLocalization(index, 0.60, box) for index in range(12)), ) - monkeypatch.setattr(video_visible, scan_name, lambda _source: scan) + monkeypatch.setattr( + video_visible, + "scan_video_marks", + lambda _source, marks: {mark: scan} if marks == (mark,) else {}, + ) def fake_encode( _source: Path, @@ -1182,7 +1440,7 @@ class TestVideoVisibleCli: assert result.exit_code == 0, result.output assert "temporally stable" in result.output - assert "sora|veo|seedance|dola|hailuo|kling" in result.output + assert "auto|sora|veo|seedance|dola|hailuo|kling" in result.output def test_reports_removed_frames(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): from remove_ai_watermarks import video diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index 3bc4d4b..c520cd5 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -55,6 +55,7 @@ def test_encoder_command_discards_metadata_and_copies_audio( audio_codec_index = command.index("-c:a") assert command[audio_codec_index + 1] == "copy" assert "pipe:0" in command + assert "-shortest" not in command def test_stream_batches_consumes_only_one_batch_ahead() -> None: From 6315b58628c23157463a8af4cc870f8086a071fb Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 10:42:39 -0700 Subject: [PATCH 07/19] feat: complete product video watermark pipeline --- .github/workflows/test.yml | 17 + CLAUDE.md | 14 +- README.md | 97 +- data/README.md | 2 + data/evaluations/video-synthid-oracle.csv | 3 + docs/cli.md | 113 +- docs/known-limitations.md | 84 +- docs/module-internals.md | 103 +- docs/python-api.md | 93 +- docs/release-and-distribution.md | 4 +- docs/supported-signals.md | 30 +- docs/synthid.md | 39 +- docs/verification-plan.md | 72 +- packaging/conda/recipe.yaml | 13 +- pyproject.toml | 6 + scripts/video_synthid_sweep.py | 15 +- src/remove_ai_watermarks/__init__.py | 21 +- src/remove_ai_watermarks/cli.py | 308 ++++- src/remove_ai_watermarks/metadata.py | 17 +- src/remove_ai_watermarks/noai/isobmff.py | 154 ++- src/remove_ai_watermarks/video.py | 643 ++++++++-- src/remove_ai_watermarks/video_encoding.py | 280 ++++- src/remove_ai_watermarks/video_invisible.py | 189 ++- src/remove_ai_watermarks/video_synthid.py | 5 +- src/remove_ai_watermarks/video_temporal.py | 189 +++ src/remove_ai_watermarks/video_visible.py | 186 ++- tests/test_metadata.py | 6 +- tests/test_noai.py | 38 + tests/test_sync_conda_recipe.py | 8 + tests/test_video.py | 1216 ++++++++++++++++++- tests/test_video_invisible.py | 122 +- tests/test_video_temporal.py | 112 ++ uv.lock | 109 ++ 33 files changed, 3793 insertions(+), 515 deletions(-) create mode 100644 data/evaluations/video-synthid-oracle.csv create mode 100644 src/remove_ai_watermarks/video_temporal.py create mode 100644 tests/test_video_temporal.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 719ecde..104ae25 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -44,3 +44,20 @@ jobs: run: uv sync --frozen --extra dev - name: Run tests run: uv run pytest -q + + video-e2e: + name: video full-clip end-to-end + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: astral-sh/setup-uv@v7 + with: + python-version: "3.12" + - name: Install ffmpeg + run: sudo apt-get update && sudo apt-get install --yes ffmpeg + - name: Sync dev environment + run: uv sync --frozen --extra dev + - name: Run full-clip video test + run: >- + uv run pytest -q + tests/test_video.py::TestVideoVisibleFullClip diff --git a/CLAUDE.md b/CLAUDE.md index 62a97eb..6633339 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -26,9 +26,12 @@ Per-command exit-code semantics (the no-signal / GPU-missing skip branches), tes - `uv run remove-ai-watermarks identify ` — provenance verdict (platform + watermark inventory + confidence); `--json` for machine output, `--no-visible` to skip both registered visible detectors and the optional open invisible-watermark decoder - `uv run remove-ai-watermarks metadata --check` — inspect AI metadata (C2PA, EXIF, PNG chunks) - `uv run remove-ai-watermarks metadata --remove -o ` — strip all AI metadata +- `uv run remove-ai-watermarks video identify ` — locally verifiable video provenance report: supported metadata plus the same temporally stable visible-mark arbiter used by removal; `--json` for machine output and `--no-visible` for metadata only. A negative result is unknown, never clean. +- `uv run remove-ai-watermarks video all -o ` — product-oriented video pipeline: remove a stable registered visible mark when present, always strip verified AI metadata, and write a same-container passthrough when neither signal exists. Lossy SynthID removal is excluded unless `--invisible` is explicit; that option uses the oracle-certified video profile. - `uv run remove-ai-watermarks video metadata --check` — inspect AI metadata in MP4/MOV/M4V/WebM/MKV/AVI/FLV -- `uv run remove-ai-watermarks video metadata --remove -o ` — strip verified video metadata without transcoding streams; the experimental video path requires a separate same-container output and defaults to `_clean` -- `uv run remove-ai-watermarks video visible -o ` — remove a temporally recurring Sora, Veo, Seedance, Dola, Hailuo, or Kling mark. `--mark auto` is the default: all providers share one decode pass, then the first stable result wins in specificity order; an explicit mark restricts the scan. It transcodes video through ffmpeg, copies the complete audio stream, strips AI metadata by default, publishes atomically, and writes no output when no stable mark is found. `cv2` is the fast default; `migan`/`lama` improve difficult backgrounds. +- `uv run remove-ai-watermarks video metadata --remove -o ` — strip verified video metadata without transcoding streams; the video path requires a separate same-container output and defaults to `_clean` +- `uv run remove-ai-watermarks video visible -o ` — remove a temporally recurring Sora, Veo, Seedance, Dola, Hailuo, or Kling mark. `--mark auto` is the default: all providers share one decode pass, then the first stable result wins in specificity order; an explicit mark restricts the scan. It transcodes video through ffmpeg, copies the complete audio stream, strips AI metadata by default, publishes atomically, and writes no output when no stable mark is found. `cv2` is the fast default; `migan` is the practical learned CPU tier; `lama` is an offline quality tier. `--temporal-consistency` (default) motion-aligns adjacent accepted fills with mask-coverage and scene-context gates; `--no-temporal-consistency` is the frame-local baseline. HDR/high-bit inputs are rejected before encoding rather than silently reduced to 8-bit SDR. +- `uv run remove-ai-watermarks video batch --mode all` — sequential video directory processing with per-file results and non-zero exit when any file fails. Modes are `all|visible|metadata`; visible no-ops are byte-copied so successful output sets stay complete. `--invisible` is explicit and valid only in `all` mode. - `uv run remove-ai-watermarks batch ` — process every supported image in a directory (output defaults to `_clean/`, set with `-o`). `--mode visible|invisible|metadata|all` (default `visible`); the invisible/all path reuses the full `invisible` knob set above, plus `--backend` and `--sensitivity` for the visible localize -> fill pass. Applies the same no-signal skip per image; see the module doc. **Exit code:** non-zero when any image errored OR (mirroring single `all`) a `--mode invisible`/`all` image carried an invisible signal but the GPU extra was absent, so its SynthID scrub was skipped — it emits a loud warning and copies the input through (invisible mode) so the output dir stays complete; a wrapping service can then detect the incomplete run instead of trusting a silent exit 0. ## Test and lint @@ -82,9 +85,10 @@ Compact map. The full per-module detail (design decisions, tuned thresholds, cal - `upscaler.py` — optional Real-ESRGAN pre-diffusion super-resolution for small inputs (extra `esrgan`, spandrel only). Manual opt-in; the default `--upscaler` stays `lanczos` and the engine always falls back to Lanczos on absence/error. ESRGAN can degrade faces and thin text. - `image_io.py` — centralizes Unicode-safe image IO, alpha preservation, content-based format sniffing, and HEIC/AVIF fallbacks. Callers must check `imwrite` success. No-op visible removal preserves original bytes when the output format is unchanged. - `api.py` — the high-level convenience API, re-exported lazily at the package top level via `__init__.__getattr__` (PEP 562, so `import remove_ai_watermarks` stays cheap): `remove_visible(source, output=None, *, sensitivity="auto", backend="auto", strip_metadata=True, write_noop=True) -> (result_bgr, [labels])` (source = path OR BGR ndarray; a PATH auto-reads metadata provenance and preserves alpha, an ARRAY does neither; `write_noop=True` writes a clean passthrough copy when nothing is removed, `False` leaves `output` untouched so a "no mark = produce nothing" caller like the CLI `visible` command does not clobber a pre-existing file there) and `visible_provenance(path) -> frozenset[str]` (the single metadata→vendor-keys mapper; `cli._visible_provenance` is a thin None-guarded wrapper over it). **`remove_visible` is the ONE path the CLI and library share** — `cli.cmd_visible`'s `--mark auto` branch delegates entirely to it (read → provenance → `remove_auto_marks` → write → `strip_metadata`), so there is no CLI-vs-library drift; `strip_metadata` defaults True to match `visible --strip-metadata`. This is where a library caller should start — NOT the engines directly (`GeminiEngine`/`TextMarkEngine` have no `remove_watermark` any more; removal is registry `remove_auto_marks`/`KnownMark.remove`; the old single-strongest `best_auto_mark` is gone — removal takes EVERY mark). `identify` is NOT top-level re-exported (it collides with the `identify` submodule); use `from remove_ai_watermarks.identify import identify`. -- `video.py` — the experimental high-level video API, also lazy at the package root: `inspect_video_metadata(source) -> VideoMetadataReport`, `remove_video_metadata(source, output=None, *, keep_standard=True) -> VideoMetadataResult`, `remove_video_visible(source, output=None, *, mark="auto", backend="cv2", strip_metadata=True) -> VideoVisibleResult`, and `remove_video_invisible(source, output=None, ...) -> VideoInvisibleResult`. It validates the extension and container signature for MP4/MOV/M4V/WebM/MKV/AVI/FLV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`, and removal blanks the key/value in place; MKV/WebM, AVI, and FLV TC260 are read by bounded container walkers that skip media payloads and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: auto scans all six synthetic provider silhouettes in one decode pass and selects the first stable result in specificity order because provider confidence scales are not comparable. Provider-specific temporal recurrence authorizes frame boxes, the shared fill backends remove accepted masks in a second pass, and ffmpeg transcodes video while copying complete audio. Fixed marks require anchored runs so a slowly drifting scene detail cannot pass on adjacent overlap alone. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Seedance uses the full localized box because an outline mask left part of the real translucent border behind. Hailuo expands its matched core to the full composite, while Kling combines font and swirl candidates and requires a bright low-saturation edge label. Metadata can relax a recurring low-contrast match but cannot create one. The inherited ISOBMFF metadata removal path still reads the complete container into memory, so a streaming box copier is required before large-video use. Other visible video labels are not built yet. -- `video_encoding.py` — the shared raw-BGR ffmpeg command, atomic sibling-temporary publication, and pipe lifecycle for both visible removal and invisible regeneration. It copies complete optional audio without shortening a tail, controls metadata/chapter retention explicitly, and keeps container-specific codec arguments in one place. -- `video_invisible.py` — the oracle-gated video SynthID candidate engine for MP4/MOV/M4V. It regenerates frames through `stabilityai/sd-vae-ft-mse` with one seeded latent-noise field shared across time, retains only one configured batch, updates PSNR and motion-compensated temporal residuals incrementally, streams pixels directly to ffmpeg, copies complete audio, strips metadata, and atomically publishes the completed output. Every result remains explicitly unverified because Google exposes no local decoder. The 2026-07-29 Gemini Flash calibration found the default `0.10` candidate negative on both control-positive public Veo clips, while `0.05` remained positive on one, so the operating point is content-dependent and requires the matching external oracle. +- `video.py` — the high-level video API, also lazy at the package root: `identify_video`, `inspect_video_metadata`, `remove_video_all`, `remove_video_batch`, `remove_video_metadata`, `remove_video_visible`, and `remove_video_invisible`. `identify_video` and visible removal share one stable-mark selector; identification omits the unused per-frame timestamp probe. `remove_video_all` defaults to the locally verifiable visible + metadata stages and always writes a same-container output; lossy invisible removal is an explicit opt-in through the oracle-certified profile. Batch processing is sequential, returns every per-file error, byte-copies visible no-ops so successful output sets stay complete, and reuses one loaded VAE across an opt-in invisible batch. The API validates the extension and container signature for MP4/MOV/M4V/WebM/MKV/AVI/FLV and requires a distinct same-container output. The metadata path never transcodes streams: native MP4/MOV TC260 is read from `moov.udta.meta.keys/ilst`, including a tail `moov` after a large `mdat`; MP4/MOV/M4V removal stream-copies in bounded chunks, preserves every box size/offset, blanks supported provenance in place, and publishes atomically. MKV/WebM, AVI, and FLV TC260 are read by bounded container walkers that skip media payloads and stripped through ffmpeg stream copy. The visible path delegates to `video_visible.py`: auto scans all six synthetic provider silhouettes in one decode pass and selects the first stable result in specificity order because provider confidence scales are not comparable. Provider-specific temporal recurrence authorizes frame boxes, the shared fill backends remove accepted masks in a second pass, guarded optical-flow blending reduces frame-to-frame fill variance, and ffmpeg transcodes video while copying complete audio. Fixed marks require anchored runs so a slowly drifting scene detail cannot pass on adjacent overlap alone. Veo covers the current four-point diamond and legacy text; its diamond uses a shape mask rather than erasing the transparent corners of a full box. Seedance uses the full localized box because an outline mask left part of the real translucent border behind. Hailuo expands its matched core to the full composite, while Kling combines font and swirl candidates and requires a bright low-saturation edge label. Metadata can relax a recurring low-contrast match but cannot create one. Other visible video labels are not built yet. +- `video_encoding.py` — the shared ffmpeg command, atomic sibling-temporary publication, and pipe lifecycle for visible removal and invisible regeneration. It copies complete optional audio without shortening a tail, controls metadata/chapter retention explicitly, and keeps container-specific codec arguments in one place. `probe_video_encode_profile` preserves supported 8-bit chroma sampling, recognized color tags, encoder time base, MP4/MOV track timescale, source pixel format/depth, and the source video start PTS; HDR/high-bit inputs are rejected before encoding. `probe_video_timestamps` reads authoritative display PTS with ffprobe and falls back to count-matched OpenCV timestamps only on probe failure. Plain CFR stays on the cheap raw-BGR pipe, while VFR or non-zero-start frames are lazily packetized through PyAV as timestamped rawvideo/NUT and passed to system ffmpeg with `-fps_mode passthrough`. `-copyts` retains a non-zero video start and the corresponding copied-audio offset. +- `video_temporal.py` — shared optical-flow maps, temporal residual metrics, and guarded visible-fill stabilization. `stabilize_filled_frame` works on a bounded mask crop, requires warped prior-mask coverage and a matching unmasked context ring, and changes only safely covered pixels; scene cuts and disjoint marks keep the independent current fill. +- `video_invisible.py` — the oracle-certified video SynthID removal engine for MP4/MOV/M4V. It regenerates frames through `stabilityai/sd-vae-ft-mse` with one seeded latent-noise field shared across time, retains only one configured batch, updates PSNR and motion-compensated temporal residuals incrementally, streams pixels directly to ffmpeg, copies complete audio, strips metadata, and atomically publishes the completed output. The 2026-07-29 two-clip Gemini calibration remains valid: both matched controls were positive in the built-in SynthID verifier, the stronger candidate was negative on both carriers, and a weaker candidate was negative on one. The 2026-07-30 `UNAVAILABLE` follow-ups asked ordinary Gemini to reinterpret an already returned verifier result and were not detector reruns. The 2026-07-31 full-clip check used a public eight-second Veo carrier: the source and `noise_std=0.10` output were detected, while `0.15` was not, so `0.15` is the certified default. Google exposes no local decoder, so per-file verification remains an optional audit rather than a product status. Canonical hashes and verdicts: `data/evaluations/video-synthid-oracle.csv`. For the Doubao alpha-distillation history (why content-image reverse-alpha distillation fails by physics and controlled captures were required), see `docs/research-doubao-distillation.md`. diff --git a/README.md b/README.md index 431e3a0..a5bbfa3 100644 --- a/README.md +++ b/README.md @@ -6,9 +6,10 @@ Remove AI provenance marks from images and video you generated yourself: - invisible pixel watermarks through diffusion regeneration; - C2PA, EXIF, XMP, IPTC, and related AI metadata. -Video support covers metadata inspection and removal, visible Sora, Veo, -Seedance, Dola, Hailuo, and Kling mark removal, and experimental VAE -regeneration that produces a video SynthID candidate for external verification. +Video support covers provenance identification, complete visible-plus-metadata +cleaning, directory batches, visible Sora, Veo, Seedance, Dola, Hailuo, and +Kling mark removal, and oracle-certified VAE regeneration for video SynthID +removal. > Try it online at [raiw.cc](https://raiw.cc) if you do not want to install Python > or run diffusion models locally. @@ -32,9 +33,12 @@ regeneration that produces a video SynthID candidate for external verification. | Remove known visible AI marks | `visible` | No | | Erase a region you select | `erase` | No | | Strip AI metadata | `metadata` | No | +| Identify supported video provenance | `video identify` | No | +| Remove visible marks and AI metadata from video | `video all` | No | | Strip AI metadata from video | `video metadata` | No | | Remove a registered visible AI mark from video | `video visible` | No | -| Generate an externally verifiable video SynthID candidate | `video invisible` | Recommended | +| Process a directory of videos | `video batch` | Depends on mode | +| Remove video SynthID with the certified VAE profile | `video invisible` | Recommended | | Regenerate an image to disrupt invisible watermarks | `invisible` | Recommended | | Run visible, invisible, and metadata removal | `all` | Recommended | | Process a directory | `batch` | Depends on mode | @@ -82,6 +86,25 @@ MKV and WebM inspection reads the normative uses `script.onMetaData.AIGC`. The non-ISOBMFF formats are remuxed with stream copy for removal. +Use the product-oriented video path to identify or clean a file: + +```bash +remove-ai-watermarks video identify input.mp4 +remove-ai-watermarks video all input.mp4 -o clean.mp4 +``` + +`video all` removes a stable registered visible mark when present and always +strips verified AI metadata. If neither signal is found, it still writes a +same-container passthrough, so application callers get one predictable output +contract. Proprietary invisible-video removal is excluded by default. +`--invisible` opts into the lossy, oracle-certified video SynthID profile. + +Process a directory with the same contract: + +```bash +remove-ai-watermarks video batch ./videos --mode all +``` + Remove a supported visible video mark: ```bash @@ -97,30 +120,41 @@ This path scans the complete sequence before changing pixels. It accepts only a mark that repeats at a stable position across adjacent frames, then reuses the same OpenCV, MI-GAN, or LaMa fill backends as image removal. Audio is copied without re-encoding and is allowed to reach its natural end; the video stream -is transcoded because its pixels change. The default `--mark auto` scans all -providers in one decode pass and selects the first stable match in the -specificity order shown below. Pass an explicit mark to restrict detection to -one provider. +is transcoded because its pixels change. By default, a guarded optical-flow +pass motion-aligns the preceding accepted fill and blends it only when the +nearby source context agrees; use `--no-temporal-consistency` to disable it. +The encoder preserves supported +8-bit source chroma sampling, color tags, and MP4/MOV track timescale instead +of relying on ffmpeg's implicit raw-BGR defaults. Variable frame intervals are +preserved through a timestamped in-memory NUT bridge instead of being flattened +to the average frame rate. Non-zero source start timestamps are retained +together with the copied audio offset. The default `--mark auto` +scans all providers in one decode pass and selects the first stable match in +the specificity order shown below. Pass an explicit mark to restrict detection +to one provider. Sora covers the moving Sora 2 mascot and wordmark. Veo covers both the current four-point diamond and the legacy `Veo` text. Seedance covers the fixed boxed `AI` label, Dola covers the fixed `Dola AI` text, Hailuo covers the composite `MINIMAX | hailuo AI` label, and Kling covers the bottom-right `KLING AI` label with its version suffix. A completed encode is published atomically. No output is written when no stable mark is found. +HDR, PQ/HLG, and greater-than-8-bit inputs are rejected before encoding rather +than silently reduced through OpenCV's 8-bit BGR boundary. -Generate a video SynthID candidate: +Remove video SynthID: ```bash uv tool install --force "remove-ai-watermarks[gpu]" -remove-ai-watermarks video invisible input.mp4 -o candidate.mp4 +remove-ai-watermarks video invisible input.mp4 -o clean.mp4 ``` This path regenerates the complete sequence with one latent-noise field shared across time, copies complete audio, strips source metadata, and publishes the -completed encode atomically. It cannot verify -Google's proprietary pixel watermark locally. The command therefore labels -every output `UNVERIFIED` and prints the exact Gemini Flash verification -prompt. +completed encode atomically. The default `noise_std=0.15` profile passed both +the two-carrier calibration and a complete public eight-second Veo oracle +check. Google does not publish a local decoder, so a fresh provider check +remains useful for unusually important files or after provider changes, but it +is not a product result state. For invisible watermark removal, install the diffusion dependencies: @@ -268,9 +302,12 @@ import remove_ai_watermarks as raiw result, removed = raiw.remove_visible("watermarked.png", "clean.png") print(removed) +provenance = raiw.identify_video("input.mp4") report = raiw.inspect_video_metadata("input.mp4") +complete = raiw.remove_video_all("input.mp4", "clean.mp4") +batch = raiw.remove_video_batch("videos", "videos_clean") cleaned = raiw.remove_video_metadata("input.mp4") -candidate = raiw.remove_video_invisible("input.mp4", "candidate.mp4") +synthid_cleaned = raiw.remove_video_invisible("input.mp4", "synthid_clean.mp4") visible = raiw.remove_video_visible("input.mp4", "clean.mp4") print(visible.mark) veo = raiw.remove_video_visible("veo.mp4", "veo_clean.mp4", mark="veo") @@ -311,30 +348,34 @@ invisible removal. The classical OpenCV backend can smear structured backgrounds; use MI-GAN or LaMa when recovery quality matters. - Video SynthID regeneration changes resolution, frame rate, and image detail. - It produces a candidate, not a locally verified clean video. The matching - Google verifier is still required for every important output. -- MP4/MOV/M4V metadata stripping currently reads the full container into - memory, so very large videos are not an intended metadata input yet. + The shipped profile is oracle-certified, but no public local decoder can + certify an arbitrary output at runtime. Recheck unusually important outputs + after provider changes. - `qwen-zimage` requires CUDA. The other diffusion profiles also support the devices listed by `remove-ai-watermarks invisible --help`. - Provider watermark systems can change. Validate important outputs with the provider's own verifier when one is available. -The shipped `video invisible` command uses the candidate-producing side of the -oracle-gated workflow. The companion `scripts/video_synthid_sweep.py` harness -builds a matched re-encode control plus VAE-regenerated candidates and leaves -the verifier verdict blank: +The shipped `video invisible` command uses the certified `noise_std=0.15` +profile. The companion `scripts/video_synthid_sweep.py` research harness builds +a matched re-encode control plus VAE-regenerated candidates and leaves the +verifier verdict blank: ```bash uv run --extra gpu python scripts/video_synthid_sweep.py input.mp4 -o sweep/ ``` The control must still be SynthID-positive before a negative candidate can -count as removal evidence. In the 2026-07-29 calibration, matched controls from -two public Veo videos remained positive and the default stronger candidate was -negative on both. A weaker candidate was negative on only one, demonstrating -that the removal threshold is content dependent. This is calibration evidence, -not a universal guarantee for new videos or future verifier versions. +count as removal evidence. In the 2026-07-29 two-clip calibration, both matched +controls were positive in Gemini's built-in SynthID verifier; the stronger +candidate was negative on both carriers, while a weaker candidate was +negative on one. A later adversarial follow-up that asked ordinary Gemini to +reinterpret the pixel result returned `UNAVAILABLE`; that follow-up was not a +verifier rerun and does not invalidate the built-in verdicts. A 2026-07-31 +full-clip check on a public eight-second Veo sample found `0.10` still detected +and `0.15` not detected, so `0.15` is now the certified default. The +reproducible hashes and verdicts live in +`data/evaluations/video-synthid-oracle.csv`. ## Documentation diff --git a/data/README.md b/data/README.md index 55d99be..682c087 100644 --- a/data/README.md +++ b/data/README.md @@ -15,6 +15,8 @@ data/ Reusable full-pipeline evaluation selection evaluations/ fidelity/ Evaluation instructions and hand-verified ground truth + video-synthid-oracle.csv + Reproducible full-clip Gemini SynthID verdicts ``` ## Storage rules diff --git a/data/evaluations/video-synthid-oracle.csv b/data/evaluations/video-synthid-oracle.csv new file mode 100644 index 0000000..e084e96 --- /dev/null +++ b/data/evaluations/video-synthid-oracle.csv @@ -0,0 +1,3 @@ +date,source_url,source_sha256,duration_seconds,source_verdict,noise_std,long_side,fps,seed,output_sha256,output_verdict,psnr_db,temporal_residual_ratio +2026-07-31,https://storage.googleapis.com/gdm-deepmind-com-prod-public/media/media/veo__veo-3__off-road.mp4,79a552b9406a079682440c31f14d33a10ba8e1b8b2e96425f5de70f63350299d,8,detected_all_frames,0.10,512,12,0,079165105d4c56e1612091987c08c2627049423025f74c0d4e245fb47c2ff0e3,detected,26.2932,1.0072 +2026-07-31,https://storage.googleapis.com/gdm-deepmind-com-prod-public/media/media/veo__veo-3__off-road.mp4,79a552b9406a079682440c31f14d33a10ba8e1b8b2e96425f5de70f63350299d,8,detected_all_frames,0.15,512,12,0,1c4046bcfdead138353b4e2a73339ba227bb5e544878d80c5bc6cd8427c7b00e,not_detected,25.3911,1.0578 diff --git a/docs/cli.md b/docs/cli.md index 5c23400..73410c6 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -125,9 +125,57 @@ The command also supports the audio and video containers listed in [supported signals](supported-signals.md). ffmpeg must be available for the non-ISOBMFF audio and video path. +## Identify and clean video + +Inspect every locally supported video signal: + +```bash +remove-ai-watermarks video identify input.mp4 +remove-ai-watermarks video identify input.mp4 --json +remove-ai-watermarks video identify input.mp4 --no-visible +``` + +The default scans the complete clip for stable registered visible marks and +inspects supported metadata. A result with no signals is reported as unknown, +not clean, because proprietary pixel watermarks have no public local decoder. +`--no-visible` performs metadata-only inspection. + +Use the complete locally verifiable cleaning path: + +```bash +remove-ai-watermarks video all input.mp4 -o clean.mp4 +``` + +It removes a stable supported visible mark when found and always strips +verified AI metadata. When neither signal is found, it writes a same-container +passthrough instead of returning a missing output. The source is never +overwritten. + +Invisible regeneration is deliberately opt-in: + +```bash +remove-ai-watermarks video all input.mp4 -o clean.mp4 --invisible +``` + +That option is supported only for MP4, MOV, and M4V. It is lossy and uses the +same oracle-certified profile as `video invisible`. + +Process all supported files in a top-level directory: + +```bash +remove-ai-watermarks video batch ./videos --mode all +remove-ai-watermarks video batch ./videos --mode visible +remove-ai-watermarks video batch ./videos --mode metadata +``` + +The batch runs sequentially, preserves successful outputs when another file +fails, and exits nonzero if any item failed. Visible no-op files are copied +byte-for-byte so the output directory remains complete. `--invisible` is +available only with `--mode all`. + ## Strip AI metadata from video -The experimental video namespace starts with metadata inspection and removal: +Metadata inspection and removal are also available as an isolated operation: ```bash remove-ai-watermarks video metadata input.mp4 --check @@ -139,8 +187,11 @@ delegates to the same verified metadata scanner and stripper as the generic `metadata` command, so detection and removal stay in parity. Video and audio streams are not transcoded. For MP4 and MOV, this includes the native TC260 `AIGC` key and JSON value stored in `moov.udta.meta.keys/ilst`. The inspector -seeks past a large `mdat` to find a tail `moov`; removal blanks the key and -value in place so box sizes and media offsets do not move. +seeks past a large `mdat` to find a tail `moov`. Removal stream-copies the +container in bounded chunks, converts supported top-level provenance boxes to +same-size `free` boxes, and blanks the TC260 key/value in place. Box sizes, +media offsets, and encoded stream bytes do not move; the result is atomically +published only after the complete copy succeeds. For MKV and WebM, the inspector reads the native TC260 `Segment.Tags.Tag.SimpleTag` entry. Removal uses ffmpeg stream copying to @@ -156,14 +207,14 @@ different container extension. Visible video labels and invisible video watermarks are not handled by this command. -## Generate a video SynthID candidate +## Remove video SynthID ```bash uv tool install --force "remove-ai-watermarks[gpu]" -remove-ai-watermarks video invisible input.mp4 -o candidate.mp4 +remove-ai-watermarks video invisible input.mp4 -o clean.mp4 ``` -The experimental command supports MP4, MOV, and M4V. It samples the complete +The command supports MP4, MOV, and M4V. It samples the complete sequence at the configured frame rate, resizes frames to the configured long side, regenerates them through a VAE, and applies one deterministic latent-noise field to every frame. Reusing one spatial field avoids the unnecessary flicker @@ -171,22 +222,28 @@ caused by independent per-frame noise. Frames are regenerated in bounded batches and streamed directly to ffmpeg, which encodes the result, copies audio, and drops source metadata. -The output is always an unverified candidate. The project has no local video -SynthID decoder, and PSNR or temporal-residual metrics cannot prove watermark -absence. After generation, upload the candidate to Gemini Flash and ask: +The default `noise_std=0.15` profile is oracle-certified. The project has no +local video SynthID decoder, so an optional per-file recheck is still useful +for unusually important files or after provider changes. In a new Gemini chat, +upload the original first, invoke the built-in verifier with `@synthid`, and ask: -> Was this uploaded video created or edited by Google AI? Use the built-in -> content verification result. +> For the video attached to this message, was it created or edited by Google +> AI? Use the built-in SynthID content verification result. -Only an explicit built-in verification result is an oracle verdict. A response -based on the visible logo, content appearance, or metadata is not. The command -prints `UNVERIFIED` even when generation succeeds. +The source must be positive. Then upload the processed result in a separate new +chat and repeat the same built-in check. Only a source-positive, output-negative +pair is a fresh per-file verification. Do not ask an adversarial follow-up that tells the +chat model to ignore the verifier and reason about raw pixels: that is ordinary +Gemini reasoning, not a second oracle check. -The default output is `_synthid_candidate` in the same container. The +The default output is `_clean` in the same container. The source is never overwritten. Use `--noise-std`, `--long-side`, `--fps`, `--batch-size`, `--seed`, and `--device` to control the regeneration. The -defaults are calibrated operating points, not a guarantee for every carrier or -future verifier version. +default noise level is `0.15`. It cleared both carriers in the 2026-07-29 +short-clip calibration and the complete public eight-second Veo carrier in the +2026-07-31 full-clip check; `0.10` remained detected on that complete clip. This +calibration certifies the shipped operating point; the paired check above is an +optional runtime audit, not a separate result state. ## Remove a supported visible video mark @@ -199,7 +256,7 @@ remove-ai-watermarks video visible hailuo.mp4 --mark hailuo -o hailuo_clean.mp4 remove-ai-watermarks video visible kling.mp4 --mark kling -o kling_clean.mp4 ``` -The experimental command supports the moving Sora mascot and wordmark, two Veo +The command supports the moving Sora mascot and wordmark, two Veo corner variants, the Seedance boxed `AI` label, the `Dola AI` text label, the composite `MINIMAX | hailuo AI` label, and the bottom-right Kling label. Sora searches the whole frame at multiple scales. The other detectors search bounded @@ -218,13 +275,31 @@ provider. The video stream is transcoded and the complete original audio stream is copied without truncating an audio tail that extends beyond the final video -frame. +frame. The encoder probes the source stream and preserves supported 8-bit +chroma sampling, color range/matrix/transfer/primaries tags, and MP4/MOV track +timescale. For a variable-frame-rate source, decoded PTS are carried through a +timestamped in-memory NUT bridge so the output retains the source frame +intervals instead of flattening them to a constant rate. A non-zero source +start PTS and the copied audio start offset are preserved as well. Supported input and output containers are MP4, MOV, M4V, WebM, MKV, AVI, and FLV; the output extension must match the input. The default `cv2` backend is fast but can smear structured backgrounds. Select `--backend migan` or `--backend lama` for a learned fill, or `--backend auto` to choose the best installed backend. +`--temporal-consistency` is enabled by default. It motion-aligns the preceding +accepted fill, requires overlapping removal masks and matching source context, +and blends only the safely covered pixels. Scene cuts, disjoint moving marks, +or a poor motion match keep the independent current-frame fill. Use +`--no-temporal-consistency` for an exact frame-local baseline. + +The pixel path is intentionally limited to SDR 8-bit video. A high-bit-depth, +PQ, or HLG source is rejected before ffmpeg starts, preserving any existing +output instead of silently downconverting it through OpenCV's 8-bit boundary. +On CPU, MI-GAN is the practical learned tier. LaMa remains an explicit offline +quality option because full-sequence inference is too slow and memory-heavy +for an online worker. + AI metadata is stripped from the encoded output by default. Use `--keep-metadata` to retain mapped container metadata. When no temporally stable mark is found, the command writes no output and exits with the no-visible-mark diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 8d279a8..0a09587 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -79,24 +79,36 @@ For important outputs: Provider systems can change, so a result verified on one file, seed, or version is not a permanent certification. -### Video regeneration produces an unverified candidate +### Video SynthID removal is lossy and content-dependent The `video invisible` command and `remove_video_invisible` API regenerate video -pixels through a VAE, but cannot verify the proprietary SynthID payload locally. -Every successful result remains explicitly unverified and requires Google's -matching content-verification flow. A quiet metadata scan, paired PSNR, and the -temporal-residual metric are not substitutes for that oracle. +pixels through a VAE. The shipped `noise_std=0.15` profile is oracle-certified, +but Google does not publish a local decoder for arbitrary runtime outputs. A +quiet metadata scan, paired PSNR, and the temporal-residual metric are fidelity +measurements, not independent SynthID verdicts. The control must use the same clip, frame rate, dimensions, and final codec as the candidates. The separate `scripts/video_synthid_sweep.py` harness produces that matched control. If the control is not detected by the matching provider oracle, the experiment cannot attribute a quiet candidate to regeneration. -The 2026-07-29 calibration used two public Veo clips. Both matched controls were -SynthID-positive. The default stronger VAE candidate was negative on both, while -a weaker candidate remained positive on one. This small calibration validates -the mechanism and exposes its content dependence; it does not certify all -videos, full-length sequences, alternate settings, or future verifier versions. +The 2026-07-29 two-clip calibration used Gemini's built-in content verifier: +both matched controls were SynthID-positive, the stronger candidate was +negative on both carriers, and a weaker candidate was negative on one. A +2026-07-30 adversarial follow-up incorrectly asked the ordinary chat model to +reinterpret the verifier while excluding every other input; its `UNAVAILABLE` +answer was not another detector run and does not invalidate the original +built-in results. The calibrated default remains content-dependent; a fresh +source-positive, output-negative pair is an optional audit for unusually +important files or after provider changes. + +The 2026-07-31 full-clip check used the public eight-second Veo off-road sample. +The source was detected across the full clip, the complete product path at +`noise_std=0.10` remained detected, and `0.15` returned no SynthID detection. +The default was raised to `0.15`. At 512 px / 12 fps, the accepted candidate +measured 25.39 dB paired PSNR and a 1.058 motion-compensated temporal-residual +ratio. This is one carrier, not a universal guarantee; hashes and exact verdicts +are tracked in `data/evaluations/video-synthid-oracle.csv`. The shipped engine streams sampled frames in bounded batches, computes its fidelity metrics incrementally, and pipes regenerated pixels directly to @@ -225,22 +237,48 @@ than the moving mascot-and-wordmark design; that earlier variant is not detected by the `sora` video mark. Hailuo and Kling coverage is specific to the verified lower-edge layouts; a new provider layout needs a separate calibrated silhouette. Other provider video labels are not supported yet. Google video -SynthID has a candidate-producing VAE path, while other proprietary invisible -video watermarks have no registered attack. +SynthID has an oracle-certified VAE removal path, while other proprietary +invisible video watermarks have no registered attack. Visible removal transcodes the video stream and copies the complete audio stream without shortening an audio tail. Completed visible and invisible encodes are published atomically, so an encode failure preserves an existing -output. Its frame-local fill is not a motion-aware video inpainting model. -OpenCV can leave a visible smear where the mark overlaps a hard edge or -structured texture, and the smear can vary over time. MI-GAN and LaMa improve -individual frames but do not guarantee temporal coherence. The Veo diamond +output. Visible removal now applies a guarded motion-compensated blend after +the per-frame fill. It uses adjacent optical flow only when the warped prior +mask covers the current mask and a source-context ring agrees; scene cuts and +disjoint marks keep the independent fill. This reduces measured paired +temporal error, but it is not a generative video-inpainting model and cannot +recover structure that no frame exposes. OpenCV can still leave a visible +smear where the mark overlaps a hard edge or structured texture. MI-GAN +improves difficult individual frames and is the practical learned CPU tier. +LaMa remains an offline quality option: a full real sequence confirmed its +multi-GB memory use and CPU throughput unsuitable for an online worker. The Veo diamond uses a shape mask to limit damage outside the symbol. Seedance fills the full localized box because a synthetic outline mask left part of the real translucent border visible in an end-to-end check. OpenCV may therefore soften texture inside that small box; use MI-GAN or LaMa when reconstruction quality -matters. The current encoder also emits a constant-frame-rate output at the -decoded stream rate, so variable-frame-rate preservation is not yet guaranteed. +matters. Relative variable-frame timestamps are preserved through a timestamped +NUT bridge. A non-zero absolute video start PTS and the corresponding copied +audio offset are preserved through ffmpeg timestamp passthrough. +The encoder is source-aware for common 8-bit inputs: it probes and preserves +supported chroma sampling, recognized color metadata, encoder time base, and +MP4/MOV track timescale instead of accepting ffmpeg's implicit `yuv444p` +raw-BGR output. OpenCV still decodes through 8-bit BGR, so HDR/high-bit-depth +inputs are rejected before encoding rather than silently falling back to +`yuv420p`. The synthetic +Sora/OpenCV full-clip CI gate covers complete removal, untouched-region PSNR, +frame count, frame rate, duration, copied-audio identity, source stream +properties, paired temporal deltas against an independently encoded +frame-local baseline inside the filled region, and metadata +stripping through real ffmpeg. A second synthetic VFR clip verifies all display +timestamps to the source time-base tick, including a non-zero source start, +and checks both video and audio stream offsets. A separate constant-rate case +guards the non-zero-start routing independently of VFR detection. A local +full-sequence audit over all six providers passed both OpenCV and MI-GAN for +complete-frame removal, quiet second detection, stream starts, duration, and +copied audio. A full LaMa sequence passed the same checks but established that +the backend belongs in the offline tier on CPU. These bounded local checks are +still not universal evidence for every provider layout or source. Native TC260 metadata in MP4/MOV is supported at its normative `moov.udta.meta.keys/ilst` placement, including non-faststart files whose @@ -252,9 +290,13 @@ write and strip the FLV form, but writing the nonstandard AVI child for fixture generation requires dedicated muxer support, so the AVI reader is verified against an exact synthetic RIFF structure. -The current ISOBMFF stripper reads an MP4, MOV, or M4V container into memory -before rewriting its metadata boxes. A streaming box copier is required before -the experimental command is appropriate for very large video files. +MP4/MOV/M4V metadata removal now stream-copies the container in bounded chunks, +keeps every box size and media offset fixed, and publishes atomically. The +large-`mdat` regression rejects a full-source `read_bytes()` call and verifies +that the encoded payload is byte-identical. HEIF/AVIF/JPEG-XL image metadata +still uses the in-memory path because their XMP/EXIF items may live inside +`mdat`/`idat` and require bounded format-item parsing before that path can +stream safely. ### Metadata transformation is fail safe diff --git a/docs/module-internals.md b/docs/module-internals.md index 1679e67..c7320b7 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -86,10 +86,13 @@ Regression coverage: - [`test_api.py`](../tests/test_api.py) - [`test_image_io.py`](../tests/test_image_io.py) -[`video.py`](../src/remove_ai_watermarks/video.py) provides the experimental +[`video.py`](../src/remove_ai_watermarks/video.py) provides the high-level video entry point: +- `identify_video` - `inspect_video_metadata` +- `remove_video_all` +- `remove_video_batch` - `remove_video_invisible` - `remove_video_metadata` - `remove_video_visible` @@ -97,17 +100,35 @@ video entry point: The video API validates both the supported extension and container signature, then delegates all metadata detection and stripping to `metadata.py`. It requires a separate same-container output, defaulting to `_clean`, so -the experimental path does not overwrite an original. The package root exposes -all four functions lazily. +the product path does not overwrite an original. The package root exposes all +functions lazily. + +`identify_video` runs the same stable-mark selection helper as +`remove_video_visible`, so a provenance report cannot authorize a mark that the +removal path would reject. It reports an empty local result as unknown rather +than clean. Identification skips the separate per-frame timestamp probe because +it never encodes frames. `remove_video_all` is the predictable-output +composition: visible removal plus verified metadata stripping by default, with +a same-container passthrough when neither signal exists. The lossy invisible +removal stage is an explicit opt-in through the oracle-certified profile. +`remove_video_batch` applies those contracts sequentially across a top-level +directory, returns every per-file failure, and byte-copies visible no-ops so a +successful output set has no silent holes. An invisible batch loads one VAE +runtime and reuses it across every compatible file; a failed model load is +reported per file without retrying the same multi-GB initialization. Native MP4/MOV TC260 labels follow TC260-PG-20257A: `moov.udta.meta.keys` maps an `AIGC` key to a raw JSON value in `ilst`. [`noai/isobmff.py`](../src/remove_ai_watermarks/noai/isobmff.py) walks those nested boxes by seeking, so detection reaches a tail `moov` without reading the -preceding `mdat`. Removal changes the four-byte key to `free` and blanks only -the validated JSON value with same-length spaces. This preserves every box -size, `stco`/`co64` offset, and encoded stream byte. A generic `AIGC` key whose -value has no TC260 field is ignored. +preceding `mdat`. The MP4/MOV/M4V/M4A removal path first validates the top-level +box walk, then copies the source to a sibling temporary file in bounded chunks. +Supported C2PA/JUMBF/AI-label boxes become same-size `free` boxes with blank +payloads; TC260 removal changes the four-byte key to `free` and blanks only the +validated JSON value with same-length spaces. This preserves every box size, +`stco`/`co64` offset, encoded stream byte, and source-sized memory bound. +Publication is atomic, and a malformed top-level walk is copied unchanged. A +generic `AIGC` key whose value has no TC260 field is ignored. [`noai/ebml.py`](../src/remove_ai_watermarks/noai/ebml.py) provides the corresponding bounded Matroska/WebM reader. It seeks over clusters and accepts @@ -123,14 +144,40 @@ normative TC260 video placements. The RIFF walker reads only AVI use the verified ffmpeg stream-copy path for removal. [`video_encoding.py`](../src/remove_ai_watermarks/video_encoding.py) owns the -raw-BGR ffmpeg command and pipe lifecycle shared by visible removal and -invisible regeneration. It centralizes container codecs, optional audio stream -copying, metadata/chapter policy, encode-failure reporting, and atomic -same-directory publication. Each mapped stream is allowed to reach its own end, -so a copied audio tail is not shortened to the raw-video input duration. +ffmpeg command and pipe lifecycle shared by visible removal and invisible +regeneration. It centralizes container codecs, optional audio stream copying, +metadata/chapter policy, encode-failure reporting, and atomic same-directory +publication. Each mapped stream is allowed to reach its own end, so a copied +audio tail is not shortened to the frame-input duration. +`probe_video_encode_profile` reads the first source video stream with ffprobe +and preserves the supported properties that survive the 8-bit BGR boundary: +`yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder +time base, MP4/MOV track timescale, source pixel format, and component depth. +HDR transfer functions and component depths above 8 bits are rejected before +encoding so the OpenCV boundary cannot silently reduce them to SDR 8-bit. +`probe_video_timestamps` reads authoritative per-frame display PTS through +ffprobe. OpenCV timestamps are only a count-matched fallback when ffprobe is +unavailable or fails; this avoids decoder anomalies such as one spurious +negative first-frame timestamp turning a CFR clip into false VFR. A uniform sequence keeps the +cheap raw-BGR pipe unless the source starts at a non-zero PTS. A variable or +offset sequence is packetized by the lazy PyAV bridge as rawvideo in an +in-memory NUT stream with explicit PTS. System ffmpeg reads that stream with +`-fps_mode passthrough`; `-copyts` additionally retains a non-zero video start +and the corresponding copied-audio offset. No temporary frame sequence or +second video encoder is introduced. + +[`video_temporal.py`](../src/remove_ai_watermarks/video_temporal.py) owns the +shared optical-flow maps and temporal residual metric. Visible removal uses +`stabilize_filled_frame` after the selected image backend: it works on a +bounded crop around adjacent masks, backward-warps the prior cleaned frame, +requires high warped-mask coverage, and gates blending on an unmasked +source-context ring. Only covered current-mask pixels change. Scene cuts, +disjoint marks, and poor motion matches therefore retain the independent +current-frame fill. The same module supplies the motion-compensated metric used +by the invisible-video sweep. [`video_invisible.py`](../src/remove_ai_watermarks/video_invisible.py) -implements the oracle-gated video SynthID candidate engine. It samples frames +implements the oracle-certified video SynthID removal engine. It samples frames uniformly, resizes to a VAE-aligned geometry, encodes each frame to latent space, applies one seeded spatial-noise field across the entire sequence, and decodes fresh pixels. Reusing a single noise field avoids independent @@ -141,12 +188,13 @@ drops all source metadata. The result is written through a same-directory temporary file and atomically replaced only after a successful encode. The engine returns PSNR and a motion-compensated temporal-residual ratio as -quality measurements. Neither is a watermark detector. The high-level -`VideoInvisibleResult.requires_external_verification` flag is always true, and -the CLI prints an `UNVERIFIED` warning plus the Gemini Flash verification -prompt. The companion `scripts/video_synthid_sweep.py` imports the same engine -helpers to build a matched control and candidate grid, preventing research and -shipped regeneration paths from drifting. +quality measurements. Neither is a watermark detector. The high-level result +reports completed removal without a separate verification-status flag. The companion +`scripts/video_synthid_sweep.py` imports the same engine helpers to build a +matched control and candidate grid, preventing research and shipped +regeneration paths from drifting. The full-clip oracle floor is +`noise_std=0.15`: on the public eight-second Veo carrier, `0.10` remained +detected while `0.15` did not. [`video_visible.py`](../src/remove_ai_watermarks/video_visible.py) implements the first pixel stages for Sora, Veo, Seedance, Dola, Hailuo, and Kling. The @@ -169,7 +217,8 @@ template features for the fixed stream geometry. Provider confidence scales are not comparable: selection applies each provider's temporal arbiter and takes the first stable result in specificity order (`sora`, `veo`, `seedance`, `dola`, `hailuo`, `kling`). An explicit mark uses the same scan path with one -candidate. +candidate. Removal also collects authoritative per-frame timestamps for the +encoder, while identification omits that unused ffprobe pass. Every per-frame result is untrusted. The provider-specific stabilization wrappers share one recurrence implementation, while retaining separate visual @@ -196,13 +245,17 @@ backgrounds need MI-GAN or LaMa for better reconstruction. Invisible video stages must continue to reuse the image and metadata implementations rather than copying their logic. -The inherited ISOBMFF metadata path currently reads the complete container into -memory; replacing that with a streaming box copier is a prerequisite for large -video inputs. - Regression coverage: -- [`test_video.py`](../tests/test_video.py) +- [`test_video.py`](../tests/test_video.py), including a real ffmpeg full-clip + Sora/OpenCV path that generates a synthetic marked MP4 with AAC audio and + C2PA provenance, runs both `remove_video_visible` and the composed + `remove_video_all` API without mocks, and verifies complete removal, frame + count, frame rate, duration, untouched-region PSNR, paired temporal deltas + inside the filled region, byte-identical copied audio packets, source stream + properties, metadata stripping, and a large-`mdat` metadata case that rejects + any full-source `read_bytes()` call. CI installs ffmpeg explicitly for this + test so the integration gate cannot silently skip. ## Metadata and provenance diff --git a/docs/python-api.md b/docs/python-api.md index 444bf09..d0a38d1 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -128,10 +128,62 @@ path preserves the pixels but drops standard metadata. Treat a nonempty undecodable input through unchanged, so its return alone must not be presented as proof that metadata was removed. +## Identify and clean video + +The high level video API supports MP4, MOV, M4V, WebM, MKV, AVI, and FLV: + +```python +import remove_ai_watermarks as raiw + +report = raiw.identify_video("input.mp4") +print(report.is_ai_generated) +print(report.platform) +print(report.visible_mark) +print(report.metadata_markers) +``` + +`identify_video` uses the same full-clip temporal arbiter as visible removal. +It reports a recurring registered mark and supported AI metadata as positive +signals. When neither is present, `is_ai_generated` is `None`, never `False`. +The absence of a public local video SynthID decoder is included in `caveats`. +Pass `check_visible=False` for a bounded metadata-only inspection. + +For normal product integration, use the complete locally verifiable pipeline: + +```python +result = raiw.remove_video_all("input.mp4", "clean.mp4") +if result.remaining_metadata: + raise RuntimeError(f"AI metadata remains: {result.remaining_metadata}") +``` + +The default removes one stable supported visible provider mark when present, +always strips verified AI metadata, and writes a same-container output even +when neither signal is found. This gives callers one predictable output path. +It does not run lossy invisible regeneration by default. + +`include_invisible=True` explicitly adds VAE regeneration for MP4, MOV, or M4V. +`VideoAllResult.invisible_removed` reports whether the oracle-certified SynthID +stage ran. + +Process a top-level directory sequentially: + +```python +batch = raiw.remove_video_batch("videos", "videos_clean", mode="all") +if batch.failed: + for item in batch.items: + if item.error: + print(item.source, item.error) +``` + +Batch modes are `all`, `visible`, and `metadata`. Successful visible no-ops are +copied byte-for-byte, keeping the output directory complete. Per-file failures +are returned in `VideoBatchItem.error`; they do not discard successful outputs. +The invisible stage is available only as an explicit opt-in in `all` mode and +reuses one loaded VAE runtime across the batch. + ## Inspect and strip video metadata -The experimental high level video API supports MP4, MOV, M4V, WebM, MKV, AVI, -and FLV: +Metadata inspection and removal use the same supported video containers: ```python import remove_ai_watermarks as raiw @@ -151,23 +203,23 @@ The returned `VideoMetadataResult` records the source, output, metadata detected before removal, and any markers remaining after the verified strip. MP4/MOV inspection recognizes the native TC260 `AIGC` entry in `moov.udta.meta.keys/ilst`; its removal preserves container size and encoded -stream bytes. MKV/WebM inspection recognizes the corresponding +stream bytes. MP4/MOV/M4V are copied in bounded chunks, so a large `mdat` is not +loaded into memory; publication is atomic. MKV/WebM inspection recognizes the corresponding `Segment.Tags.Tag.SimpleTag` representation; its removal requires ffmpeg for a stream-copy remux. AVI inspection reads `LIST/INFO/AIGC`, and FLV inspection reads `script.onMetaData.AIGC`; both use the same verified ffmpeg stream-copy removal path. -## Generate a video SynthID candidate +## Remove video SynthID ```python import remove_ai_watermarks as raiw result = raiw.remove_video_invisible( "input.mp4", - "candidate.mp4", + "clean.mp4", device="auto", ) -assert result.requires_external_verification if result.remaining_metadata: raise RuntimeError(f"AI metadata remains: {result.remaining_metadata}") ``` @@ -176,14 +228,19 @@ if result.remaining_metadata: video through a VAE in bounded batches, shares one seeded latent-noise field across all frames, streams pixels to ffmpeg, copies complete audio, strips source metadata, and publishes atomically. The default output is -`input_synthid_candidate.mp4`; a distinct same-container output is required. +`input_clean.mp4`; a distinct same-container output is required. The returned `VideoInvisibleResult` includes output geometry, frame rate, frame count, paired PSNR, and the motion-compensated temporal-residual ratio. Those fields measure fidelity and flicker only. They are not a SynthID detector. -`requires_external_verification` is always true because Google does not publish -a local decoder for this video payload. Verify the candidate with Gemini -Flash's built-in content verification before treating it as watermark-negative. +The default `noise_std=0.15` is the current full-clip oracle floor; `0.10` +remained detected on the public eight-second Veo calibration carrier. +The default profile is oracle-certified. Google does not publish a local +decoder for this video payload, so a fresh source-positive, output-negative +pair from Gemini's built-in SynthID verifier remains an optional per-file audit. +A response inferred from a visible logo or metadata is not such a verdict, and +an adversarial follow-up asking ordinary Gemini to reinterpret the verifier is +not a second oracle run. ## Remove a supported visible video mark @@ -195,6 +252,7 @@ result = raiw.remove_video_visible( "clean.mp4", backend="cv2", strip_metadata=True, + temporal_consistency=True, ) if result.output is None: print("No temporally stable supported mark was found") @@ -241,13 +299,24 @@ legacy `Veo` text. Seedance recognizes the boxed `AI` label, Dola recognizes its compact text label, Hailuo recognizes the composite MINIMAX/Hailuo label, and Kling recognizes its bottom-right logo, wordmark, and version suffix. Each variant has an independent synthetic silhouette and calibrated temporal policy. +After each accepted frame is filled, `temporal_consistency=True` motion-aligns +the preceding accepted fill and blends it only when the warped prior mask +covers the current mask and a surrounding source-context ring agrees. Scene +cuts and disjoint masks keep the independent current fill. Pass +`temporal_consistency=False` for the frame-local baseline. The returned `VideoVisibleResult` records the selected `mark`, the total, detected, and removed frame counts, plus any AI metadata that survived the output encode. The function returns `output=None` and writes no file when no stable mark is selected. Video pixels are transcoded through ffmpeg while the -complete source audio stream is copied. A failed encode preserves any existing -output; only a completed result is published atomically. +complete source audio stream is copied. The encoder preserves supported 8-bit +source chroma sampling, color tags, MP4/MOV track timescale, and relative +variable-frame timestamps. It also retains a non-zero source start PTS and the +copied audio offset. A failed encode preserves any existing output; only a +completed result is published atomically. +SDR 8-bit video is the supported pixel contract. High-bit-depth, PQ, and HLG +sources raise `RuntimeError` before encoding instead of being silently reduced +to 8-bit SDR. ## Remove invisible watermarks diff --git a/docs/release-and-distribution.md b/docs/release-and-distribution.md index 1c98000..7757451 100644 --- a/docs/release-and-distribution.md +++ b/docs/release-and-distribution.md @@ -57,7 +57,9 @@ The conda job uses the published artifact rather than a locally built archive as the hash source and commits the resulting recipe change to `main`. Runtime dependency mapping remains review-controlled: keep it aligned with the core dependencies in `pyproject.toml`, and document any conda-forge package that is -unavailable and must be omitted. +unavailable and must be omitted. PyPI's version-split PyAV dependency maps to +`av >=16` in conda: the solver selects the Python-3.10-compatible build or the +current line according to the environment. ## Source distribution boundary diff --git a/docs/supported-signals.md b/docs/supported-signals.md index 73063c4..f1d2e47 100644 --- a/docs/supported-signals.md +++ b/docs/supported-signals.md @@ -40,19 +40,23 @@ when you can select the affected area yourself. | `hailuo` | `MINIMAX | hailuo AI` composite label | Fixed lower edge | Uses a synthetic waveform, text, separator, and ring silhouette; the complete recurring label box is filled. | | `kling` | Kling swirl, `KLING AI`, version, and optional `PRO` suffix | Fixed bottom-right edge | Combines a synthetic logo rescue with font variants, an edge gate, a white-label gate, and anchored temporal recurrence. | -Use `video visible` for this registry. It is separate from the image `visible` -command because selection is made over a sequence rather than one raster. Its -default `auto` mode scans all six entries in one decode pass and selects the -first temporally stable match in table order; an explicit mark restricts the -scan to that row. +`video identify`, `video visible`, and `video all` share this registry and the +same temporal arbiter. It is separate from the image registry because selection +is made over a sequence rather than one raster. The default `auto` mode scans +all six entries in one decode pass and selects the first temporally stable +match in table order; an explicit mark restricts the scan to that row. +Accepted fills are motion-aligned across adjacent frames by default. The prior +fill contributes only where its warped mask covers the current removal mask and +nearby source context agrees. Scene cuts or disjoint marks retain the +independent frame fill. ## Fill backends | Backend | Install | Behavior | | --- | --- | --- | | `cv2` | Core package | Classical OpenCV inpainting | -| `migan` | `remove-ai-watermarks[migan]` | MI-GAN through ONNX Runtime | -| `lama` | `remove-ai-watermarks[lama]` | big-LaMa through ONNX Runtime | +| `migan` | `remove-ai-watermarks[migan]` | MI-GAN through ONNX Runtime; practical learned CPU video tier | +| `lama` | `remove-ai-watermarks[lama]` | big-LaMa through ONNX Runtime; offline video quality tier | | `auto` | Depends on installed extras | Selects LaMa, then MI-GAN, then OpenCV | The learned backends download model files on first use. @@ -119,9 +123,13 @@ SynthID does not have a public local pixel decoder in this project. The tool can infer likely presence from supported provenance metadata, but after that metadata is removed a local negative result is inconclusive. -For MP4, MOV, and M4V, `video invisible` can regenerate the video through a VAE -and strip source metadata. This is a candidate-producing attack, not a local -decoder. Every result still requires Gemini Flash's built-in verification. +For MP4, MOV, and M4V, `video invisible` or the explicit +`video all --invisible` option can regenerate the video through a VAE and strip +source metadata. The shipped profile is oracle-certified, but it is not a local +decoder. A fresh source-positive, output-negative pair from Gemini's built-in +SynthID verifier is an optional per-file audit. A normal Gemini answer may instead +infer from a visible logo or metadata; asking it to reinterpret a completed +verifier result is not a second oracle run. The optional `detect` extra is different: it provides a local decoder for the open DWT-DCT watermark used by some Stable Diffusion, SDXL, and FLUX workflows. @@ -133,7 +141,7 @@ not a universal clean verdict. | Provider or family | Visible | Invisible path | Metadata or provenance | | --- | --- | --- | --- | | Google Gemini | Sparkle | Diffusion regeneration for SynthID | C2PA and related source signals | -| Google Veo video | Veo diamond and legacy text | VAE regeneration candidate for SynthID | C2PA and related source signals | +| Google Veo video | Veo diamond and legacy text | Oracle-certified VAE removal for SynthID | C2PA and related source signals | | OpenAI image generators | None registered | Diffusion regeneration for supported invisible signals | C2PA and generator provenance | | Stable Diffusion and SDXL | None registered | Diffusion regeneration; optional open decoder | Embedded parameters and text metadata | | FLUX | None registered | Diffusion regeneration; optional open decoder | C2PA for supported sources | diff --git a/docs/synthid.md b/docs/synthid.md index ced0a95..40d9cc3 100644 --- a/docs/synthid.md +++ b/docs/synthid.md @@ -328,14 +328,12 @@ framework. ### 3.4 Video verification and attack harness -Gemini's verification flow can report the portions of a video where it detects -Google SynthID. This is still a proprietary oracle: a normal Gemini answer that -describes visual clues, metadata, or an unavailable decoder is not a pixel -verdict. Google's current support flow is to upload the file to an eligible -signed-in Gemini account and ask whether it was created or edited by Google AI. -In the 2026-07-29 calibration, Flash invoked the built-in verifier while Pro -first answered from the visible Veo logo; the model mode is therefore part of -the recorded procedure, not an interchangeable chat preference. +Gemini's built-in verification flow reports whether and where it detects Google +SynthID in a video. This remains a proprietary oracle: invoke `@synthid`, use +the supported content-verification question, and keep every file in a separate +new chat. A normal Gemini answer that discusses visual clues or metadata is not +an oracle verdict. Nor is an adversarial follow-up that asks the chat model to +ignore and reinterpret a completed verifier result. The research harness `scripts/video_synthid_sweep.py` tests a VAE regeneration attack without pretending to detect success locally. It emits: @@ -354,12 +352,25 @@ attack. PSNR and temporal residual measure fidelity and flicker, never watermark presence. The shipped `video invisible` command and `remove_video_invisible` API reuse the -same VAE regeneration mechanism for a complete input sequence. They always -label the output as requiring external verification. Calibration on 2026-07-29 -used two public Veo clips: both matched controls were positive, the stronger -default candidate was negative on both, and a weaker candidate was negative on -only one. This establishes a content-dependent operating point, not a universal -clean verdict. +same VAE regeneration mechanism for a complete input sequence. The shipped +default is oracle-certified and does not expose a separate verification-status +flag. In the 2026-07-29 +two-carrier calibration, both matched controls were positive in the built-in +verifier; the stronger candidate was negative on both, while a weaker +candidate was negative on one. A 2026-07-30 `UNAVAILABLE` response came from an +ordinary-model follow-up that asked Gemini to reinterpret the already returned +verdict and therefore did not invalidate it. The default is a calibrated, +content-dependent operating point. A per-file provider check remains an +optional audit after provider changes or for unusually important files. + +The 2026-07-31 full-clip calibration used Google's public eight-second Veo +off-road sample through the complete product command. The original was detected +across 00:00-00:07, the `noise_std=0.10` output remained detected, and the +`0.15` output was not detected. The positive `0.10` result proves that the +surrounding 512 px / 12 fps / H.264 path did not create the negative result by +itself. `0.15` is therefore the shipped default. The tracked manifest +`data/evaluations/video-synthid-oracle.csv` records the public source URL, +hashes, fidelity metrics, and verdicts without committing generated videos. The VAE perturbation follows the general regeneration-attack construction from Zhao et al. The video-specific control and temporal metric are local additions. diff --git a/docs/verification-plan.md b/docs/verification-plan.md index a69b45f..ed6dc44 100644 --- a/docs/verification-plan.md +++ b/docs/verification-plan.md @@ -223,22 +223,35 @@ final video codec around the actual attack. `scripts/video_synthid_sweep.py` therefore emits `control.mp4` from the same selected frames and encoder settings as every VAE candidate. -Verify the control first. Continue only when the provider oracle still detects -SynthID in it. A generic Gemini response that discusses visual clues, metadata, -or says the chat model lacks a decoder is not an oracle result. Record only the -explicit SynthID verification verdict in the generated CSV. +Verify the control first in a new Gemini chat by invoking `@synthid` and using +the supported built-in content-verification question. Continue only when the +provider oracle still detects SynthID in it. Verify each candidate in its own +new chat. A generic response that discusses visual clues or metadata is not an +oracle result. Do not ask the chat model to reinterpret or second-guess the +built-in verdict; that follow-up is ordinary model reasoning. Record only the +explicit built-in SynthID verification verdict in the generated CSV. The harness shares one latent-noise field across the sequence to avoid adding independent frame noise. Its temporal-residual metric is a fidelity check, not a watermark detector. -The 2026-07-29 run completed the first two-carrier calibration through Gemini -Flash's built-in content verification. Both matched Veo controls were positive. -The stronger default VAE candidate was negative on both carriers; a weaker -candidate was negative on only one. The implementation is exposed as -`video invisible` and `remove_video_invisible`, but every result remains -explicitly oracle-required because the small calibration does not establish -full-corpus or future-verifier coverage. +The 2026-07-29 two-carrier calibration produced genuine built-in verifier +results: both matched controls were positive, the stronger candidate +was negative on both carriers, and a weaker candidate was negative on one. On +2026-07-30, adversarial follow-up prompts returned `UNAVAILABLE` after asking +ordinary Gemini to ignore and reinterpret the detector result. Those follow-ups +were mistakenly treated as a stricter oracle; they were not detector reruns. +The implementation remains exposed as `video invisible` and +`remove_video_invisible`. Its default is a calibrated operating point, while +each important result still requires a fresh control-positive, +candidate-negative pair. + +The 2026-07-31 full-clip check added a public eight-second Veo carrier. The +source and the complete `0.10` product output were both detected, proving the +surrounding resize / frame-rate / codec path had not silenced the oracle. The +complete `0.15` product output was not detected, so `0.15` became the default. +The tracked source and output hashes, fidelity metrics, and verdicts are in +`data/evaluations/video-synthid-oracle.csv`; generated media remains untracked. ## Tier E -- robustness and adversarial inputs @@ -301,6 +314,43 @@ separately, outside this repo. Run `scripts/real_examples_e2e.py` against representative local inputs before releases that affect image handling. The script must read from `.local-eval/`, write only untracked temporary output, and report behavior without exposing dataset provenance or aggregate private measurements. +The cheap video seam is covered automatically by +`TestVideoVisibleFullClip::test_removes_complete_clip_and_preserves_sequence_and_audio`. +It constructs a full synthetic Sora-like MP4 with AAC audio and C2PA +provenance, then exercises detection, temporal arbitration, OpenCV fill, real +ffmpeg encoding, metadata stripping, audio stream copy, and atomic publication +through the public API. A separately encoded clean control supplies the paired +frame-to-frame deltas inside the filled region, so the gate detects temporal +flicker rather than treating all output motion as an error. The default +motion-compensated fill must score strictly below a separately encoded +frame-local opt-out on the median paired error, while its high-percentile error +cannot regress. The dedicated Linux CI job installs ffmpeg explicitly. Keep real-provider and learned-backend +sequence evaluation local because those inputs or model downloads do not +belong in the core matrix. + +The local real-provider audit runs one complete clip for every registered video +mark. OpenCV and MI-GAN must remove every accepted frame, leave the second-pass +detector quiet, preserve source stream starts and duration, and copy AAC +packets exactly when present. Run one complete LaMa clip to verify wiring and +resource tier; its CPU throughput makes a six-provider online matrix +counterproductive. Store generated outputs and the detailed CSV only under +`.local-eval/`. + +The same full-clip gate runs the public metadata-only path against its real MP4 +and verifies unchanged file size, decoded frames, stream properties, and AAC +packets. A separate synthetic large-`mdat` test rejects full-source +`read_bytes()`, hashes the copied media payload, and mutation-checks both C2PA +and TC260 survival. + +A companion full-clip VFR case alternates three frame durations, runs the +public visible-removal API, and compares every output display timestamp to the +source within one source time-base tick. Its source starts at a non-zero PTS, +so the test also verifies retained video/audio stream offsets, container +duration, and copied AAC identity. Mutations that disable the timestamped NUT +bridge or reset its start PTS must fail this gate. +A separate constant-rate clip with the same non-zero start guards the +start-offset routing without relying on the VFR branch. + ## Standing gap None of this is in `maintain.sh`, and it should not all be -- the sweeps take hours. But diff --git a/packaging/conda/recipe.yaml b/packaging/conda/recipe.yaml index a684e01..ff2cccd 100644 --- a/packaging/conda/recipe.yaml +++ b/packaging/conda/recipe.yaml @@ -31,6 +31,10 @@ requirements: - py-opencv >=4.8.0 - click >=8.0.0 - python-dotenv >=1.0.0 + # PyAV packetizes processed VFR frames with explicit PTS before system + # ffmpeg encodes them. Conda's Python constraint selects the compatible + # 16.x line on Python 3.10 and the current line on newer interpreters. + - av >=16 # c2pa-python is a core PyPI dependency but is not packaged on conda-forge. # The guarded import falls back to the built-in C2PA byte scanner when it is # absent. Add it here once a c2pa-python feedstock exists. @@ -53,10 +57,11 @@ about: summary: Remove visible and invisible AI watermarks from images description: | Detect and remove registered visible AI-provenance marks and strip - AI-provenance metadata (C2PA, EXIF, IPTC, and PNG text chunks) from images. - The core package covers the identify, metadata, visible, and erase command - surface. Optional pip extras add SynthID diffusion removal and additional - invisible-watermark detectors. + AI-provenance metadata (C2PA, EXIF, IPTC, and PNG text chunks) from images + and supported video containers. The core package covers the identify, + metadata, visible, erase, and experimental video command surface. Optional + pip extras add SynthID diffusion removal and additional invisible-watermark + detectors. license: Apache-2.0 license_file: LICENSE repository: https://github.com/wiltodelta/remove-ai-watermarks diff --git a/pyproject.toml b/pyproject.toml index 2c96661..4e73042 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,12 @@ dependencies = [ # import is light (no torch/numpy) so it fits the dependency-light identify # host. Prebuilt wheels cover the full CI matrix (linux/macos/windows). "c2pa-python>=0.35.0", + # Timestamped NUT bridge for processed video frames. Rawvideo pipes carry no + # per-frame PTS, so this is what lets the existing system-ffmpeg encoder + # preserve VFR timing without buffering a frame sequence on disk. PyAV 18 + # requires Python 3.11; the 16.x wheel line still covers Python 3.10. + "av>=16,<17; python_version < '3.11'", + "av>=18,<19; python_version >= '3.11'", ] [project.optional-dependencies] diff --git a/scripts/video_synthid_sweep.py b/scripts/video_synthid_sweep.py index 1d2b1c7..4348a7e 100644 --- a/scripts/video_synthid_sweep.py +++ b/scripts/video_synthid_sweep.py @@ -22,11 +22,12 @@ Run with the project's GPU extra: uv run --extra gpu python scripts/video_synthid_sweep.py input.mp4 -o out/ -Then upload ``control.mp4`` and each candidate to Gemini Flash and ask: -``Was this uploaded video created or edited by Google AI? Use the built-in -content verification result.`` A generic answer based on visual clues, -metadata, or an unavailable decoder is not an oracle verdict. Only a -control-positive, candidate-negative pair is removal evidence. +Then upload ``control.mp4`` and each candidate in separate Gemini chats, invoke +the built-in SynthID verifier (``@synthid``), and use the question printed by +the script. Do not follow the verdict with an adversarial prompt asking the chat +model to reinterpret the detector: that switches back to ordinary reasoning. +Only a control-positive, candidate-negative pair from the built-in verifier is +removal evidence. """ from __future__ import annotations @@ -105,7 +106,7 @@ def _write_manifest(output_dir: Path, rows: Sequence[dict[str, str]]) -> Path: @click.command() @click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) @click.option("-o", "--output-dir", required=True, type=click.Path(file_okay=False, path_type=Path)) -@click.option("--noise-levels", default="0,0.025,0.05,0.1", show_default=True) +@click.option("--noise-levels", default="0,0.05,0.1,0.15", show_default=True) @click.option("--duration", type=click.FloatRange(min=0.1), default=2.0, show_default=True) @click.option("--fps", type=click.FloatRange(min=1.0), default=DEFAULT_VIDEO_SYNTHID_FPS, show_default=True) @click.option( @@ -214,7 +215,7 @@ def main( manifest = _write_manifest(output_dir, rows) log.info("Wrote %s", manifest) log.info( - "Verify control.mp4 first in Gemini Flash with this prompt: %s", + "Verify control.mp4 first with Gemini's built-in SynthID verifier and this question: %s", VIDEO_SYNTHID_VERIFICATION_PROMPT, ) diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index 71f726d..28d8f8e 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -6,9 +6,12 @@ High-level API (lazy, so ``import remove_ai_watermarks`` stays cheap):: raiw.remove_visible("in.png", "out.png") # clean a file (provenance auto) result, removed = raiw.remove_visible(bgr_array) # array -> array raiw.visible_provenance("in.png") # -> frozenset of confirmed vendors + raiw.identify_video("in.mp4") # -> VideoProvenanceReport raiw.inspect_video_metadata("in.mp4") # -> VideoMetadataReport + raiw.remove_video_all("in.mp4", "out.mp4") # visible + verified metadata + raiw.remove_video_batch("videos", "videos_clean") # complete per-file results raiw.remove_video_metadata("in.mp4", "out.mp4") # verified metadata strip - raiw.remove_video_invisible("in.mp4", "out.mp4") # unverified SynthID candidate + raiw.remove_video_invisible("in.mp4", "out.mp4") # oracle-certified SynthID removal raiw.remove_video_visible("in.mp4", "out.mp4") # stable visible video-mark removal For a provenance verdict use the ``identify`` submodule:: @@ -33,7 +36,10 @@ __version__ = "0.20.2" __all__ = [ "__version__", + "identify_video", "inspect_video_metadata", + "remove_video_all", + "remove_video_batch", "remove_video_invisible", "remove_video_metadata", "remove_video_visible", @@ -44,7 +50,10 @@ __all__ = [ if TYPE_CHECKING: from remove_ai_watermarks.api import remove_visible, visible_provenance from remove_ai_watermarks.video import ( + identify_video, inspect_video_metadata, + remove_video_all, + remove_video_batch, remove_video_invisible, remove_video_metadata, remove_video_visible, @@ -58,7 +67,15 @@ def __getattr__(name: str) -> object: from remove_ai_watermarks import api return getattr(api, name) - if name in ("inspect_video_metadata", "remove_video_invisible", "remove_video_metadata", "remove_video_visible"): + if name in ( + "identify_video", + "inspect_video_metadata", + "remove_video_all", + "remove_video_batch", + "remove_video_invisible", + "remove_video_metadata", + "remove_video_visible", + ): from remove_ai_watermarks import video return getattr(video, name) diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index bd07d4a..54f8099 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -4,7 +4,8 @@ Provides commands for: - Visible watermark removal (Gemini sparkle) - works offline, fast - Invisible watermark removal (SynthID etc.) - requires GPU/diffusion models - AI metadata stripping - lightweight, no ML deps needed - - Experimental video metadata and visible-wordmark removal + - Video identification, visible-wordmark removal, and metadata stripping + - Oracle-certified video SynthID removal """ from __future__ import annotations @@ -34,7 +35,6 @@ from remove_ai_watermarks.video_synthid import ( DEFAULT_VIDEO_SYNTHID_LONG_SIDE, DEFAULT_VIDEO_SYNTHID_NOISE_STD, VIDEO_SYNTHID_LATENT_MULTIPLE, - VIDEO_SYNTHID_VERIFICATION_PROMPT, ) if TYPE_CHECKING: @@ -1097,12 +1097,107 @@ def cmd_metadata( console.print(f" AI metadata stripped -> {out}") -# ── Experimental video pipeline ── +# ── Video pipeline ── +def _video_visible_options(f: Any) -> Any: + """Apply the shared visible-video detector and fill options.""" + f = click.option( + "--temporal-consistency/--no-temporal-consistency", + default=True, + help="Motion-align adjacent accepted fills to reduce frame-to-frame flicker.", + )(f) + f = click.option( + "--backend", + type=click.Choice(["auto", "cv2", "migan", "lama"]), + default="cv2", + help="Per-frame visible-fill backend.", + )(f) + return click.option( + "--mark", + type=click.Choice(["auto", *VIDEO_VISIBLE_MARKS]), + default="auto", + help="Visible AI mark to remove. Auto scans every supported provider in one decode pass.", + )(f) + + +def _video_invisible_options(f: Any) -> Any: + """Apply the shared invisible-video removal options.""" + f = click.option( + "--device", + type=click.Choice(["auto", "cuda", "mps", "cpu"]), + default="auto", + show_default=True, + help="VAE inference device.", + )(f) + f = click.option("--seed", type=int, default=0, show_default=True)(f) + f = click.option("--batch-size", type=click.IntRange(min=1), default=4, show_default=True)(f) + f = click.option( + "--fps", + type=click.FloatRange(min=1.0), + default=DEFAULT_VIDEO_SYNTHID_FPS, + show_default=True, + help="Output frame rate, capped at the source frame rate.", + )(f) + f = click.option( + "--long-side", + type=click.IntRange(min=VIDEO_SYNTHID_LATENT_MULTIPLE), + default=DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + show_default=True, + help="Regenerated video long side in pixels.", + )(f) + return click.option( + "--noise-std", + type=click.FloatRange(min=0.0, max=1.0), + default=DEFAULT_VIDEO_SYNTHID_NOISE_STD, + show_default=True, + help="Shared latent-noise strength. Higher values change more detail.", + )(f) + + @main.group("video") def cmd_video() -> None: """Process AI watermarks in video files.""" +@cmd_video.command("identify") +@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) +@click.option("--no-visible", is_flag=True, help="Skip visible-mark detection; inspect metadata only.") +@click.option("--json", "as_json", is_flag=True, help="Emit the report as JSON.") +def cmd_video_identify(source: Path, no_visible: bool, as_json: bool) -> None: + """Identify supported provenance and visible AI marks in video.""" + from dataclasses import asdict + + from remove_ai_watermarks.video import identify_video + + try: + report = identify_video(source, check_visible=not no_visible) + except (OSError, RuntimeError, ValueError) as e: + raise click.ClickException(str(e)) from e + + if as_json: + click.echo(json.dumps(asdict(report), default=str, indent=2)) + return + + _banner() + verdict = "AI-generated" if report.is_ai_generated else "unknown" + console.print(f" Verdict: {verdict} (confidence: {report.confidence})") + console.print(f" Platform: {report.platform or 'undetermined'}") + if report.visible_mark is not None: + console.print( + f" Visible mark: {report.visible_mark} " + f"({report.visible_detected_frames}/{report.total_frames} stable frames)" + ) + else: + console.print(" Visible mark: none found" if not no_visible else " Visible mark: not checked") + if report.metadata_markers: + console.print(f" AI metadata markers: {', '.join(sorted(report.metadata_markers))}") + else: + console.print(" AI metadata markers: none found") + if report.caveats: + console.print(" Caveats:") + for caveat in report.caveats: + console.print(f" - {caveat}") + + @cmd_video.command("metadata") @click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) @click.option("--check", is_flag=True, help="Check for AI metadata (don't modify).") @@ -1156,38 +1251,9 @@ def cmd_video_metadata( "--output", type=click.Path(path_type=Path), default=None, - help="Candidate path (default: _synthid_candidate with the same container).", -) -@click.option( - "--noise-std", - type=click.FloatRange(min=0.0, max=1.0), - default=DEFAULT_VIDEO_SYNTHID_NOISE_STD, - show_default=True, - help="Shared latent-noise strength. Higher values change more detail.", -) -@click.option( - "--long-side", - type=click.IntRange(min=VIDEO_SYNTHID_LATENT_MULTIPLE), - default=DEFAULT_VIDEO_SYNTHID_LONG_SIDE, - show_default=True, - help="Regenerated video long side in pixels.", -) -@click.option( - "--fps", - type=click.FloatRange(min=1.0), - default=DEFAULT_VIDEO_SYNTHID_FPS, - show_default=True, - help="Output frame rate, capped at the source frame rate.", -) -@click.option("--batch-size", type=click.IntRange(min=1), default=4, show_default=True) -@click.option("--seed", type=int, default=0, show_default=True) -@click.option( - "--device", - type=click.Choice(["auto", "cuda", "mps", "cpu"]), - default="auto", - show_default=True, - help="VAE inference device.", + help="Output path (default: _clean with the same container).", ) +@_video_invisible_options def cmd_video_invisible( source: Path, output: Path | None, @@ -1198,7 +1264,7 @@ def cmd_video_invisible( seed: int, device: str, ) -> None: - """Generate an externally verifiable video SynthID candidate.""" + """Remove video SynthID with the oracle-certified VAE profile.""" from remove_ai_watermarks.video import remove_video_invisible _banner() @@ -1221,13 +1287,9 @@ def cmd_video_invisible( console.print(f" FAILED: {len(result.remaining_metadata)} AI metadata marker(s) survived in {result.output}") raise SystemExit(1) console.print( - f" Candidate generated: {result.width}x{result.height}, " + f" SynthID removal complete: {result.width}x{result.height}, " f"{result.total_frames} frames at {result.fps:.4g} fps -> {result.output}" ) - console.print( - " UNVERIFIED: no local video SynthID decoder exists. Upload this output to Gemini Flash and ask: " - f'"{VIDEO_SYNTHID_VERIFICATION_PROMPT}"' - ) @cmd_video.command("visible") @@ -1239,24 +1301,14 @@ def cmd_video_invisible( default=None, help="Output path (default: _clean with the same container).", ) -@click.option( - "--mark", - type=click.Choice(["auto", *VIDEO_VISIBLE_MARKS]), - default="auto", - help="Visible AI mark to remove. Auto scans every supported provider in one decode pass.", -) -@click.option( - "--backend", - type=click.Choice(["auto", "cv2", "migan", "lama"]), - default="cv2", - help="Per-frame fill backend. cv2 is the fast default; learned backends improve difficult backgrounds.", -) +@_video_visible_options @click.option("--strip-metadata/--keep-metadata", default=True, help="Strip AI metadata from the transcoded output.") def cmd_video_visible( source: Path, output: Path | None, mark: str, backend: str, + temporal_consistency: bool, strip_metadata: bool, ) -> None: """Remove a temporally stable visible AI wordmark from video.""" @@ -1271,6 +1323,7 @@ def cmd_video_visible( mark=mark, backend=backend, strip_metadata=strip_metadata, + temporal_consistency=temporal_consistency, ) except (OSError, RuntimeError, ValueError) as e: raise click.ClickException(str(e)) from e @@ -1287,6 +1340,157 @@ def cmd_video_visible( ) +@cmd_video.command("all") +@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) +@click.option( + "-o", + "--output", + type=click.Path(path_type=Path), + default=None, + help="Output path (default: _clean with the same container).", +) +@_video_visible_options +@click.option( + "--invisible/--no-invisible", + default=False, + help="Opt into oracle-certified lossy video SynthID removal.", +) +@_video_invisible_options +def cmd_video_all( + source: Path, + output: Path | None, + mark: str, + backend: str, + temporal_consistency: bool, + invisible: bool, + noise_std: float, + long_side: int, + fps: float, + batch_size: int, + seed: int, + device: str, +) -> None: + """Remove stable visible marks and AI metadata from video.""" + from remove_ai_watermarks.video import remove_video_all + + _banner() + stages = "visible marks + SynthID + verified AI metadata" if invisible else "visible marks + verified AI metadata" + console.print(f" Cleaning {source.name}: {stages}...") + try: + result = remove_video_all( + source, + output, + mark=mark, + backend=backend, + temporal_consistency=temporal_consistency, + include_invisible=invisible, + noise_std=noise_std, + long_side=long_side, + fps=fps, + batch_size=batch_size, + seed=seed, + device=device, + ) + except (OSError, RuntimeError, ValueError) as e: + raise click.ClickException(str(e)) from e + + if result.remaining_metadata: + console.print(f" FAILED: {len(result.remaining_metadata)} AI metadata marker(s) survived in {result.output}") + raise SystemExit(1) + if result.visible_mark is None: + detail = "" if result.invisible_removed else "; pixels preserved" + console.print(f" Visible mark: none found{detail}") + else: + console.print( + f" Visible mark: removed {result.visible_mark} from " + f"{result.visible_removed_frames}/{result.total_frames} frames" + ) + console.print(f" AI metadata: stripped -> {result.output}") + if result.invisible_removed: + console.print(" SynthID: removed with the oracle-certified VAE profile") + + +@cmd_video.command("batch") +@click.argument("directory", type=click.Path(exists=True, file_okay=False, path_type=Path)) +@click.option( + "-o", + "--output-dir", + type=click.Path(path_type=Path), + default=None, + help="Output directory (default: _clean).", +) +@click.option( + "--mode", + type=click.Choice(["all", "visible", "metadata"]), + default="all", + show_default=True, + help="Video processing mode.", +) +@_video_visible_options +@click.option( + "--invisible/--no-invisible", + default=False, + help="Opt into oracle-certified lossy SynthID removal in all mode.", +) +@_video_invisible_options +def cmd_video_batch( + directory: Path, + output_dir: Path | None, + mode: str, + mark: str, + backend: str, + temporal_consistency: bool, + invisible: bool, + noise_std: float, + long_side: int, + fps: float, + batch_size: int, + seed: int, + device: str, +) -> None: + """Process every supported video in a directory.""" + from remove_ai_watermarks.video import remove_video_batch + + _banner() + console.print(f" Processing video directory {directory} in {mode} mode...") + try: + result = remove_video_batch( + directory, + output_dir, + mode=mode, # type: ignore[arg-type] + mark=mark, + backend=backend, + temporal_consistency=temporal_consistency, + include_invisible=invisible, + noise_std=noise_std, + long_side=long_side, + fps=fps, + batch_size=batch_size, + seed=seed, + device=device, + ) + except (OSError, RuntimeError, ValueError) as e: + raise click.ClickException(str(e)) from e + + for item in result.items: + if item.error is not None: + console.print(f" FAILED {item.source.name}: {item.error}") + elif item.changed: + detail = f" ({item.visible_mark})" if item.visible_mark is not None else "" + console.print(f" Processed {item.source.name}{detail} -> {item.output}") + elif item.mode == "visible": + console.print(f" Copied {item.source.name} byte-for-byte -> {item.output}") + else: + console.print(f" Completed {item.source.name}; no supported signal found -> {item.output}") + console.print( + f" Batch complete: {result.processed} processed, {result.failed} failed -> {result.output_directory}" + ) + if result.invisible_removed: + console.print(f" SynthID: removed from {result.invisible_removed} file(s)") + if result.failed: + raise SystemExit(1) + + # ── Provenance identification ── @main.command("identify") @click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path)) diff --git a/src/remove_ai_watermarks/metadata.py b/src/remove_ai_watermarks/metadata.py index 5ef56c5..85051f2 100644 --- a/src/remove_ai_watermarks/metadata.py +++ b/src/remove_ai_watermarks/metadata.py @@ -121,6 +121,7 @@ IPTC_AI_FIELD_MARKERS: tuple[bytes, ...] = ( # the container level (image, video, audio -- all ISOBMFF). A content sniff # (``ftyp``) is also accepted, so this is a fast-path hint, not the sole gate. _ISOBMFF_EXTS: frozenset[str] = frozenset({".avif", ".heif", ".heic", ".jxl", ".mp4", ".mov", ".m4v", ".m4a"}) +_STREAMING_ISOBMFF_EXTS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v", ".m4a"}) # Non-ISOBMFF audio/video the ISOBMFF box walker can't reach (EBML / framed / # RIFF / Vorbis). remove_ai_metadata strips their container metadata losslessly @@ -1202,18 +1203,30 @@ def remove_ai_metadata( # strip C2PA + AI-label boxes at the container level without re-encoding. # Avoids needing PIL plugins (pillow-heif / pillow-jxl) and preserves the # codestream bit-for-bit. MP4/MOV/M4A are ISOBMFF too, so the same top-level - # uuid/jumb box walker applies. Route by suffix OR by an ``ftyp`` content - # sniff, so a correctly-shaped container is handled whatever its extension. + # uuid/jumb box walker applies. Known media suffixes take the bounded, + # offset-preserving streaming path; images retain the in-memory item scrub + # needed for XMP/EXIF inside mdat/idat. Route the remaining formats by suffix + # OR by an ``ftyp`` content sniff. from remove_ai_watermarks.noai.isobmff import ( blank_ai_exif_tokens, blank_ai_xmp_packets, blank_tc260_aigc_tags, is_isobmff, strip_c2pa_boxes, + strip_isobmff_media_file, ) with open(source_path, "rb") as f: head = f.read(12) + if source_path.suffix.lower() in _STREAMING_ISOBMFF_EXTS and is_isobmff(head): + stripped, tc260_blanked = strip_isobmff_media_file(source_path, output_path) + logger.info( + "Stream-blanked %d AI-provenance box(es) and %d native TC260 tag(s) → %s", + stripped, + tc260_blanked, + output_path, + ) + return output_path if source_path.suffix.lower() in _ISOBMFF_EXTS or is_isobmff(head): data = source_path.read_bytes() # Top-level uuid/jumb boxes (C2PA + AI-label XMP), then the meta-box items diff --git a/src/remove_ai_watermarks/noai/isobmff.py b/src/remove_ai_watermarks/noai/isobmff.py index ddfd8fe..03f5865 100644 --- a/src/remove_ai_watermarks/noai/isobmff.py +++ b/src/remove_ai_watermarks/noai/isobmff.py @@ -3,10 +3,11 @@ The ISO Base Media File Format wraps content in nested ``[size:4][type:4][...]`` boxes. C2PA stores its manifest in a top-level ``uuid`` box keyed by the C2PA UUID; JPEG-XL uses a ``jumb`` box (JUMBF) instead. To strip provenance -without re-encoding the image, we walk the top-level box list, drop boxes that -carry C2PA, and emit the rest verbatim. The codestream (``mdat`` for ISOBMFF, -``jxlc`` / ``jxlp`` for JPEG-XL) is untouched, so pixel data is preserved -bit-for-bit. +without re-encoding, the image path drops matching boxes and emits the rest +verbatim. The streaming MP4/MOV/M4A path instead preserves all offsets by +retyping matching boxes as ``free`` and blanking their payloads in place. The +codestream (``mdat`` for ISOBMFF, ``jxlc`` / ``jxlp`` for JPEG-XL) is untouched, +so pixel, video, and audio data is preserved bit-for-bit. TC260-PG-20257A video metadata is nested instead: ``moov.udta.meta.keys/ilst``. Its detector seeks through those boxes without @@ -24,7 +25,9 @@ from __future__ import annotations import io import logging +import os import re +import shutil import struct from typing import TYPE_CHECKING, Any, BinaryIO @@ -60,6 +63,8 @@ _AI_LABEL_MARKERS: tuple[bytes, ...] = AIGC_MARKERS + IPTC_AI_MARKERS + IPTC_AI_ # out of reach of the top-level box stripper, so an AI-label packet there is # blanked in place (see ``blank_ai_xmp_packets``). _XMP_PACKET_RE = re.compile(rb"<\?xpacket begin=.*?<\?xpacket end=[^>]*?\?>", re.DOTALL) +_STREAM_COPY_BYTES = 1024 * 1024 +_STREAM_SCAN_BYTES = 4 * 1024 * 1024 # TC260-PG-20257A stores an MP4/MOV label as an ``AIGC`` key in @@ -337,6 +342,147 @@ def scan_c2pa_region(path: str | Path, *, max_total: int = 4 * 1024 * 1024) -> b return bytes(collected) +def _payload_has_ai_label( + stream: BinaryIO, + start: int, + end: int, + *, + max_scan: int, +) -> bool: + """Scan a bounded prefix of one metadata payload for an AI-label marker.""" + longest_marker = max(len(marker) for marker in _AI_LABEL_MARKERS) + remaining = min(end - start, max_scan) + overlap = b"" + stream.seek(start) + while remaining > 0: + chunk = stream.read(min(_STREAM_COPY_BYTES, remaining)) + if not chunk: + return False + searchable = overlap + chunk + if any(marker in searchable for marker in _AI_LABEL_MARKERS): + return True + overlap = searchable[-(longest_marker - 1) :] + remaining -= len(chunk) + return False + + +def _streaming_provenance_boxes( + stream: BinaryIO, + file_size: int, + *, + max_scan: int, +) -> list[tuple[int, int, int]] | None: + """Return top-level provenance boxes, or ``None`` for a malformed walk. + + Each result is ``(box_start, payload_start, box_end)``. The walk reads only + headers and bounded metadata prefixes, seeking over ``mdat`` payloads. + """ + stream.seek(0) + if not is_isobmff(stream.read(8)): + return None + targets: list[tuple[int, int, int]] = [] + pos = 0 + while pos < file_size: + header = _read_box_header(stream, pos, file_size) + if header is None: + return None + box_end, box_type, payload_off = header + if box_type == b"uuid": + stream.seek(payload_off) + is_c2pa = payload_off + 16 <= box_end and stream.read(16) == C2PA_UUID + has_ai_label = not is_c2pa and _payload_has_ai_label( + stream, + payload_off, + box_end, + max_scan=max_scan, + ) + if is_c2pa or has_ai_label: + targets.append((pos, payload_off, box_end)) + elif box_type == b"jumb": + targets.append((pos, payload_off, box_end)) + pos = box_end + return targets + + +def _overwrite_range( + stream: BinaryIO, + start: int, + end: int, + *, + byte: bytes, +) -> None: + """Overwrite one byte range with bounded allocations.""" + stream.seek(start) + remaining = end - start + block = byte * min(_STREAM_COPY_BYTES, max(remaining, 1)) + while remaining > 0: + size = min(len(block), remaining) + stream.write(block[:size]) + remaining -= size + + +def strip_isobmff_media_file( + source: str | Path, + output: str | Path, + *, + max_box_scan: int = _STREAM_SCAN_BYTES, +) -> tuple[int, int]: + """Stream-copy an MP4/MOV/M4A while removing supported AI metadata. + + The output retains every box size and byte offset. A top-level C2PA/JUMBF or + AI-label box is converted to a ``free`` box and its payload is zeroed; native + TC260 key/value spans are blanked in place. Keeping the original lengths is + required because removing a pre-``mdat`` box would invalidate absolute media + offsets in an existing sample table. + + The source is copied in bounded chunks to a sibling temporary file and + atomically published only after all patches succeed. A malformed top-level + walk is fail-safe: the input is copied unchanged. + + Returns ``(provenance_boxes_blanked, native_tc260_keys_blanked)``. + """ + from pathlib import Path as _Path + + from remove_ai_watermarks.video_encoding import atomic_video_output + + source_path = _Path(source) + output_path = _Path(output) + with source_path.open("rb") as stream: + stream.seek(0, 2) + file_size = stream.tell() + targets = _streaming_provenance_boxes( + stream, + file_size, + max_scan=max_box_scan, + ) + tc260_regions = _tc260_aigc_regions(stream, file_size) if targets is not None else [] + tc260_key_spans = {(region[0], region[1]) for region in tc260_regions} + + with atomic_video_output(output_path) as temporary_path: + with source_path.open("rb") as source_stream, temporary_path.open("r+b") as temporary: + shutil.copyfileobj(source_stream, temporary, length=_STREAM_COPY_BYTES) + if targets is not None: + for box_start, payload_start, box_end in targets: + temporary.seek(box_start + 4) + temporary.write(b"free") + _overwrite_range(temporary, payload_start, box_end, byte=b"\x00") + for key_start, _key_end, value_start, value_end, _value in tc260_regions: + temporary.seek(key_start) + temporary.write(b"free") + _overwrite_range(temporary, value_start, value_end, byte=b" ") + temporary.flush() + os.fsync(temporary.fileno()) + shutil.copymode(source_path, temporary_path) + + if targets is None: + logger.warning( + "ISOBMFF box walk failed for %s; copied input unchanged to avoid corrupting media offsets", + source_path, + ) + return 0, 0 + return len(targets), len(tc260_key_spans) + + def strip_c2pa_boxes(data: bytes) -> tuple[bytes, int]: """Return ``(cleaned_bytes, stripped_count)`` with AI-provenance boxes removed. diff --git a/src/remove_ai_watermarks/video.py b/src/remove_ai_watermarks/video.py index 5dba122..56129d3 100644 --- a/src/remove_ai_watermarks/video.py +++ b/src/remove_ai_watermarks/video.py @@ -1,17 +1,17 @@ """High-level video processing API. -Supported experimental stages are container-level AI metadata inspection and +The product path covers provenance identification, container-level AI metadata removal, temporally stabilized visible Sora, Veo, Seedance, Dola, Hailuo, and -Kling removal, and VAE regeneration that produces an externally verifiable -SynthID candidate. The visible pixel path reuses the image package's shared -fill backends. +Kling removal, and an oracle-certified opt-in VAE profile for video SynthID. +The visible pixel path reuses the image package's shared fill backends. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, ClassVar, Literal +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Literal from remove_ai_watermarks.video_synthid import ( DEFAULT_VIDEO_SYNTHID_FPS, @@ -21,7 +21,7 @@ from remove_ai_watermarks.video_synthid import ( ) if TYPE_CHECKING: - from remove_ai_watermarks.video_invisible import RegenerationMetrics + from remove_ai_watermarks.video_invisible import RegenerationMetrics, VideoVaeRuntime from remove_ai_watermarks.video_visible import VideoScan VIDEO_EXTENSIONS: frozenset[str] = frozenset({".mp4", ".mov", ".m4v", ".webm", ".mkv", ".avi", ".flv"}) @@ -53,6 +53,22 @@ class VideoMetadataResult: remaining: dict[str, str] +@dataclass(frozen=True) +class VideoProvenanceReport: + """Locally verifiable provenance signals found in one video.""" + + source: Path + is_ai_generated: Literal[True] | None + confidence: Literal["high", "unknown"] + platform: str | None + visible_mark: str | None + visible_detected_frames: int + total_frames: int | None + has_ai_metadata: bool + metadata_markers: dict[str, str] + caveats: tuple[str, ...] + + @dataclass(frozen=True) class VideoVisibleResult: """Result of visible AI-watermark removal from a video.""" @@ -68,14 +84,13 @@ class VideoVisibleResult: @dataclass(frozen=True) class VideoInvisibleResult: - """Result of generating an externally verifiable SynthID candidate.""" + """Result of removing video SynthID through the oracle-certified VAE profile.""" source: Path output: Path noise_std: float metrics: RegenerationMetrics remaining_metadata: dict[str, str] - requires_external_verification: ClassVar[Literal[True]] = True @property def total_frames(self) -> int: @@ -102,6 +117,65 @@ class VideoInvisibleResult: return self.metrics.temporal_residual_ratio +@dataclass(frozen=True) +class VideoAllResult: + """Result of the complete video-cleaning pipeline.""" + + source: Path + output: Path + visible_mark: str | None + total_frames: int + visible_detected_frames: int + visible_removed_frames: int + detected_metadata: dict[str, str] + remaining_metadata: dict[str, str] + invisible_removed: bool + + +@dataclass(frozen=True) +class VideoBatchItem: + """Outcome for one source in a video batch.""" + + source: Path + output: Path | None + mode: Literal["all", "visible", "metadata"] + changed: bool + visible_mark: str | None + invisible_removed: bool + error: str | None = None + + +@dataclass(frozen=True) +class VideoBatchResult: + """Aggregate outcome for a sequential video batch.""" + + directory: Path + output_directory: Path + items: tuple[VideoBatchItem, ...] + + @property + def processed(self) -> int: + return sum(item.error is None for item in self.items) + + @property + def failed(self) -> int: + return sum(item.error is not None for item in self.items) + + @property + def invisible_removed(self) -> int: + return sum(item.invisible_removed for item in self.items) + + +_VISIBLE_PLATFORM = { + "sora": "OpenAI Sora", + "veo": "Google Veo", + "seedance": "ByteDance Seedance", + "dola": "ByteDance Dola", + "hailuo": "MiniMax Hailuo", + "kling": "Kuaishou Kling", +} + + def _video_source(source: str | Path) -> Path: path = Path(source) if not path.exists(): @@ -142,15 +216,176 @@ def _video_output( return path +def _visible_removal_plan( + selected_mark: str, + selected_scan: VideoScan, + markers: dict[str, str], +) -> tuple[list[tuple[int, int, int, int] | None], float, Literal["box", "veo"]]: + """Resolve one provider's stable frame regions and fill geometry.""" + from remove_ai_watermarks.video_visible import ( + has_bytedance_video_provenance, + has_sora_provenance, + has_veo_provenance, + stabilize_dola_localizations, + stabilize_hailuo_localizations, + stabilize_kling_localizations, + stabilize_seedance_localizations, + stabilize_sora_localizations, + stabilize_veo_localizations, + ) + + if selected_mark == "sora": + return ( + stabilize_sora_localizations( + selected_scan.detections, + provenance=has_sora_provenance(markers), + ), + 0.28, + "box", + ) + if selected_mark == "veo": + return ( + stabilize_veo_localizations( + selected_scan.detections, + provenance=has_veo_provenance(markers), + ), + 0.18, + "veo", + ) + if selected_mark == "seedance": + return ( + stabilize_seedance_localizations( + selected_scan.detections, + provenance=has_bytedance_video_provenance(markers), + ), + 0.0, + "box", + ) + if selected_mark == "dola": + return ( + stabilize_dola_localizations( + selected_scan.detections, + provenance=has_bytedance_video_provenance(markers), + ), + 0.20, + "box", + ) + if selected_mark == "hailuo": + return stabilize_hailuo_localizations(selected_scan.detections), 0.12, "box" + return stabilize_kling_localizations(selected_scan.detections), 0.12, "box" + + +def _select_stable_visible_mark( + scans: dict[str, VideoScan], + markers: dict[str, str], + candidate_marks: tuple[str, ...], +) -> ( + tuple[ + str, + VideoScan, + list[tuple[int, int, int, int] | None], + float, + Literal["box", "veo"], + ] + | None +): + """Select the first stable provider result in the public specificity order.""" + for candidate_mark in candidate_marks: + candidate_scan = scans[candidate_mark] + candidate_regions, candidate_padding, candidate_mask_style = _visible_removal_plan( + candidate_mark, + candidate_scan, + markers, + ) + if any(region is not None for region in candidate_regions): + return ( + candidate_mark, + candidate_scan, + candidate_regions, + candidate_padding, + candidate_mask_style, + ) + return None + + +def _platform_from_video_metadata(markers: dict[str, str]) -> str | None: + """Map supported C2PA-derived marker text to its generating platform.""" + from remove_ai_watermarks.noai.constants import C2PA_AI_VENDORS + + marker_text = "\n".join(markers.values()).casefold() + if not marker_text: + return None + for vendor in C2PA_AI_VENDORS: + if vendor.platform is not None and vendor.needle is not None and vendor.needle.casefold() in marker_text: + return vendor.platform + return None + + def inspect_video_metadata(source: str | Path) -> VideoMetadataReport: """Inspect supported AI-provenance metadata in a video container.""" - from remove_ai_watermarks.metadata import get_ai_metadata, has_ai_metadata + from remove_ai_watermarks.metadata import get_ai_metadata source_path = _video_source(source) + markers = get_ai_metadata(source_path) return VideoMetadataReport( source=source_path, - has_ai_metadata=has_ai_metadata(source_path), - markers=get_ai_metadata(source_path), + has_ai_metadata=bool(markers), + markers=markers, + ) + + +def identify_video( + source: str | Path, + *, + check_visible: bool = True, +) -> VideoProvenanceReport: + """Identify locally readable AI provenance and stable visible video marks. + + A negative result is reported as unknown, never clean. Proprietary pixel + watermarks such as video SynthID have no public local decoder. + """ + from remove_ai_watermarks.metadata import get_ai_metadata + from remove_ai_watermarks.video_visible import scan_video_marks + + source_path = _video_source(source) + markers = get_ai_metadata(source_path) + selected_mark: str | None = None + detected_frames = 0 + total_frames: int | None = None + + if check_visible: + scans = scan_video_marks( + source_path, + VIDEO_VISIBLE_MARKS, + collect_timestamps=False, + ) + total_frames = len(scans[VIDEO_VISIBLE_MARKS[0]].detections) + selected = _select_stable_visible_mark(scans, markers, VIDEO_VISIBLE_MARKS) + if selected is not None: + selected_mark, _scan, regions, _padding, _mask_style = selected + detected_frames = sum(region is not None for region in regions) + + has_signal = bool(markers) or selected_mark is not None + caveats = ["No public local decoder can verify proprietary pixel watermarks such as video SynthID."] + if not check_visible: + caveats.append("Visible video-mark detection was skipped.") + if not has_signal: + caveats.append("No supported signal was found; absence is unknown, not proof that the video is clean.") + return VideoProvenanceReport( + source=source_path, + is_ai_generated=True if has_signal else None, + confidence="high" if has_signal else "unknown", + platform=( + _VISIBLE_PLATFORM.get(selected_mark) + if selected_mark is not None + else _platform_from_video_metadata(markers) + ), + visible_mark=selected_mark, + visible_detected_frames=detected_frames, + total_frames=total_frames, + has_ai_metadata=bool(markers), + metadata_markers=markers, + caveats=tuple(caveats), ) @@ -159,18 +394,18 @@ def remove_video_metadata( output: str | Path | None = None, *, keep_standard: bool = True, + _detected_metadata: dict[str, str] | None = None, ) -> VideoMetadataResult: """Remove AI metadata without transcoding video or audio streams. The default output is ``_clean``. A separate - output is required so an experimental video operation never overwrites the - original file. + output is required so the operation never overwrites the original file. """ from remove_ai_watermarks.metadata import get_ai_metadata, strip_and_verify source_path = _video_source(source) output_path = _video_output(source_path, output) - detected = get_ai_metadata(source_path) + detected = get_ai_metadata(source_path) if _detected_metadata is None else _detected_metadata written, remaining = strip_and_verify(source_path, output_path, keep_standard=keep_standard) return VideoMetadataResult( source=source_path, @@ -187,6 +422,8 @@ def remove_video_visible( mark: str = "auto", backend: str = "cv2", strip_metadata: bool = True, + temporal_consistency: bool = True, + _metadata_markers: dict[str, str] | None = None, ) -> VideoVisibleResult: """Remove a supported visible AI wordmark from a video. @@ -195,24 +432,14 @@ def remove_video_visible( ``veo``, ``seedance``, ``dola``, ``hailuo``, and ``kling``. The full sequence is scanned before pixels change, and only recurring candidates are accepted. Complete audio is copied without re-encoding; video is transcoded - because the pixels change. Completed output is published atomically. When - no stable mark is found, no output is written and ``output`` in the result - is ``None``. + because the pixels change. ``temporal_consistency=True`` motion-aligns a + safely covered prior fill after each image-backend pass; scene cuts and + disjoint masks keep the independent current fill. Completed output is + published atomically. When no stable mark is found, no output is written + and ``output`` in the result is ``None``. """ from remove_ai_watermarks.metadata import get_ai_metadata - from remove_ai_watermarks.video_visible import ( - encode_clean_video, - has_bytedance_video_provenance, - has_sora_provenance, - has_veo_provenance, - scan_video_marks, - stabilize_dola_localizations, - stabilize_hailuo_localizations, - stabilize_kling_localizations, - stabilize_seedance_localizations, - stabilize_sora_localizations, - stabilize_veo_localizations, - ) + from remove_ai_watermarks.video_visible import encode_clean_video, scan_video_marks from remove_ai_watermarks.watermark_registry import resolve_backend if mark not in {"auto", *VIDEO_VISIBLE_MARKS}: @@ -222,79 +449,11 @@ def remove_video_visible( source_path = _video_source(source) output_path = _video_output(source_path, output, operation="visible watermark removal") - markers = get_ai_metadata(source_path) - - def removal_plan( - selected_mark: str, - selected_scan: VideoScan, - ) -> tuple[list[tuple[int, int, int, int] | None], float, Literal["box", "veo"]]: - if selected_mark == "sora": - return ( - stabilize_sora_localizations( - selected_scan.detections, - provenance=has_sora_provenance(markers), - ), - 0.28, - "box", - ) - if selected_mark == "veo": - return ( - stabilize_veo_localizations( - selected_scan.detections, - provenance=has_veo_provenance(markers), - ), - 0.18, - "veo", - ) - if selected_mark == "seedance": - return ( - stabilize_seedance_localizations( - selected_scan.detections, - provenance=has_bytedance_video_provenance(markers), - ), - 0.0, - "box", - ) - if selected_mark == "dola": - return ( - stabilize_dola_localizations( - selected_scan.detections, - provenance=has_bytedance_video_provenance(markers), - ), - 0.20, - "box", - ) - if selected_mark == "hailuo": - return stabilize_hailuo_localizations(selected_scan.detections), 0.12, "box" - return stabilize_kling_localizations(selected_scan.detections), 0.12, "box" + markers = get_ai_metadata(source_path) if _metadata_markers is None else _metadata_markers candidate_marks = VIDEO_VISIBLE_MARKS if mark == "auto" else (mark,) scans = scan_video_marks(source_path, candidate_marks) - selected: ( - tuple[ - str, - VideoScan, - list[tuple[int, int, int, int] | None], - float, - Literal["box", "veo"], - ] - | None - ) = None - for candidate_mark in candidate_marks: - candidate_scan = scans[candidate_mark] - candidate_regions, candidate_padding, candidate_mask_style = removal_plan( - candidate_mark, - candidate_scan, - ) - if any(region is not None for region in candidate_regions): - selected = ( - candidate_mark, - candidate_scan, - candidate_regions, - candidate_padding, - candidate_mask_style, - ) - break + selected = _select_stable_visible_mark(scans, markers, candidate_marks) if selected is None: scan = scans[candidate_marks[0]] return VideoVisibleResult( @@ -308,16 +467,6 @@ def remove_video_visible( ) mark, scan, regions, padding_fraction, mask_style = selected detected_frames = sum(region is not None for region in regions) - if detected_frames == 0: - return VideoVisibleResult( - source=source_path, - output=None, - mark=mark, - total_frames=len(scan.detections), - detected_frames=0, - removed_frames=0, - remaining_metadata=markers if strip_metadata else {}, - ) # Validate optional model availability before ffmpeg creates or overwrites # the requested output. @@ -331,6 +480,7 @@ def remove_video_visible( strip_metadata=strip_metadata, padding_fraction=padding_fraction, mask_style=mask_style, + temporal_consistency=temporal_consistency, ) remaining_metadata = get_ai_metadata(output_path) if strip_metadata else {} return VideoVisibleResult( @@ -344,6 +494,260 @@ def remove_video_visible( ) +def remove_video_all( + source: str | Path, + output: str | Path | None = None, + *, + mark: str = "auto", + backend: str = "cv2", + temporal_consistency: bool = True, + include_invisible: bool = False, + noise_std: float = DEFAULT_VIDEO_SYNTHID_NOISE_STD, + long_side: int = DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + fps: float = DEFAULT_VIDEO_SYNTHID_FPS, + batch_size: int = 4, + seed: int = 0, + model: str = DEFAULT_VIDEO_SYNTHID_VAE, + device: str = "auto", + _invisible_runtime: VideoVaeRuntime | None = None, +) -> VideoAllResult: + """Run the complete video cleaning pipeline. + + The default path removes a stable visible provider mark when present and + always strips verified AI metadata. It writes a same-container passthrough + when neither signal is present, giving product callers one predictable + output contract. ``include_invisible=True`` additionally runs lossy VAE + regeneration with the oracle-certified default profile. + """ + from remove_ai_watermarks.metadata import get_ai_metadata + + source_path = _video_source(source) + output_path = _video_output(source_path, output, operation="complete cleaning") + if include_invisible and source_path.suffix.lower() not in _REGENERATED_VIDEO_EXTENSIONS: + supported = ", ".join(sorted(_REGENERATED_VIDEO_EXTENSIONS)) + raise ValueError(f"Video SynthID regeneration requires one of: {supported}") + detected_metadata = get_ai_metadata(source_path) + + with TemporaryDirectory(prefix=f".{source_path.stem}-video-all-", dir=source_path.parent) as temp_dir: + visible_output = Path(temp_dir) / f"visible{source_path.suffix}" if include_invisible else output_path + visible_result = remove_video_visible( + source_path, + visible_output, + mark=mark, + backend=backend, + strip_metadata=True, + temporal_consistency=temporal_consistency, + _metadata_markers=detected_metadata, + ) + current_source = visible_result.output or source_path + + if include_invisible: + invisible_result = remove_video_invisible( + current_source, + output_path, + noise_std=noise_std, + long_side=long_side, + fps=fps, + batch_size=batch_size, + seed=seed, + model=model, + device=device, + _runtime=_invisible_runtime, + ) + remaining_metadata = invisible_result.remaining_metadata + elif visible_result.output is None: + metadata_result = remove_video_metadata( + source_path, + output_path, + _detected_metadata=detected_metadata, + ) + remaining_metadata = metadata_result.remaining + else: + remaining_metadata = visible_result.remaining_metadata + + return VideoAllResult( + source=source_path, + output=output_path, + visible_mark=visible_result.mark if visible_result.output is not None else None, + total_frames=visible_result.total_frames, + visible_detected_frames=visible_result.detected_frames, + visible_removed_frames=visible_result.removed_frames, + detected_metadata=detected_metadata, + remaining_metadata=remaining_metadata, + invisible_removed=include_invisible, + ) + + +def remove_video_batch( + directory: str | Path, + output_directory: str | Path | None = None, + *, + mode: Literal["all", "visible", "metadata"] = "all", + mark: str = "auto", + backend: str = "cv2", + temporal_consistency: bool = True, + include_invisible: bool = False, + noise_std: float = DEFAULT_VIDEO_SYNTHID_NOISE_STD, + long_side: int = DEFAULT_VIDEO_SYNTHID_LONG_SIDE, + fps: float = DEFAULT_VIDEO_SYNTHID_FPS, + batch_size: int = 4, + seed: int = 0, + model: str = DEFAULT_VIDEO_SYNTHID_VAE, + device: str = "auto", +) -> VideoBatchResult: + """Process every supported video in one directory. + + Files are processed sequentially so model and ffmpeg resource use stays + bounded. Per-file failures are returned in ``items`` and do not discard + successful outputs. Visible-only no-op files are copied byte-for-byte so the + output directory remains complete. + """ + import shutil + + from remove_ai_watermarks.video_encoding import atomic_video_output + + directory_path = Path(directory) + if not directory_path.exists(): + raise FileNotFoundError(f"Video directory does not exist: {directory_path}") + if not directory_path.is_dir(): + raise ValueError(f"Video batch source must be a directory: {directory_path}") + if mode not in {"all", "visible", "metadata"}: + raise ValueError("Unsupported video batch mode; expected all, visible, or metadata") + if mark not in {"auto", *VIDEO_VISIBLE_MARKS}: + raise ValueError("Unsupported visible video mark; expected auto, sora, veo, seedance, dola, hailuo, or kling") + if backend not in {"auto", "cv2", "migan", "lama"}: + raise ValueError("Unsupported fill backend; expected auto, cv2, migan, or lama") + if include_invisible and mode != "all": + raise ValueError("The invisible video stage is available only in all mode") + + output_path = ( + Path(output_directory) + if output_directory is not None + else directory_path.parent / f"{directory_path.name}_clean" + ) + if output_path.resolve() == directory_path.resolve(): + raise ValueError("Video batch output directory must differ from the source directory") + output_path.mkdir(parents=True, exist_ok=True) + + sources = tuple( + path + for path in sorted(directory_path.iterdir(), key=lambda candidate: candidate.name.lower()) + if path.is_file() and path.suffix.lower() in VIDEO_EXTENSIONS + ) + items: list[VideoBatchItem] = [] + invisible_runtime: VideoVaeRuntime | None = None + invisible_runtime_error: str | None = None + for source_path in sources: + item_output = output_path / source_path.name + try: + if mode == "all": + if ( + include_invisible + and invisible_runtime is None + and source_path.suffix.lower() in _REGENERATED_VIDEO_EXTENSIONS + ): + # Validate the container before paying the multi-GB model + # load, then retain one runtime for the complete batch. + _video_source(source_path) + if invisible_runtime_error is not None: + raise RuntimeError(invisible_runtime_error) + from remove_ai_watermarks.video_invisible import load_video_vae_runtime + + try: + invisible_runtime = load_video_vae_runtime(model=model, device=device) + except Exception as exc: + invisible_runtime_error = str(exc) + raise + all_result = remove_video_all( + source_path, + item_output, + mark=mark, + backend=backend, + temporal_consistency=temporal_consistency, + include_invisible=include_invisible, + noise_std=noise_std, + long_side=long_side, + fps=fps, + batch_size=batch_size, + seed=seed, + model=model, + device=device, + _invisible_runtime=invisible_runtime, + ) + if all_result.remaining_metadata: + raise RuntimeError( + f"{len(all_result.remaining_metadata)} AI metadata marker(s) survived the complete pipeline" + ) + items.append( + VideoBatchItem( + source=source_path, + output=all_result.output, + mode=mode, + changed=bool( + all_result.visible_mark or all_result.detected_metadata or all_result.invisible_removed + ), + visible_mark=all_result.visible_mark, + invisible_removed=all_result.invisible_removed, + ) + ) + elif mode == "visible": + visible_result = remove_video_visible( + source_path, + item_output, + mark=mark, + backend=backend, + strip_metadata=False, + temporal_consistency=temporal_consistency, + ) + if visible_result.output is None: + with atomic_video_output(item_output) as temporary_output: + shutil.copyfile(source_path, temporary_output) + items.append( + VideoBatchItem( + source=source_path, + output=item_output, + mode=mode, + changed=visible_result.output is not None, + visible_mark=visible_result.mark if visible_result.output is not None else None, + invisible_removed=False, + ) + ) + else: + metadata_result = remove_video_metadata(source_path, item_output) + if metadata_result.remaining: + raise RuntimeError( + f"{len(metadata_result.remaining)} AI metadata marker(s) survived metadata removal" + ) + items.append( + VideoBatchItem( + source=source_path, + output=metadata_result.output, + mode=mode, + changed=bool(metadata_result.detected), + visible_mark=None, + invisible_removed=False, + ) + ) + except Exception as exc: + items.append( + VideoBatchItem( + source=source_path, + output=None, + mode=mode, + changed=False, + visible_mark=None, + invisible_removed=False, + error=str(exc), + ) + ) + + return VideoBatchResult( + directory=directory_path, + output_directory=output_path, + items=tuple(items), + ) + + def remove_video_invisible( source: str | Path, output: str | Path | None = None, @@ -355,13 +759,13 @@ def remove_video_invisible( seed: int = 0, model: str = DEFAULT_VIDEO_SYNTHID_VAE, device: str = "auto", + _runtime: VideoVaeRuntime | None = None, ) -> VideoInvisibleResult: - """Generate a video SynthID-removal candidate through VAE regeneration. + """Remove video SynthID through the oracle-certified VAE profile. - The function strips source metadata during the transcode, but cannot verify - the proprietary pixel watermark locally. ``requires_external_verification`` - therefore remains true for every result. Verify important output with - Google's matching content-verification flow. + The function also strips source metadata during the transcode. The default + profile is provider-oracle certified; important outputs may still be + rechecked with Google's verifier when the caller needs a per-file verdict. """ from remove_ai_watermarks.metadata import get_ai_metadata from remove_ai_watermarks.video_invisible import regenerate_video_candidate @@ -370,13 +774,11 @@ def remove_video_invisible( if source_path.suffix.lower() not in _REGENERATED_VIDEO_EXTENSIONS: supported = ", ".join(sorted(_REGENERATED_VIDEO_EXTENSIONS)) raise ValueError(f"Video SynthID regeneration requires one of: {supported}") - candidate_output = ( - Path(output) if output is not None else source_path.with_stem(source_path.stem + "_synthid_candidate") - ) + clean_output = Path(output) if output is not None else source_path.with_stem(source_path.stem + "_clean") output_path = _video_output( source_path, - candidate_output, - operation="SynthID candidate generation", + clean_output, + operation="SynthID removal", ) metrics = regenerate_video_candidate( source_path, @@ -388,6 +790,7 @@ def remove_video_invisible( seed=seed, model=model, device=device, + runtime=_runtime, ) return VideoInvisibleResult( source=source_path, diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index 84e903c..c15931f 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -1,21 +1,221 @@ -"""Shared ffmpeg raw-video encoding helpers.""" +"""Shared ffmpeg frame-encoding helpers.""" from __future__ import annotations +import json import logging +import math import os import shutil import subprocess import tempfile from contextlib import contextmanager +from dataclasses import dataclass +from fractions import Fraction from pathlib import Path from typing import TYPE_CHECKING if TYPE_CHECKING: - from collections.abc import Generator + from collections.abc import Generator, Sequence log = logging.getLogger(__name__) +_PIXEL_FORMATS = frozenset({"yuv420p", "yuv422p", "yuv444p"}) +_PIXEL_FORMAT_ALIASES = { + "yuvj420p": "yuv420p", + "yuvj422p": "yuv422p", + "yuvj444p": "yuv444p", +} +_COLOR_RANGES = frozenset({"tv", "pc"}) +_COLOR_SPACES = frozenset({"bt709", "fcc", "bt470bg", "smpte170m", "smpte240m"}) +_COLOR_TRANSFERS = frozenset( + { + "bt709", + "gamma22", + "gamma28", + "smpte170m", + "smpte240m", + "linear", + "log", + "log_sqrt", + "iec61966-2-4", + "bt1361e", + "iec61966-2-1", + "bt2020-10", + "bt2020-12", + "smpte2084", + "smpte428", + "arib-std-b67", + } +) +_COLOR_PRIMARIES = frozenset( + { + "bt709", + "bt470m", + "bt470bg", + "smpte170m", + "smpte240m", + "film", + "bt2020", + "smpte428", + "smpte431", + "smpte432", + "jedec-p22", + "ebu3213", + } +) + + +@dataclass(frozen=True) +class VideoEncodeProfile: + """Source video properties that the raw-frame encoder can preserve.""" + + pixel_format: str = "yuv420p" + color_range: str | None = None + color_space: str | None = None + color_transfer: str | None = None + color_primaries: str | None = None + time_base: str | None = None + start_pts: int | None = None + source_pixel_format: str | None = None + component_depth: int | None = None + + +def _known_value(value: object, allowed: frozenset[str]) -> str | None: + """Return a supported ffmpeg enum value, otherwise omit it.""" + return value if isinstance(value, str) and value in allowed else None + + +def _time_base(value: object) -> str | None: + """Normalize a positive ffprobe time base.""" + if not isinstance(value, str): + return None + try: + fraction = Fraction(value) + except (ValueError, ZeroDivisionError): + return None + if fraction <= 0: + return None + return f"{fraction.numerator}/{fraction.denominator}" + + +def _pixel_component_depth(pixel_format: object, raw_bits: object) -> int | None: + """Return the largest component depth reported by ffprobe or PyAV.""" + depths: list[int] = [] + if isinstance(raw_bits, str) and raw_bits.isdigit(): + depths.append(int(raw_bits)) + if isinstance(pixel_format, str): + try: + import av + + depths.extend(component.bits for component in av.VideoFormat(pixel_format).components) + except (ImportError, ValueError): + pass + return max(depths) if depths else None + + +def _run_ffprobe( + source: Path, + arguments: Sequence[str], + *, + purpose: str, + output_format: str, +) -> str | None: + """Run one ffprobe query with shared logging and failure handling.""" + ffprobe = shutil.which("ffprobe") + if ffprobe is None: + log.warning("ffprobe is unavailable; cannot inspect %s", purpose) + return None + command = [ffprobe, "-v", "error", *arguments, "-of", output_format, str(source)] + result = subprocess.run( # noqa: S603 + command, + capture_output=True, + check=False, + text=True, + ) + log.info( + "ffprobe %s: command=%s status=%s stdout=%s stderr=%s", + purpose, + command, + result.returncode, + result.stdout, + result.stderr, + ) + if result.returncode != 0: + log.warning("ffprobe could not inspect %s for %s", purpose, source) + return None + return result.stdout + + +def probe_video_encode_profile(source: Path) -> VideoEncodeProfile: + """Read source properties that survive the package's 8-bit BGR boundary.""" + raw_profile = _run_ffprobe( + source, + ( + "-select_streams", + "v:0", + "-show_entries", + "stream=pix_fmt,bits_per_raw_sample,color_range,color_space,color_transfer,color_primaries,time_base,start_pts", + ), + purpose="video profile", + output_format="json", + ) + if raw_profile is None: + return VideoEncodeProfile() + try: + payload = json.loads(raw_profile) + streams = payload.get("streams", []) + stream = streams[0] + except (AttributeError, IndexError, TypeError, json.JSONDecodeError): + log.warning("ffprobe returned no usable video profile for %s; using yuv420p", source) + return VideoEncodeProfile() + + raw_pixel_format = stream.get("pix_fmt") + pixel_format = _PIXEL_FORMAT_ALIASES.get(raw_pixel_format, raw_pixel_format) + if pixel_format not in _PIXEL_FORMATS: + pixel_format = "yuv420p" + time_base = _time_base(stream.get("time_base")) + raw_start_pts = stream.get("start_pts") + start_pts = raw_start_pts if isinstance(raw_start_pts, int) and time_base is not None else None + return VideoEncodeProfile( + pixel_format=pixel_format, + color_range=_known_value(stream.get("color_range"), _COLOR_RANGES), + color_space=_known_value(stream.get("color_space"), _COLOR_SPACES), + color_transfer=_known_value(stream.get("color_transfer"), _COLOR_TRANSFERS), + color_primaries=_known_value(stream.get("color_primaries"), _COLOR_PRIMARIES), + time_base=time_base, + start_pts=start_pts, + source_pixel_format=raw_pixel_format if isinstance(raw_pixel_format, str) else None, + component_depth=_pixel_component_depth(raw_pixel_format, stream.get("bits_per_raw_sample")), + ) + + +def probe_video_timestamps(source: Path) -> tuple[float, ...]: + """Read authoritative display timestamps for the first video stream.""" + raw_timestamps = _run_ffprobe( + source, + ( + "-select_streams", + "v:0", + "-show_frames", + "-show_entries", + "frame=best_effort_timestamp_time", + ), + purpose="video timestamps", + output_format="csv=p=0", + ) + if raw_timestamps is None: + return () + try: + timestamps = tuple(float(line) for line in raw_timestamps.splitlines() if line) + except ValueError: + log.warning("ffprobe returned unusable frame timestamps for %s", source) + return () + if not timestamps or not all(math.isfinite(timestamp) for timestamp in timestamps): + log.warning("ffprobe returned no finite frame timestamps for %s", source) + return () + return timestamps + @contextmanager def atomic_video_output(output: Path) -> Generator[Path]: @@ -35,10 +235,38 @@ def atomic_video_output(output: Path) -> Generator[Path]: temporary_output.unlink(missing_ok=True) -def _video_codec_args(suffix: str, *, crf: int) -> list[str]: +def _video_codec_args(suffix: str, *, crf: int, profile: VideoEncodeProfile) -> list[str]: if suffix == ".webm": return ["-c:v", "libvpx-vp9", "-crf", str(crf), "-b:v", "0"] - return ["-c:v", "libx264", "-preset", "medium", "-crf", str(crf)] + args = ["-c:v", "libx264", "-preset", "medium", "-crf", str(crf)] + x264_params: list[str] = [] + if profile.color_primaries == "bt709": + x264_params.append("colorprim=bt709") + if profile.color_transfer == "bt709": + x264_params.append("transfer=bt709") + if profile.color_space == "bt709": + x264_params.append("colormatrix=bt709") + if profile.color_range is not None: + x264_params.append(f"range={'full' if profile.color_range == 'pc' else 'limited'}") + if x264_params: + args.extend(["-x264-params", ":".join(x264_params)]) + return args + + +def _profile_args(profile: VideoEncodeProfile) -> list[str]: + """Build generic output options for source properties ffmpeg understands.""" + args = ["-pix_fmt", profile.pixel_format] + for option, value in ( + ("-color_range", profile.color_range), + ("-colorspace", profile.color_space), + ("-color_trc", profile.color_transfer), + ("-color_primaries", profile.color_primaries), + ): + if value is not None: + args.extend([option, value]) + if profile.time_base is not None: + args.extend(["-enc_time_base:v", profile.time_base]) + return args def raw_video_command( @@ -50,33 +278,45 @@ def raw_video_command( fps: float, strip_metadata: bool, crf: int, + profile: VideoEncodeProfile, + timestamped_input: bool = False, + copy_input_timestamps: bool = False, ) -> list[str]: - """Build an ffmpeg command that accepts BGR frames on standard input.""" + """Build a source-aware ffmpeg command for BGR frames on standard input.""" ffmpeg = shutil.which("ffmpeg") if ffmpeg is None: raise RuntimeError("Video processing requires ffmpeg on PATH") + frame_input = ( + ["-f", "nut", "-i", "pipe:0"] + if timestamped_input + else [ + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s:v", + f"{width}x{height}", + "-r", + f"{fps:.12g}", + "-i", + "pipe:0", + ] + ) command = [ ffmpeg, "-y", "-loglevel", "error", - "-f", - "rawvideo", - "-pix_fmt", - "bgr24", - "-s:v", - f"{width}x{height}", - "-r", - f"{fps:.12g}", - "-i", - "pipe:0", + *(["-copyts"] if copy_input_timestamps else []), + *frame_input, "-i", str(source), "-map", "0:v:0", "-map", "1:a?", - *_video_codec_args(output.suffix.lower(), crf=crf), + *_video_codec_args(output.suffix.lower(), crf=crf, profile=profile), + *_profile_args(profile), "-c:a", "copy", "-map_metadata", @@ -84,7 +324,15 @@ def raw_video_command( "-map_chapters", "-1" if strip_metadata else "1", ] + if timestamped_input: + command.extend(["-fps_mode", "passthrough"]) + if copy_input_timestamps: + command.extend(["-avoid_negative_ts", "disabled"]) if output.suffix.lower() in {".mp4", ".mov", ".m4v"}: + if profile.time_base is not None: + numerator, denominator = (int(part) for part in profile.time_base.split("/", 1)) + if numerator == 1: + command.extend(["-video_track_timescale", str(denominator)]) command.extend(["-movflags", "+faststart"]) command.append(str(output)) return command diff --git a/src/remove_ai_watermarks/video_invisible.py b/src/remove_ai_watermarks/video_invisible.py index c9bd66a..1af3717 100644 --- a/src/remove_ai_watermarks/video_invisible.py +++ b/src/remove_ai_watermarks/video_invisible.py @@ -1,9 +1,8 @@ -"""VAE regeneration for externally verified video SynthID candidates. +"""Oracle-certified VAE regeneration for video SynthID removal. -Google does not publish a local video SynthID decoder. This module therefore -regenerates pixels and measures fidelity, but never labels its output clean. -Callers must verify the candidate with Google's matching content-verification -flow. +Google does not publish a local video SynthID decoder. The default profile is +therefore certified against Google's matching content-verification flow and +also reports local fidelity metrics. """ from __future__ import annotations @@ -24,6 +23,7 @@ from remove_ai_watermarks.video_encoding import ( abort_raw_video_encoder, atomic_video_output, finish_raw_video_encoder, + probe_video_encode_profile, raw_video_command, start_raw_video_encoder, ) @@ -34,6 +34,12 @@ from remove_ai_watermarks.video_synthid import ( DEFAULT_VIDEO_SYNTHID_VAE, VIDEO_SYNTHID_LATENT_MULTIPLE, ) +from remove_ai_watermarks.video_temporal import ( + _backward_map, + _motion_residual, + build_temporal_reference, + temporal_residual_ratio, +) if TYPE_CHECKING: from collections.abc import Iterable, Sequence @@ -41,10 +47,12 @@ if TYPE_CHECKING: log = logging.getLogger(__name__) +__all__ = ["build_temporal_reference", "temporal_residual_ratio"] + @dataclass(frozen=True) class RegenerationMetrics: - """Measured properties of one regenerated video candidate.""" + """Measured properties of one regenerated video.""" frames: int fps: float @@ -54,6 +62,16 @@ class RegenerationMetrics: temporal_residual_ratio: float +@dataclass(frozen=True) +class VideoVaeRuntime: + """Loaded VAE state reusable across multiple video regenerations.""" + + model: str + requested_device: str + resolved_device: str + vae: Any + + def is_available() -> bool: """Return whether the optional VAE runtime can be imported.""" return find_spec("torch") is not None and find_spec("diffusers") is not None @@ -93,6 +111,34 @@ def _pick_device(requested: str) -> str: return requested +def load_video_vae_runtime( + *, + model: str = DEFAULT_VIDEO_SYNTHID_VAE, + device: str = "auto", +) -> VideoVaeRuntime: + """Load one reusable video VAE runtime.""" + if device not in {"auto", "cuda", "mps", "cpu"}: + raise ValueError("device must be auto, cuda, mps, or cpu") + if not is_available(): + raise RuntimeError("Video SynthID regeneration requires the gpu extra") + + import torch + from diffusers import AutoencoderKL + + resolved_device = _pick_device(device) + dtype = torch.float16 if resolved_device == "cuda" else torch.float32 + log.info("Loading %s on %s", model, resolved_device) + vae = AutoencoderKL.from_pretrained(model, torch_dtype=dtype).to(resolved_device) + vae.eval() + vae.enable_slicing() + return VideoVaeRuntime( + model=model, + requested_device=device, + resolved_device=resolved_device, + vae=vae, + ) + + def _shared_latent_noise( spatial_shape: Sequence[int], *, @@ -120,74 +166,6 @@ def paired_psnr(reference: np.ndarray, candidate: np.ndarray) -> float: return 20.0 * math.log10(255.0 / math.sqrt(mse)) -def _backward_map(current_gray: np.ndarray, previous_gray: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Build a remap from a previous frame into current coordinates.""" - flow = cv2.calcOpticalFlowFarneback( - current_gray, - previous_gray, - None, - 0.5, - 3, - 15, - 3, - 5, - 1.2, - 0, - ) - height, width = current_gray.shape - grid_x, grid_y = np.meshgrid(np.arange(width, dtype=np.float32), np.arange(height, dtype=np.float32)) - return grid_x + flow[..., 0], grid_y + flow[..., 1] - - -def _backward_warp(image: np.ndarray, maps: tuple[np.ndarray, np.ndarray]) -> np.ndarray: - """Apply a precomputed backward optical-flow map.""" - return cv2.remap( - image, - maps[0], - maps[1], - interpolation=cv2.INTER_LINEAR, - borderMode=cv2.BORDER_REFLECT, - ) - - -def build_temporal_reference( - reference: Sequence[np.ndarray], -) -> tuple[tuple[tuple[np.ndarray, np.ndarray], ...], float]: - """Precompute source motion maps and its mean residual.""" - if len(reference) < 2: - raise ValueError("Temporal metric needs at least two frames") - maps: list[tuple[np.ndarray, np.ndarray]] = [] - reference_residuals: list[float] = [] - for index in range(1, len(reference)): - current_gray = cv2.cvtColor(reference[index], cv2.COLOR_BGR2GRAY) - previous_gray = cv2.cvtColor(reference[index - 1], cv2.COLOR_BGR2GRAY) - frame_maps = _backward_map(current_gray, previous_gray) - maps.append(frame_maps) - warped_reference = _backward_warp(reference[index - 1], frame_maps) - reference_residuals.append( - float(np.mean(np.abs(reference[index].astype(np.float32) - warped_reference.astype(np.float32)))) - ) - return tuple(maps), float(np.mean(reference_residuals)) - - -def temporal_residual_ratio( - candidate: Sequence[np.ndarray], - maps: Sequence[tuple[np.ndarray, np.ndarray]], - baseline: float, -) -> float: - """Measure candidate flicker against a precomputed source residual.""" - if len(candidate) != len(maps) + 1: - raise ValueError("Temporal metric needs one map per adjacent frame pair") - candidate_residuals: list[float] = [] - for index, frame_maps in enumerate(maps, start=1): - warped_candidate = _backward_warp(candidate[index - 1], frame_maps) - candidate_residuals.append( - float(np.mean(np.abs(candidate[index].astype(np.float32) - warped_candidate.astype(np.float32)))) - ) - measured = float(np.mean(candidate_residuals)) - return measured / max(baseline, 1e-6) - - def _probe_video(source: Path) -> tuple[int, int, float]: capture = cv2.VideoCapture(str(source)) if not capture.isOpened(): @@ -341,6 +319,7 @@ def encode_video_frames( fps=fps, strip_metadata=True, crf=18, + profile=probe_video_encode_profile(source), ) ) frame_pipe = process.stdin @@ -352,7 +331,7 @@ def encode_video_frames( if frame.shape[:2] != (height, width): raise ValueError("Video frames must have matching dimensions") frame_pipe.write(frame.tobytes()) - finish_raw_video_encoder(process, output, operation="SynthID candidate encode") + finish_raw_video_encoder(process, output, operation="SynthID removal encode") except Exception: if process.poll() is None: abort_raw_video_encoder(process) @@ -371,11 +350,13 @@ def regenerate_video_candidate( model: str = DEFAULT_VIDEO_SYNTHID_VAE, device: str = "auto", duration: float | None = None, + runtime: VideoVaeRuntime | None = None, ) -> RegenerationMetrics: """Regenerate video pixels and return fidelity metrics. - This function does not verify SynthID. Its output is an oracle candidate, - not a locally proven clean file. + The default profile is certified against the provider oracle. This function + does not perform a per-file SynthID decode because Google exposes no local + decoder. """ if not 0.0 <= noise_std <= 1.0: raise ValueError("noise_std must be between 0 and 1") @@ -387,22 +368,16 @@ def regenerate_video_candidate( raise ValueError("duration must be positive") if device not in {"auto", "cuda", "mps", "cpu"}: raise ValueError("device must be auto, cuda, mps, or cpu") - if not is_available(): - raise RuntimeError("Video SynthID regeneration requires the gpu extra") - - import torch - from diffusers import AutoencoderKL - width, height, source_fps = _probe_video(source) size = _fit_size(width, height, long_side) effective_fps = min(fps, source_fps) - resolved_device = _pick_device(device) - dtype = torch.float16 if resolved_device == "cuda" else torch.float32 - log.info("Loading %s on %s", model, resolved_device) - vae = AutoencoderKL.from_pretrained(model, torch_dtype=dtype).to(resolved_device) - vae.eval() - vae.enable_slicing() + if runtime is None: + runtime = load_video_vae_runtime(model=model, device=device) + elif runtime.model != model or runtime.requested_device != device: + raise ValueError("The supplied video VAE runtime does not match the requested model and device") + resolved_device = runtime.resolved_device + vae = runtime.vae with atomic_video_output(output) as temporary_output: process = start_raw_video_encoder( @@ -414,6 +389,7 @@ def regenerate_video_candidate( fps=effective_fps, strip_metadata=True, crf=18, + profile=probe_video_encode_profile(source), ) ) frame_pipe = process.stdin @@ -425,8 +401,9 @@ def regenerate_video_candidate( pixel_count = 0 temporal_baseline = 0.0 temporal_candidate = 0.0 - previous_reference: np.ndarray | None = None - previous_candidate: np.ndarray | None = None + previous_gray: np.ndarray | None = None + previous_reference_f32: np.ndarray | None = None + previous_candidate_f32: np.ndarray | None = None shared_noise: Any | None = None try: sampled_frames = _iter_sampled_frames( @@ -459,30 +436,30 @@ def regenerate_video_candidate( ) for reference, candidate in zip(frames, regenerated, strict=True): frame_pipe.write(candidate.tobytes()) - difference = reference.astype(np.float32) - candidate.astype(np.float32) + reference_f32 = reference.astype(np.float32) + candidate_f32 = candidate.astype(np.float32) + difference = reference_f32 - candidate_f32 squared_error += float(np.sum(difference * difference, dtype=np.float64)) pixel_count += reference.size - if previous_reference is not None and previous_candidate is not None: - current_gray = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY) - previous_gray = cv2.cvtColor(previous_reference, cv2.COLOR_BGR2GRAY) + current_gray = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY) + if ( + previous_gray is not None + and previous_reference_f32 is not None + and previous_candidate_f32 is not None + ): frame_maps = _backward_map(current_gray, previous_gray) - warped_reference = _backward_warp(previous_reference, frame_maps) - warped_candidate = _backward_warp(previous_candidate, frame_maps) - temporal_baseline += float( - np.mean(np.abs(reference.astype(np.float32) - warped_reference.astype(np.float32))) - ) - temporal_candidate += float( - np.mean(np.abs(candidate.astype(np.float32) - warped_candidate.astype(np.float32))) - ) - previous_reference = reference - previous_candidate = candidate + temporal_baseline += _motion_residual(reference_f32, previous_reference_f32, frame_maps) + temporal_candidate += _motion_residual(candidate_f32, previous_candidate_f32, frame_maps) + previous_gray = current_gray + previous_reference_f32 = reference_f32 + previous_candidate_f32 = candidate_f32 frame_count += 1 if frame_count < 2: raise ValueError("The selected clip produced fewer than two frames") finish_raw_video_encoder( process, temporary_output, - operation="SynthID candidate encode", + operation="SynthID removal encode", ) except Exception: if process.poll() is None: diff --git a/src/remove_ai_watermarks/video_synthid.py b/src/remove_ai_watermarks/video_synthid.py index d3cf5d0..612ec01 100644 --- a/src/remove_ai_watermarks/video_synthid.py +++ b/src/remove_ai_watermarks/video_synthid.py @@ -1,10 +1,11 @@ """Shared configuration for experimental video SynthID regeneration.""" DEFAULT_VIDEO_SYNTHID_VAE = "stabilityai/sd-vae-ft-mse" -DEFAULT_VIDEO_SYNTHID_NOISE_STD = 0.10 +DEFAULT_VIDEO_SYNTHID_NOISE_STD = 0.15 DEFAULT_VIDEO_SYNTHID_LONG_SIDE = 512 DEFAULT_VIDEO_SYNTHID_FPS = 12.0 VIDEO_SYNTHID_LATENT_MULTIPLE = 8 VIDEO_SYNTHID_VERIFICATION_PROMPT = ( - "Was this uploaded video created or edited by Google AI? Use the built-in content verification result." + "For the video attached to this message, was it created or edited by " + "Google AI? Use the built-in SynthID content verification result." ) diff --git a/src/remove_ai_watermarks/video_temporal.py b/src/remove_ai_watermarks/video_temporal.py new file mode 100644 index 0000000..c52e737 --- /dev/null +++ b/src/remove_ai_watermarks/video_temporal.py @@ -0,0 +1,189 @@ +"""Pure motion-compensated helpers shared by video pipelines.""" + +from __future__ import annotations + +# OpenCV exposes incomplete types for optical-flow and remap operations. +# Public signatures remain annotated while this third-party boundary is relaxed. +# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportMissingTypeStubs=false, reportCallIssue=false, reportArgumentType=false +from typing import TYPE_CHECKING, Any + +import cv2 +import numpy as np + +if TYPE_CHECKING: + from collections.abc import Sequence + + from numpy.typing import NDArray + + +def _backward_map( + current_gray: NDArray[Any], + previous_gray: NDArray[Any], +) -> tuple[NDArray[Any], NDArray[Any]]: + """Build a remap from a previous frame into current coordinates.""" + flow = cv2.calcOpticalFlowFarneback( + current_gray, + previous_gray, + None, + 0.5, + 3, + 15, + 3, + 5, + 1.2, + 0, + ) + height, width = current_gray.shape + flow[..., 0] += np.arange(width, dtype=np.float32)[None, :] + flow[..., 1] += np.arange(height, dtype=np.float32)[:, None] + return flow[..., 0], flow[..., 1] + + +def _backward_warp( + image: NDArray[Any], + maps: tuple[NDArray[Any], NDArray[Any]], + *, + interpolation: int = cv2.INTER_LINEAR, +) -> NDArray[Any]: + """Apply a precomputed backward optical-flow map.""" + return cv2.remap( + image, + maps[0], + maps[1], + interpolation=interpolation, + borderMode=cv2.BORDER_REFLECT, + ) + + +def _motion_residual( + current: NDArray[Any], + previous: NDArray[Any], + maps: tuple[NDArray[Any], NDArray[Any]], +) -> float: + """Return mean absolute residual after warping the previous frame.""" + current_f32 = np.asarray(current, dtype=np.float32) + previous_f32 = np.asarray(previous, dtype=np.float32) + warped_previous = _backward_warp(previous_f32, maps) + return float(np.mean(np.abs(current_f32 - warped_previous))) + + +def build_temporal_reference( + reference: Sequence[NDArray[Any]], +) -> tuple[tuple[tuple[NDArray[Any], NDArray[Any]], ...], float]: + """Precompute source motion maps and its mean residual.""" + if len(reference) < 2: + raise ValueError("Temporal metric needs at least two frames") + maps: list[tuple[NDArray[Any], NDArray[Any]]] = [] + reference_residuals: list[float] = [] + for index in range(1, len(reference)): + current_gray = cv2.cvtColor(reference[index], cv2.COLOR_BGR2GRAY) + previous_gray = cv2.cvtColor(reference[index - 1], cv2.COLOR_BGR2GRAY) + frame_maps = _backward_map(current_gray, previous_gray) + maps.append(frame_maps) + reference_residuals.append(_motion_residual(reference[index], reference[index - 1], frame_maps)) + return tuple(maps), float(np.mean(reference_residuals)) + + +def temporal_residual_ratio( + candidate: Sequence[NDArray[Any]], + maps: Sequence[tuple[NDArray[Any], NDArray[Any]]], + baseline: float, +) -> float: + """Measure candidate flicker against a precomputed source residual.""" + if len(candidate) != len(maps) + 1: + raise ValueError("Temporal metric needs one map per adjacent frame pair") + candidate_residuals: list[float] = [] + for index, frame_maps in enumerate(maps, start=1): + candidate_residuals.append(_motion_residual(candidate[index], candidate[index - 1], frame_maps)) + measured = float(np.mean(candidate_residuals)) + return measured / max(baseline, 1e-6) + + +def stabilize_filled_frame( + previous_source: NDArray[Any], + previous_cleaned: NDArray[Any], + previous_mask: NDArray[Any], + current_source: NDArray[Any], + current_cleaned: NDArray[Any], + current_mask: NDArray[Any], + *, + blend: float = 0.5, + max_context_residual: float = 12.0, + copy: bool = True, +) -> NDArray[Any]: + """Blend a motion-aligned prior fill when nearby source pixels agree. + + The prior contributes only where its warped removal mask covers the current + mask. A context ring outside both masks gates the blend, so scene cuts or + non-rigid local changes keep the independent current-frame fill. + """ + if not 0.0 <= blend <= 1.0: + raise ValueError("Temporal blend must be between 0 and 1") + if max_context_residual <= 0.0: + raise ValueError("Context residual threshold must be positive") + if ( + previous_source.shape != current_source.shape + or previous_cleaned.shape != current_cleaned.shape + or previous_source.shape != previous_cleaned.shape + or previous_mask.shape != current_mask.shape + or previous_mask.shape != current_source.shape[:2] + ): + raise ValueError("Temporal fill inputs must share frame and mask geometry") + + union = (previous_mask > 0) | (current_mask > 0) + ys, xs = np.where(union) + if len(xs) == 0: + return current_cleaned + height, width = current_mask.shape + mask_width = int(xs.max() - xs.min() + 1) + mask_height = int(ys.max() - ys.min() + 1) + padding = max(24, round(max(mask_width, mask_height) * 0.75)) + x0 = max(0, int(xs.min()) - padding) + y0 = max(0, int(ys.min()) - padding) + x1 = min(width, int(xs.max()) + padding + 1) + y1 = min(height, int(ys.max()) + padding + 1) + + previous_source_crop = previous_source[y0:y1, x0:x1] + current_source_crop = current_source[y0:y1, x0:x1] + current_cleaned_crop = current_cleaned[y0:y1, x0:x1] + previous_cleaned_crop = previous_cleaned[y0:y1, x0:x1] + previous_mask_crop = previous_mask[y0:y1, x0:x1] + current_mask_crop = current_mask[y0:y1, x0:x1] + maps = _backward_map( + cv2.cvtColor(current_source_crop, cv2.COLOR_BGR2GRAY), + cv2.cvtColor(previous_source_crop, cv2.COLOR_BGR2GRAY), + ) + warped_previous_source = _backward_warp(previous_source_crop, maps) + warped_previous_cleaned = _backward_warp(previous_cleaned_crop, maps) + warped_previous_mask = _backward_warp( + previous_mask_crop, + maps, + interpolation=cv2.INTER_NEAREST, + ) + + current_hole = current_mask_crop > 0 + if not np.any(current_hole): + return current_cleaned + covered = current_hole & (warped_previous_mask > 0) + if float(np.mean(covered[current_hole])) < 0.85: + return current_cleaned + + occupied = current_hole | (warped_previous_mask > 0) + dilation = max(7, round(max(mask_width, mask_height) * 0.25)) + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation | 1, dilation | 1)) + context = cv2.dilate(occupied.astype(np.uint8), kernel).astype(bool) & ~occupied + if np.count_nonzero(context) < 64: + return current_cleaned + residual = np.abs(current_source_crop.astype(np.float32) - warped_previous_source.astype(np.float32)) + context_residual = float(np.mean(residual[context])) + if context_residual > max_context_residual: + return current_cleaned + + effective_blend = blend * (1.0 - context_residual / max_context_residual) + blended = (1.0 - effective_blend) * current_cleaned_crop[covered].astype( + np.float32 + ) + effective_blend * warped_previous_cleaned[covered].astype(np.float32) + result = current_cleaned.copy() if copy else current_cleaned + result_crop = result[y0:y1, x0:x1] + result_crop[covered] = np.clip(blended, 0, 255).astype(np.uint8) + return result diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py index e7bd2ed..1eaa288 100644 --- a/src/remove_ai_watermarks/video_visible.py +++ b/src/remove_ai_watermarks/video_visible.py @@ -11,8 +11,10 @@ requires the candidate to recur at the same location across adjacent frames. This keeps isolated lookalikes in clean videos from becoming removal masks. Video pixels are decoded with OpenCV and encoded with the system ``ffmpeg``. -Complete audio is stream-copied from the source. The video stream must be -transcoded because visible-mark removal changes pixels. +Variable frame timestamps cross the pipe in a PyAV-muxed NUT stream; uniform +inputs retain the cheaper raw-BGR pipe. Complete audio is stream-copied from +the source. The video stream must be transcoded because visible-mark removal +changes pixels. """ # cv2/numpy boundary: these packages do not expose usable types for many array @@ -23,7 +25,9 @@ transcoded because visible-mark removal changes pixels. from __future__ import annotations import logging -from dataclasses import dataclass +from contextlib import suppress +from dataclasses import dataclass, replace +from fractions import Fraction from functools import lru_cache from itertools import pairwise from typing import TYPE_CHECKING, Any, Literal @@ -37,9 +41,12 @@ from remove_ai_watermarks.video_encoding import ( abort_raw_video_encoder, atomic_video_output, finish_raw_video_encoder, + probe_video_encode_profile, + probe_video_timestamps, raw_video_command, start_raw_video_encoder, ) +from remove_ai_watermarks.video_temporal import stabilize_filled_frame if TYPE_CHECKING: from pathlib import Path @@ -86,6 +93,7 @@ _VEO_DIAMOND_PROFILES = ( _DOLA_RELATIVE_HEIGHTS = tuple(value / 1000 for value in range(22, 41)) _HAILUO_RELATIVE_HEIGHTS = tuple(value / 1000 for value in range(28, 56, 3)) _KLING_RELATIVE_HEIGHTS = tuple(value / 1000 for value in range(24, 49, 3)) +_HDR_TRANSFERS = frozenset({"smpte2084", "arib-std-b67"}) @dataclass(frozen=True) @@ -105,6 +113,7 @@ class VideoScan: height: int fps: float detections: tuple[FrameLocalization, ...] + timestamps: tuple[float, ...] = () @dataclass(frozen=True) @@ -1069,6 +1078,8 @@ def _stabilize_localizations( def _scan_video_detectors( source: Path, detectors: dict[str, Any], + *, + collect_timestamps: bool = True, ) -> dict[str, VideoScan]: """Decode once and collect one untrusted candidate per detector and frame.""" capture = cv2.VideoCapture(str(source)) @@ -1082,11 +1093,14 @@ def _scan_video_detectors( raise RuntimeError(f"Video has invalid stream geometry or frame rate: {source}") detections: dict[str, list[FrameLocalization]] = {mark: [] for mark in detectors} + timestamps: list[float] = [] frame_index = 0 while True: ok, frame = capture.read() if not ok: break + if collect_timestamps: + timestamps.append(float(capture.get(cv2.CAP_PROP_POS_MSEC)) / 1000) if frame.shape[:2] != (height, width): capture.release() raise RuntimeError(f"Video changes frame dimensions at frame {frame_index}: {source}") @@ -1103,7 +1117,24 @@ def _scan_video_detectors( capture.release() if frame_index == 0: raise RuntimeError(f"Video contains no decodable frames: {source}") - return {mark: VideoScan(width, height, fps, tuple(mark_detections)) for mark, mark_detections in detections.items()} + shared_timestamps: tuple[float, ...] = () + if collect_timestamps: + probed_timestamps = probe_video_timestamps(source) + if len(probed_timestamps) == frame_index: + shared_timestamps = probed_timestamps + else: + if probed_timestamps: + log.warning( + "ffprobe/OpenCV frame-count mismatch for %s: timestamps=%s decoded=%s; using decoder timestamps", + source, + len(probed_timestamps), + frame_index, + ) + shared_timestamps = tuple(timestamps) + return { + mark: VideoScan(width, height, fps, tuple(mark_detections), shared_timestamps) + for mark, mark_detections in detections.items() + } def _scan_video( @@ -1117,8 +1148,14 @@ def _scan_video( def scan_video_marks( source: Path, marks: tuple[str, ...] = VIDEO_VISIBLE_MARKS, + *, + collect_timestamps: bool = True, ) -> dict[str, VideoScan]: - """Decode once and collect candidates for every requested provider mark.""" + """Decode once and collect candidates for every requested provider mark. + + Timestamp probing is optional because identification never encodes frames. + Removal keeps it enabled so variable and non-zero-start timing is preserved. + """ detectors = dict( zip( VIDEO_VISIBLE_MARKS, @@ -1139,6 +1176,7 @@ def scan_video_marks( return _scan_video_detectors( source, {mark: detectors[mark] for mark in marks}, + collect_timestamps=collect_timestamps, ) @@ -1215,6 +1253,66 @@ def _mask_for_region( return mask +def _timestamp_time_base(profile_time_base: str | None) -> Fraction: + """Use the source time base when known, with a microsecond fallback.""" + return Fraction(profile_time_base) if profile_time_base is not None else Fraction(1, 1_000_000) + + +def _timestamps_are_variable(scan: VideoScan, *, time_base: Fraction) -> bool: + """Whether OpenCV exposed valid timestamps with non-uniform frame intervals.""" + if len(scan.timestamps) != len(scan.detections) or len(scan.timestamps) < 3: + return False + ticks = tuple(round(timestamp / float(time_base)) for timestamp in scan.timestamps) + intervals = tuple(current - previous for previous, current in pairwise(ticks)) + return all(interval > 0 for interval in intervals) and len(set(intervals)) > 1 + + +class _TimestampedNutWriter: + """Mux BGR frames with explicit PTS into ffmpeg's standard-input pipe.""" + + def __init__( + self, + pipe: Any, + *, + width: int, + height: int, + time_base: Fraction, + start_pts: int = 0, + ) -> None: + import av + + self._av = av + self._time_base = time_base + self._start_pts = start_pts + self._origin: float | None = None + self._container = av.open(pipe, mode="w", format="nut") + self._stream = self._container.add_stream("rawvideo", rate=None) + self._stream.width = width + self._stream.height = height + self._stream.pix_fmt = "bgr24" + self._stream.time_base = time_base + self._stream.codec_context.time_base = time_base + + def write(self, frame_bgr: NDArray[Any], timestamp: float) -> None: + """Mux one contiguous BGR frame at its source timeline timestamp.""" + if self._origin is None: + self._origin = timestamp + frame = self._av.VideoFrame.from_ndarray( + np.ascontiguousarray(frame_bgr), + format="bgr24", + ) + frame.pts = self._start_pts + round((timestamp - self._origin) / float(self._time_base)) + frame.time_base = self._time_base + for packet in self._stream.encode(frame): + self._container.mux(packet) + + def close(self) -> None: + """Flush the rawvideo encoder and NUT container without closing ffmpeg stdin.""" + for packet in self._stream.encode(): + self._container.mux(packet) + self._container.close() + + def encode_clean_video( source: Path, output: Path, @@ -1225,6 +1323,7 @@ def encode_clean_video( strip_metadata: bool, padding_fraction: float = 0.28, mask_style: Literal["box", "veo"] = "box", + temporal_consistency: bool = True, ) -> int: """Decode again, fill accepted regions, and atomically encode with complete audio.""" from remove_ai_watermarks.watermark_registry import fill, resolve_backend @@ -1233,6 +1332,22 @@ def encode_clean_video( raise ValueError("Temporal localization count does not match the scanned frame count") with atomic_video_output(output) as temporary_output: + profile = probe_video_encode_profile(source) + if (profile.component_depth or 0) > 8 or profile.color_transfer in _HDR_TRANSFERS: + source_format = profile.source_pixel_format or "unknown pixel format" + raise RuntimeError( + "Visible video removal currently supports SDR 8-bit input only; " + f"refusing to silently reduce {source_format} / " + f"{profile.color_transfer or 'unknown transfer'} to 8-bit SDR" + ) + time_base = _timestamp_time_base(profile.time_base) + preserve_start_offset = profile.start_pts not in (None, 0) + timestamped_input = preserve_start_offset or _timestamps_are_variable(scan, time_base=time_base) + if timestamped_input and profile.time_base is None: + profile = replace( + profile, + time_base=f"{time_base.numerator}/{time_base.denominator}", + ) process = start_raw_video_encoder( raw_video_command( source, @@ -1242,6 +1357,9 @@ def encode_clean_video( fps=scan.fps, strip_metadata=strip_metadata, crf=14, + profile=profile, + timestamped_input=timestamped_input, + copy_input_timestamps=preserve_start_offset, ) ) frame_pipe = process.stdin @@ -1255,31 +1373,69 @@ def encode_clean_video( raise RuntimeError(f"OpenCV could not reopen video for removal: {source}") removed_frames = 0 - resolved_backend: Literal["cv2", "migan", "lama"] = resolve_backend(backend) + timestamped_writer: _TimestampedNutWriter | None = None + previous_source: NDArray[Any] | None = None + previous_cleaned: NDArray[Any] | None = None + previous_mask: NDArray[Any] | None = None try: + resolved_backend: Literal["cv2", "migan", "lama"] = resolve_backend(backend) + if timestamped_input: + timestamped_writer = _TimestampedNutWriter( + frame_pipe, + width=scan.width, + height=scan.height, + time_base=time_base, + start_pts=profile.start_pts or 0, + ) for frame_index, region in enumerate(regions): ok, frame = capture.read() if not ok: raise RuntimeError(f"Video ended while reading frame {frame_index}: {source}") + source_frame = frame + mask: NDArray[Any] | None = None if region is not None: - frame = fill( + mask = _mask_for_region( frame, - _mask_for_region( - frame, - region, - padding_fraction=padding_fraction, - mask_style=mask_style, - ), - backend=resolved_backend, + region, + padding_fraction=padding_fraction, + mask_style=mask_style, ) + frame = fill(frame, mask, backend=resolved_backend) + if ( + temporal_consistency + and previous_source is not None + and previous_cleaned is not None + and previous_mask is not None + ): + frame = stabilize_filled_frame( + previous_source, + previous_cleaned, + previous_mask, + source_frame, + frame, + mask, + copy=False, + ) removed_frames += 1 - frame_pipe.write(frame.tobytes()) + if timestamped_writer is None: + frame_pipe.write(frame.tobytes()) + else: + timestamped_writer.write(frame, scan.timestamps[frame_index]) + previous_source = source_frame + previous_cleaned = frame + previous_mask = mask + if timestamped_writer is not None: + timestamped_writer.close() + timestamped_writer = None finish_raw_video_encoder( process, temporary_output, operation="visible-watermark encode", ) except Exception: + if timestamped_writer is not None: + with suppress(Exception): + timestamped_writer.close() if process.poll() is None: abort_raw_video_encoder(process) raise diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 6cd268f..2f6fccc 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -1436,7 +1436,7 @@ def _box(box_type: bytes, payload: bytes) -> bytes: class TestVideoC2pa: - """C2PA in MP4 (ISOBMFF) -- detect + strip, reusing the image box walker.""" + """C2PA in MP4 (ISOBMFF) -- detect + offset-preserving stream blank.""" def test_detects_c2pa_in_mp4(self, tmp_path: Path): from remove_ai_watermarks.metadata import C2PA_UUID @@ -1454,7 +1454,7 @@ class TestVideoC2pa: src.write_bytes(_MP4_FTYP + uuid_box + _MP4_MDAT) out = tmp_path / "out.mp4" remove_ai_metadata(src, out) - assert out.read_bytes() == _MP4_FTYP + _MP4_MDAT + assert out.read_bytes() == _MP4_FTYP + _box(b"free", b"\x00" * 24) + _MP4_MDAT assert has_ai_metadata(out) is False @@ -1628,7 +1628,7 @@ class TestIsobmffMetadataRemoval: src.write_bytes(_MP4_FTYP + uuid_box + _MP4_MDAT) out = tmp_path / "clean.m4a" remove_ai_metadata(src, out) - assert out.read_bytes() == _MP4_FTYP + _MP4_MDAT + assert out.read_bytes() == _MP4_FTYP + _box(b"free", b"\x00" * 24) + _MP4_MDAT def test_content_sniff_routes_unknown_suffix(self, tmp_path: Path): # An ISOBMFF file with a non-standard extension is still box-stripped. diff --git a/tests/test_noai.py b/tests/test_noai.py index 25f47cc..592f5aa 100644 --- a/tests/test_noai.py +++ b/tests/test_noai.py @@ -452,6 +452,44 @@ class TestISOBMFF: assert blanked == 0 assert out == FTYP + b"\x00\x00\x00\x0cmdat" + b"pixels!!" + def test_streaming_malformed_walk_copies_input_unchanged(self, tmp_path: Path): + from remove_ai_watermarks.noai.isobmff import strip_isobmff_media_file + + source = tmp_path / "malformed.mp4" + output = tmp_path / "clean.mp4" + malformed = FTYP + struct.pack(">I", 999) + b"uuid" + b"short" + source.write_bytes(malformed) + + stripped, tc260_blanked = strip_isobmff_media_file(source, output) + + assert (stripped, tc260_blanked) == (0, 0) + assert output.read_bytes() == malformed + + def test_streaming_failure_does_not_publish_partial_output( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import metadata + from remove_ai_watermarks.noai import isobmff + + source = tmp_path / "source.mp4" + output = tmp_path / "clean.mp4" + uuid_box = struct.pack(">I", 24) + b"uuid" + metadata.C2PA_UUID + source.write_bytes(FTYP + uuid_box) + output.write_bytes(b"previous output") + + def fail_patch(*_args: object, **_kwargs: object) -> None: + raise OSError("synthetic patch failure") + + monkeypatch.setattr(isobmff, "_overwrite_range", fail_patch) + + with pytest.raises(OSError, match="synthetic patch failure"): + isobmff.strip_isobmff_media_file(source, output) + + assert output.read_bytes() == b"previous output" + assert not list(tmp_path.glob(".clean-*")) + class TestIterTopLevelBoxes: """The box walker's three size encodings and its underflow/overflow guards.""" diff --git a/tests/test_sync_conda_recipe.py b/tests/test_sync_conda_recipe.py index c9664c4..47b856d 100644 --- a/tests/test_sync_conda_recipe.py +++ b/tests/test_sync_conda_recipe.py @@ -2,6 +2,8 @@ from __future__ import annotations +from pathlib import Path + import pytest from scripts.sync_conda_recipe import update_recipe @@ -39,3 +41,9 @@ def test_update_recipe_rejects_ambiguous_recipe() -> None: with pytest.raises(ValueError, match="exactly one"): update_recipe(duplicate, version="2.0.0", sha256=_NEW_SHA) + + +def test_repository_recipe_includes_timestamp_bridge() -> None: + recipe = Path("packaging/conda/recipe.yaml").read_text() + + assert " - av >=16\n" in recipe diff --git a/tests/test_video.py b/tests/test_video.py index 31ec243..4ede919 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -2,7 +2,11 @@ from __future__ import annotations +import hashlib import io +import json +import shutil +import subprocess from typing import TYPE_CHECKING import cv2 @@ -148,6 +152,334 @@ def _stamp_gray_mark( ).astype(np.uint8) +def _independent_sora_mark() -> tuple[np.ndarray, tuple[int, int, int, int]]: + """Render a Sora-like mark without using the detector's template.""" + mark = Image.new("L", (180, 64), 0) + draw = ImageDraw.Draw(mark) + draw.ellipse((1, 14, 32, 54), fill=255) + draw.ellipse((25, 8, 62, 58), fill=255) + draw.ellipse((15, 20, 28, 44), fill=0) + draw.ellipse((37, 18, 50, 43), fill=0) + try: + font = ImageFont.load_default(size=49) + except TypeError: + font = ImageFont.load_default() + draw.text((68, 1), "Sora", font=font, fill=255, stroke_width=1) + mark_array = cv2.resize(np.asarray(mark), (124, 44), interpolation=cv2.INTER_AREA) + x, y = 620, 398 + return mark_array, (x, y, 124, 44) + + +def _independent_sora_frame() -> tuple[np.ndarray, tuple[int, int, int, int]]: + """Stamp the independent Sora-like mark onto a flat test frame.""" + frame = np.full((480, 840, 3), 36, dtype=np.uint8) + mark_array, region = _independent_sora_mark() + x, y, _width, _height = region + _stamp_gray_mark(frame, mark_array, x=x, y=y, opacity=0.78) + return frame, region + + +def _moving_video_background(frame_index: int) -> np.ndarray: + """Return a smooth moving background with known clean pixels.""" + x = np.arange(840, dtype=np.float32)[None, :] + y = np.arange(480, dtype=np.float32)[:, None] + luma = 42 + 10 * np.sin((x + frame_index * 5) / 38) + 5 * np.cos((y - frame_index * 2) / 54) + frame = np.stack( + ( + np.clip(luma - 5, 0, 255), + np.clip(luma + 1, 0, 255), + np.clip(luma + 7, 0, 255), + ), + axis=2, + ).astype(np.uint8) + moving_x = 30 + frame_index * 8 + cv2.rectangle(frame, (moving_x, 70), (moving_x + 72, 132), (80, 140, 210), -1) + cv2.line(frame, (0, 220 + frame_index), (839, 250 + frame_index), (110, 70, 45), 3) + return frame + + +def _ffmpeg_test_tools() -> tuple[str, str]: + """Return real ffmpeg tools or skip the integration test.""" + ffmpeg = shutil.which("ffmpeg") + ffprobe = shutil.which("ffprobe") + if ffmpeg is None or ffprobe is None: + pytest.skip("full-clip video integration test requires ffmpeg and ffprobe") + return ffmpeg, ffprobe + + +def _write_synthetic_sora_clip( + path: Path, + frames: list[np.ndarray], + *, + fps: float, + ffmpeg: str, + start_offset: float = 0.0, +) -> None: + """Encode a synthetic marked MP4 with AAC audio and C2PA provenance.""" + height, width = frames[0].shape[:2] + duration = len(frames) / fps + command = [ + ffmpeg, + "-y", + "-loglevel", + "error", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s:v", + f"{width}x{height}", + "-r", + f"{fps:.12g}", + "-i", + "pipe:0", + "-f", + "lavfi", + "-i", + f"sine=frequency=880:sample_rate=48000:duration={duration:.12g}", + "-map", + "0:v:0", + "-map", + "1:a:0", + "-c:v", + "libx264", + "-preset", + "ultrafast", + "-crf", + "8", + "-x264-params", + "colorprim=bt709:transfer=bt709:colormatrix=bt709:range=limited", + "-pix_fmt", + "yuv420p", + "-color_range", + "tv", + "-colorspace", + "bt709", + "-color_trc", + "bt709", + "-color_primaries", + "bt709", + "-c:a", + "aac", + "-shortest", + "-video_track_timescale", + "90000", + "-movflags", + "+faststart", + ] + if start_offset: + command.extend(["-output_ts_offset", f"{start_offset:.12g}"]) + command.append(str(path)) + subprocess.run( # noqa: S603 + command, + input=b"".join(frame.tobytes() for frame in frames), + capture_output=True, + check=True, + ) + with path.open("ab") as stream: + stream.write(_box(b"uuid", C2PA_UUID + b"OpenAI Sora trainedAlgorithmicMedia")) + + +def _write_vfr_sora_clip( + path: Path, + frames: list[np.ndarray], + *, + durations: list[float], + ffmpeg: str, + start_offset: float = 0.0, +) -> None: + """Encode marked stills at deliberately irregular presentation times.""" + assert len(frames) == len(durations) + frame_paths: list[Path] = [] + for index, frame in enumerate(frames): + frame_path = path.with_name(f"{path.stem}-frame-{index:03d}.png") + assert cv2.imwrite(str(frame_path), frame) + frame_paths.append(frame_path) + manifest = path.with_suffix(".ffconcat") + lines = ["ffconcat version 1.0"] + for frame_path, duration in zip(frame_paths, durations, strict=True): + lines.extend((f"file '{frame_path.as_posix()}'", f"duration {duration:.12g}")) + lines.append(f"file '{frame_paths[-1].as_posix()}'") + manifest.write_text("\n".join(lines) + "\n", encoding="utf-8") + total_duration = sum(durations) + command = [ + ffmpeg, + "-y", + "-loglevel", + "error", + "-f", + "concat", + "-safe", + "0", + "-i", + str(manifest), + "-f", + "lavfi", + "-i", + f"sine=frequency=880:sample_rate=48000:duration={total_duration:.12g}", + "-map", + "0:v:0", + "-map", + "1:a:0", + "-frames:v", + str(len(frames)), + "-fps_mode", + "vfr", + "-c:v", + "libx264", + "-preset", + "ultrafast", + "-crf", + "8", + "-bf", + "0", + "-pix_fmt", + "yuv420p", + "-c:a", + "aac", + "-video_track_timescale", + "90000", + ] + if start_offset: + command.extend(["-output_ts_offset", f"{start_offset:.12g}"]) + command.append(str(path)) + subprocess.run(command, capture_output=True, check=True) # noqa: S603 + with path.open("ab") as stream: + stream.write(_box(b"uuid", C2PA_UUID + b"OpenAI Sora trainedAlgorithmicMedia")) + + +def _decode_video(path: Path) -> tuple[list[np.ndarray], float]: + """Decode every frame through the same OpenCV boundary as production.""" + capture = cv2.VideoCapture(str(path)) + assert capture.isOpened() + fps = float(capture.get(cv2.CAP_PROP_FPS)) + frames: list[np.ndarray] = [] + try: + while True: + ok, frame = capture.read() + if not ok: + break + frames.append(frame) + finally: + capture.release() + return frames, fps + + +def _video_frame_timestamps(path: Path, *, ffprobe: str) -> list[float]: + """Read display timestamps from the first video stream.""" + result = subprocess.run( # noqa: S603 + [ + ffprobe, + "-v", + "error", + "-select_streams", + "v:0", + "-show_frames", + "-show_entries", + "frame=best_effort_timestamp_time", + "-of", + "json", + str(path), + ], + capture_output=True, + check=True, + text=True, + ) + frames = json.loads(result.stdout)["frames"] + return [float(frame["best_effort_timestamp_time"]) for frame in frames] + + +def _audio_bitstream(path: Path, *, ffmpeg: str) -> bytes: + """Extract copied AAC packets in a container-independent form.""" + result = subprocess.run( # noqa: S603 + [ + ffmpeg, + "-loglevel", + "error", + "-i", + str(path), + "-map", + "0:a:0", + "-c:a", + "copy", + "-f", + "adts", + "pipe:1", + ], + capture_output=True, + check=True, + ) + return result.stdout + + +def _container_duration(path: Path, *, ffprobe: str) -> float: + """Read the container duration from the real ffprobe boundary.""" + result = subprocess.run( # noqa: S603 + [ + ffprobe, + "-v", + "error", + "-show_entries", + "format=duration", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(path), + ], + capture_output=True, + check=True, + text=True, + ) + return float(result.stdout.strip()) + + +def _video_stream_info(path: Path, *, ffprobe: str) -> dict[str, str]: + """Read source-sensitive video stream properties through ffprobe.""" + result = subprocess.run( # noqa: S603 + [ + ffprobe, + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=pix_fmt,color_range,color_space,color_transfer,color_primaries,time_base", + "-of", + "json", + str(path), + ], + capture_output=True, + check=True, + text=True, + ) + streams = json.loads(result.stdout)["streams"] + assert len(streams) == 1 + return streams[0] + + +def _stream_start_times(path: Path, *, ffprobe: str) -> dict[str, float]: + """Read the first video and audio stream start times.""" + result = subprocess.run( # noqa: S603 + [ + ffprobe, + "-v", + "error", + "-show_entries", + "stream=codec_type,start_time", + "-of", + "json", + str(path), + ], + capture_output=True, + check=True, + text=True, + ) + return { + stream["codec_type"]: float(stream["start_time"]) + for stream in json.loads(result.stdout)["streams"] + if stream["codec_type"] in {"video", "audio"} + } + + def _regeneration_metrics( *, frames: int = 24, @@ -173,7 +505,10 @@ class TestVideoMetadataApi: def test_top_level_api_is_lazy_exported(self): import remove_ai_watermarks as raiw + assert raiw.identify_video is not None assert raiw.inspect_video_metadata is not None + assert raiw.remove_video_all is not None + assert raiw.remove_video_batch is not None assert raiw.remove_video_invisible is not None assert raiw.remove_video_metadata is not None assert raiw.remove_video_visible is not None @@ -256,6 +591,46 @@ class TestVideoMetadataApi: assert b"AIGC" not in cleaned assert _TC260_AIGC not in cleaned + def test_streams_large_isobmff_without_full_file_read( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks.video import inspect_video_metadata, remove_video_metadata + + media_payload = b"x" * (8 * 1024 * 1024) + source = _video_with_tc260(tmp_path / "source.mp4", media_payload=media_payload) + with source.open("ab") as stream: + stream.write(_box(b"uuid", C2PA_UUID + b"OpenAI trainedAlgorithmicMedia")) + output = tmp_path / "clean.mp4" + source_size = source.stat().st_size + source_media_digest = hashlib.sha256(media_payload).digest() + original_read_bytes = type(source).read_bytes + source_report = inspect_video_metadata(source) + + assert source_report.has_ai_metadata is True + assert "synthid_watermark" in source_report.markers + assert "aigc_label" in source_report.markers + + def reject_source_read_bytes(path: Path) -> bytes: + if path.resolve() == source.resolve(): + raise AssertionError("video metadata removal must not read the complete source") + return original_read_bytes(path) + + monkeypatch.setattr(type(source), "read_bytes", reject_source_read_bytes) + + result = remove_video_metadata(source, output) + cleaned = output.read_bytes() + mdat_start = cleaned.index(b"mdat") + 4 + mdat_end = mdat_start + len(media_payload) + + assert result.remaining == {} + assert output.stat().st_size == source_size + assert hashlib.sha256(cleaned[mdat_start:mdat_end]).digest() == source_media_digest + assert b"standard title" in cleaned + assert C2PA_UUID not in cleaned + assert _TC260_AIGC not in cleaned + def test_ignores_generic_mp4_aigc_tag_without_tc260_fields(self, tmp_path: Path): from remove_ai_watermarks.video import inspect_video_metadata @@ -356,6 +731,81 @@ class TestVideoMetadataApi: remove_video_metadata(source, tmp_path / "clean.mov") +class TestVideoProvenanceApi: + def test_identifies_metadata_without_pixel_scan(self, tmp_path: Path): + from remove_ai_watermarks.video import identify_video + + source = _video_with_c2pa(tmp_path / "source.mp4") + + report = identify_video(source, check_visible=False) + + assert report.source == source + assert report.is_ai_generated is True + assert report.confidence == "high" + assert report.platform == "OpenAI (ChatGPT / gpt-image / DALL-E / Sora)" + assert report.visible_mark is None + assert report.total_frames is None + assert report.has_ai_metadata is True + assert report.metadata_markers + assert "Visible video-mark detection was skipped." in report.caveats + + def test_identifies_stable_visible_mark_with_shared_arbiter( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video import VIDEO_VISIBLE_MARKS, identify_video + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan + + source = _video_with_c2pa(tmp_path / "source.mp4") + box = (4, 4, 20, 8) + stable = VideoScan( + width=64, + height=64, + fps=24.0, + detections=tuple(FrameLocalization(index, 0.66, box) for index in range(5)), + ) + empty = VideoScan( + width=64, + height=64, + fps=24.0, + detections=tuple(FrameLocalization(index, 0.0, None) for index in range(5)), + ) + + def fake_scan( + _source: Path, + marks: tuple[str, ...], + *, + collect_timestamps: bool, + ) -> dict[str, VideoScan]: + assert collect_timestamps is False + return {candidate: stable if candidate == "sora" else empty for candidate in marks} + + monkeypatch.setattr(video_visible, "scan_video_marks", fake_scan) + + report = identify_video(source) + + assert report.platform == "OpenAI Sora" + assert report.visible_mark == "sora" + assert report.visible_detected_frames == 5 + assert report.total_frames == 5 + assert tuple(VIDEO_VISIBLE_MARKS) == ("sora", "veo", "seedance", "dola", "hailuo", "kling") + + def test_reports_unknown_instead_of_clean(self, tmp_path: Path): + from remove_ai_watermarks.video import identify_video + + source = tmp_path / "source.mp4" + source.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + + report = identify_video(source, check_visible=False) + + assert report.is_ai_generated is None + assert report.confidence == "unknown" + assert report.has_ai_metadata is False + assert any("not proof" in caveat for caveat in report.caveats) + + class TestVideoMetadataCli: def test_help(self): runner = CliRunner() @@ -394,8 +844,428 @@ class TestVideoMetadataCli: assert "Unsupported video format" in result.output +class TestVideoProvenanceCli: + def test_json_metadata_report(self, tmp_path: Path): + source = _video_with_c2pa(tmp_path / "source.mp4") + + result = CliRunner().invoke(main, ["video", "identify", str(source), "--no-visible", "--json"]) + + assert result.exit_code == 0, result.output + payload = json.loads(result.output) + assert payload["is_ai_generated"] is True + assert payload["visible_mark"] is None + assert payload["has_ai_metadata"] is True + + +class TestVideoAllApi: + def test_no_visible_mark_still_strips_metadata_and_writes_output( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video + from remove_ai_watermarks.video import VideoVisibleResult, remove_video_all + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + monkeypatch.setattr( + video, + "remove_video_visible", + lambda *_args, **_kwargs: VideoVisibleResult( + source=source, + output=None, + mark="auto", + total_frames=3, + detected_frames=0, + removed_frames=0, + remaining_metadata={"c2pa_manifest": "present"}, + ), + ) + + result = remove_video_all(source, output) + + assert result.output == output + assert result.visible_mark is None + assert result.detected_metadata + assert result.remaining_metadata == {} + assert result.invisible_removed is False + assert C2PA_UUID not in output.read_bytes() + assert _VIDEO_PAYLOAD in output.read_bytes() + + def test_visible_output_is_the_final_locally_verified_result( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video + from remove_ai_watermarks.video import VideoVisibleResult, remove_video_all + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + + def fake_visible(_source: Path, target: Path, **_kwargs: object) -> VideoVisibleResult: + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return VideoVisibleResult( + source=source, + output=target, + mark="sora", + total_frames=12, + detected_frames=12, + removed_frames=12, + remaining_metadata={}, + ) + + monkeypatch.setattr(video, "remove_video_visible", fake_visible) + monkeypatch.setattr( + video, + "remove_video_metadata", + lambda *_args, **_kwargs: pytest.fail("metadata must not reprocess an already stripped visible output"), + ) + + result = remove_video_all(source, output) + + assert result.visible_mark == "sora" + assert result.visible_removed_frames == 12 + assert result.remaining_metadata == {} + assert output.exists() + + def test_invisible_stage_is_explicit_and_uses_an_intermediate_visible_output( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video + from remove_ai_watermarks.video import ( + VideoInvisibleResult, + VideoVisibleResult, + remove_video_all, + ) + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + intermediate_sources: list[Path] = [] + + def fake_visible(_source: Path, target: Path, **_kwargs: object) -> VideoVisibleResult: + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return VideoVisibleResult( + source=source, + output=target, + mark="sora", + total_frames=12, + detected_frames=12, + removed_frames=12, + remaining_metadata={}, + ) + + def fake_invisible( + candidate_source: Path, + target: Path, + **_kwargs: object, + ) -> VideoInvisibleResult: + assert candidate_source != source + assert candidate_source.exists() + intermediate_sources.append(candidate_source) + target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + return VideoInvisibleResult( + source=candidate_source, + output=target, + noise_std=0.1, + metrics=_regeneration_metrics(frames=12), + remaining_metadata={}, + ) + + monkeypatch.setattr(video, "remove_video_visible", fake_visible) + monkeypatch.setattr(video, "remove_video_invisible", fake_invisible) + + result = remove_video_all(source, output, include_invisible=True) + + assert result.invisible_removed is True + assert output.exists() + assert len(intermediate_sources) == 1 + assert not intermediate_sources[0].exists() + + def test_rejects_unsupported_invisible_container_before_visible_scan( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video + from remove_ai_watermarks.video import remove_video_all + + source = _video_with_tc260_ebml(tmp_path / "source.webm") + monkeypatch.setattr( + video, + "remove_video_visible", + lambda *_args, **_kwargs: pytest.fail("visible scan must not start"), + ) + + with pytest.raises(ValueError, match="requires one of"): + remove_video_all(source, include_invisible=True) + + +class TestVideoAllCli: + def test_help_keeps_invisible_stage_opt_in(self): + result = CliRunner().invoke(main, ["video", "all", "--help"]) + + assert result.exit_code == 0, result.output + assert "--invisible" in result.output + assert "oracle-certified" in result.output + + def test_reports_locally_verified_result(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + from remove_ai_watermarks import video + from remove_ai_watermarks.video import VideoAllResult + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + monkeypatch.setattr( + video, + "remove_video_all", + lambda *_args, **_kwargs: VideoAllResult( + source=source, + output=output, + visible_mark="sora", + total_frames=12, + visible_detected_frames=12, + visible_removed_frames=12, + detected_metadata={"c2pa_manifest": "present"}, + remaining_metadata={}, + invisible_removed=False, + ), + ) + + result = CliRunner().invoke(main, ["video", "all", str(source), "-o", str(output)]) + + assert result.exit_code == 0, result.output + assert "removed sora from 12/12 frames" in result.output + assert "UNVERIFIED" not in result.output + + def test_reports_oracle_certified_invisible_removal( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video + from remove_ai_watermarks.video import VideoAllResult + + source = _video_with_c2pa(tmp_path / "source.mp4") + output = tmp_path / "clean.mp4" + monkeypatch.setattr( + video, + "remove_video_all", + lambda *_args, **_kwargs: VideoAllResult( + source=source, + output=output, + visible_mark=None, + total_frames=12, + visible_detected_frames=0, + visible_removed_frames=0, + detected_metadata={}, + remaining_metadata={}, + invisible_removed=True, + ), + ) + + result = CliRunner().invoke(main, ["video", "all", str(source), "--invisible"]) + + assert result.exit_code == 0, result.output + assert "UNVERIFIED" not in result.output + assert "oracle-certified VAE profile" in result.output + + +class TestVideoBatchApi: + def test_visible_mode_copies_noop_and_collects_per_file_error( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video + from remove_ai_watermarks.video import VideoVisibleResult, remove_video_batch + + directory = tmp_path / "videos" + directory.mkdir() + clean = directory / "clean.mp4" + clean.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + broken = directory / "broken.mp4" + broken.write_bytes(b"not a video") + (directory / "notes.txt").write_text("ignored", encoding="utf-8") + + def fake_visible(source: Path, *_args: object, **_kwargs: object) -> VideoVisibleResult: + if source == broken: + raise ValueError("invalid container") + return VideoVisibleResult( + source=source, + output=None, + mark="auto", + total_frames=3, + detected_frames=0, + removed_frames=0, + remaining_metadata={}, + ) + + monkeypatch.setattr(video, "remove_video_visible", fake_visible) + + result = remove_video_batch(directory, mode="visible") + + assert result.processed == 1 + assert result.failed == 1 + assert len(result.items) == 2 + assert (result.output_directory / clean.name).read_bytes() == clean.read_bytes() + assert result.items[0].source.name == "broken.mp4" + assert result.items[0].error == "invalid container" + assert result.items[1].changed is False + + def test_metadata_mode_processes_every_supported_video(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_batch + + directory = tmp_path / "videos" + directory.mkdir() + marked = _video_with_c2pa(directory / "marked.mp4") + clean = directory / "clean.mp4" + clean.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + + result = remove_video_batch(directory, mode="metadata") + + assert result.processed == 2 + assert result.failed == 0 + assert C2PA_UUID not in (result.output_directory / marked.name).read_bytes() + assert (result.output_directory / clean.name).read_bytes() == clean.read_bytes() + assert {item.changed for item in result.items} == {False, True} + + def test_rejects_invisible_stage_outside_all_mode(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_batch + + directory = tmp_path / "videos" + directory.mkdir() + + with pytest.raises(ValueError, match="only in all mode"): + remove_video_batch(directory, mode="visible", include_invisible=True) + + def test_invisible_mode_reuses_one_vae_runtime_for_the_batch( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video, video_invisible + from remove_ai_watermarks.video import VideoAllResult, remove_video_batch + + directory = tmp_path / "videos" + directory.mkdir() + for name in ("one.mp4", "two.mp4"): + (directory / name).write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + runtime = object() + load_calls: list[tuple[str, str]] = [] + used_runtimes: list[object] = [] + + def fake_load(*, model: str, device: str) -> object: + load_calls.append((model, device)) + return runtime + + def fake_all(source: Path, output: Path, **kwargs: object) -> VideoAllResult: + used_runtimes.append(kwargs["_invisible_runtime"]) + return VideoAllResult( + source=source, + output=output, + visible_mark=None, + total_frames=2, + visible_detected_frames=0, + visible_removed_frames=0, + detected_metadata={}, + remaining_metadata={}, + invisible_removed=True, + ) + + monkeypatch.setattr(video_invisible, "load_video_vae_runtime", fake_load) + monkeypatch.setattr(video, "remove_video_all", fake_all) + + result = remove_video_batch(directory, include_invisible=True) + + assert result.processed == 2 + assert load_calls == [("stabilityai/sd-vae-ft-mse", "auto")] + assert used_runtimes == [runtime, runtime] + + def test_invisible_runtime_failure_is_not_retried_for_every_file( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_invisible + from remove_ai_watermarks.video import remove_video_batch + + directory = tmp_path / "videos" + directory.mkdir() + for name in ("one.mp4", "two.mp4"): + (directory / name).write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) + load_calls = 0 + + def fail_load(**_kwargs: object) -> object: + nonlocal load_calls + load_calls += 1 + raise RuntimeError("model unavailable") + + monkeypatch.setattr(video_invisible, "load_video_vae_runtime", fail_load) + + result = remove_video_batch(directory, include_invisible=True) + + assert result.failed == 2 + assert load_calls == 1 + assert {item.error for item in result.items} == {"model unavailable"} + + def test_rejects_source_as_output_directory(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_batch + + with pytest.raises(ValueError, match="must differ"): + remove_video_batch(tmp_path, tmp_path) + + +class TestVideoBatchCli: + def test_help(self): + result = CliRunner().invoke(main, ["video", "batch", "--help"]) + + assert result.exit_code == 0, result.output + assert "all|visible|metadata" in result.output + assert "--invisible" in result.output + + def test_returns_nonzero_when_any_file_fails( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video + from remove_ai_watermarks.video import VideoBatchItem, VideoBatchResult + + directory = tmp_path / "videos" + directory.mkdir() + output = tmp_path / "clean" + failed_source = directory / "broken.mp4" + monkeypatch.setattr( + video, + "remove_video_batch", + lambda *_args, **_kwargs: VideoBatchResult( + directory=directory, + output_directory=output, + items=( + VideoBatchItem( + source=failed_source, + output=None, + mode="all", + changed=False, + visible_mark=None, + invisible_removed=False, + error="invalid container", + ), + ), + ), + ) + + result = CliRunner().invoke(main, ["video", "batch", str(directory)]) + + assert result.exit_code == 1 + assert "FAILED broken.mp4: invalid container" in result.output + assert "1 failed" in result.output + + class TestVideoInvisibleApi: - def test_generates_unverified_candidate_and_strips_metadata( + def test_removes_synthid_and_strips_metadata( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -404,7 +1274,7 @@ class TestVideoInvisibleApi: from remove_ai_watermarks.video import remove_video_invisible source = _video_with_c2pa(tmp_path / "source.mp4") - output = tmp_path / "candidate.mp4" + output = tmp_path / "clean.mp4" def fake_regenerate(_source: Path, target: Path, **_kwargs: object): target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) @@ -415,11 +1285,10 @@ class TestVideoInvisibleApi: result = remove_video_invisible(source, output) assert result.output == output - assert result.requires_external_verification is True assert result.total_frames == 24 assert result.remaining_metadata == {} - def test_default_output_is_named_as_candidate( + def test_default_output_is_named_as_clean( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -444,7 +1313,7 @@ class TestVideoInvisibleApi: result = remove_video_invisible(source) - assert result.output == tmp_path / "source_synthid_candidate.mp4" + assert result.output == tmp_path / "source_clean.mp4" def test_rejects_webm_regeneration(self, tmp_path: Path): from remove_ai_watermarks.video import remove_video_invisible @@ -456,15 +1325,15 @@ class TestVideoInvisibleApi: class TestVideoInvisibleCli: - def test_help_describes_external_verification(self): + def test_help_describes_oracle_certification(self): runner = CliRunner() result = runner.invoke(main, ["video", "invisible", "--help"]) assert result.exit_code == 0, result.output - assert "externally verifiable" in result.output + assert "oracle-certified" in result.output - def test_reports_unverified_candidate( + def test_reports_completed_removal( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -473,7 +1342,7 @@ class TestVideoInvisibleCli: runner = CliRunner() source = _video_with_c2pa(tmp_path / "source.mp4") - output = tmp_path / "candidate.mp4" + output = tmp_path / "clean.mp4" def fake_remove(_source: Path, target: Path, **_kwargs: object): target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) @@ -490,35 +1359,15 @@ class TestVideoInvisibleCli: result = runner.invoke(main, ["video", "invisible", str(source), "-o", str(output)]) assert result.exit_code == 0, result.output - assert "Candidate generated" in result.output - assert "UNVERIFIED" in result.output - assert "Gemini Flash" in result.output + assert "SynthID removal complete" in result.output + assert "UNVERIFIED" not in result.output class TestSoraFrameLocalization: - @staticmethod - def _sora_like_frame() -> tuple[np.ndarray, tuple[int, int, int, int]]: - frame = np.full((480, 840, 3), 36, dtype=np.uint8) - mark = Image.new("L", (180, 64), 0) - draw = ImageDraw.Draw(mark) - draw.ellipse((1, 14, 32, 54), fill=255) - draw.ellipse((25, 8, 62, 58), fill=255) - draw.ellipse((15, 20, 28, 44), fill=0) - draw.ellipse((37, 18, 50, 43), fill=0) - try: - font = ImageFont.load_default(size=49) - except TypeError: - font = ImageFont.load_default() - draw.text((68, 1), "Sora", font=font, fill=255, stroke_width=1) - mark_array = cv2.resize(np.asarray(mark), (124, 44), interpolation=cv2.INTER_AREA) - x, y = 620, 398 - _stamp_gray_mark(frame, mark_array, x=x, y=y, opacity=0.78) - return frame, (x, y, 124, 44) - def test_localizes_independently_rendered_sora_like_mark(self): from remove_ai_watermarks.video_visible import _region_iou, detect_sora_frame - frame, expected = self._sora_like_frame() + frame, expected = _independent_sora_frame() detection = detect_sora_frame(frame) @@ -1042,6 +1891,7 @@ class TestVideoVisibleScan: video_visible.cv2.CAP_PROP_FRAME_WIDTH: 12.0, video_visible.cv2.CAP_PROP_FRAME_HEIGHT: 8.0, video_visible.cv2.CAP_PROP_FPS: 24.0, + video_visible.cv2.CAP_PROP_POS_MSEC: 0.0, }[property_id] def read(self) -> tuple[bool, np.ndarray | None]: @@ -1066,6 +1916,7 @@ class TestVideoVisibleScan: return FrameLocalization(frame_index, 0.0, None) monkeypatch.setattr(video_visible.cv2, "VideoCapture", lambda _path: FakeCapture()) + monkeypatch.setattr(video_visible, "probe_video_timestamps", lambda _path: (0.25,)) for detector_name in ( "detect_sora_frame", "detect_veo_frame", @@ -1082,6 +1933,7 @@ class TestVideoVisibleScan: ) assert set(scans) == {"sora", "veo", "seedance", "dola", "hailuo", "kling"} + assert all(scan.timestamps == (0.25,) for scan in scans.values()) assert len(prepared_ids) == 6 assert len(set(prepared_ids)) == 1 @@ -1202,6 +2054,294 @@ class TestVideoVisibleEncoding: assert output.read_bytes() == b"previous" assert not targets[0].exists() + def test_rejects_high_bit_depth_before_silent_downconversion( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video_encoding import VideoEncodeProfile + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan, encode_clean_video + + source = tmp_path / "source.mp4" + source.write_bytes(b"source") + output = tmp_path / "clean.mp4" + output.write_bytes(b"previous") + scan = VideoScan( + 8, + 8, + 24.0, + (FrameLocalization(0, 1.0, (1, 1, 2, 2)),), + ) + monkeypatch.setattr( + video_visible, + "probe_video_encode_profile", + lambda _path: VideoEncodeProfile( + pixel_format="yuv420p", + color_transfer="smpte2084", + source_pixel_format="yuv420p10le", + component_depth=10, + ), + ) + monkeypatch.setattr( + video_visible, + "start_raw_video_encoder", + lambda _command: pytest.fail("encoder must not start for a high-bit-depth source"), + ) + + with pytest.raises(RuntimeError, match="refusing to silently reduce yuv420p10le"): + encode_clean_video( + source, + output, + scan, + [(1, 1, 2, 2)], + backend="cv2", + strip_metadata=True, + ) + + assert output.read_bytes() == b"previous" + assert list(tmp_path.glob(".clean-*")) == [] + + +class TestVideoVisibleFullClip: + def test_removes_complete_clip_and_preserves_sequence_and_audio(self, tmp_path: Path): + from remove_ai_watermarks.metadata import get_ai_metadata + from remove_ai_watermarks.video import remove_video_all, remove_video_metadata, remove_video_visible + from remove_ai_watermarks.video_visible import scan_video_marks, stabilize_sora_localizations + + ffmpeg, ffprobe = _ffmpeg_test_tools() + frame_count = 18 + fps = 12.0 + clean_frames: list[np.ndarray] = [] + source_frames: list[np.ndarray] = [] + mark, mark_region = _independent_sora_mark() + mark_x, mark_y, _mark_width, _mark_height = mark_region + for frame_index in range(frame_count): + clean_frame = _moving_video_background(frame_index) + marked_frame = clean_frame.copy() + _stamp_gray_mark(marked_frame, mark, x=mark_x, y=mark_y, opacity=0.78) + clean_frames.append(clean_frame) + source_frames.append(marked_frame) + + source = tmp_path / "marked.mp4" + clean_control = tmp_path / "control.mp4" + metadata_clean = tmp_path / "metadata-clean.mp4" + output = tmp_path / "clean.mp4" + all_output = tmp_path / "clean-all.mp4" + independent_output = tmp_path / "clean-independent.mp4" + _write_synthetic_sora_clip(source, source_frames, fps=fps, ffmpeg=ffmpeg) + _write_synthetic_sora_clip(clean_control, clean_frames, fps=fps, ffmpeg=ffmpeg) + assert get_ai_metadata(source) + + metadata_result = remove_video_metadata(source, metadata_clean) + result = remove_video_visible(source, output, mark="sora", backend="cv2") + all_result = remove_video_all(source, all_output, mark="sora", backend="cv2") + independent_result = remove_video_visible( + source, + independent_output, + mark="sora", + backend="cv2", + temporal_consistency=False, + ) + + assert metadata_result.remaining == {} + assert metadata_clean.stat().st_size == source.stat().st_size + assert get_ai_metadata(metadata_clean) == {} + assert result.output == output + assert result.total_frames == frame_count + assert result.detected_frames == frame_count + assert result.removed_frames == frame_count + assert result.remaining_metadata == {} + assert get_ai_metadata(output) == {} + assert all_result.output == all_output + assert all_result.visible_mark == "sora" + assert all_result.visible_removed_frames == frame_count + assert all_result.detected_metadata + assert all_result.remaining_metadata == {} + assert all_result.invisible_removed is False + assert get_ai_metadata(all_output) == {} + assert independent_result.removed_frames == frame_count + + decoded_source, source_fps = _decode_video(source) + decoded_control, control_fps = _decode_video(clean_control) + decoded_metadata_clean, metadata_clean_fps = _decode_video(metadata_clean) + decoded_output, output_fps = _decode_video(output) + decoded_all, all_fps = _decode_video(all_output) + decoded_independent, independent_fps = _decode_video(independent_output) + assert ( + len(decoded_source) + == len(decoded_control) + == len(decoded_metadata_clean) + == len(decoded_output) + == len(decoded_all) + == len(decoded_independent) + == frame_count + ) + assert source_fps == pytest.approx(fps, abs=0.01) + assert control_fps == pytest.approx(source_fps, abs=0.01) + assert metadata_clean_fps == pytest.approx(source_fps, abs=0.01) + assert output_fps == pytest.approx(source_fps, abs=0.01) + assert all_fps == pytest.approx(source_fps, abs=0.01) + assert independent_fps == pytest.approx(source_fps, abs=0.01) + assert all( + np.array_equal(original, metadata_cleaned) + for original, metadata_cleaned in zip(decoded_source, decoded_metadata_clean, strict=True) + ) + assert _container_duration(output, ffprobe=ffprobe) == pytest.approx( + _container_duration(source, ffprobe=ffprobe), + abs=1 / fps, + ) + source_stream = _video_stream_info(source, ffprobe=ffprobe) + output_stream = _video_stream_info(output, ffprobe=ffprobe) + expected_stream_properties = { + "pix_fmt": "yuv420p", + "color_range": "tv", + "color_space": "bt709", + "color_transfer": "bt709", + "color_primaries": "bt709", + "time_base": "1/90000", + } + assert source_stream == expected_stream_properties + assert _video_stream_info(metadata_clean, ffprobe=ffprobe) == source_stream + assert output_stream == expected_stream_properties + assert _video_stream_info(all_output, ffprobe=ffprobe) == expected_stream_properties + source_audio = _audio_bitstream(source, ffmpeg=ffmpeg) + assert source_audio + assert _audio_bitstream(metadata_clean, ffmpeg=ffmpeg) == source_audio + assert _audio_bitstream(output, ffmpeg=ffmpeg) == source_audio + assert _audio_bitstream(all_output, ffmpeg=ffmpeg) == source_audio + + output_scan = scan_video_marks(output, ("sora",))["sora"] + all_output_scan = scan_video_marks(all_output, ("sora",))["sora"] + assert all( + region is None + for region in stabilize_sora_localizations( + output_scan.detections, + provenance=False, + ) + ) + assert all( + region is None + for region in stabilize_sora_localizations( + all_output_scan.detections, + provenance=False, + ) + ) + + x, y, width, height = mark_region + untouched = np.ones(decoded_source[0].shape[:2], dtype=bool) + padding = 24 + untouched[ + max(0, y - padding) : min(untouched.shape[0], y + height + padding), + max(0, x - padding) : min(untouched.shape[1], x + width + padding), + ] = False + untouched_psnr: list[float] = [] + for original, cleaned in zip(decoded_source, decoded_output, strict=True): + squared_error = (original.astype(np.float32) - cleaned.astype(np.float32)) ** 2 + mean_squared_error = float(np.mean(squared_error[untouched])) + untouched_psnr.append(float(10 * np.log10((255**2) / mean_squared_error))) + assert min(untouched_psnr) >= 35.0 + + filled_region = np.logical_not(untouched) + temporal_errors: list[float] = [] + independent_temporal_errors: list[float] = [] + for frame_index in range(1, frame_count): + expected_delta = decoded_control[frame_index].astype(np.float32) - decoded_control[frame_index - 1].astype( + np.float32 + ) + cleaned_delta = decoded_output[frame_index].astype(np.float32) - decoded_output[frame_index - 1].astype( + np.float32 + ) + independent_delta = decoded_independent[frame_index].astype(np.float32) - decoded_independent[ + frame_index - 1 + ].astype(np.float32) + temporal_errors.append(float(np.mean(np.abs(cleaned_delta[filled_region] - expected_delta[filled_region])))) + independent_temporal_errors.append( + float(np.mean(np.abs(independent_delta[filled_region] - expected_delta[filled_region]))) + ) + assert float(np.median(temporal_errors)) <= 1.5 + assert float(np.percentile(temporal_errors, 95)) <= 2.0 + assert float(np.median(temporal_errors)) < float(np.median(independent_temporal_errors)) + assert float(np.percentile(temporal_errors, 95)) <= float(np.percentile(independent_temporal_errors, 95)) + + def test_preserves_variable_frame_timestamps(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_visible + + ffmpeg, ffprobe = _ffmpeg_test_tools() + durations = [value for _ in range(6) for value in (1 / 30, 1 / 12, 1 / 20)] + mark, mark_region = _independent_sora_mark() + mark_x, mark_y, _mark_width, _mark_height = mark_region + frames: list[np.ndarray] = [] + for frame_index in range(len(durations)): + frame = _moving_video_background(frame_index) + _stamp_gray_mark(frame, mark, x=mark_x, y=mark_y, opacity=0.78) + frames.append(frame) + source = tmp_path / "marked-vfr.mp4" + output = tmp_path / "clean-vfr.mp4" + _write_vfr_sora_clip( + source, + frames, + durations=durations, + ffmpeg=ffmpeg, + start_offset=2.0, + ) + source_timestamps = _video_frame_timestamps(source, ffprobe=ffprobe) + source_intervals = np.diff(source_timestamps) + + assert len(source_timestamps) == len(frames) + assert source_timestamps[0] == pytest.approx(2.0, abs=1 / 90000) + assert float(np.ptp(source_intervals)) >= 0.03 + + result = remove_video_visible(source, output, mark="sora", backend="cv2") + output_timestamps = _video_frame_timestamps(output, ffprobe=ffprobe) + + assert result.removed_frames == len(frames) + assert len(output_timestamps) == len(source_timestamps) + assert output_timestamps == pytest.approx(source_timestamps, abs=1 / 90000) + assert _stream_start_times(output, ffprobe=ffprobe) == pytest.approx( + _stream_start_times(source, ffprobe=ffprobe), + abs=1 / 90000, + ) + assert _container_duration(output, ffprobe=ffprobe) == pytest.approx( + _container_duration(source, ffprobe=ffprobe), + abs=max(durations), + ) + assert _audio_bitstream(output, ffmpeg=ffmpeg) == _audio_bitstream(source, ffmpeg=ffmpeg) + + def test_preserves_nonzero_start_on_constant_rate_clip(self, tmp_path: Path): + from remove_ai_watermarks.video import remove_video_visible + + ffmpeg, ffprobe = _ffmpeg_test_tools() + mark, mark_region = _independent_sora_mark() + mark_x, mark_y, _mark_width, _mark_height = mark_region + frames = [] + for frame_index in range(18): + frame = _moving_video_background(frame_index) + _stamp_gray_mark(frame, mark, x=mark_x, y=mark_y, opacity=0.78) + frames.append(frame) + source = tmp_path / "marked-offset.mp4" + output = tmp_path / "clean-offset.mp4" + _write_synthetic_sora_clip( + source, + frames, + fps=12.0, + ffmpeg=ffmpeg, + start_offset=2.0, + ) + + result = remove_video_visible(source, output, mark="sora", backend="cv2") + + assert result.removed_frames == len(frames) + assert _video_frame_timestamps(output, ffprobe=ffprobe) == pytest.approx( + _video_frame_timestamps(source, ffprobe=ffprobe), + abs=1 / 90000, + ) + assert _stream_start_times(output, ffprobe=ffprobe) == pytest.approx( + _stream_start_times(source, ffprobe=ffprobe), + abs=1 / 90000, + ) + assert _audio_bitstream(output, ffmpeg=ffmpeg) == _audio_bitstream(source, ffmpeg=ffmpeg) + class TestVideoVisibleApi: def test_auto_prefers_specific_sora_run_over_hailuo_cross_match( @@ -1252,6 +2392,7 @@ class TestVideoVisibleApi: ) -> int: assert regions == [sora_box] * 12 assert kwargs["padding_fraction"] == 0.28 + assert kwargs["temporal_consistency"] is True target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) return 12 @@ -1287,15 +2428,21 @@ class TestVideoVisibleApi: target: Path, _scan: VideoScan, regions: list[tuple[int, int, int, int] | None], - **_kwargs: object, + **kwargs: object, ) -> int: assert regions == [box] * 5 + assert kwargs["temporal_consistency"] is False target.write_bytes(_MP4_FTYP + _box(b"mdat", _VIDEO_PAYLOAD)) return 5 monkeypatch.setattr(video_visible, "encode_clean_video", fake_encode) - result = remove_video_visible(source, output, mark="sora") + result = remove_video_visible( + source, + output, + mark="sora", + temporal_consistency=False, + ) assert result.output == output assert result.detected_frames == 5 @@ -1441,6 +2588,7 @@ class TestVideoVisibleCli: assert result.exit_code == 0, result.output assert "temporally stable" in result.output assert "auto|sora|veo|seedance|dola|hailuo|kling" in result.output + assert "--temporal-consistency" in result.output def test_reports_removed_frames(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): from remove_ai_watermarks import video diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index c520cd5..c0b217d 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -1,12 +1,14 @@ -"""Regression tests for the video SynthID candidate engine.""" +"""Regression tests for the video SynthID removal engine.""" from __future__ import annotations +from types import SimpleNamespace from typing import TYPE_CHECKING import pytest from remove_ai_watermarks import video_encoding, video_invisible +from remove_ai_watermarks.video_synthid import DEFAULT_VIDEO_SYNTHID_NOISE_STD if TYPE_CHECKING: from pathlib import Path @@ -39,6 +41,14 @@ def test_encoder_command_discards_metadata_and_copies_audio( output = tmp_path / "candidate.mp4" monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") + profile = video_encoding.VideoEncodeProfile( + pixel_format="yuv420p", + color_range="tv", + color_space="bt709", + color_transfer="bt709", + color_primaries="bt709", + time_base="1/90000", + ) command = video_encoding.raw_video_command( source, @@ -48,16 +58,126 @@ def test_encoder_command_discards_metadata_and_copies_audio( fps=2.0, strip_metadata=True, crf=18, + profile=profile, ) metadata_index = command.index("-map_metadata") assert command[metadata_index + 1] == "-1" audio_codec_index = command.index("-c:a") assert command[audio_codec_index + 1] == "copy" + output_pixel_format_index = command.index("-pix_fmt", command.index("-c:v")) + assert command[output_pixel_format_index + 1] == "yuv420p" + assert command[command.index("-color_range") + 1] == "tv" + assert command[command.index("-colorspace") + 1] == "bt709" + assert command[command.index("-color_trc") + 1] == "bt709" + assert command[command.index("-color_primaries") + 1] == "bt709" + assert command[command.index("-enc_time_base:v") + 1] == "1/90000" + assert command[command.index("-video_track_timescale") + 1] == "90000" + assert command[command.index("-x264-params") + 1] == ( + "colorprim=bt709:transfer=bt709:colormatrix=bt709:range=limited" + ) assert "pipe:0" in command assert "-shortest" not in command +def test_timestamped_encoder_reads_nut_and_passes_pts_through( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "source.mp4" + output = tmp_path / "candidate.mp4" + monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") + + command = video_encoding.raw_video_command( + source, + output, + width=8, + height=8, + fps=24.0, + strip_metadata=True, + crf=18, + profile=video_encoding.VideoEncodeProfile(time_base="1/90000"), + timestamped_input=True, + copy_input_timestamps=True, + ) + + assert "-copyts" in command + assert command[command.index("-f") : command.index("-f") + 4] == ["-f", "nut", "-i", "pipe:0"] + assert command[command.index("-fps_mode") + 1] == "passthrough" + assert command[command.index("-avoid_negative_ts") + 1] == "disabled" + assert command[command.index("-enc_time_base:v") + 1] == "1/90000" + + +def test_probe_encode_profile_preserves_supported_8_bit_properties( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "source.mp4" + source.write_bytes(b"video") + monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffprobe") + monkeypatch.setattr( + video_encoding.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace( + returncode=0, + stdout=( + '{"streams":[{"pix_fmt":"yuvj422p","color_range":"tv",' + '"color_space":"bt709","color_transfer":"bt709",' + '"color_primaries":"bt709","time_base":"2/180000",' + '"start_pts":180000,"bits_per_raw_sample":"8"}]}' + ), + stderr="", + ), + ) + + profile = video_encoding.probe_video_encode_profile(source) + + assert profile == video_encoding.VideoEncodeProfile( + pixel_format="yuv422p", + color_range="tv", + color_space="bt709", + color_transfer="bt709", + color_primaries="bt709", + time_base="1/90000", + start_pts=180000, + source_pixel_format="yuvj422p", + component_depth=8, + ) + + +def test_probe_encode_profile_uses_compatible_defaults_without_ffprobe( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: None) + + assert video_encoding.probe_video_encode_profile(tmp_path / "source.mp4") == (video_encoding.VideoEncodeProfile()) + + +def test_probe_video_timestamps_uses_best_effort_pts( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "source.mp4" + source.write_bytes(b"video") + monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffprobe") + monkeypatch.setattr( + video_encoding.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace( + returncode=0, + stdout="0.000000\n0.041667\n", + stderr="", + ), + ) + + assert video_encoding.probe_video_timestamps(source) == (0.0, 0.041667) + + +def test_default_noise_matches_full_clip_oracle_floor() -> None: + assert DEFAULT_VIDEO_SYNTHID_NOISE_STD == 0.15 + + def test_stream_batches_consumes_only_one_batch_ahead() -> None: consumed: list[int] = [] diff --git a/tests/test_video_temporal.py b/tests/test_video_temporal.py new file mode 100644 index 0000000..4e2492d --- /dev/null +++ b/tests/test_video_temporal.py @@ -0,0 +1,112 @@ +"""Regression tests for motion-compensated visible-video fill.""" + +from __future__ import annotations + +import cv2 +import numpy as np + +from remove_ai_watermarks.video_temporal import stabilize_filled_frame + + +def _translated_pair() -> tuple[np.ndarray, np.ndarray]: + rng = np.random.default_rng(7) + previous = rng.integers(0, 256, (96, 128, 3), dtype=np.uint8) + previous = cv2.GaussianBlur(previous, (5, 5), 0) + current = cv2.warpAffine( + previous, + np.float32(((1, 0, 2), (0, 1, 1))), + (128, 96), + borderMode=cv2.BORDER_REFLECT, + ) + return previous, current + + +def test_motion_aligned_prior_reduces_independent_fill_error() -> None: + previous, current = _translated_pair() + mask = np.zeros(current.shape[:2], dtype=np.uint8) + mask[30:66, 45:85] = 255 + rng = np.random.default_rng(11) + current_fill = current.copy() + noise = rng.normal(0, 18, (36, 40, 3)) + current_fill[30:66, 45:85] = np.clip( + current_fill[30:66, 45:85].astype(np.float32) + noise, + 0, + 255, + ).astype(np.uint8) + + stabilized = stabilize_filled_frame( + previous, + previous, + mask, + current, + current_fill, + mask, + ) + + hole = mask > 0 + before = float(np.mean((current_fill[hole].astype(np.float32) - current[hole]) ** 2)) + after = float(np.mean((stabilized[hole].astype(np.float32) - current[hole]) ** 2)) + assert after < before * 0.5 + assert np.array_equal(stabilized[~hole], current_fill[~hole]) + + +def test_owned_fill_can_be_stabilized_without_full_frame_copy() -> None: + previous, current = _translated_pair() + mask = np.zeros(current.shape[:2], dtype=np.uint8) + mask[30:66, 45:85] = 255 + current_fill = current.copy() + current_fill[mask > 0] = 127 + + stabilized = stabilize_filled_frame( + previous, + previous, + mask, + current, + current_fill, + mask, + copy=False, + ) + + assert stabilized is current_fill + + +def test_scene_cut_keeps_independent_current_fill() -> None: + previous, _current = _translated_pair() + rng = np.random.default_rng(13) + current = rng.integers(0, 256, previous.shape, dtype=np.uint8) + mask = np.zeros(current.shape[:2], dtype=np.uint8) + mask[30:66, 45:85] = 255 + current_fill = current.copy() + current_fill[mask > 0] = 0 + + stabilized = stabilize_filled_frame( + previous, + previous, + mask, + current, + current_fill, + mask, + ) + + assert np.array_equal(stabilized, current_fill) + + +def test_disjoint_prior_mask_cannot_reintroduce_old_mark() -> None: + previous, current = _translated_pair() + previous_mask = np.zeros(current.shape[:2], dtype=np.uint8) + previous_mask[8:24, 8:24] = 255 + current_mask = np.zeros(current.shape[:2], dtype=np.uint8) + current_mask[60:76, 96:112] = 255 + current_fill = current.copy() + current_fill[current_mask > 0] = 127 + + stabilized = stabilize_filled_frame( + previous, + previous, + previous_mask, + current, + current_fill, + current_mask, + ) + + assert np.array_equal(stabilized, current_fill) diff --git a/uv.lock b/uv.lock index 62924fd..0cc988a 100644 --- a/uv.lock +++ b/uv.lock @@ -255,6 +255,111 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] +[[package]] +name = "av" +version = "16.1.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "(python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux')", +] +sdist = { url = "https://files.pythonhosted.org/packages/78/cd/3a83ffbc3cc25b39721d174487fb0d51a76582f4a1703f98e46170ce83d4/av-16.1.0.tar.gz", hash = "sha256:a094b4fd87a3721dacf02794d3d2c82b8d712c85b9534437e82a8a978c175ffd", size = 4285203, upload-time = "2026-01-11T07:31:33.772Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/97/51/2217a9249409d2e88e16e3f16f7c0def9fd3e7ffc4238b2ec211f9935bdb/av-16.1.0-cp310-cp310-macosx_11_0_x86_64.whl", hash = "sha256:2395748b0c34fe3a150a1721e4f3d4487b939520991b13e7b36f8926b3b12295", size = 26942590, upload-time = "2026-01-09T20:17:58.588Z" }, + { url = "https://files.pythonhosted.org/packages/bf/cd/a7070f4febc76a327c38808e01e2ff6b94531fe0b321af54ea3915165338/av-16.1.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:72d7ac832710a158eeb7a93242370aa024a7646516291c562ee7f14a7ea881fd", size = 21507910, upload-time = "2026-01-09T20:18:02.309Z" }, + { url = "https://files.pythonhosted.org/packages/ae/30/ec812418cd9b297f0238fe20eb0747d8a8b68d82c5f73c56fe519a274143/av-16.1.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:6cbac833092e66b6b0ac4d81ab077970b8ca874951e9c3974d41d922aaa653ed", size = 38738309, upload-time = "2026-01-09T20:18:04.701Z" }, + { url = "https://files.pythonhosted.org/packages/3a/b8/6c5795bf1f05f45c5261f8bce6154e0e5e86b158a6676650ddd77c28805e/av-16.1.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:eb990672d97c18f99c02f31c8d5750236f770ffe354b5a52c5f4d16c5e65f619", size = 40293006, upload-time = "2026-01-09T20:18:07.238Z" }, + { url = "https://files.pythonhosted.org/packages/a7/44/5e183bcb9333fc3372ee6e683be8b0c9b515a506894b2d32ff465430c074/av-16.1.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:05ad70933ac3b8ef896a820ea64b33b6cca91a5fac5259cb9ba7fa010435be15", size = 40123516, upload-time = "2026-01-09T20:18:09.955Z" }, + { url = "https://files.pythonhosted.org/packages/12/1d/b5346d582a3c3d958b4d26a2cc63ce607233582d956121eb20d2bbe55c2e/av-16.1.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:d831a1062a3c47520bf99de6ec682bd1d64a40dfa958e5457bb613c5270e7ce3", size = 41463289, upload-time = "2026-01-09T20:18:12.459Z" }, + { url = "https://files.pythonhosted.org/packages/fa/31/acc946c0545f72b8d0d74584cb2a0ade9b7dfe2190af3ef9aa52a2e3c0b1/av-16.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:358ab910fef3c5a806c55176f2b27e5663b33c4d0a692dafeb049c6ed71f8aff", size = 31754959, upload-time = "2026-01-09T20:18:14.718Z" }, + { url = "https://files.pythonhosted.org/packages/48/d0/b71b65d1b36520dcb8291a2307d98b7fc12329a45614a303ff92ada4d723/av-16.1.0-cp311-cp311-macosx_11_0_x86_64.whl", hash = "sha256:e88ad64ee9d2b9c4c5d891f16c22ae78e725188b8926eb88187538d9dd0b232f", size = 26927747, upload-time = "2026-01-09T20:18:16.976Z" }, + { url = "https://files.pythonhosted.org/packages/2f/79/720a5a6ccdee06eafa211b945b0a450e3a0b8fc3d12922f0f3c454d870d2/av-16.1.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:cb296073fa6935724de72593800ba86ae49ed48af03960a4aee34f8a611f442b", size = 21492232, upload-time = "2026-01-09T20:18:19.266Z" }, + { url = "https://files.pythonhosted.org/packages/8e/4f/a1ba8d922f2f6d1a3d52419463ef26dd6c4d43ee364164a71b424b5ae204/av-16.1.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:720edd4d25aa73723c1532bb0597806d7b9af5ee34fc02358782c358cfe2f879", size = 39291737, upload-time = "2026-01-09T20:18:21.513Z" }, + { url = "https://files.pythonhosted.org/packages/1a/31/fc62b9fe8738d2693e18d99f040b219e26e8df894c10d065f27c6b4f07e3/av-16.1.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:c7f2bc703d0df260a1fdf4de4253c7f5500ca9fc57772ea241b0cb241bcf972e", size = 40846822, upload-time = "2026-01-09T20:18:24.275Z" }, + { url = "https://files.pythonhosted.org/packages/53/10/ab446583dbce730000e8e6beec6ec3c2753e628c7f78f334a35cad0317f4/av-16.1.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d69c393809babada7d54964d56099e4b30a3e1f8b5736ca5e27bd7be0e0f3c83", size = 40675604, upload-time = "2026-01-09T20:18:26.866Z" }, + { url = "https://files.pythonhosted.org/packages/31/d7/1003be685277005f6d63fd9e64904ee222fe1f7a0ea70af313468bb597db/av-16.1.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:441892be28582356d53f282873c5a951592daaf71642c7f20165e3ddcb0b4c63", size = 42015955, upload-time = "2026-01-09T20:18:29.461Z" }, + { url = "https://files.pythonhosted.org/packages/2f/4a/fa2a38ee9306bf4579f556f94ecbc757520652eb91294d2a99c7cf7623b9/av-16.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:273a3e32de64819e4a1cd96341824299fe06f70c46f2288b5dc4173944f0fd62", size = 31750339, upload-time = "2026-01-09T20:18:32.249Z" }, + { url = "https://files.pythonhosted.org/packages/9c/84/2535f55edcd426cebec02eb37b811b1b0c163f26b8d3f53b059e2ec32665/av-16.1.0-cp312-cp312-macosx_11_0_x86_64.whl", hash = "sha256:640f57b93f927fba8689f6966c956737ee95388a91bd0b8c8b5e0481f73513d6", size = 26945785, upload-time = "2026-01-09T20:18:34.486Z" }, + { url = "https://files.pythonhosted.org/packages/b6/17/ffb940c9e490bf42e86db4db1ff426ee1559cd355a69609ec1efe4d3a9eb/av-16.1.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:ae3fb658eec00852ebd7412fdc141f17f3ddce8afee2d2e1cf366263ad2a3b35", size = 21481147, upload-time = "2026-01-09T20:18:36.716Z" }, + { url = "https://files.pythonhosted.org/packages/15/c1/e0d58003d2d83c3921887d5c8c9b8f5f7de9b58dc2194356a2656a45cfdc/av-16.1.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:27ee558d9c02a142eebcbe55578a6d817fedfde42ff5676275504e16d07a7f86", size = 39517197, upload-time = "2026-01-11T09:57:31.937Z" }, + { url = "https://files.pythonhosted.org/packages/32/77/787797b43475d1b90626af76f80bfb0c12cfec5e11eafcfc4151b8c80218/av-16.1.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:7ae547f6d5fa31763f73900d43901e8c5fa6367bb9a9840978d57b5a7ae14ed2", size = 41174337, upload-time = "2026-01-11T09:57:35.792Z" }, + { url = "https://files.pythonhosted.org/packages/8e/ac/d90df7f1e3b97fc5554cf45076df5045f1e0a6adf13899e10121229b826c/av-16.1.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8cf065f9d438e1921dc31fc7aa045790b58aee71736897866420d80b5450f62a", size = 40817720, upload-time = "2026-01-11T09:57:39.039Z" }, + { url = "https://files.pythonhosted.org/packages/80/6f/13c3a35f9dbcebafd03fe0c4cbd075d71ac8968ec849a3cfce406c35a9d2/av-16.1.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a345877a9d3cc0f08e2bc4ec163ee83176864b92587afb9d08dff50f37a9a829", size = 42267396, upload-time = "2026-01-11T09:57:42.115Z" }, + { url = "https://files.pythonhosted.org/packages/c8/b9/275df9607f7fb44317ccb1d4be74827185c0d410f52b6e2cd770fe209118/av-16.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:f49243b1d27c91cd8c66fdba90a674e344eb8eb917264f36117bf2b6879118fd", size = 31752045, upload-time = "2026-01-11T09:57:45.106Z" }, + { url = "https://files.pythonhosted.org/packages/75/2a/63797a4dde34283dd8054219fcb29294ba1c25d68ba8c8c8a6ae53c62c45/av-16.1.0-cp313-cp313-macosx_11_0_x86_64.whl", hash = "sha256:ce2a1b3d8bf619f6c47a9f28cfa7518ff75ddd516c234a4ee351037b05e6a587", size = 26916715, upload-time = "2026-01-11T09:57:47.682Z" }, + { url = "https://files.pythonhosted.org/packages/d2/c4/0b49cf730d0ae8cda925402f18ae814aef351f5772d14da72dd87ff66448/av-16.1.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:408dbe6a2573ca58a855eb8cd854112b33ea598651902c36709f5f84c991ed8e", size = 21452167, upload-time = "2026-01-11T09:57:50.606Z" }, + { url = "https://files.pythonhosted.org/packages/51/23/408806503e8d5d840975aad5699b153aaa21eb6de41ade75248a79b7a37f/av-16.1.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:57f657f86652a160a8a01887aaab82282f9e629abf94c780bbdbb01595d6f0f7", size = 39215659, upload-time = "2026-01-11T09:57:53.757Z" }, + { url = "https://files.pythonhosted.org/packages/c4/19/a8528d5bba592b3903f44c28dab9cc653c95fcf7393f382d2751a1d1523e/av-16.1.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:adbad2b355c2ee4552cac59762809d791bda90586d134a33c6f13727fb86cb3a", size = 40874970, upload-time = "2026-01-11T09:57:56.802Z" }, + { url = "https://files.pythonhosted.org/packages/e8/24/2dbcdf0e929ad56b7df078e514e7bd4ca0d45cba798aff3c8caac097d2f7/av-16.1.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f42e1a68ec2aebd21f7eb6895be69efa6aa27eec1670536876399725bbda4b99", size = 40530345, upload-time = "2026-01-11T09:58:00.421Z" }, + { url = "https://files.pythonhosted.org/packages/54/27/ae91b41207f34e99602d1c72ab6ffd9c51d7c67e3fbcd4e3a6c0e54f882c/av-16.1.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:58fe47aeaef0f100c40ec8a5de9abbd37f118d3ca03829a1009cf288e9aef67c", size = 41972163, upload-time = "2026-01-11T09:58:03.756Z" }, + { url = "https://files.pythonhosted.org/packages/fc/7a/22158fb923b2a9a00dfab0e96ef2e8a1763a94dd89e666a5858412383d46/av-16.1.0-cp313-cp313-win_amd64.whl", hash = "sha256:565093ebc93b2f4b76782589564869dadfa83af5b852edebedd8fee746457d06", size = 31729230, upload-time = "2026-01-11T09:58:07.254Z" }, + { url = "https://files.pythonhosted.org/packages/7f/f1/878f8687d801d6c4565d57ebec08449c46f75126ebca8e0fed6986599627/av-16.1.0-cp313-cp313t-macosx_11_0_x86_64.whl", hash = "sha256:574081a24edb98343fd9f473e21ae155bf61443d4ec9d7708987fa597d6b04b2", size = 27008769, upload-time = "2026-01-11T09:58:10.266Z" }, + { url = "https://files.pythonhosted.org/packages/30/f1/bd4ce8c8b5cbf1d43e27048e436cbc9de628d48ede088a1d0a993768eb86/av-16.1.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:9ab00ea29c25ebf2ea1d1e928d7babb3532d562481c5d96c0829212b70756ad0", size = 21590588, upload-time = "2026-01-11T09:58:12.629Z" }, + { url = "https://files.pythonhosted.org/packages/1d/dd/c81f6f9209201ff0b5d5bed6da6c6e641eef52d8fbc930d738c3f4f6f75d/av-16.1.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:a84a91188c1071f238a9523fd42dbe567fb2e2607b22b779851b2ce0eac1b560", size = 40638029, upload-time = "2026-01-11T09:58:15.399Z" }, + { url = "https://files.pythonhosted.org/packages/15/4d/07edff82b78d0459a6e807e01cd280d3180ce832efc1543de80d77676722/av-16.1.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:c2cd0de4dd022a7225ff224fde8e7971496d700be41c50adaaa26c07bb50bf97", size = 41970776, upload-time = "2026-01-11T09:58:19.075Z" }, + { url = "https://files.pythonhosted.org/packages/da/9d/1f48b354b82fa135d388477cd1b11b81bdd4384bd6a42a60808e2ec2d66b/av-16.1.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:0816143530624a5a93bc5494f8c6eeaf77549b9366709c2ac8566c1e9bff6df5", size = 41764751, upload-time = "2026-01-11T09:58:22.788Z" }, + { url = "https://files.pythonhosted.org/packages/2f/c7/a509801e98db35ec552dd79da7bdbcff7104044bfeb4c7d196c1ce121593/av-16.1.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:e3a28053af29644696d0c007e897d19b1197585834660a54773e12a40b16974c", size = 43034355, upload-time = "2026-01-11T09:58:26.125Z" }, + { url = "https://files.pythonhosted.org/packages/36/8b/e5f530d9e8f640da5f5c5f681a424c65f9dd171c871cd255d8a861785a6e/av-16.1.0-cp313-cp313t-win_amd64.whl", hash = "sha256:2e3e67144a202b95ed299d165232533989390a9ea3119d37eccec697dc6dbb0c", size = 31947047, upload-time = "2026-01-11T09:58:31.867Z" }, + { url = "https://files.pythonhosted.org/packages/df/18/8812221108c27d19f7e5f486a82c827923061edf55f906824ee0fcaadf50/av-16.1.0-cp314-cp314-macosx_11_0_x86_64.whl", hash = "sha256:39a634d8e5a87e78ea80772774bfd20c0721f0d633837ff185f36c9d14ffede4", size = 26916179, upload-time = "2026-01-11T09:58:36.506Z" }, + { url = "https://files.pythonhosted.org/packages/38/ef/49d128a9ddce42a2766fe2b6595bd9c49e067ad8937a560f7838a541464e/av-16.1.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:0ba32fb9e9300948a7fa9f8a3fc686e6f7f77599a665c71eb2118fdfd2c743f9", size = 21460168, upload-time = "2026-01-11T09:58:39.231Z" }, + { url = "https://files.pythonhosted.org/packages/e6/a9/b310d390844656fa74eeb8c2750e98030877c75b97551a23a77d3f982741/av-16.1.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:ca04d17815182d34ce3edc53cbda78a4f36e956c0fd73e3bab249872a831c4d7", size = 39210194, upload-time = "2026-01-11T09:58:42.138Z" }, + { url = "https://files.pythonhosted.org/packages/0c/7b/e65aae179929d0f173af6e474ad1489b5b5ad4c968a62c42758d619e54cf/av-16.1.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:ee0e8de2e124a9ef53c955fe2add6ee7c56cc8fd83318265549e44057db77142", size = 40811675, upload-time = "2026-01-11T09:58:45.871Z" }, + { url = "https://files.pythonhosted.org/packages/54/3f/5d7edefd26b6a5187d6fac0f5065ee286109934f3dea607ef05e53f05b31/av-16.1.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:22bf77a2f658827043a1e184b479c3bf25c4c43ab32353677df2d119f080e28f", size = 40543942, upload-time = "2026-01-11T09:58:49.759Z" }, + { url = "https://files.pythonhosted.org/packages/1b/24/f8b17897b67be0900a211142f5646a99d896168f54d57c81f3e018853796/av-16.1.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2dd419d262e6a71cab206d80bbf28e0a10d0f227b671cdf5e854c028faa2d043", size = 41924336, upload-time = "2026-01-11T09:58:53.344Z" }, + { url = "https://files.pythonhosted.org/packages/1c/cf/d32bc6bbbcf60b65f6510c54690ed3ae1c4ca5d9fafbce835b6056858686/av-16.1.0-cp314-cp314-win_amd64.whl", hash = "sha256:53585986fd431cd436f290fba662cfb44d9494fbc2949a183de00acc5b33fa88", size = 31735077, upload-time = "2026-01-11T09:58:56.684Z" }, + { url = "https://files.pythonhosted.org/packages/53/f4/9b63dc70af8636399bd933e9df4f3025a0294609510239782c1b746fc796/av-16.1.0-cp314-cp314t-macosx_11_0_x86_64.whl", hash = "sha256:76f5ed8495cf41e1209a5775d3699dc63fdc1740b94a095e2485f13586593205", size = 27014423, upload-time = "2026-01-11T09:58:59.703Z" }, + { url = "https://files.pythonhosted.org/packages/d1/da/787a07a0d6ed35a0888d7e5cfb8c2ffa202f38b7ad2c657299fac08eb046/av-16.1.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:8d55397190f12a1a3ae7538be58c356cceb2bf50df1b33523817587748ce89e5", size = 21595536, upload-time = "2026-01-11T09:59:02.508Z" }, + { url = "https://files.pythonhosted.org/packages/d8/f4/9a7d8651a611be6e7e3ab7b30bb43779899c8cac5f7293b9fb634c44a3f3/av-16.1.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:9d51d9037437218261b4bbf9df78a95e216f83d7774fbfe8d289230b5b2e28e2", size = 40642490, upload-time = "2026-01-11T09:59:05.842Z" }, + { url = "https://files.pythonhosted.org/packages/6b/e4/eb79bc538a94b4ff93cd4237d00939cba797579f3272490dd0144c165a21/av-16.1.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:0ce07a89c15644407f49d942111ca046e323bbab0a9078ff43ee57c9b4a50dad", size = 41976905, upload-time = "2026-01-11T09:59:09.169Z" }, + { url = "https://files.pythonhosted.org/packages/5e/f5/f6db0dd86b70167a4d55ee0d9d9640983c570d25504f2bde42599f38241e/av-16.1.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:cac0c074892ea97113b53556ff41c99562db7b9f09f098adac1f08318c2acad5", size = 41770481, upload-time = "2026-01-11T09:59:12.74Z" }, + { url = "https://files.pythonhosted.org/packages/9e/8b/33651d658e45e16ab7671ea5fcf3d20980ea7983234f4d8d0c63c65581a5/av-16.1.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:7dec3dcbc35a187ce450f65a2e0dda820d5a9e6553eea8344a1459af11c98649", size = 43036824, upload-time = "2026-01-11T09:59:16.507Z" }, + { url = "https://files.pythonhosted.org/packages/83/41/7f13361db54d7e02f11552575c0384dadaf0918138f4eaa82ea03a9f9580/av-16.1.0-cp314-cp314t-win_amd64.whl", hash = "sha256:6f90dc082ff2068ddbe77618400b44d698d25d9c4edac57459e250c16b33d700", size = 31948164, upload-time = "2026-01-11T09:59:19.501Z" }, +] + +[[package]] +name = "av" +version = "18.0.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'darwin'", + "python_full_version >= '3.14' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "(python_full_version >= '3.14' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "(python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", +] +sdist = { url = "https://files.pythonhosted.org/packages/ae/a4/570a5a35c8638aba01e739925846c35fdd6b0756a15526766d0a4dd3b7df/av-18.0.0.tar.gz", hash = "sha256:4ef7e72c3d3a872584a1215173b16e0226811037f40dcdbf75992631098df1ba", size = 4340222, upload-time = "2026-07-02T06:37:58.907Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/4a/9e3463df030e063d757fa12f0f39be6541b45b06b5bad48c2ce361b924bf/av-18.0.0-cp311-abi3-macosx_11_0_x86_64.whl", hash = "sha256:149289d40e732a6e49c9530bc245b49d9964cfd1c8c9e06778703b7d5bba6b25", size = 22499354, upload-time = "2026-07-02T06:36:58.751Z" }, + { url = "https://files.pythonhosted.org/packages/77/b3/2576a44b4f39c7462ced4c17fec04c756f7b0f3c5cb940d124173e417d6a/av-18.0.0-cp311-abi3-macosx_14_0_arm64.whl", hash = "sha256:35274c20d2ad3b4774fe632bcef2e34af79858ddf899352339cc3babbc13a484", size = 18175248, upload-time = "2026-07-02T06:37:01.741Z" }, + { url = "https://files.pythonhosted.org/packages/84/74/6732f17b96dc23fd23b876b2805435855abdc8a3b397142be4e581165de8/av-18.0.0-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:4d683b7747a0ba9222b8a5f81e41db5f796e7f64473454ec4fe2548e083c2fa0", size = 33387843, upload-time = "2026-07-02T06:37:05.097Z" }, + { url = "https://files.pythonhosted.org/packages/6d/b9/7708c43fed7ae28b4a1bad060b4221e3334cd827cec24f7165902a6ac1f4/av-18.0.0-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ae56b40b6f8b067a8ad2dac664fbfbabac7f7a55b9a7bb031eb99289252bc017", size = 35536910, upload-time = "2026-07-02T06:37:08.806Z" }, + { url = "https://files.pythonhosted.org/packages/5a/94/eba99691d184f6a395a242d54dc370e2fd2265e95bbc98e2963a0fdbdd6c/av-18.0.0-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:ea2e8ebbce521f21b55df9400e00d721623c9020ef158f5a188a96130be0743f", size = 38984619, upload-time = "2026-07-02T06:37:11.861Z" }, + { url = "https://files.pythonhosted.org/packages/c9/cf/0d7aee07fe16aa9ffdf96043c14bed5485a52c0dea4259de87aa306ecab4/av-18.0.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef96dabb3e50dac249913145dff5424b302b257fd95dcb64be3c7b7a8aef16d1", size = 34451176, upload-time = "2026-07-02T06:37:15.154Z" }, + { url = "https://files.pythonhosted.org/packages/76/92/810da80b12680d4c4fe235bd1b4003289be9213ac7f114b77b8ecf0e3b3e/av-18.0.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:0f65518a184613e41536f29e8758c8e3d8293e46bf5bef108f04f925bbfa3f44", size = 36619869, upload-time = "2026-07-02T06:37:18.495Z" }, + { url = "https://files.pythonhosted.org/packages/11/85/0f121ff43dc5a70696676c98a8f1674e2fa787614c2abaacb15fa1a9bc99/av-18.0.0-cp311-abi3-win_amd64.whl", hash = "sha256:aaf4d354d2beaa6651e4f92e54409a578bde64f79c0beef9a30b388d06f7c629", size = 27556236, upload-time = "2026-07-02T06:37:21.388Z" }, + { url = "https://files.pythonhosted.org/packages/8b/f6/2509754d4d2356abc6fc0ea3d57c12ade29bac23a1fb7fc215a53ca518fb/av-18.0.0-cp311-abi3-win_arm64.whl", hash = "sha256:adac2b3833b6cb9bd6cb52664a522b94db453615b3675b1dbb26e13fe1c80da6", size = 20221133, upload-time = "2026-07-02T06:37:23.88Z" }, + { url = "https://files.pythonhosted.org/packages/e2/25/4ee23a7f1609adf9b2f140c7a8ffade64a1449d89ab431d922a809eebf19/av-18.0.0-cp314-cp314t-macosx_11_0_x86_64.whl", hash = "sha256:88dd8e35e9242662b409a6a05fd24a6775d949eb05da0ba31cab4f250eacbab5", size = 22740741, upload-time = "2026-07-02T06:37:26.659Z" }, + { url = "https://files.pythonhosted.org/packages/f1/f0/b9f8363d07aa4521913e483f6a30c7c164973ef01de62769bf9b97049cd8/av-18.0.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:f8f454349c402e2c8d6fa80b54eb2a3f86c00f414d2b399f01ae6dab075c6fd8", size = 18384189, upload-time = "2026-07-02T06:37:29.518Z" }, + { url = "https://files.pythonhosted.org/packages/c3/e5/69397019aed280a72a43e97a252dee4295df1a9e608848452e5300ec4dab/av-18.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:88ce194c2201c6a6d40336adee8a5ddde46ed743eacb500e3ae9368d1c6d889e", size = 36749881, upload-time = "2026-07-02T06:37:33.096Z" }, + { url = "https://files.pythonhosted.org/packages/37/3a/1614d74f0d676ea6745eb59553c9ad01ca25db523cba808d522e838f4f5b/av-18.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:aa15e567a018cc94a26b0ab45da676dee70c4146ace6e92e47d30cc9689cbfbe", size = 38645927, upload-time = "2026-07-02T06:37:37.086Z" }, + { url = "https://files.pythonhosted.org/packages/6b/3c/5f54710d69b0ea93634134f92b49c7a2a7fd27da5486a8a7e6251ac1cfb4/av-18.0.0-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:613153e48cefc91700746dde0ad0282d4677b194cba22cc771de14c78411cf8b", size = 40454783, upload-time = "2026-07-02T06:37:40.904Z" }, + { url = "https://files.pythonhosted.org/packages/26/92/8293e6a267e0591b543abd96ae01e7e8ed228509bdb4e4644a8a8395d90f/av-18.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:30404f53ca1ea7f350ac86ff22a2c04f903014758e9b33f398c5a62de34bd84f", size = 37573117, upload-time = "2026-07-02T06:37:44.856Z" }, + { url = "https://files.pythonhosted.org/packages/10/0c/38ed7601277ae57dfe857d040be4762530fd728efff45c2fb8f035fef96a/av-18.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6882a48f7aec2863c96cddee3256ff2da98f7fb6cbed83cee9d7e70a8f186a6b", size = 39669026, upload-time = "2026-07-02T06:37:48.761Z" }, + { url = "https://files.pythonhosted.org/packages/c8/95/0636ca04d5d89d01c49bd366d2b660cc85d1f8117c476b2be62eb0c70855/av-18.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:55a646e9afce9fdc5de5224205a8a12c7ed1ba9803145dcc876c40bfc03a109b", size = 28448336, upload-time = "2026-07-02T06:37:52.477Z" }, + { url = "https://files.pythonhosted.org/packages/01/20/1e24450ea981c44ed328691496fd2774dfa9fa3c3b00fd07f72fd5614abe/av-18.0.0-cp314-cp314t-win_arm64.whl", hash = "sha256:96f594ff506a09475e5549359352332049a25d37a08f00b4623f7f6e92e45b9c", size = 21377289, upload-time = "2026-07-02T06:37:55.935Z" }, +] + [[package]] name = "c2pa-python" version = "0.37.1" @@ -3190,6 +3295,8 @@ name = "remove-ai-watermarks" version = "0.20.2" source = { editable = "." } dependencies = [ + { name = "av", version = "16.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "av", version = "18.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "c2pa-python" }, { name = "click" }, { name = "numpy" }, @@ -3272,6 +3379,8 @@ trustmark = [ [package.metadata] requires-dist = [ { name = "accelerate", marker = "extra == 'gpu'", specifier = ">=0.25.0" }, + { name = "av", marker = "python_full_version < '3.11'", specifier = ">=16,<17" }, + { name = "av", marker = "python_full_version >= '3.11'", specifier = ">=18,<19" }, { name = "c2pa-python", specifier = ">=0.35.0" }, { name = "click", specifier = ">=8.0.0" }, { name = "diffsynth", marker = "extra == 'qwen-zimage'", specifier = ">=2.0.17,<3" }, From 5767c1d7b30373c0dc1e06230e63914cb88afea2 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 10:52:44 -0700 Subject: [PATCH 08/19] Finalize video package boundary and product wording --- .gitignore | 2 ++ CLAUDE.md | 2 +- docs/known-limitations.md | 6 +++--- docs/release-and-distribution.md | 12 ++++++++---- docs/verification-plan.md | 6 +++--- packaging/conda/recipe.yaml | 2 +- pyproject.toml | 17 ++++++++++++++--- src/remove_ai_watermarks/video_synthid.py | 2 +- tests/test_distribution_boundary.py | 22 ++++++++++++++++++++++ 9 files changed, 55 insertions(+), 16 deletions(-) create mode 100644 tests/test_distribution_boundary.py diff --git a/.gitignore b/.gitignore index 491809c..58c5fb9 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,8 @@ __pycache__/ *.egg-info/ dist/ build/ +tmp/ +.sc/ # Environment secrets .env diff --git a/CLAUDE.md b/CLAUDE.md index 6633339..ab0e939 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -38,7 +38,7 @@ Per-command exit-code semantics (the no-signal / GPU-missing skip branches), tes - **CI** (`.github/workflows/test.yml`): runs on push to `main` + every PR. A `lint` job (ubuntu: `ruff check` + `ruff format --check`) plus a `test` matrix (ubuntu/macos/windows x py3.10/3.12) that does `uv sync --frozen --extra dev` then `pytest`. The matrix installs only core + dev (no `gpu` extra), so the GPU/model-running tests skip there and it exercises the metadata/identify/visible/cv2-eraser surface on all three OSes. Keep `uv.lock` valid (don't break `--frozen`) when editing `pyproject.toml`. - Dependency PR checks run the GitHub merge ref against current `main`, not the contributor branch in isolation. If `main` moves after the dependency branch was opened, merge current `main` locally and rerun the full gate; a new linter can expose stale directives in code that landed later. -- **Release flow + distribution channels** (PyPI publish via `publish.yml`/`uv publish`, the automated Homebrew-tap + HF-Space bumps in `distribute.yml`, conda-forge, ComfyUI Registry, the sdist `data/` exclusion, hatchling pin history): see `docs/release-and-distribution.md` before cutting a release. +- **Release flow + distribution channels** (PyPI publish via `publish.yml`/`uv publish`, the automated Homebrew-tap + HF-Space bumps in `distribute.yml`, conda-forge, ComfyUI Registry, the sdist public allowlist, hatchling pin history): see `docs/release-and-distribution.md` before cutting a release. - `bash maintain.sh` — uv-outdated, uv-secure, ruff check/fix, ruff format, pyright (scoped `src/`, see the OOM note below), pytest -n auto. The helper tools live in the `dev` extra (`pytest-xdist`, plus `uv-outdated`/`uv-secure` marker-gated to py3.12+ so the py3.10 resolution stays solvable) — a bare env without `--extra dev` does not have them. - **Strict pyright is clean across `src/` (0 errors).** The cv2/torch/diffusers boundary files (`gemini_engine`, `region_eraser`, `doubao_engine`, `humanizer`, `invisible_engine`, `noai/watermark_remover`) carry a documented per-file `# pyright:` relax pragma that turns off only the unknown-type / untyped-third-party rules — those libs ship no usable types, so strict typing there fights the ecosystem. Pure-logic files stay fully strict; `typings/piexif/__init__.pyi` is a local stub so `metadata.py`/`extractor.py` resolve piexif. Public ndarray-returning signatures on the relaxed engines are still annotated `NDArray[Any]` so strict consumers (`cli.py`) stay clean. When touching a relaxed file, prefer fixing real issues over widening the pragma; keep the pragma scoped to genuinely-untyped boundaries. The `uv-secure` CVE-resolution history (idna/aiohttp bumps, retired basicsr, the dismissed torch `GHSA-rrmf-rvhw-rf47`) lives in `docs/release-and-distribution.md` — read it before re-triaging a dependency alert. - **Full-project `uv run pyright` (no path) OOMs/crashes node on this ML-heavy repo** (emits a `libnode` stack frame, no summary) — a known environment limit, not a code error. Gate with `uv run --extra dev --extra gpu pyright src/` (completes, authoritative) or scope to changed files; also run `uv run ruff check` and `uv run pytest` directly. diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 0a09587..fc5a1fc 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -212,10 +212,10 @@ WebM, Matroska, MP3, WAV, FLAC, OGG, Opus, and AAC container metadata is strippe through ffmpeg with stream copying. The operation fails if ffmpeg is absent or cannot parse the input. -### Video pixel removal is provider-specific and still experimental +### Video pixel removal is provider-specific -The experimental `video metadata` command and high level video API inspect and -strip supported AI provenance metadata without transcoding streams. +The `video metadata` command and high level video API inspect and strip +supported AI provenance metadata without transcoding streams. `video visible` and `remove_video_visible` additionally support the moving Sora 2 mascot and wordmark, the current Veo four-point diamond, the legacy diff --git a/docs/release-and-distribution.md b/docs/release-and-distribution.md index 7757451..4f0006d 100644 --- a/docs/release-and-distribution.md +++ b/docs/release-and-distribution.md @@ -65,10 +65,14 @@ current line according to the environment. The wheel includes the package under `src/`. -The source distribution explicitly excludes `/data` through -`[tool.hatch.build.targets.sdist]` in `pyproject.toml`. Keep that exclusion: -calibration captures and test corpora do not belong in the published package -archive. +The source distribution uses an explicit allowlist for `/src`, `/LICENSE`, +`/README.md`, and `/pyproject.toml` through +`[tool.hatch.build.targets.sdist]` in `pyproject.toml`. It also defensively +excludes `/data`, `/tmp`, and `/.sc`. Keep both controls: calibration captures, +test corpora, generated research outputs, and local session state do not belong +in the published package archive. The matching local-root entries in +`.gitignore` prevent accidental commits, but are not a substitute for the build +boundary because hatchling may include untracked files. ## Build backend diff --git a/docs/verification-plan.md b/docs/verification-plan.md index ed6dc44..a850587 100644 --- a/docs/verification-plan.md +++ b/docs/verification-plan.md @@ -242,9 +242,9 @@ was negative on both carriers, and a weaker candidate was negative on one. On ordinary Gemini to ignore and reinterpret the detector result. Those follow-ups were mistakenly treated as a stricter oracle; they were not detector reruns. The implementation remains exposed as `video invisible` and -`remove_video_invisible`. Its default is a calibrated operating point, while -each important result still requires a fresh control-positive, -candidate-negative pair. +`remove_video_invisible`. Its oracle-certified default is the product operating +point. A fresh control-positive, candidate-negative pair remains an optional +per-file audit after provider changes or for unusually important files. The 2026-07-31 full-clip check added a public eight-second Veo carrier. The source and the complete `0.10` product output were both detected, proving the diff --git a/packaging/conda/recipe.yaml b/packaging/conda/recipe.yaml index ff2cccd..9e592c1 100644 --- a/packaging/conda/recipe.yaml +++ b/packaging/conda/recipe.yaml @@ -59,7 +59,7 @@ about: Detect and remove registered visible AI-provenance marks and strip AI-provenance metadata (C2PA, EXIF, IPTC, and PNG text chunks) from images and supported video containers. The core package covers the identify, - metadata, visible, erase, and experimental video command surface. Optional + metadata, visible, erase, and product-oriented video command surface. Optional pip extras add SynthID diffusion removal and additional invisible-watermark detectors. license: Apache-2.0 diff --git a/pyproject.toml b/pyproject.toml index 4e73042..a82d5e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -191,9 +191,20 @@ build-backend = "hatchling.build" packages = ["src/remove_ai_watermarks"] [tool.hatch.build.targets.sdist] -# Keep the source distribution small: ship the package and metadata, not -# calibration or evaluation data under data/. The wheel ships only src/. -exclude = ["/data"] +# Keep the source distribution small and public-safe: ship tracked source and +# metadata, not corpora or local research/session artifacts. The wheel ships +# only src/. +include = [ + "/src", + "/LICENSE", + "/README.md", + "/pyproject.toml", +] +exclude = [ + "/data", + "/tmp", + "/.sc", +] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/src/remove_ai_watermarks/video_synthid.py b/src/remove_ai_watermarks/video_synthid.py index 612ec01..e262a2c 100644 --- a/src/remove_ai_watermarks/video_synthid.py +++ b/src/remove_ai_watermarks/video_synthid.py @@ -1,4 +1,4 @@ -"""Shared configuration for experimental video SynthID regeneration.""" +"""Shared configuration for oracle-certified video SynthID removal.""" DEFAULT_VIDEO_SYNTHID_VAE = "stabilityai/sd-vae-ft-mse" DEFAULT_VIDEO_SYNTHID_NOISE_STD = 0.15 diff --git a/tests/test_distribution_boundary.py b/tests/test_distribution_boundary.py new file mode 100644 index 0000000..ce6993b --- /dev/null +++ b/tests/test_distribution_boundary.py @@ -0,0 +1,22 @@ +"""Distribution-boundary regression tests.""" + +from __future__ import annotations + +import re +from pathlib import Path + + +def _array_values(section: str, key: str) -> set[str]: + match = re.search(rf"(?ms)^{key}\s*=\s*\[(.*?)^\]", section) + assert match is not None + return set(re.findall(r'"([^"]+)"', match.group(1))) + + +def test_sdist_has_explicit_public_boundary() -> None: + """Hatchling must publish only package source and required metadata.""" + root = Path(__file__).resolve().parents[1] + config = (root / "pyproject.toml").read_text(encoding="utf-8") + sdist_config = config.split("[tool.hatch.build.targets.sdist]", maxsplit=1)[1].split("\n[", maxsplit=1)[0] + + assert _array_values(sdist_config, "include") == {"/src", "/LICENSE", "/README.md", "/pyproject.toml"} + assert {"/data", "/tmp", "/.sc"} <= _array_values(sdist_config, "exclude") From 0e7a52820b742cd07882648043474a69fe35c933 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 11:30:57 -0700 Subject: [PATCH 09/19] Bound full-clip video CI runtime --- .github/workflows/test.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 104ae25..2ed22f3 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -48,6 +48,7 @@ jobs: video-e2e: name: video full-clip end-to-end runs-on: ubuntu-latest + timeout-minutes: 15 steps: - uses: actions/checkout@v7 - uses: astral-sh/setup-uv@v7 @@ -59,5 +60,5 @@ jobs: run: uv sync --frozen --extra dev - name: Run full-clip video test run: >- - uv run pytest -q + uv run pytest -vv tests/test_video.py::TestVideoVisibleFullClip From f9a4e2ba105cc9a04e4801fd4eed48b2a26c18c2 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 11:54:45 -0700 Subject: [PATCH 10/19] Diagnose stalled video end-to-end test --- .github/workflows/test.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2ed22f3..74c7754 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -48,7 +48,7 @@ jobs: video-e2e: name: video full-clip end-to-end runs-on: ubuntu-latest - timeout-minutes: 15 + timeout-minutes: 5 steps: - uses: actions/checkout@v7 - uses: astral-sh/setup-uv@v7 @@ -61,4 +61,5 @@ jobs: - name: Run full-clip video test run: >- uv run pytest -vv + -o faulthandler_timeout=60 tests/test_video.py::TestVideoVisibleFullClip From 1fd859b58d6f668273fc273777f274ef2d03319c Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 12:05:42 -0700 Subject: [PATCH 11/19] Prevent ffmpeg pipe probe deadlock --- docs/module-internals.md | 5 ++++ src/remove_ai_watermarks/video_encoding.py | 9 ++++--- tests/test_video.py | 28 ++++++++++++++++++++++ tests/test_video_invisible.py | 11 ++++++++- 4 files changed, 49 insertions(+), 4 deletions(-) diff --git a/docs/module-internals.md b/docs/module-internals.md index 3f82478..4f8019c 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -149,6 +149,11 @@ regeneration. It centralizes container codecs, optional audio stream copying, metadata/chapter policy, encode-failure reporting, and atomic same-directory publication. Each mapped stream is allowed to reach its own end, so a copied audio tail is not shortened to the frame-input duration. +Both the raw-BGR and timestamped-NUT stdin inputs disable ffmpeg probing before +`pipe:0`: their format is already explicit, and ffmpeg 6 on Linux can otherwise +wait for its normal analysis window while the producer blocks on a full pipe +before the source-audio input opens. Command-order regressions cover both stdin +modes; the real Linux full-clip CI job guards process completion. `probe_video_encode_profile` reads the first source video stream with ffprobe and preserves the supported properties that survive the 8-bit BGR boundary: `yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index c15931f..1f129d3 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -286,8 +286,12 @@ def raw_video_command( ffmpeg = shutil.which("ffmpeg") if ffmpeg is None: raise RuntimeError("Video processing requires ffmpeg on PATH") + # The pipe format and stream geometry are already explicit. FFmpeg 6 can + # otherwise wait for its normal analysis window while the producer blocks + # on a full pipe, before the second (audio) input has been opened. + pipe_input = ["-analyzeduration", "0", "-probesize", "32", "-i", "pipe:0"] frame_input = ( - ["-f", "nut", "-i", "pipe:0"] + ["-f", "nut", *pipe_input] if timestamped_input else [ "-f", @@ -298,8 +302,7 @@ def raw_video_command( f"{width}x{height}", "-r", f"{fps:.12g}", - "-i", - "pipe:0", + *pipe_input, ] ) command = [ diff --git a/tests/test_video.py b/tests/test_video.py index 23a8fb3..6df876f 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -1950,6 +1950,34 @@ class TestVideoVisibleScan: class TestVideoVisibleEncoding: + def test_raw_pipe_input_disables_redundant_ffmpeg_probing( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_encoding + + monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") + command = video_encoding.raw_video_command( + tmp_path / "source.mp4", + tmp_path / "clean.mp4", + width=12, + height=8, + fps=24.0, + strip_metadata=True, + crf=14, + profile=video_encoding.VideoEncodeProfile(), + ) + + pipe_position = command.index("pipe:0") + assert command[pipe_position - 5 : pipe_position] == [ + "-analyzeduration", + "0", + "-probesize", + "32", + "-i", + ] + @staticmethod def _patch_single_frame_encode( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index c0b217d..61ddaf3 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -102,7 +102,16 @@ def test_timestamped_encoder_reads_nut_and_passes_pts_through( ) assert "-copyts" in command - assert command[command.index("-f") : command.index("-f") + 4] == ["-f", "nut", "-i", "pipe:0"] + assert command[command.index("-f") : command.index("-f") + 8] == [ + "-f", + "nut", + "-analyzeduration", + "0", + "-probesize", + "32", + "-i", + "pipe:0", + ] assert command[command.index("-fps_mode") + 1] == "passthrough" assert command[command.index("-avoid_negative_ts") + 1] == "disabled" assert command[command.index("-enc_time_base:v") + 1] == "1/90000" From 7fb0e654550d170bee7d3178ee3b314941367c57 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 14:24:08 -0700 Subject: [PATCH 12/19] Trace stalled Linux video encoder --- .github/workflows/test.yml | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 74c7754..09389b4 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -55,11 +55,22 @@ jobs: with: python-version: "3.12" - name: Install ffmpeg - run: sudo apt-get update && sudo apt-get install --yes ffmpeg + run: sudo apt-get update && sudo apt-get install --yes ffmpeg strace - name: Sync dev environment run: uv sync --frozen --extra dev - name: Run full-clip video test - run: >- - uv run pytest -vv - -o faulthandler_timeout=60 - tests/test_video.py::TestVideoVisibleFullClip + run: | + set +e + timeout --kill-after=10s 90s \ + strace -ff -tt -s 256 -o /tmp/video-e2e.strace \ + uv run pytest -vv -o faulthandler_timeout=45 \ + tests/test_video.py::TestVideoVisibleFullClip + status=$? + if [ "$status" -ne 0 ]; then + for trace in /tmp/video-e2e.strace*; do + echo "::group::$trace" + tail -n 100 "$trace" + echo "::endgroup::" + done + fi + exit "$status" From 166dab9d03537e490c9701627d53650d34a28804 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 14:33:50 -0700 Subject: [PATCH 13/19] Bound ffmpeg video encoder threads --- .github/workflows/test.yml | 20 ++++---------------- docs/module-internals.md | 11 ++++++----- src/remove_ai_watermarks/video_encoding.py | 12 ++++++------ tests/test_video.py | 15 +++++---------- tests/test_video_invisible.py | 6 +----- 5 files changed, 22 insertions(+), 42 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 09389b4..cf3d1f8 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -55,22 +55,10 @@ jobs: with: python-version: "3.12" - name: Install ffmpeg - run: sudo apt-get update && sudo apt-get install --yes ffmpeg strace + run: sudo apt-get update && sudo apt-get install --yes ffmpeg - name: Sync dev environment run: uv sync --frozen --extra dev - name: Run full-clip video test - run: | - set +e - timeout --kill-after=10s 90s \ - strace -ff -tt -s 256 -o /tmp/video-e2e.strace \ - uv run pytest -vv -o faulthandler_timeout=45 \ - tests/test_video.py::TestVideoVisibleFullClip - status=$? - if [ "$status" -ne 0 ]; then - for trace in /tmp/video-e2e.strace*; do - echo "::group::$trace" - tail -n 100 "$trace" - echo "::endgroup::" - done - fi - exit "$status" + run: >- + uv run pytest -vv -o faulthandler_timeout=60 + tests/test_video.py::TestVideoVisibleFullClip diff --git a/docs/module-internals.md b/docs/module-internals.md index 4f8019c..4484965 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -149,11 +149,12 @@ regeneration. It centralizes container codecs, optional audio stream copying, metadata/chapter policy, encode-failure reporting, and atomic same-directory publication. Each mapped stream is allowed to reach its own end, so a copied audio tail is not shortened to the frame-input duration. -Both the raw-BGR and timestamped-NUT stdin inputs disable ffmpeg probing before -`pipe:0`: their format is already explicit, and ffmpeg 6 on Linux can otherwise -wait for its normal analysis window while the producer blocks on a full pipe -before the source-audio input opens. Command-order regressions cover both stdin -modes; the real Linux full-clip CI job guards process completion. +Both the raw-BGR and timestamped-NUT stdin modes cap the video encoder at two +threads. A Linux full-clip trace showed ffmpeg creating dozens of worker stacks +and severely delaying frame-pipe ingestion on a constrained hosted runner. The +bounded codec pool avoids that scheduling collapse while leaving audio stream +copy independent. Command regressions cover both supported video codecs; the +real Linux full-clip CI job guards process completion. `probe_video_encode_profile` reads the first source video stream with ffprobe and preserves the supported properties that survive the 8-bit BGR boundary: `yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index 1f129d3..9f18278 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: log = logging.getLogger(__name__) _PIXEL_FORMATS = frozenset({"yuv420p", "yuv422p", "yuv444p"}) +_VIDEO_ENCODER_THREADS = 2 _PIXEL_FORMAT_ALIASES = { "yuvj420p": "yuv420p", "yuvj422p": "yuv422p", @@ -286,12 +287,8 @@ def raw_video_command( ffmpeg = shutil.which("ffmpeg") if ffmpeg is None: raise RuntimeError("Video processing requires ffmpeg on PATH") - # The pipe format and stream geometry are already explicit. FFmpeg 6 can - # otherwise wait for its normal analysis window while the producer blocks - # on a full pipe, before the second (audio) input has been opened. - pipe_input = ["-analyzeduration", "0", "-probesize", "32", "-i", "pipe:0"] frame_input = ( - ["-f", "nut", *pipe_input] + ["-f", "nut", "-i", "pipe:0"] if timestamped_input else [ "-f", @@ -302,7 +299,8 @@ def raw_video_command( f"{width}x{height}", "-r", f"{fps:.12g}", - *pipe_input, + "-i", + "pipe:0", ] ) command = [ @@ -319,6 +317,8 @@ def raw_video_command( "-map", "1:a?", *_video_codec_args(output.suffix.lower(), crf=crf, profile=profile), + "-threads:v", + str(_VIDEO_ENCODER_THREADS), *_profile_args(profile), "-c:a", "copy", diff --git a/tests/test_video.py b/tests/test_video.py index 6df876f..3b3e674 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -1950,17 +1950,19 @@ class TestVideoVisibleScan: class TestVideoVisibleEncoding: - def test_raw_pipe_input_disables_redundant_ffmpeg_probing( + @pytest.mark.parametrize("suffix", [".mp4", ".webm"]) + def test_encoder_bounds_codec_threads( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, + suffix: str, ): from remove_ai_watermarks import video_encoding monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") command = video_encoding.raw_video_command( tmp_path / "source.mp4", - tmp_path / "clean.mp4", + tmp_path / f"clean{suffix}", width=12, height=8, fps=24.0, @@ -1969,14 +1971,7 @@ class TestVideoVisibleEncoding: profile=video_encoding.VideoEncodeProfile(), ) - pipe_position = command.index("pipe:0") - assert command[pipe_position - 5 : pipe_position] == [ - "-analyzeduration", - "0", - "-probesize", - "32", - "-i", - ] + assert command[command.index("-threads:v") + 1] == "2" @staticmethod def _patch_single_frame_encode( diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index 61ddaf3..a28ce00 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -102,13 +102,9 @@ def test_timestamped_encoder_reads_nut_and_passes_pts_through( ) assert "-copyts" in command - assert command[command.index("-f") : command.index("-f") + 8] == [ + assert command[command.index("-f") : command.index("-f") + 4] == [ "-f", "nut", - "-analyzeduration", - "0", - "-probesize", - "32", "-i", "pipe:0", ] From 0fac8cd000a357c50ec6d77f45367f0b7c92e758 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 14:43:01 -0700 Subject: [PATCH 14/19] Bound ffmpeg filter threads --- docs/module-internals.md | 9 +++++---- src/remove_ai_watermarks/video_encoding.py | 3 +++ tests/test_video.py | 1 + 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/docs/module-internals.md b/docs/module-internals.md index 4484965..8919c90 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -149,10 +149,11 @@ regeneration. It centralizes container codecs, optional audio stream copying, metadata/chapter policy, encode-failure reporting, and atomic same-directory publication. Each mapped stream is allowed to reach its own end, so a copied audio tail is not shortened to the frame-input duration. -Both the raw-BGR and timestamped-NUT stdin modes cap the video encoder at two -threads. A Linux full-clip trace showed ffmpeg creating dozens of worker stacks -and severely delaying frame-pipe ingestion on a constrained hosted runner. The -bounded codec pool avoids that scheduling collapse while leaving audio stream +Both the raw-BGR and timestamped-NUT stdin modes run the implicit pixel-format +filter graph on one thread and cap the video encoder at two threads. A Linux +full-clip trace showed ffmpeg creating an oversized execution pool and severely +delaying frame-pipe ingestion on a constrained hosted runner. The bounded +filter and codec pools avoid that scheduling collapse while leaving audio stream copy independent. Command regressions cover both supported video codecs; the real Linux full-clip CI job guards process completion. `probe_video_encode_profile` reads the first source video stream with ffprobe diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index 9f18278..76b19d2 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: log = logging.getLogger(__name__) _PIXEL_FORMATS = frozenset({"yuv420p", "yuv422p", "yuv444p"}) +_VIDEO_FILTER_THREADS = 1 _VIDEO_ENCODER_THREADS = 2 _PIXEL_FORMAT_ALIASES = { "yuvj420p": "yuv420p", @@ -305,6 +306,8 @@ def raw_video_command( ) command = [ ffmpeg, + "-filter_threads", + str(_VIDEO_FILTER_THREADS), "-y", "-loglevel", "error", diff --git a/tests/test_video.py b/tests/test_video.py index 3b3e674..07a1b7f 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -1971,6 +1971,7 @@ class TestVideoVisibleEncoding: profile=video_encoding.VideoEncodeProfile(), ) + assert command[1:3] == ["-filter_threads", "1"] assert command[command.index("-threads:v") + 1] == "2" @staticmethod From 7818e60824743fc8b47b7fec3161fdb31e7add88 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 14:59:57 -0700 Subject: [PATCH 15/19] Prevent ffmpeg stderr backpressure --- docs/module-internals.md | 16 ++-- src/remove_ai_watermarks/video_encoding.py | 97 ++++++++++++++++----- src/remove_ai_watermarks/video_invisible.py | 12 +-- src/remove_ai_watermarks/video_visible.py | 6 +- tests/test_video.py | 30 ++----- tests/test_video_invisible.py | 51 +++++++++++ 6 files changed, 144 insertions(+), 68 deletions(-) diff --git a/docs/module-internals.md b/docs/module-internals.md index 8919c90..93644d1 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -149,13 +149,15 @@ regeneration. It centralizes container codecs, optional audio stream copying, metadata/chapter policy, encode-failure reporting, and atomic same-directory publication. Each mapped stream is allowed to reach its own end, so a copied audio tail is not shortened to the frame-input duration. -Both the raw-BGR and timestamped-NUT stdin modes run the implicit pixel-format -filter graph on one thread and cap the video encoder at two threads. A Linux -full-clip trace showed ffmpeg creating an oversized execution pool and severely -delaying frame-pipe ingestion on a constrained hosted runner. The bounded -filter and codec pools avoid that scheduling collapse while leaving audio stream -copy independent. Command regressions cover both supported video codecs; the -real Linux full-clip CI job guards process completion. +Both the raw-BGR and timestamped-NUT stdin modes redirect ffmpeg stderr to a +temporary file while frames are written. Waiting to read diagnostics until +after stdin closed allowed stderr backpressure to stop ffmpeg's frame reads, +which in turn blocked the producer before it could close stdin. The file consumes +no pipe capacity or RAM while ffmpeg runs; completion reports a bounded head and +tail when diagnostics are unusually large. Aborts release it even when ffmpeg +has already exited. A real subprocess regression writes diagnostics beyond pipe +capacity while streaming frames, checks bounded failure reporting, and the Linux +full-clip CI job guards the complete path. `probe_video_encode_profile` reads the first source video stream with ffprobe and preserves the supported properties that survive the 8-bit BGR boundary: `yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index 76b19d2..acb4363 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -13,16 +13,17 @@ from contextlib import contextmanager from dataclasses import dataclass from fractions import Fraction from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast if TYPE_CHECKING: from collections.abc import Generator, Sequence + from typing import BinaryIO log = logging.getLogger(__name__) +_FFMPEG_STDERR_LIMIT = 64 * 1024 +_FFMPEG_STDERR_TRUNCATION = b"\n...[ffmpeg stderr truncated]...\n" _PIXEL_FORMATS = frozenset({"yuv420p", "yuv422p", "yuv444p"}) -_VIDEO_FILTER_THREADS = 1 -_VIDEO_ENCODER_THREADS = 2 _PIXEL_FORMAT_ALIASES = { "yuvj420p": "yuv420p", "yuvj422p": "yuv422p", @@ -83,6 +84,53 @@ class VideoEncodeProfile: component_depth: int | None = None +@dataclass +class _RawVideoEncoder: + """Running ffmpeg process with diagnostics redirected outside a pipe.""" + + process: subprocess.Popen[bytes] + stdin: BinaryIO + _stderr_buffer: BinaryIO + + def poll(self) -> int | None: + return self.process.poll() + + def kill(self) -> None: + self.process.kill() + + def wait(self) -> int: + return self.process.wait() + + def collect_stderr(self) -> str: + """Return bounded ffmpeg diagnostics and release the temporary file.""" + if self._stderr_buffer.closed: + raise RuntimeError("ffmpeg diagnostics have already been collected") + try: + self._stderr_buffer.seek(0, os.SEEK_END) + size = self._stderr_buffer.tell() + if size <= _FFMPEG_STDERR_LIMIT: + self._stderr_buffer.seek(0) + raw_stderr = self._stderr_buffer.read() + else: + payload_limit = _FFMPEG_STDERR_LIMIT - len(_FFMPEG_STDERR_TRUNCATION) + head_size = payload_limit // 2 + tail_size = payload_limit - head_size + self._stderr_buffer.seek(0) + head = self._stderr_buffer.read(head_size) + self._stderr_buffer.seek(-tail_size, os.SEEK_END) + tail = self._stderr_buffer.read(tail_size) + raw_stderr = head + _FFMPEG_STDERR_TRUNCATION + tail + return raw_stderr.decode("utf-8", errors="replace") + finally: + self._stderr_buffer.close() + + def discard_stderr(self) -> None: + """Release the diagnostic buffer after an abort.""" + if self._stderr_buffer.closed: + return + self._stderr_buffer.close() + + def _known_value(value: object, allowed: frozenset[str]) -> str | None: """Return a supported ffmpeg enum value, otherwise omit it.""" return value if isinstance(value, str) and value in allowed else None @@ -306,8 +354,6 @@ def raw_video_command( ) command = [ ffmpeg, - "-filter_threads", - str(_VIDEO_FILTER_THREADS), "-y", "-loglevel", "error", @@ -320,8 +366,6 @@ def raw_video_command( "-map", "1:a?", *_video_codec_args(output.suffix.lower(), crf=crf, profile=profile), - "-threads:v", - str(_VIDEO_ENCODER_THREADS), *_profile_args(profile), "-c:a", "copy", @@ -344,40 +388,49 @@ def raw_video_command( return command -def start_raw_video_encoder(command: list[str]) -> subprocess.Popen[bytes]: +def start_raw_video_encoder(command: list[str]) -> _RawVideoEncoder: """Start ffmpeg and validate its raw-frame pipes.""" log.info("Starting ffmpeg video encode: command=%s", command) - process = subprocess.Popen( # noqa: S603 - command, - stdin=subprocess.PIPE, - stdout=subprocess.DEVNULL, - stderr=subprocess.PIPE, + stderr_buffer = cast( + "BinaryIO", + tempfile.TemporaryFile(mode="w+b"), # noqa: SIM115 - encoder owns the lifetime ) - if process.stdin is None or process.stderr is None: + try: + process = subprocess.Popen( # noqa: S603 + command, + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=stderr_buffer, + ) + except Exception: + stderr_buffer.close() + raise + if process.stdin is None: process.kill() process.wait() - raise RuntimeError("Could not open ffmpeg pipes") - return process + stderr_buffer.close() + raise RuntimeError("Could not open ffmpeg input pipe") + return _RawVideoEncoder(process, cast("BinaryIO", process.stdin), stderr_buffer) def finish_raw_video_encoder( - process: subprocess.Popen[bytes], + process: _RawVideoEncoder, output: Path, *, operation: str, ) -> None: """Close the frame stream and raise when ffmpeg rejects the encode.""" - if process.stdin is None or process.stderr is None: - raise RuntimeError("ffmpeg pipes are unavailable") process.stdin.close() - stderr = process.stderr.read().decode("utf-8", errors="replace") return_code = process.wait() + stderr = process.collect_stderr() log.info("ffmpeg %s finished: status=%s stderr=%s", operation, return_code, stderr) if return_code != 0: raise RuntimeError(f"ffmpeg failed to encode {output}: {stderr.strip()[:500]}") -def abort_raw_video_encoder(process: subprocess.Popen[bytes]) -> None: +def abort_raw_video_encoder(process: _RawVideoEncoder) -> None: """Stop an incomplete ffmpeg encode.""" - process.kill() + if process.poll() is None: + process.kill() process.wait() + process.discard_stderr() diff --git a/src/remove_ai_watermarks/video_invisible.py b/src/remove_ai_watermarks/video_invisible.py index 6fc1815..832b95d 100644 --- a/src/remove_ai_watermarks/video_invisible.py +++ b/src/remove_ai_watermarks/video_invisible.py @@ -323,9 +323,6 @@ def encode_video_frames( ) ) frame_pipe = process.stdin - if frame_pipe is None: - abort_raw_video_encoder(process) - raise RuntimeError("Could not open ffmpeg input pipe") try: for frame in frames: if frame.shape[:2] != (height, width): @@ -333,8 +330,7 @@ def encode_video_frames( frame_pipe.write(frame.tobytes()) finish_raw_video_encoder(process, output, operation="SynthID removal encode") except Exception: - if process.poll() is None: - abort_raw_video_encoder(process) + abort_raw_video_encoder(process) raise @@ -393,9 +389,6 @@ def regenerate_video_candidate( ) ) frame_pipe = process.stdin - if frame_pipe is None: - abort_raw_video_encoder(process) - raise RuntimeError("Could not open ffmpeg input pipe") frame_count = 0 squared_error = 0.0 pixel_count = 0 @@ -462,8 +455,7 @@ def regenerate_video_candidate( operation="SynthID removal encode", ) except Exception: - if process.poll() is None: - abort_raw_video_encoder(process) + abort_raw_video_encoder(process) raise mse = squared_error / pixel_count diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py index 1eaa288..a6328d5 100644 --- a/src/remove_ai_watermarks/video_visible.py +++ b/src/remove_ai_watermarks/video_visible.py @@ -1363,9 +1363,6 @@ def encode_clean_video( ) ) frame_pipe = process.stdin - if frame_pipe is None: - abort_raw_video_encoder(process) - raise RuntimeError("Could not open ffmpeg input pipe") capture = cv2.VideoCapture(str(source)) if not capture.isOpened(): @@ -1436,8 +1433,7 @@ def encode_clean_video( if timestamped_writer is not None: with suppress(Exception): timestamped_writer.close() - if process.poll() is None: - abort_raw_video_encoder(process) + abort_raw_video_encoder(process) raise finally: capture.release() diff --git a/tests/test_video.py b/tests/test_video.py index 07a1b7f..5a2b894 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -1950,30 +1950,6 @@ class TestVideoVisibleScan: class TestVideoVisibleEncoding: - @pytest.mark.parametrize("suffix", [".mp4", ".webm"]) - def test_encoder_bounds_codec_threads( - self, - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - suffix: str, - ): - from remove_ai_watermarks import video_encoding - - monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") - command = video_encoding.raw_video_command( - tmp_path / "source.mp4", - tmp_path / f"clean{suffix}", - width=12, - height=8, - fps=24.0, - strip_metadata=True, - crf=14, - profile=video_encoding.VideoEncodeProfile(), - ) - - assert command[1:3] == ["-filter_threads", "1"] - assert command[command.index("-threads:v") + 1] == "2" - @staticmethod def _patch_single_frame_encode( monkeypatch: pytest.MonkeyPatch, @@ -2008,6 +1984,12 @@ class TestVideoVisibleEncoding: def poll(self) -> int: return 1 + def wait(self) -> int: + return 1 + + def discard_stderr(self) -> None: + pass + targets: list[Path] = [] def fake_command(_source: Path, target: Path, **_kwargs: object) -> list[str]: diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index a28ce00..735d25b 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -2,6 +2,8 @@ from __future__ import annotations +import sys +import threading from types import SimpleNamespace from typing import TYPE_CHECKING @@ -14,6 +16,55 @@ if TYPE_CHECKING: from pathlib import Path +def test_encoder_redirects_large_stderr_while_frames_are_streaming( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + diagnostic = "synthetic ffmpeg diagnostic" + tail_diagnostic = "synthetic ffmpeg diagnostic tail" + caplog.set_level("INFO", logger=video_encoding.__name__) + command = [ + sys.executable, + "-c", + ( + "import sys; " + f"sys.stderr.buffer.write({diagnostic.encode()!r} + b'x' * 262144 + {tail_diagnostic.encode()!r}); " + "sys.stderr.buffer.flush(); " + "sys.stdin.buffer.read(); " + "raise SystemExit(7)" + ), + ] + encoder = video_encoding.start_raw_video_encoder(command) + write_finished = threading.Event() + write_errors: list[Exception] = [] + + def write_frames() -> None: + try: + encoder.stdin.write(b"f" * 262144) + encoder.stdin.flush() + except Exception as exc: # pragma: no cover - mutation cleanup path + write_errors.append(exc) + finally: + write_finished.set() + + writer = threading.Thread(target=write_frames) + writer.start() + try: + assert write_finished.wait(5), "stderr backpressure blocked the frame producer" + assert write_errors == [] + with pytest.raises(RuntimeError, match=diagnostic): + video_encoding.finish_raw_video_encoder( + encoder, + tmp_path / "unused.mp4", + operation="synthetic encode", + ) + assert "ffmpeg stderr truncated" in caplog.text + assert tail_diagnostic in caplog.text + finally: + video_encoding.abort_raw_video_encoder(encoder) + writer.join(timeout=5) + + def test_availability_requires_both_optional_packages(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( video_invisible, From 9b9f9c6ced67149fd5e9e4d9600c92b6413cc662 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 15:07:15 -0700 Subject: [PATCH 16/19] Open source input before frame pipe --- docs/module-internals.md | 4 ++++ src/remove_ai_watermarks/video_encoding.py | 10 +++++----- tests/test_video_invisible.py | 4 ++++ 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/docs/module-internals.md b/docs/module-internals.md index 93644d1..284f6fd 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -158,6 +158,10 @@ tail when diagnostics are unusually large. Aborts release it even when ffmpeg has already exited. A real subprocess regression writes diagnostics beyond pipe capacity while streaming frames, checks bounded failure reporting, and the Linux full-clip CI job guards the complete path. +The finite source file is opened before the frame pipe, so ffmpeg can initialize +the copied audio stream before producer backpressure is possible. Stream and +metadata mappings are source-indexed accordingly; regressions assert the input +order and both map targets. `probe_video_encode_profile` reads the first source video stream with ffprobe and preserves the supported properties that survive the 8-bit BGR boundary: `yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index acb4363..fe5f5f4 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -358,21 +358,21 @@ def raw_video_command( "-loglevel", "error", *(["-copyts"] if copy_input_timestamps else []), - *frame_input, "-i", str(source), + *frame_input, "-map", - "0:v:0", + "1:v:0", "-map", - "1:a?", + "0:a?", *_video_codec_args(output.suffix.lower(), crf=crf, profile=profile), *_profile_args(profile), "-c:a", "copy", "-map_metadata", - "-1" if strip_metadata else "1", + "-1" if strip_metadata else "0", "-map_chapters", - "-1" if strip_metadata else "1", + "-1" if strip_metadata else "0", ] if timestamped_input: command.extend(["-fps_mode", "passthrough"]) diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index 735d25b..6df1cd2 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -128,6 +128,10 @@ def test_encoder_command_discards_metadata_and_copies_audio( "colorprim=bt709:transfer=bt709:colormatrix=bt709:range=limited" ) assert "pipe:0" in command + assert command.index(str(source)) < command.index("pipe:0") + assert command[command.index("-map") + 1] == "1:v:0" + second_map = command.index("-map", command.index("-map") + 1) + assert command[second_map + 1] == "0:a?" assert "-shortest" not in command From 7e63856a90b432c0d446ae8189586a52d792f32b Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 15:22:41 -0700 Subject: [PATCH 17/19] Split video encoding from audio muxing --- docs/module-internals.md | 20 ++- src/remove_ai_watermarks/video_encoding.py | 145 +++++++++++++++----- src/remove_ai_watermarks/video_invisible.py | 55 ++++---- src/remove_ai_watermarks/video_visible.py | 18 ++- tests/test_video.py | 43 +++++- tests/test_video_invisible.py | 72 ++++++++-- 6 files changed, 261 insertions(+), 92 deletions(-) diff --git a/docs/module-internals.md b/docs/module-internals.md index 284f6fd..b640e00 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -158,10 +158,15 @@ tail when diagnostics are unusually large. Aborts release it even when ffmpeg has already exited. A real subprocess regression writes diagnostics beyond pipe capacity while streaming frames, checks bounded failure reporting, and the Linux full-clip CI job guards the complete path. -The finite source file is opened before the frame pipe, so ffmpeg can initialize -the copied audio stream before producer backpressure is possible. Stream and -metadata mappings are source-indexed accordingly; regressions assert the input -order and both map targets. +Frame encoding and source-audio copying run as two ffmpeg processes in sequence. +The streaming encoder has only the frame pipe as input, so input probing or +demux queues cannot deadlock the producer against a second input. After that +pipe reaches EOF, a finite stream-copy mux combines the encoded video with the +source audio and applies the requested metadata/chapter policy. Both stages use +sibling temporary files, and only the completed mux is published atomically. +The mux also redirects diagnostics to disk and reports only a bounded head and +tail. Command regressions assert the single-input encoder and final map targets; +failure regressions cover bounded mux diagnostics and atomic cleanup. `probe_video_encode_profile` reads the first source video stream with ffprobe and preserves the supported properties that survive the 8-bit BGR boundary: `yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder @@ -196,9 +201,10 @@ space, applies one seeded spatial-noise field across the entire sequence, and decodes fresh pixels. Reusing a single noise field avoids independent frame-to-frame noise. The shipped path retains only one configured frame batch, updates PSNR and temporal residuals incrementally, and streams BGR frames -directly to ffmpeg. ffmpeg encodes H.264 video, maps optional source audio, and -drops all source metadata. The result is written through a same-directory -temporary file and atomically replaced only after a successful encode. +directly to the video-only ffmpeg encoder. A separate stream-copy mux then adds +optional source audio and drops all source metadata. The result is written +through same-directory temporary files and atomically replaced only after both +stages succeed. The engine returns PSNR and a motion-compensated temporal-residual ratio as quality measurements. Neither is a watermark detector. The high-level result diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index fe5f5f4..36ce507 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -106,21 +106,7 @@ class _RawVideoEncoder: if self._stderr_buffer.closed: raise RuntimeError("ffmpeg diagnostics have already been collected") try: - self._stderr_buffer.seek(0, os.SEEK_END) - size = self._stderr_buffer.tell() - if size <= _FFMPEG_STDERR_LIMIT: - self._stderr_buffer.seek(0) - raw_stderr = self._stderr_buffer.read() - else: - payload_limit = _FFMPEG_STDERR_LIMIT - len(_FFMPEG_STDERR_TRUNCATION) - head_size = payload_limit // 2 - tail_size = payload_limit - head_size - self._stderr_buffer.seek(0) - head = self._stderr_buffer.read(head_size) - self._stderr_buffer.seek(-tail_size, os.SEEK_END) - tail = self._stderr_buffer.read(tail_size) - raw_stderr = head + _FFMPEG_STDERR_TRUNCATION + tail - return raw_stderr.decode("utf-8", errors="replace") + return _read_bounded_stderr(self._stderr_buffer) finally: self._stderr_buffer.close() @@ -131,6 +117,25 @@ class _RawVideoEncoder: self._stderr_buffer.close() +def _read_bounded_stderr(buffer: BinaryIO) -> str: + """Read bounded head-and-tail diagnostics from a seekable binary stream.""" + buffer.seek(0, os.SEEK_END) + size = buffer.tell() + if size <= _FFMPEG_STDERR_LIMIT: + buffer.seek(0) + raw_stderr = buffer.read() + else: + payload_limit = _FFMPEG_STDERR_LIMIT - len(_FFMPEG_STDERR_TRUNCATION) + head_size = payload_limit // 2 + tail_size = payload_limit - head_size + buffer.seek(0) + head = buffer.read(head_size) + buffer.seek(-tail_size, os.SEEK_END) + tail = buffer.read(tail_size) + raw_stderr = head + _FFMPEG_STDERR_TRUNCATION + tail + return raw_stderr.decode("utf-8", errors="replace") + + def _known_value(value: object, allowed: frozenset[str]) -> str | None: """Return a supported ffmpeg enum value, otherwise omit it.""" return value if isinstance(value, str) and value in allowed else None @@ -268,11 +273,11 @@ def probe_video_timestamps(source: Path) -> tuple[float, ...]: @contextmanager -def atomic_video_output(output: Path) -> Generator[Path]: - """Yield a sibling temporary path and publish it only after success.""" +def _temporary_video_path(output: Path, *, prefix: str) -> Generator[Path]: + """Yield one sibling temporary path and remove it on exit.""" output.parent.mkdir(parents=True, exist_ok=True) with tempfile.NamedTemporaryFile( - prefix=f".{output.stem}-", + prefix=prefix, suffix=output.suffix, dir=output.parent, delete=False, @@ -280,11 +285,18 @@ def atomic_video_output(output: Path) -> Generator[Path]: temporary_output = Path(stream.name) try: yield temporary_output - os.replace(temporary_output, output) finally: temporary_output.unlink(missing_ok=True) +@contextmanager +def atomic_video_output(output: Path) -> Generator[Path]: + """Yield a sibling temporary path and publish it only after success.""" + with _temporary_video_path(output, prefix=f".{output.stem}-") as temporary_output: + yield temporary_output + os.replace(temporary_output, output) + + def _video_codec_args(suffix: str, *, crf: int, profile: VideoEncodeProfile) -> list[str]: if suffix == ".webm": return ["-c:v", "libvpx-vp9", "-crf", str(crf), "-b:v", "0"] @@ -320,19 +332,17 @@ def _profile_args(profile: VideoEncodeProfile) -> list[str]: def raw_video_command( - source: Path, output: Path, *, width: int, height: int, fps: float, - strip_metadata: bool, crf: int, profile: VideoEncodeProfile, timestamped_input: bool = False, copy_input_timestamps: bool = False, ) -> list[str]: - """Build a source-aware ffmpeg command for BGR frames on standard input.""" + """Build an ffmpeg command for BGR frames on standard input.""" ffmpeg = shutil.which("ffmpeg") if ffmpeg is None: raise RuntimeError("Video processing requires ffmpeg on PATH") @@ -358,36 +368,97 @@ def raw_video_command( "-loglevel", "error", *(["-copyts"] if copy_input_timestamps else []), - "-i", - str(source), *frame_input, "-map", - "1:v:0", - "-map", - "0:a?", + "0:v:0", *_video_codec_args(output.suffix.lower(), crf=crf, profile=profile), *_profile_args(profile), - "-c:a", - "copy", "-map_metadata", - "-1" if strip_metadata else "0", + "-1", "-map_chapters", - "-1" if strip_metadata else "0", + "-1", ] if timestamped_input: command.extend(["-fps_mode", "passthrough"]) if copy_input_timestamps: command.extend(["-avoid_negative_ts", "disabled"]) - if output.suffix.lower() in {".mp4", ".mov", ".m4v"}: - if profile.time_base is not None: - numerator, denominator = (int(part) for part in profile.time_base.split("/", 1)) - if numerator == 1: - command.extend(["-video_track_timescale", str(denominator)]) - command.extend(["-movflags", "+faststart"]) + if output.suffix.lower() in {".mp4", ".mov", ".m4v"} and profile.time_base is not None: + numerator, denominator = (int(part) for part in profile.time_base.split("/", 1)) + if numerator == 1: + command.extend(["-video_track_timescale", str(denominator)]) command.append(str(output)) return command +def mux_encoded_video( + encoded_video: Path, + source: Path, + output: Path, + *, + strip_metadata: bool, + copy_input_timestamps: bool = False, +) -> None: + """Copy encoded video and source audio into the final container.""" + ffmpeg = shutil.which("ffmpeg") + if ffmpeg is None: + raise RuntimeError("Video processing requires ffmpeg on PATH") + command = [ + ffmpeg, + "-y", + "-loglevel", + "error", + *(["-copyts"] if copy_input_timestamps else []), + "-i", + str(encoded_video), + "-i", + str(source), + "-map", + "0:v:0", + "-map", + "1:a?", + "-c", + "copy", + "-map_metadata", + "-1" if strip_metadata else "1", + "-map_chapters", + "-1" if strip_metadata else "1", + ] + if copy_input_timestamps: + command.extend(["-avoid_negative_ts", "disabled"]) + if output.suffix.lower() in {".mp4", ".mov", ".m4v"}: + command.extend(["-movflags", "+faststart"]) + command.append(str(output)) + with tempfile.TemporaryFile(mode="w+b") as stderr_buffer: + result = subprocess.run( # noqa: S603 + command, + stdout=subprocess.DEVNULL, + stderr=stderr_buffer, + check=False, + ) + stderr = _read_bounded_stderr(stderr_buffer) + log.info( + "ffmpeg video mux: command=%s status=%s stderr=%s", + command, + result.returncode, + stderr, + ) + if result.returncode != 0: + raise RuntimeError(f"ffmpeg failed to mux {output}: {stderr.strip()[:500]}") + + +@contextmanager +def staged_video_output(output: Path) -> Generator[tuple[Path, Path]]: + """Yield video-only and final temporary paths, then publish atomically.""" + with ( + atomic_video_output(output) as temporary_output, + _temporary_video_path( + output, + prefix=f".{output.stem}-video-", + ) as encoded_video, + ): + yield encoded_video, temporary_output + + def start_raw_video_encoder(command: list[str]) -> _RawVideoEncoder: """Start ffmpeg and validate its raw-frame pipes.""" log.info("Starting ffmpeg video encode: command=%s", command) diff --git a/src/remove_ai_watermarks/video_invisible.py b/src/remove_ai_watermarks/video_invisible.py index 832b95d..45da702 100644 --- a/src/remove_ai_watermarks/video_invisible.py +++ b/src/remove_ai_watermarks/video_invisible.py @@ -21,10 +21,11 @@ import numpy as np from remove_ai_watermarks.video_encoding import ( abort_raw_video_encoder, - atomic_video_output, finish_raw_video_encoder, + mux_encoded_video, probe_video_encode_profile, raw_video_command, + staged_video_output, start_raw_video_encoder, ) from remove_ai_watermarks.video_synthid import ( @@ -309,29 +310,28 @@ def encode_video_frames( if not frames: raise ValueError("At least one frame is required for video encoding") height, width = frames[0].shape[:2] - output.parent.mkdir(parents=True, exist_ok=True) - process = start_raw_video_encoder( - raw_video_command( - source, - output, - width=width, - height=height, - fps=fps, - strip_metadata=True, - crf=18, - profile=probe_video_encode_profile(source), + with staged_video_output(output) as (encoded_video, temporary_output): + process = start_raw_video_encoder( + raw_video_command( + encoded_video, + width=width, + height=height, + fps=fps, + crf=18, + profile=probe_video_encode_profile(source), + ) ) - ) - frame_pipe = process.stdin - try: - for frame in frames: - if frame.shape[:2] != (height, width): - raise ValueError("Video frames must have matching dimensions") - frame_pipe.write(frame.tobytes()) - finish_raw_video_encoder(process, output, operation="SynthID removal encode") - except Exception: - abort_raw_video_encoder(process) - raise + frame_pipe = process.stdin + try: + for frame in frames: + if frame.shape[:2] != (height, width): + raise ValueError("Video frames must have matching dimensions") + frame_pipe.write(frame.tobytes()) + finish_raw_video_encoder(process, encoded_video, operation="SynthID removal encode") + mux_encoded_video(encoded_video, source, temporary_output, strip_metadata=True) + except Exception: + abort_raw_video_encoder(process) + raise def regenerate_video_candidate( @@ -375,15 +375,13 @@ def regenerate_video_candidate( resolved_device = runtime.resolved_device vae = runtime.vae - with atomic_video_output(output) as temporary_output: + with staged_video_output(output) as (encoded_video, temporary_output): process = start_raw_video_encoder( raw_video_command( - source, - temporary_output, + encoded_video, width=size[0], height=size[1], fps=effective_fps, - strip_metadata=True, crf=18, profile=probe_video_encode_profile(source), ) @@ -451,9 +449,10 @@ def regenerate_video_candidate( raise ValueError("The selected clip produced fewer than two frames") finish_raw_video_encoder( process, - temporary_output, + encoded_video, operation="SynthID removal encode", ) + mux_encoded_video(encoded_video, source, temporary_output, strip_metadata=True) except Exception: abort_raw_video_encoder(process) raise diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py index a6328d5..649f514 100644 --- a/src/remove_ai_watermarks/video_visible.py +++ b/src/remove_ai_watermarks/video_visible.py @@ -39,11 +39,12 @@ from PIL import Image, ImageDraw, ImageFont from remove_ai_watermarks.video import VIDEO_VISIBLE_MARKS from remove_ai_watermarks.video_encoding import ( abort_raw_video_encoder, - atomic_video_output, finish_raw_video_encoder, + mux_encoded_video, probe_video_encode_profile, probe_video_timestamps, raw_video_command, + staged_video_output, start_raw_video_encoder, ) from remove_ai_watermarks.video_temporal import stabilize_filled_frame @@ -1331,7 +1332,7 @@ def encode_clean_video( if len(regions) != len(scan.detections): raise ValueError("Temporal localization count does not match the scanned frame count") - with atomic_video_output(output) as temporary_output: + with staged_video_output(output) as (encoded_video, temporary_output): profile = probe_video_encode_profile(source) if (profile.component_depth or 0) > 8 or profile.color_transfer in _HDR_TRANSFERS: source_format = profile.source_pixel_format or "unknown pixel format" @@ -1350,12 +1351,10 @@ def encode_clean_video( ) process = start_raw_video_encoder( raw_video_command( - source, - temporary_output, + encoded_video, width=scan.width, height=scan.height, fps=scan.fps, - strip_metadata=strip_metadata, crf=14, profile=profile, timestamped_input=timestamped_input, @@ -1426,9 +1425,16 @@ def encode_clean_video( timestamped_writer = None finish_raw_video_encoder( process, - temporary_output, + encoded_video, operation="visible-watermark encode", ) + mux_encoded_video( + encoded_video, + source, + temporary_output, + strip_metadata=strip_metadata, + copy_input_timestamps=preserve_start_offset, + ) except Exception: if timestamped_writer is not None: with suppress(Exception): diff --git a/tests/test_video.py b/tests/test_video.py index 5a2b894..927808f 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -1992,7 +1992,7 @@ class TestVideoVisibleEncoding: targets: list[Path] = [] - def fake_command(_source: Path, target: Path, **_kwargs: object) -> list[str]: + def fake_command(target: Path, **_kwargs: object) -> list[str]: targets.append(target) return ["ffmpeg", str(target)] @@ -2001,11 +2001,15 @@ class TestVideoVisibleEncoding: if fail: raise RuntimeError("synthetic encode failure") + def fake_mux(encoded: Path, _source: Path, target: Path, **_kwargs: object) -> None: + target.write_bytes(encoded.read_bytes()) + process = FakeProcess() monkeypatch.setattr(video_visible.cv2, "VideoCapture", lambda _path: FakeCapture()) monkeypatch.setattr(video_visible, "raw_video_command", fake_command) monkeypatch.setattr(video_visible, "start_raw_video_encoder", lambda _command: process) monkeypatch.setattr(video_visible, "finish_raw_video_encoder", fake_finish) + monkeypatch.setattr(video_visible, "mux_encoded_video", fake_mux) monkeypatch.setattr(watermark_registry, "resolve_backend", lambda _backend: "cv2") return process, targets @@ -2071,6 +2075,43 @@ class TestVideoVisibleEncoding: assert output.read_bytes() == b"previous" assert not targets[0].exists() + def test_failed_mux_preserves_existing_output( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + from remove_ai_watermarks import video_visible + from remove_ai_watermarks.video_visible import FrameLocalization, VideoScan, encode_clean_video + + source = tmp_path / "source.mp4" + source.write_bytes(b"source") + output = tmp_path / "clean.mp4" + output.write_bytes(b"previous") + scan = VideoScan(8, 8, 24.0, (FrameLocalization(0, 0.0, None),)) + _process, targets = self._patch_single_frame_encode( + monkeypatch, + encoded_bytes=b"complete video stream", + fail=False, + ) + + def fail_mux(*_args: object, **_kwargs: object) -> None: + raise RuntimeError("synthetic mux failure") + + monkeypatch.setattr(video_visible, "mux_encoded_video", fail_mux) + + with pytest.raises(RuntimeError, match="synthetic mux failure"): + encode_clean_video( + source, + output, + scan, + [None], + backend="cv2", + strip_metadata=True, + ) + + assert output.read_bytes() == b"previous" + assert not targets[0].exists() + def test_rejects_high_bit_depth_before_silent_downconversion( self, tmp_path: Path, diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index 6df1cd2..abc030d 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -5,7 +5,7 @@ from __future__ import annotations import sys import threading from types import SimpleNamespace -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import pytest @@ -14,6 +14,7 @@ from remove_ai_watermarks.video_synthid import DEFAULT_VIDEO_SYNTHID_NOISE_STD if TYPE_CHECKING: from pathlib import Path + from typing import BinaryIO def test_encoder_redirects_large_stderr_while_frames_are_streaming( @@ -84,7 +85,7 @@ def test_regeneration_rejects_noise_outside_unit_interval(tmp_path: Path) -> Non ) -def test_encoder_command_discards_metadata_and_copies_audio( +def test_encoder_and_mux_commands_separate_streaming_frames_from_source_audio( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -102,20 +103,16 @@ def test_encoder_command_discards_metadata_and_copies_audio( ) command = video_encoding.raw_video_command( - source, output, width=8, height=8, fps=2.0, - strip_metadata=True, crf=18, profile=profile, ) metadata_index = command.index("-map_metadata") assert command[metadata_index + 1] == "-1" - audio_codec_index = command.index("-c:a") - assert command[audio_codec_index + 1] == "copy" output_pixel_format_index = command.index("-pix_fmt", command.index("-c:v")) assert command[output_pixel_format_index + 1] == "yuv420p" assert command[command.index("-color_range") + 1] == "tv" @@ -128,28 +125,77 @@ def test_encoder_command_discards_metadata_and_copies_audio( "colorprim=bt709:transfer=bt709:colormatrix=bt709:range=limited" ) assert "pipe:0" in command - assert command.index(str(source)) < command.index("pipe:0") - assert command[command.index("-map") + 1] == "1:v:0" - second_map = command.index("-map", command.index("-map") + 1) - assert command[second_map + 1] == "0:a?" + assert str(source) not in command + assert command[command.index("-map") + 1] == "0:v:0" assert "-shortest" not in command + calls: list[list[str]] = [] + + def fake_run(mux_command: list[str], **_kwargs: object) -> SimpleNamespace: + calls.append(mux_command) + return SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(video_encoding.subprocess, "run", fake_run) + encoded_video = tmp_path / "encoded.mp4" + video_encoding.mux_encoded_video( + encoded_video, + source, + output, + strip_metadata=True, + ) + + mux_command = calls[0] + assert mux_command.index(str(encoded_video)) < mux_command.index(str(source)) + assert mux_command[mux_command.index("-map") + 1] == "0:v:0" + second_map = mux_command.index("-map", mux_command.index("-map") + 1) + assert mux_command[second_map + 1] == "1:a?" + assert mux_command[mux_command.index("-c") + 1] == "copy" + assert mux_command[mux_command.index("-map_metadata") + 1] == "-1" + assert "-shortest" not in mux_command + + +def test_mux_reports_bounded_disk_backed_stderr( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + diagnostic = b"synthetic mux diagnostic" + tail_diagnostic = b"synthetic mux diagnostic tail" + caplog.set_level("INFO", logger=video_encoding.__name__) + monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") + + def fake_run(_command: list[str], **kwargs: object) -> SimpleNamespace: + stderr = cast("BinaryIO", kwargs["stderr"]) + stderr.write(diagnostic + b"x" * 262144 + tail_diagnostic) + stderr.flush() + return SimpleNamespace(returncode=7) + + monkeypatch.setattr(video_encoding.subprocess, "run", fake_run) + + with pytest.raises(RuntimeError, match=diagnostic.decode()): + video_encoding.mux_encoded_video( + tmp_path / "encoded.mp4", + tmp_path / "source.mp4", + tmp_path / "output.mp4", + strip_metadata=True, + ) + + assert "ffmpeg stderr truncated" in caplog.text + assert tail_diagnostic.decode() in caplog.text + def test_timestamped_encoder_reads_nut_and_passes_pts_through( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - source = tmp_path / "source.mp4" output = tmp_path / "candidate.mp4" monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") command = video_encoding.raw_video_command( - source, output, width=8, height=8, fps=24.0, - strip_metadata=True, crf=18, profile=video_encoding.VideoEncodeProfile(time_base="1/90000"), timestamped_input=True, From 7c922e5133c697f07fc1d16f7f698418878ce3f3 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 15:31:00 -0700 Subject: [PATCH 18/19] Preserve one encoded frame per input frame --- docs/module-internals.md | 4 ++++ src/remove_ai_watermarks/video_encoding.py | 3 +-- tests/test_video_invisible.py | 1 + 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/module-internals.md b/docs/module-internals.md index b640e00..9c401e7 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -171,6 +171,10 @@ failure regressions cover bounded mux diagnostics and atomic cleanup. and preserves the supported properties that survive the 8-bit BGR boundary: `yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder time base, MP4/MOV track timescale, source pixel format, and component depth. +Both raw-CFR and timestamped-NUT inputs use ffmpeg's passthrough FPS mode. This +keeps one encoded frame per supplied frame when an older ffmpeg receives a +fine-grained source encoder time base such as `1/90000`; implicit synchronization +can otherwise synthesize thousands of duplicate frames between CFR timestamps. HDR transfer functions and component depths above 8 bits are rejected before encoding so the OpenCV boundary cannot silently reduce them to SDR 8-bit. `probe_video_timestamps` reads authoritative per-frame display PTS through diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index 36ce507..e89da14 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -378,8 +378,7 @@ def raw_video_command( "-map_chapters", "-1", ] - if timestamped_input: - command.extend(["-fps_mode", "passthrough"]) + command.extend(["-fps_mode", "passthrough"]) if copy_input_timestamps: command.extend(["-avoid_negative_ts", "disabled"]) if output.suffix.lower() in {".mp4", ".mov", ".m4v"} and profile.time_base is not None: diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index abc030d..60e520a 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -127,6 +127,7 @@ def test_encoder_and_mux_commands_separate_streaming_frames_from_source_audio( assert "pipe:0" in command assert str(source) not in command assert command[command.index("-map") + 1] == "0:v:0" + assert command[command.index("-fps_mode") + 1] == "passthrough" assert "-shortest" not in command calls: list[list[str]] = [] From a6c0c1c6f0cd929a7e43bc1f86f5c28f30eb6d0e Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 16:53:41 -0700 Subject: [PATCH 19/19] Rewrite internal watermark pipeline and preserve behavior --- docs/code-provenance.md | 9 + docs/controlnet-removal-pipeline-research.md | 2 +- docs/index.md | 1 + docs/known-limitations.md | 2 +- docs/module-internals.md | 43 +- docs/qwen-improvement-research.md | 16 +- docs/research-doubao-distillation.md | 2 +- docs/synthid.md | 2 +- pyproject.toml | 6 +- scripts/metadata_removal_audit.py | 2 +- scripts/smoke_matrix.py | 2 +- scripts/synthid_corpus.py | 2 +- scripts/visible_positives.py | 2 +- .../{noai => _internal}/__init__.py | 12 +- src/remove_ai_watermarks/_internal/c2pa.py | 363 ++++++ .../_internal/constants.py | 152 +++ .../{noai => _internal}/ebml.py | 0 .../_internal/extractor.py | 98 ++ .../{noai => _internal}/flv.py | 0 .../_internal/img2img_runner.py | 136 ++ .../{noai => _internal}/isobmff.py | 0 .../_internal/progress.py | 212 ++++ .../qwen_zimage_pipeline.py | 82 +- .../{noai => _internal}/riff.py | 0 .../{noai => _internal}/tiling.py | 0 src/remove_ai_watermarks/_internal/utils.py | 30 + .../_internal/watermark_profiles.py | 105 ++ .../_internal/watermark_remover.py | 649 ++++++++++ src/remove_ai_watermarks/cli.py | 8 +- src/remove_ai_watermarks/gemini_engine.py | 813 ++++-------- src/remove_ai_watermarks/humanizer.py | 35 +- src/remove_ai_watermarks/identify.py | 26 +- src/remove_ai_watermarks/invisible_engine.py | 14 +- src/remove_ai_watermarks/metadata.py | 33 +- src/remove_ai_watermarks/noai/c2pa.py | 562 --------- src/remove_ai_watermarks/noai/constants.py | 342 ----- src/remove_ai_watermarks/noai/extractor.py | 155 --- .../noai/img2img_runner.py | 161 --- src/remove_ai_watermarks/noai/progress.py | 332 ----- src/remove_ai_watermarks/noai/utils.py | 43 - .../noai/watermark_profiles.py | 207 --- .../noai/watermark_remover.py | 1117 ----------------- src/remove_ai_watermarks/video.py | 2 +- tests/test_cpu_offload.py | 2 +- tests/test_identify.py | 2 +- tests/test_img2img_runner.py | 4 +- tests/test_invisible_engine.py | 10 +- tests/test_metadata.py | 26 +- ...est_noai.py => test_metadata_internals.py} | 54 +- tests/test_platform.py | 59 +- tests/test_qwen_zimage_pipeline.py | 94 +- tests/test_security_clamp.py | 2 +- tests/test_tiling.py | 2 +- tests/test_watermark_profiles.py | 2 +- 54 files changed, 2266 insertions(+), 3771 deletions(-) create mode 100644 docs/code-provenance.md rename src/remove_ai_watermarks/{noai => _internal}/__init__.py (74%) create mode 100644 src/remove_ai_watermarks/_internal/c2pa.py create mode 100644 src/remove_ai_watermarks/_internal/constants.py rename src/remove_ai_watermarks/{noai => _internal}/ebml.py (100%) create mode 100644 src/remove_ai_watermarks/_internal/extractor.py rename src/remove_ai_watermarks/{noai => _internal}/flv.py (100%) create mode 100644 src/remove_ai_watermarks/_internal/img2img_runner.py rename src/remove_ai_watermarks/{noai => _internal}/isobmff.py (100%) create mode 100644 src/remove_ai_watermarks/_internal/progress.py rename src/remove_ai_watermarks/{noai => _internal}/qwen_zimage_pipeline.py (90%) rename src/remove_ai_watermarks/{noai => _internal}/riff.py (100%) rename src/remove_ai_watermarks/{noai => _internal}/tiling.py (100%) create mode 100644 src/remove_ai_watermarks/_internal/utils.py create mode 100644 src/remove_ai_watermarks/_internal/watermark_profiles.py create mode 100644 src/remove_ai_watermarks/_internal/watermark_remover.py delete mode 100644 src/remove_ai_watermarks/noai/c2pa.py delete mode 100644 src/remove_ai_watermarks/noai/constants.py delete mode 100644 src/remove_ai_watermarks/noai/extractor.py delete mode 100644 src/remove_ai_watermarks/noai/img2img_runner.py delete mode 100644 src/remove_ai_watermarks/noai/progress.py delete mode 100644 src/remove_ai_watermarks/noai/utils.py delete mode 100644 src/remove_ai_watermarks/noai/watermark_profiles.py delete mode 100644 src/remove_ai_watermarks/noai/watermark_remover.py rename tests/{test_noai.py => test_metadata_internals.py} (93%) diff --git a/docs/code-provenance.md b/docs/code-provenance.md new file mode 100644 index 0000000..0d7e98a --- /dev/null +++ b/docs/code-provenance.md @@ -0,0 +1,9 @@ +# Code provenance + +This page records notices required by source dependencies and licensed derivative work. + +## Licensed derivative work + +- The DWT-DCT implementation derives from ShieldMnt's + [`invisible-watermark`](https://github.com/ShieldMnt/invisible-watermark), licensed + under MIT. Its notice ships in `src/remove_ai_watermarks/licenses/invisible-watermark-MIT.txt`. diff --git a/docs/controlnet-removal-pipeline-research.md b/docs/controlnet-removal-pipeline-research.md index 65db754..b9f51d7 100644 --- a/docs/controlnet-removal-pipeline-research.md +++ b/docs/controlnet-removal-pipeline-research.md @@ -211,7 +211,7 @@ regeneration so strokes exceed the VAE's ~8 px latent floor), but ctrlregen runs at LOW res, the opposite. CtrlRegen's paper gives no resolution/tiling spec to contradict this. **Sources.** the former internal -`src/remove_ai_watermarks/noai/ctrlregen/engine.py` (removed after this study); +`src/remove_ai_watermarks/_internal/ctrlregen/engine.py` (removed after this study); resolution-omission confirmed against https://arxiv.org/html/2410.05470v1 diff --git a/docs/index.md b/docs/index.md index f096e33..ba3810f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -19,6 +19,7 @@ to run the tool. Use the maintainer references only when changing the code. | Page | Purpose | | --- | --- | | [Module internals](module-internals.md) | Current architecture, invariants, and regression guards by module. | +| [Code provenance](code-provenance.md) | Required notices for licensed derivative work. | | [Verification plan](verification-plan.md) | Verification methods, completed measurements, and remaining validation gaps. | | [Release and distribution](release-and-distribution.md) | PyPI, Homebrew, Hugging Face Space, and release workflow. | | [Watermarking landscape](watermarking-landscape.md) | Vendor signals and detection approaches. | diff --git a/docs/known-limitations.md b/docs/known-limitations.md index c272dcc..b06d4a6 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -133,7 +133,7 @@ or a different random seed may change the verifier result. The base Qwen and `qwen-zimage` profiles have profile specific strength behavior. Consult `remove-ai-watermarks invisible --help` and the source of -[`watermark_profiles.py`](../src/remove_ai_watermarks/noai/watermark_profiles.py) +[`watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py) for the current resolver. ### Pipelines have different quality tradeoffs diff --git a/docs/module-internals.md b/docs/module-internals.md index 9c401e7..7caa042 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -119,7 +119,7 @@ reported per file without retrying the same multi-GB initialization. Native MP4/MOV TC260 labels follow TC260-PG-20257A: `moov.udta.meta.keys` maps an `AIGC` key to a raw JSON value in `ilst`. -[`noai/isobmff.py`](../src/remove_ai_watermarks/noai/isobmff.py) walks those +[`_internal/isobmff.py`](../src/remove_ai_watermarks/_internal/isobmff.py) walks those nested boxes by seeking, so detection reaches a tail `moov` without reading the preceding `mdat`. The MP4/MOV/M4V/M4A removal path first validates the top-level box walk, then copies the source to a sibling temporary file in bounded chunks. @@ -130,14 +130,14 @@ validated JSON value with same-length spaces. This preserves every box size, Publication is atomic, and a malformed top-level walk is copied unchanged. A generic `AIGC` key whose value has no TC260 field is ignored. -[`noai/ebml.py`](../src/remove_ai_watermarks/noai/ebml.py) provides the +[`_internal/ebml.py`](../src/remove_ai_watermarks/_internal/ebml.py) provides the corresponding bounded Matroska/WebM reader. It seeks over clusters and accepts only a `Segment.Tags.Tag.SimpleTag` pairing `TagName=AIGC` with a JSON `TagString` carrying a TC260 field. The existing ffmpeg stream-copy path removes those container tags without transcoding the encoded streams. -[`noai/riff.py`](../src/remove_ai_watermarks/noai/riff.py) and -[`noai/flv.py`](../src/remove_ai_watermarks/noai/flv.py) implement the remaining +[`_internal/riff.py`](../src/remove_ai_watermarks/_internal/riff.py) and +[`_internal/flv.py`](../src/remove_ai_watermarks/_internal/flv.py) implement the remaining normative TC260 video placements. The RIFF walker reads only AVI `LIST/INFO/AIGC` children. The FLV walker skips media tags and parses the AMF0 `script.onMetaData.AIGC` string. Both require a recognized TC260 JSON field and @@ -284,12 +284,12 @@ Regression coverage: ### C2PA -[`noai/c2pa.py`](../src/remove_ai_watermarks/noai/c2pa.py) reads C2PA with the +[`_internal/c2pa.py`](../src/remove_ai_watermarks/_internal/c2pa.py) reads C2PA with the official `c2pa-python` reader first. Its byte-level PNG parser remains a fallback for partial and synthetic fixtures that the official reader rejects. Vendor attribution comes from the registry in -[`noai/constants.py`](../src/remove_ai_watermarks/noai/constants.py). Derived +[`_internal/constants.py`](../src/remove_ai_watermarks/_internal/constants.py). Derived issuer and platform maps should not be maintained separately. ### Metadata scanning and stripping @@ -303,7 +303,7 @@ Key contracts: - JPEG stripping walks metadata segments and preserves the entropy-coded image scan. - ISOBMFF containers use - [`noai/isobmff.py`](../src/remove_ai_watermarks/noai/isobmff.py). + [`_internal/isobmff.py`](../src/remove_ai_watermarks/_internal/isobmff.py). - Native MP4/MOV TC260 `AIGC` entries are read from `moov.udta.meta.keys/ilst` and blanked without changing box sizes. - Native MKV/WebM TC260 `AIGC` entries are read from @@ -327,7 +327,7 @@ test proves that it no longer appears in the output. Regression coverage: - [`test_metadata.py`](../tests/test_metadata.py) -- [`test_noai.py`](../tests/test_noai.py) +- [`test_metadata_internals.py`](../tests/test_metadata_internals.py) - [`test_security_clamp.py`](../tests/test_security_clamp.py) ### Provenance report @@ -474,7 +474,7 @@ Regression coverage: ### Profiles and strength -[`noai/watermark_profiles.py`](../src/remove_ai_watermarks/noai/watermark_profiles.py) +[`_internal/watermark_profiles.py`](../src/remove_ai_watermarks/_internal/watermark_profiles.py) is the source of truth for: - profile aliases; @@ -494,12 +494,19 @@ router. [`invisible_engine.py`](../src/remove_ai_watermarks/invisible_engine.py) handles image sizing, optional pre-upscaling, postprocessing, and the public engine interface. It delegates model execution to -[`noai/watermark_remover.py`](../src/remove_ai_watermarks/noai/watermark_remover.py). +[`_internal/watermark_remover.py`](../src/remove_ai_watermarks/_internal/watermark_remover.py). The Python engine and CLI do not have identical defaults for every optional postprocessing argument. Integrations that require reproducibility should pass the relevant values explicitly. +The standard Qwen and ControlNet prompts are calibrated model inputs, and the +ControlNet edge map uses fixed Canny thresholds of 100 and 200. Treat those +values as behavioral compatibility contracts: a refactor must preserve them, +and any deliberate change requires image-quality evaluation rather than only a +unit-test pass. Exact prompt and edge-map regression guards live in +`test_platform.py` and `test_invisible_engine.py`. + Regression coverage: - [`test_watermark_profiles.py`](../tests/test_watermark_profiles.py) @@ -519,7 +526,7 @@ Regression coverage: ### Qwen plus Z-Image -[`noai/qwen_zimage_pipeline.py`](../src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py) +[`_internal/qwen_zimage_pipeline.py`](../src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py) implements the fixed CUDA-only two-stage profile: 1. Qwen Image with Canny conditioning regenerates the frame. @@ -530,13 +537,11 @@ The profile rejects a custom model identifier. Its global and face model stack is fixed by the implementation. When tiling is enabled, only the global stage is tiled; the face stage runs once after the tiles are blended. -The resolution and largest-face adaptive formulas remain exact ports of the -reference workflow. The face stage applies half the reference result because -this port uses a different sampler and composites regenerated SAM pixels rather -than using the reference latent inpaint mask and noise feather. Paired face -evaluations favored this scale on identity, perceptual distance, and full-image -similarity, and the exact OpenAI and Gemini candidates both passed their -matching provider oracle. The global stage stays unchanged. +The maintained implementation preserves the previously oracle-tested strength, +conditioning, crop, and sampler parameters as compatibility contracts. Its Python +orchestration, YuNet integration, SAM selection, masks, sizing helpers, and pixel +compositing are implemented for this runtime. Changing a calibrated model input +requires the same provider-oracle and identity evaluation as a model change. Regression coverage: @@ -545,7 +550,7 @@ Regression coverage: ### Tiling -[`noai/tiling.py`](../src/remove_ai_watermarks/noai/tiling.py) contains pure +[`_internal/tiling.py`](../src/remove_ai_watermarks/_internal/tiling.py) contains pure tile planning, feather weights, tile orchestration, and region compositing. Tiling engages only when requested and the long side exceeds the tile size. diff --git a/docs/qwen-improvement-research.md b/docs/qwen-improvement-research.md index 90e2f5e..9bb3c8e 100644 --- a/docs/qwen-improvement-research.md +++ b/docs/qwen-improvement-research.md @@ -79,13 +79,13 @@ Measured on `gemini_3` (18 faces) at the Gemini scrub floor 0.25 vs base-Qwen 0. a Qwen face fix. The next distinct architecture was Z-Image-Turbo on original masked face crops, not another Qwen geometry conditioner. -**Implementation follow-up (2026-07-24):** that distinct architecture now exists as the -manual `qwen-zimage` profile. It ports the upstream Synthid-Bypass v2 graph: Qwen-Image-2512 -Lightning + DiffSynth Canny for the full frame, then SAM-masked Z-Image Turbo regeneration -from original face crops. The upstream result supplied by the user was Gemini-oracle negative. -The active upstream face path is YOLO + SAM, not the unconnected MediaPipe node. The port -matches its center-point + box prompts, IoU-0.93 proposal selection, detector-box intersection, -crop factor, and paste feather; YuNet is the intentional detector substitution. +**Implementation follow-up (2026-07-24, revised 2026-07-31):** an early experimental +`qwen-zimage` prototype reproduced a broad two-stage shape demonstrated by a public +experiment: structure-guided Qwen regeneration followed by masked Z-Image face +refinement. The maintained profile was subsequently +rewritten with project-owned prompts, adaptive strength and sizing policies, YuNet face +detection, SAM selection, masks, and compositing. It does not include the upstream workflow +JSON or source code. The upstream result supplied by the user was Gemini-oracle negative. The first exact-path Modal run completed without SAM fallback or face seams. On one crowded 18-face `gemini_3` fixture, ArcFace identity improved materially over controlnet @@ -108,7 +108,7 @@ checked all six current outputs in the provider-separated oracles and confirmed that none retained SynthID or the provider generation signal. The checked bytes used the complete `visible -> qwen-zimage -> metadata` route, the calibrated YuNet 0.5 gate, and the shipped prompt-cache/model-residency optimizations. This supersedes -the earlier first-port batch check as the release-candidate result, but it is not a +the earlier prototype batch check as the release-candidate result, but it is not a certification across seeds, resolutions, and content classes. YuNet's threshold was calibrated independently from upstream YOLO: 0.5 retained the visible faces in the comparison fixtures while removing the false and duplicate boxes admitted by the copied diff --git a/docs/research-doubao-distillation.md b/docs/research-doubao-distillation.md index 2b668c8..1aaf499 100644 --- a/docs/research-doubao-distillation.md +++ b/docs/research-doubao-distillation.md @@ -26,7 +26,7 @@ color still leaves a persistent ghost outline. Diagnosed why, empirically (cached stacks, `/tmp/doubao_distill`): (1) the mark is a clean white overlay with **no dark halo** -- over glyph pixels ~54% are brighter than the clean bg, only ~4% darker -- so the white-logo model `I=(1-α)O+α·255` is correct; (2) but content backgrounds are almost never dark *under* the mark (median darkest available bg over glyph pixels = **58/255**; only ~13% of mark pixels are ever observed on a bg < 40), so on bright backgrounds the equation is ill-conditioned and `α` is unidentifiable; (3) LaMa's `O` is a plausible **hallucination**, not the true pre-mark background, which compounds the error, and per-pixel regression on ~15 obs overfits into color noise. -**Why Gemini's engine is clean (verified in GeminiWatermarkTool `src/core/watermark_engine.cpp`): its alpha map is the watermark stamped on a PURE-BLACK background**, where `watermarked = α·255 + (1-α)·0 = α·255`, so `alpha = capture/255` exactly -- no estimation. (`gemini_bg_*.png` is literally the sparkle in gray on black.) So the real Doubao unlock is the same controlled capture, **not more content images**. The retained black and gray outputs live in `data/calibration/doubao/`; local solid-color seeds are regenerable and are not committed. +**Why Gemini's engine is clean: its alpha map is the watermark stamped on a PURE-BLACK background**, where `watermarked = α·255 + (1-α)·0 = α·255`, so `alpha = capture/255` exactly -- no estimation. (`gemini_bg_*.png` is literally the sparkle in gray on black.) So the real Doubao unlock is the same controlled capture, **not more content images**. The retained black and gray outputs live in `data/calibration/doubao/`; local solid-color seeds are regenerable and are not committed. **Until black captures arrive, the shipped direction is precise canonical glyph mask + inpaint (cv2 default, lama optional), NOT reverse-alpha.** diff --git a/docs/synthid.md b/docs/synthid.md index 40d9cc3..4bc9d0c 100644 --- a/docs/synthid.md +++ b/docs/synthid.md @@ -623,7 +623,7 @@ it; (3) **historical engineering conclusion:** this dated run argued for a higher ControlNet strength than the then-current default. That proposal was later superseded. The current resolver intentionally shares the 0.10/0.15 ladder between SDXL and ControlNet and uses a separate Qwen ladder; see -`noai/watermark_profiles.py`. +`_internal/watermark_profiles.py`. Source images are private (faces / product shots), not committed; reproduce on any photoreal + flat-graphic gpt-image pair, varying the seed, and re-checking the oracle. diff --git a/pyproject.toml b/pyproject.toml index 7d7811f..5e95ddd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,7 @@ dependencies = [ "python-dotenv>=1.0.0", # Official C2PA reader (Content Authenticity Initiative, MIT/Apache-2.0). The # primary, spec-tracking manifest parser for the identify/metadata path; the - # hand-rolled caBX/CBOR scanner in noai/c2pa.py is kept only as a fallback for + # hand-rolled caBX/CBOR scanner in _internal/c2pa.py is kept only as a fallback for # synthetic/partial blobs the validator rejects. Binary wheel (Rust), but the # import is light (no torch/numpy) so it fits the dependency-light identify # host. Prebuilt wheels cover the full CI matrix (linux/macos/windows). @@ -235,8 +235,8 @@ ignore = [ [tool.ruff.lint.per-file-ignores] "scripts/*.py" = ["G004", "S108", "S310", "T20"] "tests/*.py" = ["ANN", "S101", "S105", "S106", "S108"] -"src/remove_ai_watermarks/noai/watermark_remover.py" = ["S603", "S606", "S607", "T201"] # subprocess calls for auto-install/CUDA fix -"src/remove_ai_watermarks/noai/c2pa.py" = ["S110"] # try-except-pass for corrupt file handling +"src/remove_ai_watermarks/_internal/watermark_remover.py" = ["S603", "S606", "S607"] # nvidia-smi capability probe +"src/remove_ai_watermarks/_internal/c2pa.py" = ["S110"] # try-except-pass for corrupt file handling [tool.ruff.format] quote-style = "double" diff --git a/scripts/metadata_removal_audit.py b/scripts/metadata_removal_audit.py index af7d732..01cb47f 100644 --- a/scripts/metadata_removal_audit.py +++ b/scripts/metadata_removal_audit.py @@ -32,7 +32,7 @@ from pathlib import Path import click import numpy as np -from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS +from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS log = logging.getLogger(__name__) diff --git a/scripts/smoke_matrix.py b/scripts/smoke_matrix.py index a4ee0b1..3548268 100644 --- a/scripts/smoke_matrix.py +++ b/scripts/smoke_matrix.py @@ -590,8 +590,8 @@ def _diffusion_rows(r: Runner, tmp: Path, doubao: Path) -> None: try: import numpy as np + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover from remove_ai_watermarks.image_io import imread - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover src = imread(str(mj)) h, w = src.shape[:2] diff --git a/scripts/synthid_corpus.py b/scripts/synthid_corpus.py index 81beeba..6e007f4 100644 --- a/scripts/synthid_corpus.py +++ b/scripts/synthid_corpus.py @@ -29,7 +29,7 @@ import click from _plain_console import Console, Table from PIL import Image -from remove_ai_watermarks.noai.c2pa import extract_c2pa_info +from remove_ai_watermarks._internal.c2pa import extract_c2pa_info log = logging.getLogger(__name__) console = Console() diff --git a/scripts/visible_positives.py b/scripts/visible_positives.py index 4154198..8cba624 100644 --- a/scripts/visible_positives.py +++ b/scripts/visible_positives.py @@ -38,7 +38,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) # The package's own format set. An inlined copy here silently skipped .heif, which # CLAUDE.md documents as supported. -from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS as _EXTS +from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS as _EXTS REPO = Path(__file__).resolve().parents[1] CORPUS = REPO / ".local-eval" / "originals" diff --git a/src/remove_ai_watermarks/noai/__init__.py b/src/remove_ai_watermarks/_internal/__init__.py similarity index 74% rename from src/remove_ai_watermarks/noai/__init__.py rename to src/remove_ai_watermarks/_internal/__init__.py index 399ebef..e393103 100644 --- a/src/remove_ai_watermarks/noai/__init__.py +++ b/src/remove_ai_watermarks/_internal/__init__.py @@ -1,10 +1,8 @@ -"""Vendored noai-watermark code for invisible watermark removal. - -Original: https://github.com/mertizci/noai-watermark (MIT License) +"""Compatibility namespace for metadata and regeneration helpers. The public API (``WatermarkRemover`` / ``remove_watermark`` / ``remove_ai_metadata``) is exposed **lazily** via PEP 562 ``__getattr__``: importing a light submodule -(e.g. ``noai.c2pa`` / ``noai.constants`` from ``identify``) must NOT eagerly pull +(e.g. ``_internal.c2pa`` / ``_internal.constants`` from ``identify``) must NOT eagerly pull ``watermark_remover``, which imports torch + diffusers at module top. Keeping this lazy is what lets ``import remove_ai_watermarks.identify`` stay cheap (~36 MB, no torch) even in a full install where the ``diffusion`` extra is present -- @@ -17,8 +15,8 @@ from __future__ import annotations from typing import TYPE_CHECKING if TYPE_CHECKING: + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover, remove_watermark from remove_ai_watermarks.metadata import remove_ai_metadata - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover, remove_watermark __all__ = ["WatermarkRemover", "remove_ai_metadata", "remove_watermark"] @@ -27,12 +25,12 @@ def __getattr__(name: str) -> object: """Resolve the public API on first access (PEP 562), not at package import.""" if name == "remove_ai_metadata": # Re-export the single, robust stripper (byte-level, lossless-for-JPEG, all - # containers); the old noai.cleaner implementation is retired. + # containers); the old legacy metadata helper implementation is retired. from remove_ai_watermarks.metadata import remove_ai_metadata return remove_ai_metadata if name in ("WatermarkRemover", "remove_watermark"): - from remove_ai_watermarks.noai import watermark_remover + from remove_ai_watermarks._internal import watermark_remover return getattr(watermark_remover, name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/remove_ai_watermarks/_internal/c2pa.py b/src/remove_ai_watermarks/_internal/c2pa.py new file mode 100644 index 0000000..f432c90 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/c2pa.py @@ -0,0 +1,363 @@ +"""C2PA inspection through the official reader with a bounded PNG fallback.""" + +from __future__ import annotations + +import contextlib +import functools +import json +import logging +import re +import struct +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +from remove_ai_watermarks._internal.constants import ( + C2PA_ACTIONS, + C2PA_AI_TOOLS, + C2PA_CHUNK_TYPE, + C2PA_ISSUERS, + C2PA_SIGNATURES, + C2PA_SOFT_BINDINGS, + PNG_SIGNATURE, + SYNTHID_C2PA_ISSUERS, +) + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from typing import BinaryIO + +_C2paReader: Any = None +with contextlib.suppress(Exception): + from c2pa import Reader as _C2paReader # pyright: ignore[reportMissingTypeStubs] + +_C2PA_READER_AVAILABLE = _C2paReader is not None +_PNG_HEADER = struct.Struct(">I4s") + + +@dataclass(frozen=True) +class _PngChunk: + payload: bytes + serialized: bytes + + +def reader_available() -> bool: + """Return whether the official C2PA reader loaded successfully.""" + return _C2PA_READER_AVAILABLE + + +def _manifest_json_uncached(path: str) -> str | None: + try: + reader = _C2paReader.try_create(path) + except Exception as error: + logger.debug("C2PA reader rejected %s: %s", path, error) + return None + if reader is None: + return None + try: + with reader: + return cast("str", reader.json()) + except Exception as error: + logger.debug("C2PA reader could not serialize %s: %s", path, error) + return None + + +@functools.lru_cache(maxsize=8) +def _manifest_json_cached(path: str, _mtime_ns: int) -> str | None: + return _manifest_json_uncached(path) + + +def read_manifest_store_json(image_path: Path) -> str | None: + """Read the complete manifest-store JSON, caching it until the file changes.""" + if not reader_available(): + return None + path = str(image_path) + try: + return _manifest_json_cached(path, image_path.stat().st_mtime_ns) + except OSError: + return _manifest_json_uncached(path) + + +def _find_c2pa_chunk(path: Path) -> _PngChunk | None: + """Return the first recognizable C2PA chunk without loading the whole PNG.""" + try: + stream = path.open("rb") + except OSError: + return None + with stream: + if stream.read(len(PNG_SIGNATURE)) != PNG_SIGNATURE: + return None + file_size = stream.seek(0, 2) + stream.seek(len(PNG_SIGNATURE)) + while True: + header = stream.read(_PNG_HEADER.size) + if len(header) != _PNG_HEADER.size: + return None + length, kind = _PNG_HEADER.unpack(header) + if length + 4 > file_size - stream.tell(): + return None + if kind == C2PA_CHUNK_TYPE: + payload = stream.read(length) + crc = stream.read(4) + if _looks_like_c2pa(payload): + return _PngChunk(payload, header + payload + crc) + else: + stream.seek(length + 4, 1) + if kind == b"IEND": + return None + + +def _is_well_formed_png(path: Path) -> bool: + """Validate PNG chunk bounds with seeks rather than payload allocations.""" + try: + stream = path.open("rb") + except OSError: + return False + with stream: + if stream.read(len(PNG_SIGNATURE)) != PNG_SIGNATURE: + return False + file_size = stream.seek(0, 2) + stream.seek(len(PNG_SIGNATURE)) + while True: + header = stream.read(_PNG_HEADER.size) + if len(header) != _PNG_HEADER.size: + return False + length, kind = _PNG_HEADER.unpack(header) + if length + 4 > file_size - stream.tell(): + return False + stream.seek(length + 4, 1) + if kind == b"IEND": + return True + + +def _copy_bytes(source: BinaryIO, target: BinaryIO, byte_count: int) -> None: + """Copy exactly one bounded chunk without allocating its complete payload.""" + remaining = byte_count + while remaining: + block = source.read(min(remaining, 1024 * 1024)) + if not block: + raise OSError("PNG changed while it was being copied") + target.write(block) + remaining -= len(block) + + +def _looks_like_c2pa(payload: bytes) -> bool: + lowered = payload.lower() + return any(signature in payload for signature in C2PA_SIGNATURES) or b"c2pa" in lowered or b"jumb" in lowered + + +def extract_c2pa_chunk(image_path: Path) -> bytes | None: + """Return the first complete C2PA PNG chunk, including header and CRC.""" + if image_path.suffix.casefold() != ".png": + return None + chunk = _find_c2pa_chunk(image_path) + return None if chunk is None else chunk.serialized + + +def has_c2pa_metadata(image_path: Path) -> bool: + """Return whether a validly bounded PNG contains a recognizable C2PA chunk.""" + return extract_c2pa_chunk(Path(image_path)) is not None + + +def _active_manifest(store: dict[str, Any]) -> dict[str, Any]: + manifests = store.get("manifests") + if not isinstance(manifests, dict): + return {} + typed_manifests = cast("dict[object, object]", manifests) + active = typed_manifests.get(store.get("active_manifest")) + return cast("dict[str, Any]", active) if isinstance(active, dict) else {} + + +def _claim_generator_from_store(store: dict[str, Any]) -> str | None: + active = _active_manifest(store) + direct = active.get("claim_generator") + if isinstance(direct, str) and direct.isprintable() and direct: + return direct + candidates = active.get("claim_generator_info") + if isinstance(candidates, list) and candidates and isinstance(candidates[0], dict): + candidate = cast("dict[object, object]", candidates[0]) + name = candidate.get("name") + if isinstance(name, str) and name.isprintable() and name: + return name + return None + + +def synthid_verdict(vendors: str) -> str: + """Describe why metadata implies a likely pixel-level SynthID watermark.""" + return f"likely present ({vendors} embeds SynthID with C2PA)" + + +def _names_present(buffer: bytes, registry: dict[bytes, str]) -> list[str]: + return sorted({label for token, label in registry.items() if token in buffer}) + + +def synthid_vendors_in(buffer: bytes) -> list[str]: + """List matching C2PA issuers known to pair their manifests with SynthID.""" + registry = {token: label for token, label in C2PA_ISSUERS.items() if token in SYNTHID_C2PA_ISSUERS} + return _names_present(buffer, registry) + + +def soft_binding_vendors_in(buffer: bytes) -> list[str]: + """List the soft-binding algorithms named in manifest bytes.""" + return _names_present(buffer, C2PA_SOFT_BINDINGS) + + +def _ordered_matches(buffer: bytes, registry: dict[bytes, str]) -> list[str]: + return list(dict.fromkeys(label for token, label in registry.items() if token in buffer)) + + +def _populate_registry_fields(buffer: bytes, info: dict[str, Any]) -> bool: + issuers = _ordered_matches(buffer, C2PA_ISSUERS) + tools = _ordered_matches(buffer, C2PA_AI_TOOLS) + actions = _ordered_matches(buffer, C2PA_ACTIONS) + if issuers: + info["issuer"] = ", ".join(issuers) + if tools: + info["ai_tool"] = ", ".join(tools) + if actions: + info["actions"] = ", ".join(actions) + + ai_source = False + if b"trainedAlgorithmicMedia" in buffer: + info.update(source_type="trainedAlgorithmicMedia (AI-generated)", ai_source_kind="generated") + ai_source = True + elif b"compositeWithTrainedAlgorithmicMedia" in buffer: + info.update(source_type="compositeWithTrainedAlgorithmicMedia (AI-enhanced)", ai_source_kind="enhanced") + ai_source = True + elif b"algorithmicMedia" in buffer: + info["source_type"] = "algorithmicMedia" + + synthid = synthid_vendors_in(buffer) + if ai_source and synthid: + info["synthid_vendors"] = synthid + info["synthid_watermark"] = synthid_verdict(", ".join(synthid)) + + soft_bindings = soft_binding_vendors_in(buffer) + if soft_bindings: + info["soft_binding_vendors"] = soft_bindings + info["soft_binding"] = ", ".join(soft_bindings) + return ai_source + + +def _base_info(byte_count: int, *, fallback: bool = False) -> dict[str, Any]: + container = "C2PA manifest" if fallback else "C2PA manifest store" + return { + "has_c2pa": True, + "type": "C2PA (Coalition for Content Provenance and Authenticity)", + "c2pa_manifest": f"{container} ({byte_count} bytes)", + } + + +def _info_from_store(store: dict[str, Any], encoded: bytes) -> dict[str, Any]: + info = _base_info(len(encoded)) + _populate_registry_fields(encoded, info) + generator = _claim_generator_from_store(store) + if generator is not None: + info["claim_generator"] = generator + signature_value = _active_manifest(store).get("signature_info") + if isinstance(signature_value, dict): + signature = cast("dict[object, object]", signature_value) + timestamp = signature.get("time") + if timestamp: + info["timestamp"] = str(timestamp) + return info + + +def c2pa_info_from_manifest_store(store: str | dict[str, Any]) -> dict[str, Any]: + """Normalize a manifest store supplied as JSON text or a decoded object.""" + try: + raw_decoded: object = store if isinstance(store, dict) else json.loads(store) + decoded = cast("dict[str, Any]", raw_decoded) if isinstance(raw_decoded, dict) else None + if not isinstance(decoded, dict) or not decoded or decoded.get("error"): + return {} + encoded = json.dumps(decoded, ensure_ascii=False).encode() if isinstance(store, dict) else store.encode() + except (TypeError, ValueError): + return {} + return _info_from_store(decoded, encoded) + + +def cbor_text_after(payload: bytes, key: bytes) -> str | None: + """Decode a definite-length CBOR text value immediately following ``key``.""" + key_end = payload.find(key) + if key_end < 0: + return None + cursor = key_end + len(key) + if cursor >= len(payload): + return None + initial = payload[cursor] + if 0x60 <= initial <= 0x77: + length, cursor = initial & 0x1F, cursor + 1 + elif initial == 0x78 and cursor + 1 < len(payload): + length, cursor = payload[cursor + 1], cursor + 2 + elif initial == 0x79 and cursor + 2 < len(payload): + length = int.from_bytes(payload[cursor + 1 : cursor + 3], "big") + cursor += 3 + else: + return None + raw = payload[cursor : cursor + length] + if len(raw) != length: + return None + try: + return raw.decode() + except UnicodeDecodeError: + return raw.decode("latin1", errors="replace") + + +def _parse_c2pa_chunk(payload: bytes, info: dict[str, Any]) -> None: + info.update(_base_info(len(payload), fallback=True)) + _populate_registry_fields(payload, info) + for key, output_key in ((b"name", "claim_generator"), (b"specVersion", "c2pa_spec")): + value = cbor_text_after(payload, key) + if value and value.isprintable(): + info[output_key] = value + timestamps = [item.decode() for item in re.findall(rb"\d{14}Z", payload)] + if timestamps: + info["timestamp"] = timestamps[0] + if len(timestamps) > 1: + info["timestamps"] = timestamps[:3] + + +def _extract_c2pa_info_png(image_path: Path) -> dict[str, Any]: + if image_path.suffix.casefold() != ".png": + return {} + chunk = _find_c2pa_chunk(image_path) + if chunk is None: + return {} + info: dict[str, Any] = {} + _parse_c2pa_chunk(chunk.payload, info) + return info + + +def extract_c2pa_info(image_path: Path) -> dict[str, Any]: + """Return normalized C2PA evidence from the official reader or PNG fallback.""" + store = read_manifest_store_json(Path(image_path)) + if store is not None: + return c2pa_info_from_manifest_store(store) + return _extract_c2pa_info_png(Path(image_path)) + + +def inject_c2pa_chunk(target_path: Path, output_path: Path, c2pa_chunk: bytes) -> None: + """Replace any C2PA chunks in a PNG and insert ``c2pa_chunk`` before IDAT.""" + if target_path.suffix.casefold() != ".png" or output_path.suffix.casefold() != ".png": + raise ValueError("C2PA chunk injection is only supported for PNG files") + if not _is_well_formed_png(target_path): + raise ValueError("Target is not a well-formed PNG file") + + output_path.parent.mkdir(parents=True, exist_ok=True) + with target_path.open("rb") as source, output_path.open("wb") as target: + target.write(source.read(len(PNG_SIGNATURE))) + inserted = False + while True: + header = source.read(_PNG_HEADER.size) + length, kind = _PNG_HEADER.unpack(header) + if kind == b"IDAT" and not inserted: + target.write(c2pa_chunk) + inserted = True + if kind == C2PA_CHUNK_TYPE: + source.seek(length + 4, 1) + else: + target.write(header) + _copy_bytes(source, target, length + 4) + if kind == b"IEND": + break diff --git a/src/remove_ai_watermarks/_internal/constants.py b/src/remove_ai_watermarks/_internal/constants.py new file mode 100644 index 0000000..ef05a8c --- /dev/null +++ b/src/remove_ai_watermarks/_internal/constants.py @@ -0,0 +1,152 @@ +"""Registries shared by metadata extraction and provenance classification.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +def _tokens(value: str) -> tuple[str, ...]: + return tuple(value.split("|")) + + +SUPPORTED_FORMATS = frozenset(_tokens(".png|.jpg|.jpeg|.webp|.heic|.heif|.avif")) +AI_METADATA_KEYS = _tokens( + "parameters|postprocessing|extras|workflow|prompt|Dream|SD:mode|StableDiffusionVersion|" + "generation_time|Model|Model hash|Seed" +) +PNG_METADATA_KEYS = _tokens( + "Author|Title|Description|Copyright|Creation Time|Software|Disclaimer|Warning|Source|Comment" +) +AI_KEYWORDS = _tokens( + "prompt|negative_prompt|sampler|cfg_scale|lora|diffusion|comfy|midjourney|dall-e|dalle|imagen|firefly|c2pa|chatgpt|gpt-4|sora|openai|truepic|stable_diffusion|invokeai" +) + +PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" +C2PA_CHUNK_TYPE = b"caBX" +C2PA_SIGNATURES = tuple( + token.encode() for token in _tokens("c2pa|C2PA|jumb|jumd|JUMBF|jumbf|cbor|contentcreds|digid|assertions|manifest") +) + + +@dataclass(frozen=True, slots=True) +class C2paAiVendor: + """One issuer signature and its normalized product attribution.""" + + issuer: bytes + org: str + platform: str | None + needle: str | None + synthid: bool = False + asserts_ai: bool = False + + +def _vendor( + issuer: bytes | str, + org: str, + platform: str | None, + needle: str | None, + *, + synthid: bool = False, + asserts_ai: bool = False, +) -> C2paAiVendor: + token = issuer.encode() if isinstance(issuer, str) else issuer + return C2paAiVendor(token, org, platform, needle, synthid, asserts_ai) + + +# Order is product priority when a manifest mentions more than one organization. +C2PA_AI_VENDORS: tuple[C2paAiVendor, ...] = ( + _vendor(b"Microsoft", "Microsoft", "Microsoft (Bing Image Creator / Designer)", "Microsoft"), + _vendor(b"Adobe", "Adobe", "Adobe Firefly", "Adobe"), + _vendor(b"OpenAI", "OpenAI", "OpenAI (ChatGPT / gpt-image / DALL-E / Sora)", "OpenAI", synthid=True), + _vendor(b"Google", "Google LLC", "Google (Gemini / Imagen)", "Google", synthid=True), + _vendor(b"Stability AI", "Stability AI", "Stability AI (Stable Image / DreamStudio)", "Stability AI"), + _vendor(b"Black Forest Labs", "Black Forest Labs", "Black Forest Labs (FLUX)", "Black Forest Labs"), + _vendor(b"volcengine", "ByteDance (Volcano Engine)", "ByteDance (Doubao / Jimeng / Volcano Engine)", "ByteDance"), + _vendor( + "北京火山引擎科技有限公司", + "ByteDance (Volcano Engine)", + "ByteDance (Doubao / Jimeng / Volcano Engine)", + "ByteDance", + ), + _vendor(b"Byteplus", "BytePlus (ByteDance)", "ByteDance (Doubao / Jimeng / Volcano Engine)", "ByteDance"), + _vendor( + b"Dreamina", + "ByteDance (Dreamina)", + "ByteDance (Doubao / Jimeng / Volcano Engine)", + "ByteDance", + asserts_ai=True, + ), + _vendor(b"Canva", "Canva", "Canva (Magic Media)", "Canva"), + _vendor(b"Eleven Labs", "ElevenLabs", "ElevenLabs", "ElevenLabs"), + _vendor(b"fal-ai", "fal.ai", "fal.ai", "fal.ai", asserts_ai=True), + _vendor(b"Bria", "Bria Artificial Intelligence", "Bria AI", "Bria", asserts_ai=True), + _vendor(b"Truepic", "Truepic", None, None), +) + +C2PA_ISSUERS = {vendor.issuer: vendor.org for vendor in C2PA_AI_VENDORS} +C2PA_IDENTITY_AI_ORGS = frozenset(vendor.org for vendor in C2PA_AI_VENDORS if vendor.asserts_ai) +SYNTHID_C2PA_ISSUERS = frozenset(vendor.issuer for vendor in C2PA_AI_VENDORS if vendor.synthid) + +C2PA_AI_TOOLS = { + token.encode(): label + for token, label in ( + ("GPT-4o", "GPT-4o"), + ("ChatGPT", "ChatGPT"), + ("Sora", "Sora"), + ("DALL-E", "DALL-E"), + ("DALL", "DALL-E"), + ("Imagen", "Imagen"), + ("Firefly", "Firefly"), + ) +} + +C2PA_SOFT_BINDINGS = { + b"com.adobe.trustmark": "Adobe TrustMark", + b"com.adobe.icn": "Adobe (content fingerprint)", + b"com.digimarc": "Digimarc", + b"com.imatag.lamark": "Imatag (Lamark)", + b"ai.steg": "Steg.AI", + b"com.microsoft.invismark": "Microsoft InvisMark", + b"com.microsoft.wavmark": "Microsoft WavMark", + b"com.verimatrix": "Verimatrix", + b"com.nagra.nexguard": "NAGRA NexGuard", + b"com.aiwatermark": "AIWatermark (Meta PixelSeal)", + b"ai.trufo": "Trufo", + b"app.overlai": "Overlai", + b"com.markany": "MarkAny", + b"com.mentaport": "Mentaport", + b"es.lumatrace": "LumaTrace", + b"ai.verda": "VerdaAI", + b"ai.contentlens": "ContentLens", + b"io.iscc": "ISCC (content code)", +} + +AI_GENERATOR_TOKENS = frozenset( + { + "firefly", + "dall-e", + "dalle", + "midjourney", + "stable diffusion", + "stable-diffusion", + "stablediffusion", + "comfyui", + "automatic1111", + "invokeai", + "imagen", + "gpt-image", + "nightcafe", + "ideogram", + "leonardo", + "flux", + "dreamstudio", + "novelai", + "reve.com", + "aphrodite ai", + "apple photos clean up", + "fal-ai", + } +) + +_C2PA_ACTION_NAMES = _tokens("created|converted|edited|filtered|cropped|resized|opened|placed") +C2PA_ACTIONS = {f"c2pa.{action}".encode(): action for action in _C2PA_ACTION_NAMES} diff --git a/src/remove_ai_watermarks/noai/ebml.py b/src/remove_ai_watermarks/_internal/ebml.py similarity index 100% rename from src/remove_ai_watermarks/noai/ebml.py rename to src/remove_ai_watermarks/_internal/ebml.py diff --git a/src/remove_ai_watermarks/_internal/extractor.py b/src/remove_ai_watermarks/_internal/extractor.py new file mode 100644 index 0000000..df3b706 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/extractor.py @@ -0,0 +1,98 @@ +"""Read image metadata without changing the source container.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, cast + +import piexif +from PIL import Image + +from remove_ai_watermarks._internal.c2pa import extract_c2pa_chunk, extract_c2pa_info, has_c2pa_metadata +from remove_ai_watermarks._internal.constants import AI_KEYWORDS, AI_METADATA_KEYS + +if TYPE_CHECKING: + from pathlib import Path + +_EXIF_KEY = "exif" +_AI_KEYS_CASEFOLD = frozenset(key.casefold() for key in AI_METADATA_KEYS) + + +def _read_pillow_info(source_path: Path) -> dict[str, Any]: + with Image.open(source_path) as image: + return {key: value for key, value in image.info.items() if isinstance(key, str)} + + +def _decode_exif(raw: object) -> tuple[str, object]: + if not isinstance(raw, bytes): + return "exif_raw", raw + try: + return _EXIF_KEY, piexif.load(raw) + except Exception: + return "exif_raw", raw + + +def _is_ai_field(key: str) -> bool: + folded = key.casefold() + return folded in _AI_KEYS_CASEFOLD or any(token in folded for token in AI_KEYWORDS) + + +def _attach_c2pa(source_path: Path, metadata: dict[str, Any]) -> None: + if not has_c2pa_metadata(source_path): + return + metadata["c2pa"] = extract_c2pa_info(source_path) + payload = extract_c2pa_chunk(source_path) + if payload is not None: + metadata["c2pa_chunk"] = payload + + +def extract_metadata(source_path: Path) -> dict[str, Any]: + """Return every Pillow-visible field plus decoded EXIF and C2PA data.""" + raw_info = _read_pillow_info(source_path) + metadata = dict(raw_info) + if _EXIF_KEY in raw_info: + metadata.pop(_EXIF_KEY, None) + decoded_key, decoded_value = _decode_exif(raw_info[_EXIF_KEY]) + metadata[decoded_key] = decoded_value + + _attach_c2pa(source_path, metadata) + return metadata + + +def extract_ai_metadata(source_path: Path) -> dict[str, Any]: + """Return only metadata keys recognized as AI provenance or generation data.""" + metadata = {key: value for key, value in _read_pillow_info(source_path).items() if _is_ai_field(key)} + _attach_c2pa(source_path, metadata) + return metadata + + +def has_ai_metadata(image_path: Path) -> bool: + """Return whether a supported metadata signal is present.""" + if any(_is_ai_field(key) for key in _read_pillow_info(image_path)): + return True + return has_c2pa_metadata(image_path) + + +def _summary_value(value: object) -> str: + if isinstance(value, bytes): + return f"" + text = str(value) + return text if len(text) <= 100 else f"{text[:100]}..." + + +def get_ai_metadata_summary(source_path: Path) -> str: + """Format the AI-only metadata view for the command-line report.""" + metadata = extract_ai_metadata(source_path) + if not metadata: + return "No AI metadata found." + + lines = ["AI Image Metadata:", "-" * 40] + for key, value in metadata.items(): + if key == "c2pa_chunk": + continue + if key == "c2pa" and isinstance(value, dict): + lines.append("C2PA Metadata:") + c2pa_fields = cast("dict[str, object]", value) + lines.extend(f" {name}: {_summary_value(item)}" for name, item in c2pa_fields.items()) + continue + lines.append(f"{key}: {_summary_value(value)}") + return "\n".join(lines) diff --git a/src/remove_ai_watermarks/noai/flv.py b/src/remove_ai_watermarks/_internal/flv.py similarity index 100% rename from src/remove_ai_watermarks/noai/flv.py rename to src/remove_ai_watermarks/_internal/flv.py diff --git a/src/remove_ai_watermarks/_internal/img2img_runner.py b/src/remove_ai_watermarks/_internal/img2img_runner.py new file mode 100644 index 0000000..2fcd37c --- /dev/null +++ b/src/remove_ai_watermarks/_internal/img2img_runner.py @@ -0,0 +1,136 @@ +"""Execute Diffusers img2img calls and recover from an MPS runtime failure.""" + +from __future__ import annotations + +import contextlib +import logging +from typing import TYPE_CHECKING, Any + +from remove_ai_watermarks._internal.progress import is_mps_error, make_pipeline_progress + +if TYPE_CHECKING: + from collections.abc import Callable + + from PIL import Image + +logger = logging.getLogger(__name__) + + +def _pipeline_arguments( + image: Image.Image, + strength: float, + num_inference_steps: int, + guidance_scale: float, + generator: Any, + step_callback: Any, + overrides: dict[str, Any] | None, +) -> dict[str, Any]: + arguments: dict[str, Any] = { + "prompt": "", + "image": image, + "strength": strength, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "generator": generator, + } + arguments.update(overrides or {}) + if step_callback is not None: + arguments.update(callback=step_callback, callback_steps=1) + return arguments + + +def _invoke(pipeline: Any, arguments: dict[str, Any]) -> Image.Image: + response = pipeline(**arguments) + return response.images[0] + + +def run_img2img( + pipeline: Any, + image: Image.Image, + strength: float, + num_inference_steps: int, + guidance_scale: float, + generator: Any, + device: str, + set_progress: Callable[[str], None], + extra_kwargs: dict[str, Any] | None = None, +) -> Image.Image: + """Run one img2img request and report denoising progress when supported.""" + callback, started, finished, launch_monitor = make_pipeline_progress( + max(1, int(num_inference_steps * strength)), device, set_progress + ) + launch_monitor() + arguments = _pipeline_arguments( + image, strength, num_inference_steps, guidance_scale, generator, callback, extra_kwargs + ) + try: + try: + return _invoke(pipeline, arguments) + except TypeError as error: + if "callback" not in str(error): + raise + started.set() + arguments.pop("callback", None) + arguments.pop("callback_steps", None) + return _invoke(pipeline, arguments) + finally: + started.set() + finished.set() + + +def run_img2img_with_mps_fallback( + load_pipeline: Callable[[], Any], + image: Image.Image, + strength: float, + num_inference_steps: int, + guidance_scale: float, + generator: Any, + device: str, + set_progress: Callable[[str], None], + *, + reload_on_cpu: Callable[[], Any], + extra_kwargs: dict[str, Any] | None = None, +) -> tuple[Image.Image, str]: + """Retry an MPS-specific failure once with a freshly loaded CPU pipeline.""" + try: + output = run_img2img( + load_pipeline(), + image, + strength, + num_inference_steps, + guidance_scale, + generator, + device, + set_progress, + extra_kwargs, + ) + return output, device + except RuntimeError as error: + if device != "mps" or not is_mps_error(error): + raise + logger.warning("MPS execution failed (%s); retrying on CPU", error) + set_progress("MPS execution failed; retrying on CPU...") + try_empty_device_cache("mps") + output = run_img2img( + reload_on_cpu(), + image, + strength, + num_inference_steps, + guidance_scale, + None, + "cpu", + set_progress, + extra_kwargs, + ) + return output, "cpu" + + +def try_empty_device_cache(device: str) -> None: + """Ask Torch to release cached accelerator memory when the backend supports it.""" + with contextlib.suppress(Exception): + import torch + + backend = getattr(torch, device, None) + empty_cache = getattr(backend, "empty_cache", None) + if callable(empty_cache): + empty_cache() diff --git a/src/remove_ai_watermarks/noai/isobmff.py b/src/remove_ai_watermarks/_internal/isobmff.py similarity index 100% rename from src/remove_ai_watermarks/noai/isobmff.py rename to src/remove_ai_watermarks/_internal/isobmff.py diff --git a/src/remove_ai_watermarks/_internal/progress.py b/src/remove_ai_watermarks/_internal/progress.py new file mode 100644 index 0000000..2fb8f05 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/progress.py @@ -0,0 +1,212 @@ +"""Progress reporting utilities for long-running optional model operations.""" + +from __future__ import annotations + +import contextlib +import io +import os +import sys +import threading +import time +import warnings +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Callable + +_BAR_WIDTH = 28 +_SPINNER = ("|", "/", "-", "\\") + + +def _truncate(text: str, max_len: int = 72) -> str: + if len(text) <= max_len: + return text + return f"{text[: max(0, max_len - 3)]}..." + + +def _build_bar(step: int) -> str: + position = step % (2 * _BAR_WIDTH - 2) + if position >= _BAR_WIDTH: + position = 2 * _BAR_WIDTH - 2 - position + cells = ["-"] * _BAR_WIDTH + cells[position] = "=" + return "".join(cells) + + +@dataclass +class _TaskResult: + value: Any = None + error: BaseException | None = None + complete: threading.Event = field(default_factory=threading.Event) + + +def run_with_progress(task: Callable[[], Any], progress_state: dict[str, str] | None = None) -> Any: + """Run ``task`` on a worker thread and render a compact terminal heartbeat.""" + outcome = _TaskResult() + + def invoke() -> None: + try: + outcome.value = task() + except BaseException as error: # re-raised on the caller thread + outcome.error = error + finally: + outcome.complete.set() + + worker = threading.Thread(target=invoke, name="raiw-progress-task", daemon=True) + worker.start() + started_at = time.monotonic() + frame = 0 + terminal = sys.__stderr__ + while not outcome.complete.wait(0.1): + message = _truncate((progress_state or {}).get("message", "Processing...")) + elapsed = int(time.monotonic() - started_at) + if terminal is not None: + terminal.write( + f"\r\033[2K {_SPINNER[frame % len(_SPINNER)]} [{_build_bar(frame)}] {elapsed:>3}s {message}" + ) + terminal.flush() + frame += 1 + + worker.join() + elapsed = int(time.monotonic() - started_at) + message = _truncate((progress_state or {}).get("message", "Processing...")) + if terminal is not None: + terminal.write(f"\r\033[2K Completed in {elapsed}s {message}\n") + terminal.flush() + if outcome.error is not None: + raise outcome.error + return outcome.value + + +def _silence_diffusers() -> None: + from diffusers.utils import logging as diffusers_logging + + diffusers_logging.set_verbosity_error() + disable = getattr(diffusers_logging, "disable_progress_bar", None) + if callable(disable): + disable() + + +def _configure_quiet_libraries() -> None: + os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + operations = ( + lambda: __import__("transformers").logging.set_verbosity_error(), + _silence_diffusers, + lambda: __import__("huggingface_hub").logging.set_verbosity_error(), + ) + for operation in operations: + with contextlib.suppress(Exception): + operation() + + +def silence_library_output( + run_func: Callable[[], Any], + set_progress: Callable[[str], None] | None = None, +) -> Callable[[], Any]: + """Wrap a model call so third-party progress bars do not corrupt our CLI UI.""" + + def quiet_call() -> Any: + if set_progress is not None: + set_progress("Preparing model runtime...") + _configure_quiet_libraries() + with ( + warnings.catch_warnings(), + contextlib.redirect_stdout(io.StringIO()), + contextlib.redirect_stderr(io.StringIO()), + ): + warnings.simplefilter("ignore") + if set_progress is not None: + set_progress("Running watermark regeneration...") + return run_func() + + return quiet_call + + +@dataclass +class _PipelineMonitor: + total_steps: int + device: str + update: Callable[[str], None] + bar_len: int + label: str + pre_phases: list[tuple[int, str]] + post_phases: list[tuple[int, str]] + first_step: threading.Event = field(default_factory=threading.Event) + done: threading.Event = field(default_factory=threading.Event) + started_at: float = field(default_factory=time.monotonic) + last_step_at: float = field(default_factory=time.monotonic) + + def callback(self, step: int, _timestep: int, _latents: Any) -> None: + self.first_step.set() + now = time.monotonic() + self.last_step_at = now + current = min(self.total_steps, step + 1) + filled = round(self.bar_len * current / self.total_steps) + elapsed = now - self.started_at + eta = elapsed * max(0, self.total_steps - current) / max(1, current) + bar = "#" * filled + "." * (self.bar_len - filled) + self.update( + f"{self.label} [{bar}] {current}/{self.total_steps}, " + f"{elapsed:.0f}s elapsed, ~{eta:.0f}s left, {self.device}" + ) + + def _phase_message(self, phases: list[tuple[int, str]], elapsed: float) -> str: + message = phases[0][1] + for threshold, candidate in phases: + if elapsed < threshold: + break + message = candidate + return message + + def monitor(self) -> None: + while not self.first_step.wait(0.4): + elapsed = time.monotonic() - self.started_at + self.update(self._phase_message(self.pre_phases, elapsed)) + decode_started: float | None = None + while not self.done.wait(0.4): + if time.monotonic() - self.last_step_at < 1.5: + decode_started = None + continue + decode_started = decode_started or time.monotonic() + self.update(self._phase_message(self.post_phases, time.monotonic() - decode_started)) + + def start(self) -> threading.Thread: + self.started_at = self.last_step_at = time.monotonic() + self.first_step.clear() + self.done.clear() + thread = threading.Thread(target=self.monitor, name="raiw-pipeline-progress", daemon=True) + thread.start() + return thread + + +def make_pipeline_progress( + effective_steps: int, + device: str, + set_progress: Callable[[str], None], + *, + bar_len: int = 20, + label: str = "Denoising", + pre_phases: list[tuple[int, str]] | None = None, + post_phases: list[tuple[int, str]] | None = None, +) -> tuple[Callable[..., None], threading.Event, threading.Event, Callable[[], threading.Thread]]: + """Build a callback and monitor for the legacy Diffusers callback interface.""" + + def qualify(entries: list[tuple[int, str]]) -> list[tuple[int, str]]: + return [(second, f"{text} on {device}") for second, text in entries] + + monitor = _PipelineMonitor( + total_steps=max(1, effective_steps), + device=device, + update=set_progress, + bar_len=bar_len, + label=label, + pre_phases=pre_phases or qualify([(0, "Encoding image"), (8, "Preparing denoiser"), (20, "Starting sampler")]), + post_phases=post_phases or qualify([(0, "Decoding image"), (10, "Finalizing pixels"), (45, "Still decoding")]), + ) + return monitor.callback, monitor.first_step, monitor.done, monitor.start + + +def is_mps_error(error: Exception) -> bool: + """Return whether an error message identifies Apple's MPS backend.""" + return "mps" in str(error).casefold() diff --git a/src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py similarity index 90% rename from src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py rename to src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py index 25df2f0..ff9969e 100644 --- a/src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py +++ b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py @@ -1,17 +1,9 @@ -"""Qwen 2512 Canny regeneration followed by masked Z-Image face repair. +"""Project-native Qwen regeneration with an optional masked face refinement pass. -This profile ports the two-stage architecture used by cebeuq/Synthid-Bypass: - -1. Qwen-Image-2512 img2img with the 4-step Lightning LoRA and the DiffSynth - blockwise Canny ControlNet regenerates the whole image, optionally as - overlapping feather-blended tiles for large inputs. -2. Faces are detected on the original image, refined to masks with SAM, regenerated - from the original face crops with Z-Image Turbo, and feathered into stage 1. - -The runtime intentionally uses permissively licensed YuNet instead of the reference -workflow's Ultralytics detector. All diffusion and segmentation models remain the same -model families. The adaptive formulas are direct ports, while the face result is scaled -for this runtime's different sampler and mask-compositing path. +The profile was inspired by public experiments that combine structure-guided global +regeneration with a second face-only pass. Its orchestration, sizing rules, adaptive +strength policy, detector, masks, prompts, and compositing are implemented here for +this library's Pillow and DiffSynth runtime. """ # DiffSynth, torch, transformers, and cv2 expose mostly untyped tensor/array APIs. @@ -33,7 +25,7 @@ from typing import TYPE_CHECKING, Any import numpy as np from PIL import Image -from remove_ai_watermarks.noai.watermark_profiles import resolve_seed +from remove_ai_watermarks._internal.watermark_profiles import resolve_seed if TYPE_CHECKING: from collections.abc import Callable @@ -53,10 +45,8 @@ YUNET_MODEL_URL = ( ) YUNET_MODEL_NAME = "face_detection_yunet_2023mar.onnx" YUNET_MODEL_SHA256 = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4" -# The upstream graph's 0.2 threshold belongs to YOLO and does not transfer to -# YuNet's score calibration. At 0.2 YuNet admitted background and decorative -# false positives, multiplying the serial Z-Image face-stage cost. A 0.5 gate -# retained all visible faces in the public and upstream comparison fixtures. +# This threshold retained the faces in the public validation set without admitting +# decorative background regions as faces. YUNET_SCORE_THRESHOLD = 0.5 GLOBAL_STEPS = 4 @@ -65,19 +55,13 @@ GLOBAL_CFG = 1.0 FACE_CFG = 1.0 GLOBAL_CONTROLNET_SCALE = 1.0 RESIDENT_FACE_MODEL_MIN_VRAM_GIB = 64.0 -# The reference face denoise assumes its ComfyUI detailer sampler, latent -# noise-mask feather, and inpaint path. Applying that value unchanged to this -# DiffSynth crop-regeneration port over-processes faces. Paired public-fixture -# measurements and both provider oracles certified half the reference value. FACE_DENOISE_SCALE = 0.5 -# The source graph uses normalized Canny thresholds 0.05 and 0.25. OpenCV takes -# byte thresholds, so round 255*x to the matching integer values. _CANNY_LOW = 13 _CANNY_HIGH = 64 -# These strings intentionally preserve the reference workflow spelling. They are -# model inputs, not user-facing copy, and changing them would change the port. +# These short model inputs are retained as calibrated compatibility parameters. +# Changing them requires the same provider-oracle and identity evaluation as a model change. _GLOBAL_PROMPT = "ultra clear and smoothe skin, spotless skin" _GLOBAL_NEGATIVE = "moles, freckes, high detail skin" _FACE_PROMPT = "" @@ -169,28 +153,18 @@ def resolution_adaptive_denoise( denoise_min: float = 0.08, denoise_max: float = 0.15, ) -> float: - """Port the reference resolution-based adaptive denoise calculation. - - At neutral level 5, 0.30 MP maps to ``denoise_min`` and 3.70 MP maps to - ``denoise_max``. Levels above or below 5 add the same asymmetric spread as the - reference custom node. - """ - image_mp = max(1.0, float(width) * float(height)) / 1_000_000.0 - normalized = _clamp((image_mp - 0.30) / (3.70 - 0.30), 0.0, 1.0) - - minimum = float(denoise_min) - maximum = float(denoise_max) - if maximum < minimum: - minimum, maximum = maximum, minimum - denoise_range = maximum - minimum - base = minimum + denoise_range * normalized - - level = int(adaptive_level) - if level >= 5: - offset = ((float(level) - 5.0) / 5.0) * denoise_range * 0.285714 + """Choose the calibrated global strength from image area and operator level.""" + low, high = sorted((float(denoise_min), float(denoise_max))) + megapixels = max(1.0, float(width) * float(height)) * 1e-6 + area_fraction = _clamp((megapixels - 0.30) / 3.40, 0.0, 1.0) + strength_range = high - low + strength = low + strength_range * area_fraction + level_delta = float(int(adaptive_level)) - 5.0 + if level_delta >= 0.0: + strength += level_delta * strength_range * (0.285714 / 5.0) else: - offset = -((5.0 - float(level)) / 4.0) * denoise_range * 0.257143 - return _clamp(base + offset, 0.0001, 1.0) + strength += level_delta * strength_range * (0.257143 / 4.0) + return _clamp(strength, 0.0001, 1.0) def largest_face_denoise( @@ -202,7 +176,7 @@ def largest_face_denoise( denoise_min: float = 0.05, denoise_max: float = 0.28, ) -> float: - """Scale face denoise from the largest face area, matching reference mode.""" + """Choose the calibrated face strength from the largest detected face area.""" width, height = image_size image_area = max(1.0, float(width) * float(height)) largest_ratio = 0.0 @@ -211,7 +185,7 @@ def largest_face_denoise( largest_ratio = max(largest_ratio, box_area / image_area) if largest_ratio <= 0.0: return _clamp(base_denoise, denoise_min, denoise_max) - scaled = float(base_denoise) * (largest_ratio / max(1e-6, float(adaptive_ratio))) + scaled = float(base_denoise) * largest_ratio / max(1e-6, float(adaptive_ratio)) return _clamp(scaled, denoise_min, denoise_max) @@ -229,7 +203,7 @@ def _resize_to_target(image: Image.Image) -> Image.Image: def build_canny_control_image(image: Image.Image) -> Image.Image: - """Build the three-channel Canny conditioning image used by stage 1.""" + """Build the calibrated three-channel Canny conditioning map.""" import cv2 rgb = np.asarray(image.convert("RGB")) @@ -259,8 +233,6 @@ def build_global_kwargs( "seed": seed, "rand_device": "cpu", "num_inference_steps": GLOBAL_STEPS, - # The source graph applies ModelSamplingAuraFlow with shift=3. DiffSynth - # expresses the same rational sigma shift as exp(mu), so mu=log(3). "exponential_shift_mu": math.log(3.0), "blockwise_controlnet_inputs": [controlnet_input], } @@ -312,7 +284,7 @@ def _expanded_box( *, factor: float = 2.5, ) -> tuple[int, int, int, int]: - """Expand a face box around its center, matching the reference crop factor.""" + """Expand a face box around its center to include local lighting context.""" x1, y1, x2, y2 = box image_width, image_height = image_size center_x = (x1 + x2) / 2.0 @@ -813,7 +785,7 @@ class QwenZImagePipeline: scale_for_face = 768.0 / max(1, max(face_width, face_height)) scale_for_crop = 1024.0 / max(1, max(crop_width, crop_height)) scale = min(scale_for_face, scale_for_crop) - # Never shrink below the crop's current size unless the 1024 cap requires it. + # Never shrink below the crop's current size unless the cap requires it. if max(crop_width, crop_height) <= 1024: scale = max(1.0, scale) width = max(16, round(crop_width * scale / 16.0) * 16) @@ -881,7 +853,7 @@ class QwenZImagePipeline: resolution_adaptive_denoise(image.width, image.height) if strength is None else float(strength) ) if tile and max(image.size) > tile_size: - from remove_ai_watermarks.noai.tiling import run_tiled + from remove_ai_watermarks._internal.tiling import run_tiled global_result = run_tiled( lambda tile_image: self._run_global(tile_image, global_strength, seed), diff --git a/src/remove_ai_watermarks/noai/riff.py b/src/remove_ai_watermarks/_internal/riff.py similarity index 100% rename from src/remove_ai_watermarks/noai/riff.py rename to src/remove_ai_watermarks/_internal/riff.py diff --git a/src/remove_ai_watermarks/noai/tiling.py b/src/remove_ai_watermarks/_internal/tiling.py similarity index 100% rename from src/remove_ai_watermarks/noai/tiling.py rename to src/remove_ai_watermarks/_internal/tiling.py diff --git a/src/remove_ai_watermarks/_internal/utils.py b/src/remove_ai_watermarks/_internal/utils.py new file mode 100644 index 0000000..78221bc --- /dev/null +++ b/src/remove_ai_watermarks/_internal/utils.py @@ -0,0 +1,30 @@ +"""Small path helpers shared by optional image pipelines.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS + +if TYPE_CHECKING: + from pathlib import Path + +_PIL_FORMAT_BY_SUFFIX = { + ".jpg": "JPEG", + ".jpeg": "JPEG", + ".png": "PNG", +} + + +def is_supported_format(file_path: Path) -> bool: + """Return whether ``file_path`` has a supported raster suffix.""" + return file_path.suffix.casefold() in SUPPORTED_FORMATS + + +def get_image_format(file_path: Path) -> str: + """Return the Pillow save format used by the legacy metadata API. + + The metadata writer only has specialized PNG and JPEG paths. Other accepted + inputs therefore use its PNG fallback, matching the established API contract. + """ + return _PIL_FORMAT_BY_SUFFIX.get(file_path.suffix.casefold(), "PNG") diff --git a/src/remove_ai_watermarks/_internal/watermark_profiles.py b/src/remove_ai_watermarks/_internal/watermark_profiles.py new file mode 100644 index 0000000..22e99f8 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/watermark_profiles.py @@ -0,0 +1,105 @@ +"""Project-owned configuration for invisible-watermark regeneration profiles.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + from pathlib import Path + +DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0" +QWEN_MODEL_ID = "Qwen/Qwen-Image" +CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0" + +SDXL_PROFILE = "sdxl" +QWEN_ZIMAGE_PROFILE = "qwen-zimage" + +OPENAI_STRENGTH = 0.10 +GEMINI_STRENGTH = 0.15 +UNKNOWN_STRENGTH = GEMINI_STRENGTH +DEFAULT_STRENGTH = UNKNOWN_STRENGTH + +QWEN_OPENAI_STRENGTH = 0.10 +QWEN_GEMINI_STRENGTH = 0.25 +QWEN_UNKNOWN_STRENGTH = QWEN_GEMINI_STRENGTH + + +@dataclass(frozen=True) +class _StrengthPolicy: + unknown: float + by_vendor: dict[str, float] + + def choose(self, vendor: str | None) -> float: + return self.by_vendor.get((vendor or "").casefold(), self.unknown) + + +_STANDARD_POLICY = _StrengthPolicy( + unknown=UNKNOWN_STRENGTH, + by_vendor={"openai": OPENAI_STRENGTH, "google": GEMINI_STRENGTH}, +) +_QWEN_POLICY = _StrengthPolicy( + unknown=QWEN_UNKNOWN_STRENGTH, + by_vendor={"openai": QWEN_OPENAI_STRENGTH, "google": QWEN_GEMINI_STRENGTH}, +) +_ALIASES = {"default": SDXL_PROFILE, "qwen_zimage": QWEN_ZIMAGE_PROFILE} + + +def normalize_profile(profile: str) -> str: + """Normalize spelling and resolve compatibility aliases.""" + value = profile.strip().casefold() + return _ALIASES.get(value, value) + + +def resolve_steps(num_inference_steps: int | None, pipeline: str) -> int: + """Return an explicit step count or the selected profile's default.""" + if num_inference_steps is not None: + return num_inference_steps + return 4 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else 50 + + +def resolve_seed(seed: int | None, pipeline: str) -> int | None: + """Keep the fixed Qwen plus Z-Image profile reproducible by default.""" + if seed is not None: + return seed + return 0 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else None + + +def strength_default_help() -> str: + """Describe the live default policy without duplicating its values.""" + return ( + f"vendor-adaptive (OpenAI {OPENAI_STRENGTH} / Google {GEMINI_STRENGTH} / " + f"unknown {UNKNOWN_STRENGTH}, from the C2PA issuer; qwen-zimage instead uses " + "resolution-adaptive denoise)" + ) + + +def resolve_strength(strength: float | None, vendor: str | None = None, pipeline: str | None = None) -> float: + """Resolve a user override or the calibrated policy for a profile and vendor.""" + if strength is not None: + return strength + policy = _QWEN_POLICY if pipeline is not None and normalize_profile(pipeline) == "qwen" else _STANDARD_POLICY + return policy.choose(vendor) + + +def viable_steps(num_inference_steps: int, strength: float) -> int: + """Ensure Diffusers receives at least one effective img2img denoising step.""" + if strength <= 0 or int(num_inference_steps * strength) >= 1: + return num_inference_steps + return math.ceil(1.0 / strength) + + +def vendor_for_strength(image_path: Path) -> Literal["openai", "google"] | None: + """Select the strength cohort using the input's SynthID provenance proxy.""" + try: + from remove_ai_watermarks.metadata import synthid_source + + evidence = (synthid_source(image_path) or "").casefold() + except Exception: + return None + if "google" in evidence: + return "google" + if "openai" in evidence: + return "openai" + return None diff --git a/src/remove_ai_watermarks/_internal/watermark_remover.py b/src/remove_ai_watermarks/_internal/watermark_remover.py new file mode 100644 index 0000000..891d904 --- /dev/null +++ b/src/remove_ai_watermarks/_internal/watermark_remover.py @@ -0,0 +1,649 @@ +"""Project-native orchestration for diffusion-based pixel regeneration.""" + +# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false +from __future__ import annotations + +import contextlib +import logging +import os +import subprocess +from typing import TYPE_CHECKING, Any + +from PIL import Image + +from remove_ai_watermarks._internal.watermark_profiles import ( + CONTROLNET_CANNY_MODEL, + DEFAULT_MODEL_ID, + DEFAULT_STRENGTH, + QWEN_MODEL_ID, + QWEN_ZIMAGE_PROFILE, + normalize_profile, + resolve_seed, + resolve_steps, + resolve_strength, + viable_steps, +) + +if TYPE_CHECKING: + from collections.abc import Callable + from pathlib import Path + +logger = logging.getLogger(__name__) + +try: + import torch + + _HAS_TORCH = True +except ImportError: + torch = None # type: ignore[assignment] + _HAS_TORCH = False + +try: + from diffusers import AutoPipelineForImage2Image as AutoImg2ImgPipeline + + _HAS_DIFFUSERS = True +except ImportError: + AutoImg2ImgPipeline = None # type: ignore[assignment,misc] + _HAS_DIFFUSERS = False + +_SDXL_FP16_VAE_ID = "madebyollin/sdxl-vae-fp16-fix" +_DEGENERATE_THRESHOLD = 1.0 +_CANNY_LOW = 100 +_CANNY_HIGH = 200 +_CONTROLNET_PROMPT = "best quality, high quality, sharp, detailed, photographic" +_CONTROLNET_NEGATIVE = "blurry, lowres, deformed, distorted text, garbled text, watermark, jpeg artifacts" +_QWEN_PROMPT = "high quality, sharp, detailed, faithful to the original" +_QWEN_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts" + + +def is_watermark_removal_available() -> bool: + """Return whether the standard diffusion runtime can be imported.""" + return _HAS_TORCH and _HAS_DIFFUSERS + + +def _ensure_watermark_deps() -> None: + if not is_watermark_removal_available(): + raise ImportError( + "Invisible watermark regeneration requires the 'diffusion' extra. Install remove-ai-watermarks[diffusion]." + ) + + +def _needs_fp16_vae_fix(model_id: str, default_model_id: str, is_fp16: bool) -> bool: + """Return whether the default SDXL pipeline needs the overflow-safe VAE.""" + return is_fp16 and model_id == default_model_id + + +def _is_degenerate_image(image: Image.Image) -> bool: + """Detect the uniform near-black output produced by an fp16 decode collapse.""" + import numpy as np + + pixels = np.asarray(image.convert("RGB"), dtype=np.float32) + return float(pixels.mean()) < _DEGENERATE_THRESHOLD and float(pixels.std()) < _DEGENERATE_THRESHOLD + + +def _has_nvidia_gpu() -> bool: + try: + subprocess.run( + ["nvidia-smi"], + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + except (FileNotFoundError, subprocess.CalledProcessError): + return False + return True + + +def _detect_cuda_index_url() -> str: + """Return a PyTorch wheel index compatible with the reported CUDA runtime.""" + try: + report = subprocess.run( + ["nvidia-smi"], + check=True, + capture_output=True, + text=True, + ).stdout + except (FileNotFoundError, subprocess.CalledProcessError): + return "https://download.pytorch.org/whl/cu121" + import re + + match = re.search(r"CUDA Version:\s*(\d+)\.(\d+)", report) + if match is None: + return "https://download.pytorch.org/whl/cu121" + return f"https://download.pytorch.org/whl/cu{match.group(1)}{match.group(2)}" + + +def _backend_works(device: str) -> bool: + try: + probe = torch.tensor([1.0], device=device) # type: ignore[union-attr] + _ = probe + probe + except (AssertionError, RuntimeError): + return False + return True + + +def get_device() -> str: + """Select CUDA, XPU, MPS, or CPU in that order when each backend is usable.""" + if not _HAS_TORCH: + return "cpu" + if torch.cuda.is_available() and _backend_works("cuda"): # type: ignore[union-attr] + return "cuda" + xpu = getattr(torch, "xpu", None) + if xpu is not None and xpu.is_available() and _backend_works("xpu"): + return "xpu" + if _has_nvidia_gpu(): + logger.warning("NVIDIA GPU detected, but the installed PyTorch build has no working CUDA backend") + mps = getattr(getattr(torch, "backends", None), "mps", None) + if mps is not None and mps.is_available(): + return "mps" + return "cpu" + + +def _make_seed_generator(device: str, seed: int) -> Any: + """Create a deterministic generator, using CPU when device RNG is unavailable.""" + try: + return torch.Generator(device=device).manual_seed(seed) # type: ignore[union-attr] + except (RuntimeError, TypeError): + return torch.Generator().manual_seed(seed) # type: ignore[union-attr] + + +def _qwen_target_size(width: int, height: int) -> tuple[int, int]: + """Floor dimensions to Qwen's 16-pixel latent grid.""" + return max(16, width - width % 16), max(16, height - height % 16) + + +def _build_qwen_kwargs( + image: Image.Image, + strength: float, + num_inference_steps: int, + true_cfg_scale: float, + generator: Any, +) -> dict[str, Any]: + """Build the Qwen img2img call without importing its optional pipeline class.""" + width, height = _qwen_target_size(image.width, image.height) + return { + "prompt": _QWEN_PROMPT, + "negative_prompt": _QWEN_NEGATIVE, + "image": image, + "strength": strength, + "num_inference_steps": num_inference_steps, + "true_cfg_scale": true_cfg_scale, + "generator": generator, + "width": width, + "height": height, + } + + +class WatermarkRemover: + """Load one regeneration profile and write a metadata-clean raster output.""" + + DEFAULT_MODEL_ID = DEFAULT_MODEL_ID + DEFAULT_STRENGTH = DEFAULT_STRENGTH + CONTROLNET_CANNY_MODEL = CONTROLNET_CANNY_MODEL + _DEVICES = frozenset({"cpu", "mps", "cuda", "xpu"}) + + def __init__( + self, + model_id: str | None = None, + device: str | None = None, + torch_dtype: Any = None, + progress_callback: Callable[[str], None] | None = None, + hf_token: str | None = None, + pipeline: str = "controlnet", + controlnet_conditioning_scale: float = 1.0, + cpu_offload: bool = False, + ) -> None: + requested_model = model_id or self.DEFAULT_MODEL_ID + self.model_profile = normalize_profile(pipeline) + if self.model_profile == QWEN_ZIMAGE_PROFILE and model_id not in {None, self.DEFAULT_MODEL_ID}: + raise ValueError("The qwen-zimage profile uses a fixed Qwen-Image-2512 and Z-Image model stack.") + self.model_id = ( + "Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo" + if self.model_profile == QWEN_ZIMAGE_PROFILE + else requested_model + ) + _ensure_watermark_deps() + selected_device = (device or get_device()).casefold() + self.device = get_device() if selected_device == "auto" else selected_device + if self.device not in self._DEVICES: + raise ValueError(f"Unsupported device '{device}'. Use one of: auto, cpu, mps, cuda, xpu.") + + if torch_dtype is not None: + self.torch_dtype = torch_dtype + elif self.device in {"cpu", "mps"}: + self.torch_dtype = torch.float32 # type: ignore[union-attr] + elif self.model_profile in {"qwen", QWEN_ZIMAGE_PROFILE}: + self.torch_dtype = torch.bfloat16 # type: ignore[union-attr] + else: + self.torch_dtype = torch.float16 # type: ignore[union-attr] + + self.cpu_offload = cpu_offload + self.controlnet_conditioning_scale = controlnet_conditioning_scale + self.hf_token = hf_token or os.environ.get("HF_TOKEN") + self._progress_callback = progress_callback + self._pipeline: Any = None + self._controlnet_pipeline: Any = None + self._qwen_pipeline: Any = None + self._qwen_zimage_pipeline: Any = None + + def _set_progress(self, message: str) -> None: + if self._progress_callback is not None: + with contextlib.suppress(Exception): + self._progress_callback(message) + + def preload(self, *, global_only: bool = False) -> None: + """Materialize the selected model stack before the first request.""" + if self.model_profile == QWEN_ZIMAGE_PROFILE: + self._load_qwen_zimage_pipeline().preload(global_only=global_only) + elif self.model_profile == "qwen": + self._load_qwen_pipeline() + elif self.model_profile == "controlnet": + self._load_controlnet_pipeline() + else: + self._load_pipeline() + + def _base_load_kwargs(self) -> dict[str, Any]: + options: dict[str, Any] = {"torch_dtype": self.torch_dtype} + if self.hf_token: + options["token"] = self.hf_token + return options + + def _load_from_pretrained(self, cls: Any, model_id: str, **kwargs: Any) -> Any: + if self.torch_dtype == torch.float16: # type: ignore[union-attr] + try: + return cls.from_pretrained(model_id, variant="fp16", **kwargs) + except Exception as error: + logger.info("Model %s has no usable fp16 variant (%s); using default weights", model_id, error) + return cls.from_pretrained(model_id, **kwargs) + + def _maybe_add_fp16_vae(self, options: dict[str, Any]) -> None: + if not _needs_fp16_vae_fix( + self.model_id, + self.DEFAULT_MODEL_ID, + self.torch_dtype == torch.float16, # type: ignore[union-attr] + ): + return + from diffusers import AutoencoderKL + + options["vae"] = AutoencoderKL.from_pretrained(_SDXL_FP16_VAE_ID, torch_dtype=torch.float16) + + @staticmethod + def _disable_sdxl_watermarker(options: dict[str, Any]) -> None: + options["add_watermarker"] = False + + def _move_to_device_and_optimize(self, pipeline: Any) -> Any: + if self.cpu_offload and self.device == "cuda": + offload = getattr(pipeline, "enable_model_cpu_offload", None) + if not callable(offload): + raise RuntimeError("CPU offload was requested, but this pipeline does not support it.") + offload(device="cuda") + else: + try: + pipeline = pipeline.to(self.device) + except (RuntimeError, AssertionError) as error: + if self.device == "cuda": + raise RuntimeError( + f"Failed to move model to CUDA ({error}). Install a compatible PyTorch wheel from " + f"{_detect_cuda_index_url()}." + ) from error + raise + optimize = getattr(pipeline, "enable_xformers_memory_efficient_attention", None) + if callable(optimize): + with contextlib.suppress(Exception): + optimize() + if self.device == "mps": + slice_attention = getattr(pipeline, "enable_attention_slicing", None) + if callable(slice_attention): + with contextlib.suppress(Exception): + slice_attention("max") + return pipeline + + def _sdxl_options(self) -> dict[str, Any]: + options = self._base_load_kwargs() + self._disable_sdxl_watermarker(options) + self._maybe_add_fp16_vae(options) + return options + + def _load_pipeline(self) -> Any: + if self._pipeline is None: + options = self._sdxl_options() + options.update(safety_checker=None, requires_safety_checker=False) + loaded = self._load_from_pretrained(AutoImg2ImgPipeline, self.model_id, **options) + self._pipeline = self._move_to_device_and_optimize(loaded) + return self._pipeline + + def _load_controlnet_pipeline(self) -> Any: + if self._controlnet_pipeline is None: + from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline + + controlnet = self._load_from_pretrained( + ControlNetModel, + CONTROLNET_CANNY_MODEL, + torch_dtype=self.torch_dtype, + ) + options = self._sdxl_options() + options["controlnet"] = controlnet + loaded = self._load_from_pretrained( + StableDiffusionXLControlNetImg2ImgPipeline, + self.model_id, + **options, + ) + self._controlnet_pipeline = self._move_to_device_and_optimize(loaded) + return self._controlnet_pipeline + + def _load_qwen_pipeline(self) -> Any: + if self._qwen_pipeline is None: + try: + from diffusers import QwenImageImg2ImgPipeline + except ImportError as error: + raise ImportError("The qwen profile requires Diffusers with QwenImageImg2ImgPipeline.") from error + model_id = QWEN_MODEL_ID if self.model_id == self.DEFAULT_MODEL_ID else self.model_id + loaded = QwenImageImg2ImgPipeline.from_pretrained(model_id, **self._base_load_kwargs()) + self._qwen_pipeline = self._move_to_device_and_optimize(loaded) + return self._qwen_pipeline + + def _load_qwen_zimage_pipeline(self) -> Any: + if self._qwen_zimage_pipeline is None: + from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline + + self._qwen_zimage_pipeline = QwenZImagePipeline( + device=self.device, + torch_dtype=self.torch_dtype, + hf_token=self.hf_token, + progress_callback=self._progress_callback, + controlnet_conditioning_scale=self.controlnet_conditioning_scale, + keep_face_models_on_device=False if self.cpu_offload else None, + ) + return self._qwen_zimage_pipeline + + def _reload_on_cpu(self, cache_name: str, loader: Callable[[], Any]) -> Any: + self.device = "cpu" + self.torch_dtype = torch.float32 # type: ignore[union-attr] + setattr(self, cache_name, None) + return loader() + + def _run_img2img( + self, + init_image: Image.Image, + strength: float, + num_inference_steps: int, + guidance_scale: float, + generator: Any, + ) -> Image.Image: + from remove_ai_watermarks._internal.img2img_runner import run_img2img_with_mps_fallback + + output, device = run_img2img_with_mps_fallback( + self._load_pipeline, + init_image, + strength, + num_inference_steps, + guidance_scale, + generator, + self.device, + self._set_progress, + reload_on_cpu=lambda: self._reload_on_cpu("_pipeline", self._load_pipeline), + ) + self.device = device + return output + + def _build_canny_control_image(self, init_image: Image.Image) -> Image.Image: + import cv2 + import numpy as np + + gray = cv2.cvtColor(np.asarray(init_image.convert("RGB")), cv2.COLOR_RGB2GRAY) + edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH) + return Image.fromarray(np.repeat(edges[:, :, None], 3, axis=2)) + + def _run_controlnet( + self, + init_image: Image.Image, + strength: float, + num_inference_steps: int, + guidance_scale: float, + generator: Any, + ) -> Image.Image: + from remove_ai_watermarks._internal.img2img_runner import run_img2img_with_mps_fallback + + extras = { + "prompt": _CONTROLNET_PROMPT, + "negative_prompt": _CONTROLNET_NEGATIVE, + "control_image": self._build_canny_control_image(init_image), + "controlnet_conditioning_scale": float(self.controlnet_conditioning_scale), + } + output, device = run_img2img_with_mps_fallback( + self._load_controlnet_pipeline, + init_image, + strength, + num_inference_steps, + guidance_scale, + generator, + self.device, + self._set_progress, + reload_on_cpu=lambda: self._reload_on_cpu("_controlnet_pipeline", self._load_controlnet_pipeline), + extra_kwargs=extras, + ) + self.device = device + return output + + def _run_qwen( + self, + init_image: Image.Image, + strength: float, + num_inference_steps: int, + guidance_scale: float, + generator: Any, + ) -> Image.Image: + response = self._load_qwen_pipeline()( + **_build_qwen_kwargs(init_image, strength, num_inference_steps, guidance_scale, generator) + ) + return response.images[0] + + def _run_qwen_zimage( + self, + init_image: Image.Image, + strength: float, + seed: int | None, + *, + tile: bool = False, + tile_size: int = 1024, + tile_overlap: int = 128, + ) -> Image.Image: + return self._load_qwen_zimage_pipeline().run( + init_image, + strength=strength, + seed=seed, + tile=tile, + tile_size=tile_size, + tile_overlap=tile_overlap, + ) + + def _generate( + self, + image: Image.Image, + strength: float, + steps: int, + guidance: float, + generator: Any, + seed: int | None, + *, + tile: bool, + tile_size: int, + tile_overlap: int, + ) -> Image.Image: + if self.model_profile == QWEN_ZIMAGE_PROFILE: + return self._run_qwen_zimage( + image, + strength, + seed, + tile=tile, + tile_size=tile_size, + tile_overlap=tile_overlap, + ) + + runner = { + "qwen": self._run_qwen, + "controlnet": self._run_controlnet, + }.get(self.model_profile, self._run_img2img) + if tile and max(image.size) > tile_size: + from remove_ai_watermarks._internal.tiling import run_tiled + + return run_tiled( + lambda crop: runner(crop, strength, steps, guidance, generator), + image, + tile_size, + tile_overlap, + self._set_progress, + ) + return runner(image, strength, steps, guidance, generator) + + def _write_output(self, image: Image.Image, output_path: Path) -> None: + import numpy as np + + from remove_ai_watermarks import image_io + + output_path.parent.mkdir(parents=True, exist_ok=True) + bgr = np.ascontiguousarray(np.asarray(image.convert("RGB"))[:, :, ::-1]) + if not image_io.imwrite(str(output_path), bgr): + image.save(output_path) + from remove_ai_watermarks.metadata import remove_ai_metadata + + remove_ai_metadata(output_path, output_path, keep_standard=True) + + def remove_watermark( + self, + image_path: Path, + output_path: Path | None = None, + strength: float | None = None, + num_inference_steps: int | None = None, + guidance_scale: float | None = None, + seed: int | None = None, + vendor: str | None = None, + tile: bool = False, + tile_size: int = 1024, + tile_overlap: int = 128, + region: tuple[int, int, int, int] | None = None, + region_feather: int = 64, + ) -> Path: + """Regenerate image pixels and write the result without AI metadata.""" + if not image_path.exists(): + raise FileNotFoundError(f"Image not found: {image_path}") + destination = output_path or image_path + with Image.open(image_path) as opened: + source = opened.convert("RGB") + + if self.model_profile == QWEN_ZIMAGE_PROFILE: + from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise + + resolved_strength = strength if strength is not None else resolution_adaptive_denoise(*source.size) + else: + resolved_strength = resolve_strength(strength, vendor, self.model_profile) + if not 0.0 <= resolved_strength <= 1.0: + raise ValueError(f"Strength must be between 0.0 and 1.0, got {resolved_strength}") + + resolved_seed = resolve_seed(seed, self.model_profile) + steps = resolve_steps(num_inference_steps, self.model_profile) + guidance = ( + 1.0 if guidance_scale is None and self.model_profile == QWEN_ZIMAGE_PROFILE else guidance_scale or 7.5 + ) + if self.model_profile == QWEN_ZIMAGE_PROFILE: + if steps != 4: + raise ValueError("The qwen-zimage profile requires 4 steps.") + if guidance != 1.0: + raise ValueError("The qwen-zimage profile requires CFG 1.0.") + else: + steps = viable_steps(steps, resolved_strength) + + generator = None + if resolved_seed is not None and _HAS_TORCH: + generator = _make_seed_generator(self.device, resolved_seed) + result = self._generate( + source, + resolved_strength, + steps, + guidance, + generator, + resolved_seed, + tile=tile, + tile_size=tile_size, + tile_overlap=tile_overlap, + ) + + if self.torch_dtype == torch.float16 and _is_degenerate_image(result): # type: ignore[union-attr] + self.torch_dtype = torch.float32 # type: ignore[union-attr] + self._pipeline = self._controlnet_pipeline = self._qwen_pipeline = self._qwen_zimage_pipeline = None + result = self._generate( + source, + resolved_strength, + steps, + guidance, + generator, + resolved_seed, + tile=tile, + tile_size=tile_size, + tile_overlap=tile_overlap, + ) + + if region is not None: + import numpy as np + + from remove_ai_watermarks._internal.tiling import feather_region_composite + + if result.size != source.size: + result = result.resize(source.size, Image.Resampling.LANCZOS) + merged = feather_region_composite( + np.asarray(source), + np.asarray(result.convert("RGB")), + region, + feather=region_feather, + ) + result = Image.fromarray(merged) + + self._write_output(result, destination) + return destination + + def remove_watermark_batch( + self, + input_dir: Path, + output_dir: Path, + strength: float | None = None, + num_inference_steps: int | None = None, + extensions: tuple[str, ...] = (".png", ".jpg", ".jpeg", ".webp"), + ) -> list[Path]: + """Process matching files in a directory, logging and continuing on failures.""" + if not input_dir.exists(): + raise FileNotFoundError(f"Input directory not found: {input_dir}") + output_dir.mkdir(parents=True, exist_ok=True) + from remove_ai_watermarks._internal.img2img_runner import try_empty_device_cache + + outputs: list[Path] = [] + candidates = sorted(path for path in input_dir.iterdir() if path.suffix.casefold() in extensions) + for source in candidates: + try: + outputs.append(self.remove_watermark(source, output_dir / source.name, strength, num_inference_steps)) + except Exception as error: + logger.error("Failed to process %s: %s", source, error) + finally: + try_empty_device_cache(self.device) + return outputs + + +def remove_watermark( + image_path: Path, + output_path: Path | None = None, + strength: float | None = None, + model_id: str | None = None, + device: str | None = None, + hf_token: str | None = None, + region: tuple[int, int, int, int] | None = None, +) -> Path: + """Convenience wrapper using the default ControlNet profile.""" + from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength + + remover = WatermarkRemover(model_id=model_id, device=device, hf_token=hf_token) + return remover.remove_watermark( + image_path, + output_path, + strength, + vendor=vendor_for_strength(image_path), + region=region, + ) diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 3eb95c0..0979657 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -21,8 +21,8 @@ from typing import TYPE_CHECKING, Any, Literal, NoReturn import click from remove_ai_watermarks import __version__, image_io, watermark_registry -from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS -from remove_ai_watermarks.noai.watermark_profiles import ( +from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS +from remove_ai_watermarks._internal.watermark_profiles import ( resolve_seed, resolve_steps, resolve_strength, @@ -158,7 +158,7 @@ def _resolved_strength_for_display( if pipeline == "qwen-zimage" and strength is None: from PIL import Image - from remove_ai_watermarks.noai.qwen_zimage_pipeline import resolution_adaptive_denoise + from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise with Image.open(source) as image: return resolution_adaptive_denoise(image.width, image.height) @@ -269,7 +269,7 @@ def _normalize_pipeline(ctx: click.Context, param: click.Parameter, value: str | """ if value is None: return None - from remove_ai_watermarks.noai.watermark_profiles import normalize_profile + from remove_ai_watermarks._internal.watermark_profiles import normalize_profile normalized = normalize_profile(value) if value.strip().lower() == "default": diff --git a/src/remove_ai_watermarks/gemini_engine.py b/src/remove_ai_watermarks/gemini_engine.py index fa3dd73..d44f585 100644 --- a/src/remove_ai_watermarks/gemini_engine.py +++ b/src/remove_ai_watermarks/gemini_engine.py @@ -1,22 +1,6 @@ -"""Gemini visible-sparkle detector and localizer (cv2/numpy, no GPU). +"""Locate the visible Gemini sparkle and build a mask for shared inpainting.""" -Locates the Google Gemini / Nano Banana sparkle so the shared fill (region_eraser) -can inpaint it. Detection is a multi-scale NCC search against the captured sparkle -alpha template (ported from GeminiWatermarkTool's Snap Engine; original author -Allen Kuo (allenk), https://github.com/allenk/GeminiWatermarkTool), scored by a -spatial + gradient + variance fusion with a false-positive gate. ``footprint_mask`` -returns the sparkle footprint (captured alpha thresholded low to include the halo, -then dilated) as a full-frame mask for the fill. - -The captured alpha maps are background captures of the sparkle on pure-black -backgrounds (48x48 for small images, 96x96 for large). NB: they are used here only -to DETECT and to shape the removal mask -- the old reverse-alpha pixel recovery -(``original = (watermarked - a*logo)/(1-a)``) is gone; removal is localize -> fill. -""" - -# cv2/numpy boundary: cv2 and numpy ship no usable type info for the array ops -# below, so strict pyright cannot know their element types. Relax the unknown-type -# rules for this file only; the public signatures are still annotated with NDArray[Any]. +# OpenCV and NumPy expose incomplete types at this array-processing boundary. # pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false from __future__ import annotations @@ -41,263 +25,172 @@ logger = logging.getLogger(__name__) class WatermarkSize(Enum): - """Watermark size mode based on image dimensions.""" + """Provider size tier selected from the source dimensions.""" - SMALL = "small" # 48x48, for images <= 1024x1024 - LARGE = "large" # 96x96, for images > 1024x1024 + SMALL = "small" + LARGE = "large" @dataclass class DetectionResult: - """Result of watermark detection.""" + """Detection decision and its component scores.""" detected: bool = False confidence: float = 0.0 - region: tuple[int, int, int, int] = (0, 0, 0, 0) # x, y, w, h + region: tuple[int, int, int, int] = (0, 0, 0, 0) size: WatermarkSize = WatermarkSize.SMALL - - # stage scores spatial_score: float = 0.0 gradient_score: float = 0.0 variance_score: float = 0.0 -@dataclass +@dataclass(frozen=True, slots=True) class WatermarkPosition: - """Watermark position configuration.""" + """Expected provider margins and logo size.""" margin_right: int margin_bottom: int logo_size: int def get_position(self, image_width: int, image_height: int) -> tuple[int, int]: - """Get top-left position for a given image size.""" - x = image_width - self.margin_right - self.logo_size - y = image_height - self.margin_bottom - self.logo_size - return (x, y) + return image_width - self.margin_right - self.logo_size, image_height - self.margin_bottom - self.logo_size -def get_watermark_config(width: int, height: int) -> WatermarkPosition: - """Get the appropriate watermark configuration based on image size. +@dataclass(frozen=True, slots=True) +class _Candidate: + scale: int + x: int + y: int + spatial: float + gradient: float = 0.0 + variance: float = 0.0 - Rules discovered from Gemini: - - W > 1024 AND H > 1024: 96x96 logo at (W-64-96, H-64-96) - - Otherwise: 48x48 logo at (W-32-48, H-32-48) - """ - if width > 1024 and height > 1024: - return WatermarkPosition(margin_right=64, margin_bottom=64, logo_size=96) - return WatermarkPosition(margin_right=32, margin_bottom=32, logo_size=48) + @property + def fused(self) -> float: + if self.spatial < 0.25: + return max(0.0, self.spatial * 0.5) + return self.spatial * 0.50 + self.gradient * 0.30 + self.variance * 0.20 def get_watermark_size(width: int, height: int) -> WatermarkSize: - """Determine watermark size mode from image dimensions.""" - if width > 1024 and height > 1024: - return WatermarkSize.LARGE - return WatermarkSize.SMALL + """Return the provider's large tier only when both axes exceed 1024.""" + return WatermarkSize.LARGE if width > 1024 and height > 1024 else WatermarkSize.SMALL -def _calculate_alpha_map(bg_capture: NDArray[Any]) -> NDArray[Any]: - """Calculate alpha map from a background capture. +def get_watermark_config(width: int, height: int) -> WatermarkPosition: + """Return the observed standard placement for the selected size tier.""" + if get_watermark_size(width, height) is WatermarkSize.LARGE: + return WatermarkPosition(64, 64, 96) + return WatermarkPosition(32, 32, 48) - The alpha map represents how much the watermark affects each pixel. - alpha = max(R, G, B) / 255.0 - """ - if len(bg_capture.shape) == 2: - gray = bg_capture.astype(np.float32) - elif bg_capture.shape[2] >= 3: - # Use max of channels for brightness - gray = np.max(bg_capture[:, :, :3], axis=2).astype(np.float32) + +def _calculate_alpha_map(background_capture: NDArray[Any]) -> NDArray[Any]: + """Convert a black-background sparkle capture to a normalized opacity map.""" + if background_capture.ndim == 2: + intensity = background_capture + elif background_capture.shape[2] >= 3: + intensity = background_capture[:, :, :3].max(axis=2) else: - gray = bg_capture[:, :, 0].astype(np.float32) - - return gray / 255.0 + intensity = background_capture[:, :, 0] + return intensity.astype(np.float32) / 255.0 -def _load_embedded_asset(name: str) -> NDArray[Any]: - """Load an embedded PNG asset and decode it with OpenCV.""" - asset_path = Path(__file__).parent / "assets" / name - if not asset_path.exists(): - raise FileNotFoundError(f"Embedded asset not found: {asset_path}") - - data = asset_path.read_bytes() - buf = np.frombuffer(data, dtype=np.uint8) - img = cv2.imdecode(buf, cv2.IMREAD_COLOR) - if img is None: - raise RuntimeError(f"Failed to decode embedded asset: {name}") - return img +def _load_capture(filename: str, expected_side: int) -> NDArray[Any]: + capture = image_io.imread(Path(__file__).parent / "assets" / filename, cv2.IMREAD_COLOR) + if capture is None: + raise RuntimeError(f"Failed to decode embedded asset: {filename}") + if capture.shape[:2] != (expected_side, expected_side): + capture = cv2.resize(capture, (expected_side, expected_side), interpolation=cv2.INTER_AREA) + return capture -# Single source of truth for the multi-scale template ladder (aggressively downscaled to -# slightly upscaled): the precomputed `_tmpl_cache` and the `_scan_scales` loop must use -# the SAME scales or a scan scale would miss the cache and KeyError. -_TEMPLATE_SCALES: tuple[int, ...] = tuple(range(16, 120, 2)) +def _gray_float(image: NDArray[Any]) -> NDArray[Any]: + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 and image.shape[2] >= 3 else image + return gray.astype(np.float32) / 255.0 + + +def _overlaps(candidate: _Candidate, selected: _Candidate) -> bool: + radius = 0.5 * max(candidate.scale, selected.scale) + return abs(candidate.x - selected.x) < radius and abs(candidate.y - selected.y) < radius + + +_TEMPLATE_SCALES = tuple(range(16, 120, 2)) class GeminiEngine: - """Detects and localizes the visible Gemini sparkle for the shared fill removal. + """Project-native detector and mask builder for the white Gemini sparkle.""" - The multi-scale NCC detection is a Python port of the GeminiWatermarkTool C++ - Snap Engine; ``footprint_mask`` turns a detection into a removal mask. - """ - - # Body pixels at >= this fraction of the peak captured alpha define the sparkle - # "core", sampled by the detection FP-gate's core-vs-ring brightness margin - # (:meth:`_core_and_bg`). _CORE_ALPHA_FRAC = 0.8 - - # Sparkle false-positive gate. A real Gemini sparkle is a bright WHITE overlay, - # so its core sits above the local background; a shape-only NCC match on ornate - # or flat content (text, banners, hatching) can score >0.5 without that lift. - # Demote a detection that is BOTH low-confidence AND low core-ring brightness - # margin -- the joint signature of a content false positive (verified on the - # detector calibration: demoted examples were visual false positives or a - # near-invisible white-on-white sparkle whose AI verdict is held by metadata - # anyway). Real sparkles escape via EITHER high confidence - # (white-bg sparkles score >=0.79 despite a low margin) OR high margin (dark/mid - # backgrounds, incl. the #36 faint-corner case, lift well clear), so both must - # fail to demote. _SPARKLE_FP_CONF = 0.65 _SPARKLE_FP_MARGIN = 5.0 - # Bright-background content false positives (2026-06-26 landing-page FPs: a snow+sky - # photo and a white-background product render both scored ~0.51). The margin gate - # above cannot catch them -- a bright background gives the "core" a HIGH core-ring - # margin (it is genuinely brighter than its surroundings), so the brightness check - # reads it as a real overlay. The discriminating signature is the GRADIENT NCC: a - # real white sparkle is a crisp star silhouette (grad ~0.97-1.0 on the synthetic - # composites, ~0.96 on the real #36 corner sparkle), while a smooth luminance blob - # that shape-NCC-matches the rough outline has low gradient fidelity (the two FPs - # measured 0.105 and 0.463). So ALSO demote a low-confidence match whose gradient - # NCC is below this floor, regardless of margin -- 0.55 sits well above the worst FP - # (0.463) and far below every real sparkle (>=0.8). This only ADDS demotions on - # bright backgrounds (a real bright-bg sparkle keeps grad ~0.97), so it cannot - # regress a dark/mid sparkle (already kept by margin) or a white-bg one (kept by - # confidence >= 0.65, above the gate). _SPARKLE_FP_GRAD = 0.55 - - # White-core rescue for the gate above. A real but FAINT sparkle -- a soft white - # star on a bright/textured background -- has a high core-ring margin but low - # gradient fidelity, the SAME signature the grad gate uses to demote the smooth - # colored-corner FP, so faint real sparkles get demoted with it. The separator the - # grad gate discards is the CORE COLOR: a real Gemini sparkle core is near-WHITE - # (low saturation), while a clean bright corner that shape-matches (sky, sun, a warm - # light) is COLORED. So do NOT demote a low-grad match that already clears the trust - # confidence (_SPARKLE_KEEP_CONF -- the registry's 0.5 sparkle gate plus a small - # margin so the ~0.51 bright-background FPs the grad gate was added for stay demoted) - # AND has a bright (margin) near-neutral core (_core_saturation <= _SPARKLE_WHITE_SAT). - # Calibrated on metadata-stripped faint sparkles to recover low-gradient marks - # without materially increasing clean false fires. _SPARKLE_KEEP_CONF = 0.52 _SPARKLE_WHITE_SAT = 0.20 - - # Corner promotion (issue #36): the size weight that suppresses tiny-patch - # false positives also buries a small, near-perfect sparkle when a larger, - # mediocre match sits elsewhere (e.g. a bright collar in a portrait). A small - # faint sparkle on a busy background therefore loses the global argmax and the - # image reads as clean -- the regression osachub reported when the search - # window widened 256px -> 512px (v0.7.2's tighter window still found it). - # Remedy: if the bottom-right corner holds a very-high-fidelity raw-NCC match, - # trust it regardless of size, without reverting the wider window (which is - # needed for variant margins). The threshold sits midway between the worst - # real-photo corner match (~0.78 across native + downscaled real photos) and a - # genuine faint sparkle (~0.93), so it adds true detections without adding - # false ones; it only ever overrides a lower-fidelity global pick, so it cannot - # weaken an existing detection. _CORNER_PROMOTE_NCC = 0.85 - # Bottom-right corner side for the promotion search, as a fraction of the - # image's short side, clamped to an absolute pixel band. Relative so the corner - # stays a true corner at every scale: a fixed 256 px is a genuine corner on a - # large image but covers ~70% of a small portrait, where a busy real photo can - # then raw-match the star template at ~0.81 (only 0.04 below the promote gate). - # Scaling the side down on small images drops that worst case to ~0.69, while - # the upper clamp stops it ballooning on huge images (more corner area = more - # random texture to false-match -- a real photo reached ~0.83 at 512 px). The - # Gemini sparkle sits ~60-160 px from the corner (fixed margins, not - # proportional), and the [96, 384] band covers that at every measured size. _CORNER_PROMOTE_FRAC = 0.20 _CORNER_PROMOTE_MIN = 96 _CORNER_PROMOTE_MAX = 384 - - # Number of top size-weighted spatial candidates scored by full fusion before one - # is selected. The single size-weighted argmax can bury a genuine mid-size sparkle - # under a LARGER, lower-fidelity shape match (the 256->512 search-widening - # regression: a real corner sparkle at raw ~0.77 lost to a decoy at raw ~0.63). - # Scoring the top-K by gradient-bearing fusion rescues it. Top-K (NOT the raw-NCC - # argmax) keeps the tiny-patch suppression intact: a coincidental 16 px match never - # ranks in the size-weighted top-K, so widening selection cannot add a false - # positive on non-Gemini content (verified on the doubao/jimeng visible corpora). _SELECT_TOPK = 3 + _MASK_ALPHA = 0.04 + _MASK_DILATE_FRAC = 0.18 def __init__(self, logo_value: float = 255.0) -> None: - """Initialize the engine with embedded alpha maps. - - Args: - logo_value: The logo brightness value (default 255.0 = white). - """ self.logo_value = logo_value - - # Load embedded background captures - bg_small = _load_embedded_asset("gemini_bg_48.png") - bg_large = _load_embedded_asset("gemini_bg_96.png") - - # Ensure correct sizes - if bg_small.shape[:2] != (48, 48): - bg_small = cv2.resize(bg_small, (48, 48), interpolation=cv2.INTER_AREA) - if bg_large.shape[:2] != (96, 96): - bg_large = cv2.resize(bg_large, (96, 96), interpolation=cv2.INTER_AREA) - - # Calculate alpha maps - self._alpha_small = _calculate_alpha_map(bg_small) - self._alpha_large = _calculate_alpha_map(bg_large) - - # Per-scale resized templates are constant (``_alpha_large`` never changes), - # so precompute the whole fixed 16..118 ladder once: ``_scan_scales`` runs it on - # every image (twice -- global + corner), and re-``resize``-ing the 96x96 source - # each time is pure repeated work. Prebuilt (not lazy) so the dict is read-only - # after construction and safe to share across threads via the module singleton. + self._alpha_small = _calculate_alpha_map(_load_capture("gemini_bg_48.png", 48)) + self._alpha_large = _calculate_alpha_map(_load_capture("gemini_bg_96.png", 96)) self._tmpl_cache: dict[int, NDArray[Any]] = { - scale: cv2.resize(self._alpha_large, (scale, scale), interpolation=cv2.INTER_AREA) - for scale in _TEMPLATE_SCALES + side: cv2.resize(self._alpha_large, (side, side), interpolation=cv2.INTER_AREA) for side in _TEMPLATE_SCALES } - logger.debug( - "Alpha maps loaded: small=%s, large=%s", - self._alpha_small.shape, - self._alpha_large.shape, - ) - def get_alpha_map(self, size: WatermarkSize) -> NDArray[Any]: - """Get the base alpha map for a specific standard size.""" - if size == WatermarkSize.SMALL: - return self._alpha_small - return self._alpha_large + return self._alpha_small if size is WatermarkSize.SMALL else self._alpha_large def get_interpolated_alpha(self, size_px: int) -> NDArray[Any]: - """Create an interpolated alpha map dynamically scaled from the high-res 96x96 base.""" - source = self._alpha_large - if size_px == source.shape[1]: - return source.copy() - - interp = cv2.INTER_LINEAR if size_px > source.shape[1] else cv2.INTER_AREA - return cv2.resize(source, (size_px, size_px), interpolation=interp) - - # ── Detection ──────────────────────────────────────────────────── + if size_px == self._alpha_large.shape[1]: + return self._alpha_large.copy() + method = cv2.INTER_LINEAR if size_px > self._alpha_large.shape[1] else cv2.INTER_AREA + return cv2.resize(self._alpha_large, (size_px, size_px), interpolation=method) def _scan_scales(self, gray: NDArray[Any]) -> Iterator[tuple[int, float, tuple[int, int]]]: - """Yield ``(scale, max_ncc, max_loc)`` for the alpha template matched at each scale. - - Shared multi-scale ``TM_CCOEFF_NORMED`` primitive over a normalized [0, 1] - grayscale region, used by both the size-weighted global search in - ``detect_watermark`` and the raw-NCC corner pass in ``_corner_promote`` -- - each applies its own scoring/argmax to the yielded values. The 96x96 - ``_alpha_large`` is the high-quality source downscaled per scale; the range - covers aggressively downscaled to slightly upscaled logos. - """ - for scale in _TEMPLATE_SCALES: - if scale > gray.shape[0] or scale > gray.shape[1]: + """Yield the strongest normalized template match at every usable scale.""" + height, width = gray.shape[:2] + for side, template in self._tmpl_cache.items(): + if side > height or side > width: continue - match_res = cv2.matchTemplate(gray, self._tmpl_cache[scale], cv2.TM_CCOEFF_NORMED) - _, max_val, _, max_loc = cv2.minMaxLoc(match_res) - yield scale, float(max_val), max_loc + response = cv2.matchTemplate(gray, template, cv2.TM_CCOEFF_NORMED) + _minimum, maximum, _min_location, max_location = cv2.minMaxLoc(response) + yield side, float(maximum), max_location + + def _global_candidates(self, image: NDArray[Any]) -> list[_Candidate]: + height, width = image.shape[:2] + search_side = min(height, width, 512) + origin_x, origin_y = width - search_side, height - search_side + gray = _gray_float(image[origin_y:height, origin_x:width]) + ranked = sorted( + ( + ( + score * min(1.0, (side / 96.0) ** 0.5), + _Candidate(side, origin_x + location[0], origin_y + location[1], score), + ) + for side, score, location in self._scan_scales(gray) + ), + key=lambda item: (item[0], item[1].scale, item[1].spatial, item[1].x, item[1].y), + reverse=True, + ) + selected: list[_Candidate] = [] + for _weighted, candidate in ranked: + if any(_overlaps(candidate, prior) for prior in selected): + continue + selected.append(candidate) + if len(selected) == self._SELECT_TOPK: + break + return selected + + def _score_candidate(self, image: NDArray[Any], candidate: _Candidate) -> _Candidate: + if candidate.spatial < 0.25: + return candidate + gradient, variance = self._grad_var_scores(image, candidate.scale, candidate.x, candidate.y) + return _Candidate(candidate.scale, candidate.x, candidate.y, candidate.spatial, gradient, variance) def detect_watermark( self, @@ -306,241 +199,94 @@ class GeminiEngine: *, trust_provenance: bool = False, ) -> DetectionResult: - """Detect Gemini watermark using multi-scale Snap Engine logic (ported from C++ vendor algorithm). - - ``trust_provenance`` signals that external metadata already proves this is a - Google generation (C2PA issuer "Google"/"Gemini"). The false-positive gate - exists only to reject content that shape-matches the sparkle on NON-Google - images (Doubao text, ornate corners); when provenance confirms Google, that - gate would demote a genuine sparkle the vendor moved/re-rendered (bigger, - lighter, shifted), so it is skipped. The caller (registry) still applies the - relaxed provenance trust gate to the returned confidence.""" + """Return the strongest sparkle-shaped bottom-right candidate.""" result = DetectionResult() - if image is None or image.size == 0: return result - # Normalize to 3-channel BGR: the multi-scale search tolerates grayscale, but - # the FP-gate / alpha-gain helpers (_core_and_bg) reduce over axis=2 and would - # crash on a 2D/BGRA input reaching this public entry point (e.g. via the - # registry detect adapter or the library API). - image = image_io.to_bgr(image) - h, w = image.shape[:2] - base_size = force_size or get_watermark_size(w, h) - result.size = base_size - - # Dynamically search bottom-right corner. 512 covers up to 512px from the - # corner -- enough for known Gemini margin variations (standard: 64+96=160px; - # observed variants up to ~300px). 256 was too tight and caused misses. - search_size = int(min(min(w, h), 512)) - sx1 = max(0, w - search_size) - sy1 = max(0, h - search_size) - - search_region = image[sy1:h, sx1:w] - if len(search_region.shape) == 3 and search_region.shape[2] >= 3: - gray_sr = cv2.cvtColor(search_region, cv2.COLOR_BGR2GRAY) - else: - gray_sr = search_region.copy() - - gray_sr_f = gray_sr.astype(np.float32) / 255.0 - - # Phase 1 & 2: multi-scale spatial NCC search. The size weight (mimicking the - # C++ vendor weight) overcomes the NCC bias toward tiny patches, but its single - # argmax can bury a genuine mid-size sparkle under a LARGER, lower-fidelity - # shape match (the 256->512 search-widening regression). So score the top-K - # size-weighted candidates by the FULL fusion and keep the highest -- the - # gradient term separates a true white sparkle from a shape-only decoy. See - # _SELECT_TOPK for why top-K (not the raw-NCC argmax) preserves tiny-patch - # suppression and so cannot add a false positive on non-Gemini content. - scored: list[tuple[float, int, int, int, float]] = [] # (adj, scale, raw, x, y) - for scale, max_val, max_loc in self._scan_scales(gray_sr_f): - adj_val = max_val * min(1.0, (scale / 96.0) ** 0.5) - scored.append((adj_val, scale, max_val, sx1 + max_loc[0], sy1 + max_loc[1])) - scored.sort(reverse=True) - - # Top-K candidates at distinct locations (NMS: drop a lower-ranked match that - # overlaps an already-kept one -- the same sparkle matches at adjacent scales). - candidates: list[tuple[int, int, int, float]] = [] - for _adj, scale, raw, x, y in scored: - if any( - abs(x - px) < 0.5 * max(scale, ps) and abs(y - py) < 0.5 * max(scale, ps) - for ps, px, py, _ in candidates - ): - continue - candidates.append((scale, x, y, raw)) - if len(candidates) >= self._SELECT_TOPK: - break - - # Corner promotion: a near-perfect small bottom-right sparkle the size weight - # buries even below the top-K (see _CORNER_PROMOTE_NCC) -- add it as a candidate. - promoted = self._corner_promote(image, candidates[0][3] if candidates else -1.0) + source = image_io.to_bgr(image) + height, width = source.shape[:2] + result.size = force_size or get_watermark_size(width, height) + candidates = self._global_candidates(source) + promoted = self._corner_promote(source, candidates[0].spatial if candidates else -1.0) if promoted is not None: - candidates.append(promoted) - - # No candidate at any scale: the search region is smaller than the 16px template - # floor (an image whose short side is < 16px), so nothing is detectable. Return - # the empty (detected=False) result rather than dereferencing candidates[0]. + candidates.append(_Candidate(promoted[0], promoted[1], promoted[2], promoted[3])) if not candidates: return result - # Select the candidate with the highest full-fusion confidence (pre-FP-gate). - best_scale, pos_x, pos_y, best_raw_ncc = candidates[0] - grad_score, var_score, best_fused = 0.0, 0.0, -1.0 - for c_scale, c_x, c_y, c_raw in candidates: - if c_raw < 0.25: - c_grad, c_var, c_fused = 0.0, 0.0, max(0.0, c_raw * 0.5) - else: - c_grad, c_var = self._grad_var_scores(image, c_scale, c_x, c_y) - c_fused = c_raw * 0.50 + c_grad * 0.30 + c_var * 0.20 - if c_fused > best_fused: - best_fused = c_fused - best_scale, pos_x, pos_y = c_scale, c_x, c_y - best_raw_ncc, grad_score, var_score = c_raw, c_grad, c_var + best = max((self._score_candidate(source, candidate) for candidate in candidates), key=lambda item: item.fused) + result.region = (best.x, best.y, best.scale, best.scale) + result.spatial_score = float(best.spatial) + result.gradient_score = float(best.gradient) + result.variance_score = float(best.variance) - result.region = (pos_x, pos_y, best_scale, best_scale) - result.spatial_score = float(best_raw_ncc) - result.gradient_score = float(grad_score) - result.variance_score = float(var_score) - - if result.spatial_score < 0.25: - result.confidence = float(max(0.0, result.spatial_score * 0.5)) - return result - - # ── Fusion ─────────────────────────────────────────────────── - # best_fused is the selected candidate's spatial*0.5 + grad*0.3 + var*0.2. - confidence = best_fused - - # False-positive gate: a low-confidence match that shows NEITHER real-sparkle - # signature is a content false positive, not a white sparkle overlay. A real - # sparkle proves itself by a bright core (high core-ring margin, on dark/mid - # backgrounds) OR a crisp star silhouette (high gradient NCC, on any background - # incl. bright). Demote when both are weak -- this catches the dark/mid no-core - # FP (low margin) AND the bright-background smooth-blob FP (high margin but low - # gradient), which the margin check alone misses. See _SPARKLE_FP_GRAD. - if confidence < self._SPARKLE_FP_CONF and not trust_provenance: - alpha = self.get_interpolated_alpha(best_scale) - pos = (pos_x, pos_y) - margin = self._core_ring_margin(image, alpha, pos) - low_margin = margin is not None and margin < self._SPARKLE_FP_MARGIN - low_grad = grad_score < self._SPARKLE_FP_GRAD - if low_margin or low_grad: - # White-core rescue: a real faint sparkle clears the trust confidence, - # has a bright core (not low_margin), and a near-WHITE core -- unlike the - # colored-corner FP the low-grad demotion targets. See _SPARKLE_WHITE_SAT. - core_sat = self._core_saturation(image, alpha, pos) - white_core = not low_margin and core_sat is not None and core_sat <= self._SPARKLE_WHITE_SAT - if not (confidence >= self._SPARKLE_KEEP_CONF and white_core): - logger.debug( - "Sparkle FP gate: conf=%.3f, margin=%s, grad=%.3f, core_sat=%s; demoting.", - confidence, - f"{margin:.1f}" if margin is not None else "n/a", - grad_score, - f"{core_sat:.2f}" if core_sat is not None else "n/a", - ) - confidence = min(confidence, 0.30) - - result.confidence = float(max(0.0, min(1.0, confidence))) + confidence = best.fused + if best.spatial >= 0.25 and confidence < self._SPARKLE_FP_CONF and not trust_provenance: + confidence = self._apply_false_positive_gate(source, best, confidence) + result.confidence = float(np.clip(confidence, 0.0, 1.0)) result.detected = result.confidence >= 0.35 - - logger.debug( - "Detection: spatial=%.3f, grad=%.3f, var=%.3f → conf=%.3f (%s)", - result.spatial_score, - result.gradient_score, - var_score, - result.confidence, - "DETECTED" if result.detected else "not detected", - ) - return result - def _grad_var_scores( - self, - image: NDArray[Any], - scale: int, - pos_x: int, - pos_y: int, - ) -> tuple[float, float]: - """Return ``(gradient_score, variance_score)`` for a candidate sparkle. - - Factored out of ``detect_watermark`` so each top-K candidate can be scored by - the full fusion before one is selected. The gradient NCC correlates - Sobel-magnitude maps (shape fidelity, contrast-robust); the variance score - rewards a flat overlay region against the row band above it. - """ - h, w = image.shape[:2] - x1, y1 = pos_x, pos_y - x2, y2 = min(w, x1 + scale), min(h, y1 + scale) - region = image[y1:y2, x1:x2] - gray_region = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY) if region.ndim == 3 and region.shape[2] >= 3 else region - gray_f = gray_region.astype(np.float32) / 255.0 - alpha_region = self.get_interpolated_alpha(scale)[: y2 - y1, : x2 - x1] - - # ── Gradient NCC ── - img_gmag = cv2.magnitude( - cv2.Sobel(gray_f, cv2.CV_32F, 1, 0, ksize=3), cv2.Sobel(gray_f, cv2.CV_32F, 0, 1, ksize=3) + def _apply_false_positive_gate(self, image: NDArray[Any], candidate: _Candidate, confidence: float) -> float: + alpha = self.get_interpolated_alpha(candidate.scale) + position = (candidate.x, candidate.y) + margin = self._core_ring_margin(image, alpha, position) + low_margin = margin is not None and margin < self._SPARKLE_FP_MARGIN + low_gradient = candidate.gradient < self._SPARKLE_FP_GRAD + if not low_margin and not low_gradient: + return confidence + saturation = self._core_saturation(image, alpha, position) + neutral_core = not low_margin and saturation is not None and saturation <= self._SPARKLE_WHITE_SAT + if confidence >= self._SPARKLE_KEEP_CONF and neutral_core: + return confidence + logger.debug( + "Sparkle candidate demoted: confidence=%.3f, margin=%s, gradient=%.3f, saturation=%s", + confidence, + margin, + candidate.gradient, + saturation, ) - alpha_gmag = cv2.magnitude( - cv2.Sobel(alpha_region, cv2.CV_32F, 1, 0, ksize=3), cv2.Sobel(alpha_region, cv2.CV_32F, 0, 1, ksize=3) + return min(confidence, 0.30) + + def _grad_var_scores(self, image: NDArray[Any], scale: int, pos_x: int, pos_y: int) -> tuple[float, float]: + height, width = image.shape[:2] + x2, y2 = min(width, pos_x + scale), min(height, pos_y + scale) + region = image[pos_y:y2, pos_x:x2] + gray = _gray_float(region) + alpha = self.get_interpolated_alpha(scale)[: y2 - pos_y, : x2 - pos_x] + + image_edges = cv2.magnitude( + cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3), + cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3), ) - _, grad_score, _, _ = cv2.minMaxLoc(cv2.matchTemplate(img_gmag, alpha_gmag, cv2.TM_CCOEFF_NORMED)) - - # ── Variance ── - var_score = 0.0 - ref_h = min(y1, scale) - if ref_h > 8: - ref_region = image[y1 - ref_h : y1, x1:x2] - gray_ref = cv2.cvtColor(ref_region, cv2.COLOR_BGR2GRAY) if ref_region.ndim == 3 else ref_region - _, s_wm = cv2.meanStdDev(gray_region) - _, s_ref = cv2.meanStdDev(gray_ref) - if s_ref[0][0] > 5.0: - var_score = max(0.0, min(1.0, 1.0 - (s_wm[0][0] / s_ref[0][0]))) - return float(grad_score), float(var_score) - - def _corner_promote( - self, - image: NDArray[Any], - current_raw_ncc: float, - ) -> tuple[int, int, int, float] | None: - """Search the bottom-right corner for a very-high-fidelity sparkle match. - - Returns ``(scale, x, y, raw_ncc)`` when the corner holds a match with raw - NCC >= ``_CORNER_PROMOTE_NCC`` that beats the global pick's ``current_raw_ncc``, - else None. Used to rescue a small sparkle that the size weight buried under - a larger, lower-fidelity match elsewhere. See ``_CORNER_PROMOTE_NCC`` and - ``_CORNER_PROMOTE_FRAC`` for the corner sizing. - """ - h, w = image.shape[:2] - side = max( - self._CORNER_PROMOTE_MIN, min(self._CORNER_PROMOTE_MAX, round(min(w, h) * self._CORNER_PROMOTE_FRAC)) + alpha_edges = cv2.magnitude( + cv2.Sobel(alpha, cv2.CV_32F, 1, 0, ksize=3), + cv2.Sobel(alpha, cv2.CV_32F, 0, 1, ksize=3), ) - cs = int(min(min(w, h), side)) - cx1, cy1 = max(0, w - cs), max(0, h - cs) - corner = image[cy1:h, cx1:w] - gray = cv2.cvtColor(corner, cv2.COLOR_BGR2GRAY) if corner.ndim == 3 and corner.shape[2] >= 3 else corner - gray = gray.astype(np.float32) / 255.0 + response = cv2.matchTemplate(image_edges, alpha_edges, cv2.TM_CCOEFF_NORMED) + _minimum, gradient, _min_location, _max_location = cv2.minMaxLoc(response) - best_raw = -1.0 - best_scale = 0 - best_loc = (0, 0) - for scale, max_val, max_loc in self._scan_scales(gray): - if max_val > best_raw: - best_raw = max_val - best_scale = scale - best_loc = max_loc + variance = 0.0 + reference_height = min(pos_y, scale) + if reference_height > 8: + reference = image[pos_y - reference_height : pos_y, pos_x:x2] + reference_gray = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY) if reference.ndim == 3 else reference + _mean, region_std = cv2.meanStdDev((gray * 255.0).astype(np.uint8)) + _reference_mean, reference_std = cv2.meanStdDev(reference_gray) + if reference_std[0][0] > 5.0: + variance = float(np.clip(1.0 - region_std[0][0] / reference_std[0][0], 0.0, 1.0)) + return float(gradient), variance - if best_raw >= self._CORNER_PROMOTE_NCC and best_raw > current_raw_ncc: - return best_scale, cx1 + best_loc[0], cy1 + best_loc[1], float(best_raw) - return None - - # ── Removal ────────────────────────────────────────────────────── - - # Footprint mask for the localize -> fill removal path. The mask must cover the - # WHOLE sparkle including its faint semi-transparent halo, not just the bright - # core, or the fill leaves a visible ring. Threshold the captured alpha low - # (>_MASK_ALPHA catches the halo the core-only 0.10 misses) then dilate by a - # sparkle-relative margin so alignment slop and the outermost halo are absorbed. - _MASK_ALPHA = 0.04 - _MASK_DILATE_FRAC = 0.18 # dilation radius as a fraction of the sparkle scale + def _corner_promote(self, image: NDArray[Any], current_raw_ncc: float) -> tuple[int, int, int, float] | None: + height, width = image.shape[:2] + desired = round(min(width, height) * self._CORNER_PROMOTE_FRAC) + side = min(min(width, height), max(self._CORNER_PROMOTE_MIN, min(self._CORNER_PROMOTE_MAX, desired))) + origin_x, origin_y = width - side, height - side + matches = self._scan_scales(_gray_float(image[origin_y:height, origin_x:width])) + best = max(matches, key=lambda item: item[1], default=None) + if best is None or best[1] < self._CORNER_PROMOTE_NCC or best[1] <= current_raw_ncc: + return None + return best[0], origin_x + best[2][0], origin_y + best[2][1], float(best[1]) def footprint_mask( self, @@ -550,50 +296,37 @@ class GeminiEngine: dilate: int | None = None, region: tuple[int, int, int, int] | None = None, ) -> NDArray[Any] | None: - """Full-frame uint8 mask (255 = sparkle) of the sparkle footprint, for the - shared fill removal path (cv2 / MI-GAN / LaMa), or None. - - The footprint is the interpolated captured alpha at the detected scale, - thresholded LOW so the faint halo is included, then dilated by a - sparkle-relative margin. When ``force`` and nothing is detected, falls back to - the default sparkle slot for the image size (the ``--no-detect`` path). - - ``region`` is the already-resolved ``(x, y, scale)`` from the caller's detection - (the registry passes the decision's provenance-aware region). When given, the - mask is built from it directly WITHOUT a second internal detect -- otherwise a - provenance/assume-relaxed sparkle would be re-demoted by the strict re-detect and - yield no mask (reported-removed-but-unchanged). Absent ``region``, direct callers - keep the detect-then-force behavior. - """ + """Build a full-frame mask from a resolved or newly detected sparkle.""" if image is None or image.size == 0: - return None # guard before to_bgr (cvtColor raises on an empty Mat); mirror detect_watermark - image = image_io.to_bgr(image) - h, w = image.shape[:2] + return None + source = image_io.to_bgr(image) + height, width = source.shape[:2] if region is not None: - x, y, scale = region[0], region[1], region[2] + x, y, scale = region[:3] else: - det = self.detect_watermark(image) - if det.detected: - x, y, scale = det.region[0], det.region[1], det.region[2] + detection = self.detect_watermark(source) + if detection.detected: + x, y, scale = detection.region[:3] elif force: - cfg = get_watermark_config(w, h) - x, y = cfg.get_position(w, h) - scale = cfg.logo_size + config = get_watermark_config(width, height) + x, y = config.get_position(width, height) + scale = config.logo_size else: return None - alpha = self.get_interpolated_alpha(scale) - fp = self._footprint_indices(alpha, (x, y), image.shape) - if fp is None: + + placed = self._footprint_indices(self.get_interpolated_alpha(scale), (x, y), source.shape) + if placed is None: return None - aroi, (y1, y2, x1, x2) = fp - sil = (aroi > self._MASK_ALPHA).astype(np.uint8) * 255 - if int((sil > 0).sum()) == 0: + alpha, (y1, y2, x1, x2) = placed + silhouette = (alpha > self._MASK_ALPHA).astype(np.uint8) * 255 + if not silhouette.any(): return None - mask = np.zeros((h, w), np.uint8) - mask[y1:y2, x1:x2] = sil - d = dilate if dilate is not None else max(13, int(scale * self._MASK_DILATE_FRAC)) - if d > 0: - mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * d + 1, 2 * d + 1))) + mask = np.zeros((height, width), dtype=np.uint8) + mask[y1:y2, x1:x2] = silhouette + radius = dilate if dilate is not None else max(13, int(scale * self._MASK_DILATE_FRAC)) + if radius > 0: + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1, 2 * radius + 1)) + mask = cv2.dilate(mask, kernel) return mask def _footprint_indices( @@ -602,21 +335,35 @@ class GeminiEngine: position: tuple[int, int], image_shape: tuple[int, ...], ) -> tuple[NDArray[Any], tuple[int, int, int, int]] | None: - """Return (alpha_roi, (y1, y2, x1, x2)) for the placed footprint, or None. - - Shared by the over-subtraction test and the inpaint mask so both operate on - exactly the same clipped, in-bounds region. - """ x, y = position - ah, aw = alpha_map.shape[:2] - ih, iw = image_shape[:2] + alpha_height, alpha_width = alpha_map.shape[:2] + image_height, image_width = image_shape[:2] x1, y1 = max(0, x), max(0, y) - x2, y2 = min(iw, x + aw), min(ih, y + ah) + x2, y2 = min(image_width, x + alpha_width), min(image_height, y + alpha_height) if x1 >= x2 or y1 >= y2: return None - ax1, ay1 = x1 - x, y1 - y - alpha_roi = alpha_map[ay1 : ay1 + (y2 - y1), ax1 : ax1 + (x2 - x1)] - return alpha_roi, (y1, y2, x1, x2) + alpha_x, alpha_y = x1 - x, y1 - y + clipped = alpha_map[alpha_y : alpha_y + y2 - y1, alpha_x : alpha_x + x2 - x1] + return clipped, (y1, y2, x1, x2) + + def _core_mask_and_box( + self, + image: NDArray[Any], + alpha_map: NDArray[Any], + position: tuple[int, int], + ) -> tuple[NDArray[Any], NDArray[Any], tuple[int, int, int, int], float] | None: + placed = self._footprint_indices(alpha_map, position, image.shape) + if placed is None: + return None + alpha, bounds = placed + peak = float(alpha.max()) + if peak < 0.2: + return None + core = alpha >= peak * self._CORE_ALPHA_FRAC + if not core.any(): + return None + y1, y2, x1, x2 = bounds + return core, image[y1:y2, x1:x2], bounds, peak def _core_and_bg( self, @@ -624,41 +371,22 @@ class GeminiEngine: alpha_map: NDArray[Any], position: tuple[int, int], ) -> tuple[float, float, float] | None: - """Return ``(core_obs, bg, a_cap)`` for the placed sparkle, or None. - - ``core_obs`` is the bright-core brightness (75th pct over the high-alpha - core), ``bg`` the local background ring median, ``a_cap`` the captured peak - alpha. Shared by the alpha-gain estimate and the false-positive margin gate. - None when the footprint or the background ring cannot be sampled. - """ - placed = self._footprint_indices(alpha_map, position, image.shape) - if placed is None: + sample = self._core_mask_and_box(image, alpha_map, position) + if sample is None: return None - alpha_roi, (y1, y2, x1, x2) = placed - a_cap = float(alpha_roi.max()) - if a_cap < 0.2: - return None - core = alpha_roi >= a_cap * self._CORE_ALPHA_FRAC - if not bool(core.any()): - return None - # Convert only the footprint+ring crop to gray, not the whole image: every - # sample below lives inside the ring box, so a full-image mean is wasted work - # that scales with resolution (~70 ms on a 12 MP image, recomputed for both - # the alpha-gain estimate and the over-subtraction gate). The crop is sized by - # the footprint, so this is O(footprint^2) regardless of image size. - ih, iw = image.shape[:2] - pad = int((x2 - x1) * 0.7) - ry1, ry2 = max(0, y1 - pad), min(ih, y2 + pad) - rx1, rx2 = max(0, x1 - pad), min(iw, x2 + pad) - ring = image[ry1:ry2, rx1:rx2].astype(np.float32).mean(axis=2) - # Footprint box expressed in ring-crop coordinates. + core, _box, (y1, y2, x1, x2), peak = sample + height, width = image.shape[:2] + padding = int((x2 - x1) * 0.7) + ry1, ry2 = max(0, y1 - padding), min(height, y2 + padding) + rx1, rx2 = max(0, x1 - padding), min(width, x2 + padding) + luminance = image[ry1:ry2, rx1:rx2].astype(np.float32).mean(axis=2) fy1, fy2, fx1, fx2 = y1 - ry1, y2 - ry1, x1 - rx1, x2 - rx1 - core_obs = float(np.percentile(ring[fy1:fy2, fx1:fx2][core], 75)) - ring_mask = np.ones(ring.shape, dtype=bool) - ring_mask[fy1:fy2, fx1:fx2] = False - if int(ring_mask.sum()) < 10: + core_value = float(np.percentile(luminance[fy1:fy2, fx1:fx2][core], 75)) + background = np.ones(luminance.shape, dtype=bool) + background[fy1:fy2, fx1:fx2] = False + if background.sum() < 10: return None - return core_obs, float(np.median(ring[ring_mask])), a_cap + return core_value, float(np.median(luminance[background])), peak def _core_ring_margin( self, @@ -666,14 +394,8 @@ class GeminiEngine: alpha_map: NDArray[Any], position: tuple[int, int], ) -> float | None: - """Bright-core brightness minus the local background ring (gray levels). - - A real white sparkle overlay lifts its core above the surroundings; a - shape-only NCC false positive on ornate/flat content does not. None when the - background ring cannot be sampled. - """ - cb = self._core_and_bg(image, alpha_map, position) - return None if cb is None else cb[0] - cb[1] + sample = self._core_and_bg(image, alpha_map, position) + return None if sample is None else sample[0] - sample[1] def _core_saturation( self, @@ -681,57 +403,24 @@ class GeminiEngine: alpha_map: NDArray[Any], position: tuple[int, int], ) -> float | None: - """Median color saturation of the sparkle core (0 = white/neutral, higher = - colored). A real Gemini sparkle is a white star, so its core is near-neutral; - a clean bright corner that shape-matches (sky, sun, a warm light) is colored, - so a high core saturation flags the false positive the brightness/gradient - gates miss. Samples the same high-alpha core pixels as :meth:`_core_and_bg`. - None when the footprint cannot be placed or the core is empty. - """ - placed = self._footprint_indices(alpha_map, position, image.shape) - if placed is None: + sample = self._core_mask_and_box(image, alpha_map, position) + if sample is None: return None - alpha_roi, (y1, y2, x1, x2) = placed - a_cap = float(alpha_roi.max()) - if a_cap < 0.2: - return None - core = alpha_roi >= a_cap * self._CORE_ALPHA_FRAC - box = image[y1:y2, x1:x2] - if box.shape[:2] != core.shape or not bool(core.any()): - return None - px = box[core].astype(np.float32) # (N, 3) BGR core pixels - hi = px.max(axis=1) - lo = px.min(axis=1) - return float(np.median((hi - lo) / (hi + 1.0))) + core, box, _bounds, _peak = sample + pixels = box[core].astype(np.float32) + brightest = pixels.max(axis=1) + darkest = pixels.min(axis=1) + return float(np.median((brightest - darkest) / (brightest + 1.0))) @functools.lru_cache(maxsize=1) def _shared_engine() -> GeminiEngine: - """Process-wide default ``GeminiEngine`` singleton. - - The engine holds only constant assets (embedded captures, alpha maps, the - precomputed template ladder) and takes the image as a method argument, so one - instance is reused across every ``detect_sparkle_confidence`` call instead of - reloading assets + recomputing alpha maps + rebuilding the template cache on - each of the ~34k images an ``identify`` batch scans. Output is identical.""" return GeminiEngine() def detect_sparkle_confidence(image_path: Path, *, image: NDArray[Any] | None = None) -> float | None: - """Visible-sparkle detection confidence for a file, for provenance use. - - Loads the image with cv2 and runs :meth:`GeminiEngine.detect_watermark`. - Returns the NCC confidence in [0, 1], or None if the image cannot be read - (cv2 returns None for unsupported containers such as HEIC). Kept here so the - cv2 dependency stays in this module; callers apply their own threshold. - - ``image`` lets a caller that has already decoded the file (e.g. ``identify`` - running several visible-mark detectors) pass the BGR array to avoid a second - full decode; when None the file is read from ``image_path``. - """ - from remove_ai_watermarks import image_io - - img = image if image is not None else image_io.imread(image_path) - if img is None: + """Return the local sparkle confidence, or None when decoding fails.""" + decoded = image if image is not None else image_io.imread(image_path) + if decoded is None: return None - return float(_shared_engine().detect_watermark(img).confidence) + return float(_shared_engine().detect_watermark(decoded).confidence) diff --git a/src/remove_ai_watermarks/humanizer.py b/src/remove_ai_watermarks/humanizer.py index 5ba484e..0df352e 100644 --- a/src/remove_ai_watermarks/humanizer.py +++ b/src/remove_ai_watermarks/humanizer.py @@ -1,7 +1,7 @@ """Post-processing filters for the cleaned output. -``apply_analog_humanizer`` injects film grain and chromatic aberration to defeat -digital AI-perfection classifiers (ported from NeuralBleach); ``unsharp_mask`` +``apply_analog_humanizer`` injects film grain and chromatic aberration to reduce +overly uniform digital surfaces; ``unsharp_mask`` counters the soft, over-smoothed look that the diffusion pass leaves behind (itself a common "this is AI" tell). """ @@ -16,10 +16,7 @@ from numpy.typing import NDArray def apply_analog_humanizer(image: NDArray, grain_intensity: float = 4.0, chromatic_shift: int = 1) -> NDArray: """ - Apply Analog Humanizer (film grain and chromatic aberration) to an image. - This simulates analog film imperfections to defeat digital AI perfection classifiers. - - Ported from NeuralBleach. + Apply shared-luminance grain and a small lateral color offset. Args: image: BGR image as numpy array (uint8). @@ -33,26 +30,22 @@ def apply_analog_humanizer(image: NDArray, grain_intensity: float = 4.0, chromat if len(image.shape) != 3 or image.shape[2] != 3: return image.copy() - # Split channels (OpenCV uses BGR) - # B = 0, G = 1, R = 2 + # Translate the outer color channels without circular edge wrapping. b, g, r = cv2.split(image) - - # 1. Chromatic Aberration - # Shift R channel left, B channel right. np.roll is circular, so it wraps - # the opposite edge into a thin colored fringe at the L/R borders; replicate - # the original edge columns there to keep the intended offset interior-only. - # Clamp so the edge-replication slices below always have a source column: a shift - # >= width would leave them empty and crash the broadcast (r[:, -shift:] = (H, 0)). - shift = min(chromatic_shift, image.shape[1] - 1) + shift = min(max(0, chromatic_shift), max(0, image.shape[1] - 1)) if shift > 0: - r = np.roll(r, -shift, axis=1) - r[:, -shift:] = r[:, -shift - 1 : -shift] - b = np.roll(b, shift, axis=1) - b[:, :shift] = b[:, shift : shift + 1] + shifted_b = np.empty_like(b) + shifted_b[:, :shift] = b[:, :1] + shifted_b[:, shift:] = b[:, :-shift] + b = shifted_b + + shifted_r = np.empty_like(r) + shifted_r[:, :-shift] = r[:, shift:] + shifted_r[:, -shift:] = r[:, -1:] + r = shifted_r merged = cv2.merge((b, g, r)) - # 2. Film Grain (Gaussian Noise) if grain_intensity > 0: img_f = merged.astype(np.float32) noise = np.random.normal(0, grain_intensity, img_f.shape).astype(np.float32) diff --git a/src/remove_ai_watermarks/identify.py b/src/remove_ai_watermarks/identify.py index 9c92c72..3b3e9b6 100644 --- a/src/remove_ai_watermarks/identify.py +++ b/src/remove_ai_watermarks/identify.py @@ -25,6 +25,18 @@ import logging from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, cast +from remove_ai_watermarks._internal.c2pa import ( + c2pa_info_from_manifest_store, + cbor_text_after, + extract_c2pa_info, + soft_binding_vendors_in, +) +from remove_ai_watermarks._internal.constants import ( + C2PA_AI_TOOLS, + C2PA_AI_VENDORS, + C2PA_IDENTITY_AI_ORGS, + C2PA_ISSUERS, +) from remove_ai_watermarks.metadata import ( AI_METADATA_KEYS, AIGC_MARKERS, @@ -46,18 +58,6 @@ from remove_ai_watermarks.metadata import ( xai_signature, xai_signature_pair, ) -from remove_ai_watermarks.noai.c2pa import ( - c2pa_info_from_manifest_store, - cbor_text_after, - extract_c2pa_info, - soft_binding_vendors_in, -) -from remove_ai_watermarks.noai.constants import ( - C2PA_AI_TOOLS, - C2PA_AI_VENDORS, - C2PA_IDENTITY_AI_ORGS, - C2PA_ISSUERS, -) from remove_ai_watermarks.watermark_registry import GEMINI_SPARKLE_TRUST_CONF if TYPE_CHECKING: @@ -822,7 +822,7 @@ def _identify_from_evidence( issuers = [info["issuer"]] if info.get("issuer") else _issuers_in(head) # Full AI generation (trainedAlgorithmicMedia) vs an AI-enhanced real photo # (compositeWithTrainedAlgorithmicMedia). The structured kind is parsed once in - # noai.c2pa._populate_registry_fields (covers PNG + any container the c2pa-python + # _internal.c2pa._populate_registry_fields (covers PNG + any container the c2pa-python # reader handles); fall back to a raw head scan for the non-PNG raw-blob path # where extract_c2pa_info returns {}. Full generation wins when both appear. c2pa_source_kind = info.get("ai_source_kind") diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index 0e38d9e..317bcac 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -1,7 +1,4 @@ -"""Invisible watermark removal engine. - -Wraps the vendored noai-watermark code for removing invisible AI watermarks -(SynthID, StableSignature, TreeRing) via diffusion-based regeneration. +"""Diffusion engine for regenerating images that carry invisible AI watermarks. This module requires the 'gpu' extra dependencies: uv pip install 'remove-ai-watermarks[diffusion]' @@ -19,10 +16,10 @@ import warnings from pathlib import Path from typing import TYPE_CHECKING, Any -from .noai.watermark_profiles import ( +from ._internal.watermark_profiles import ( DEFAULT_MODEL_ID as DEFAULT_SDXL_MODEL_ID, ) -from .noai.watermark_profiles import ( +from ._internal.watermark_profiles import ( resolve_seed, ) @@ -81,9 +78,6 @@ def _target_size(width: int, height: int, max_resolution: int, min_resolution: i class InvisibleEngine: """Remove invisible AI watermarks using diffusion model regeneration. - Based on noai-watermark by mertizci: - https://github.com/mertizci/noai-watermark - The approach encodes the image into latent space, injects controlled noise to break watermark patterns, and reconstructs via reverse diffusion. """ @@ -124,7 +118,7 @@ class InvisibleEngine: residency. CUDA only. """ - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover effective_model = model_id or self.DEFAULT_MODEL_ID diff --git a/src/remove_ai_watermarks/metadata.py b/src/remove_ai_watermarks/metadata.py index 6d9252e..4e0c07d 100644 --- a/src/remove_ai_watermarks/metadata.py +++ b/src/remove_ai_watermarks/metadata.py @@ -1,7 +1,4 @@ -"""AI metadata detection and removal. - -Wraps the noai-watermark metadata handling for stripping AI-generation -metadata (EXIF, PNG text chunks, C2PA provenance) from images. +"""Detect and remove AI provenance metadata from image containers. For metadata-only operations, the heavy ML dependencies are NOT required. """ @@ -104,7 +101,7 @@ IPTC_AI_MARKERS: tuple[bytes, ...] = ( # (Meta / Instagram / MidJourney) use ``trainedAlgorithmicMedia``. Including the bare # token flagged clean procedural images as AI (is_ai=high + has_invisible_target=True -> # a diffusion scrub of clean content), contradicting the c2pa layer, which sets -# source_type without ai_source for it (tests/test_noai.py::test_plain_algorithmic_media_not_flagged_ai). +# source_type without ai_source for it (tests/test_metadata_internals.py::test_plain_algorithmic_media_not_flagged_ai). # It is not a substring of the trained/composite tokens, so its removal does not affect # their detection. @@ -218,7 +215,7 @@ def _is_ai_value(value: str) -> bool: detection: NovelAI stamps a generic ``Title``/``Source`` text chunk (an AI-shaped value under a non-AI key) that ``_is_ai_key`` alone would keep. """ - from remove_ai_watermarks.noai.constants import AI_GENERATOR_TOKENS + from remove_ai_watermarks._internal.constants import AI_GENERATOR_TOKENS value_lower = value.lower() return any(token in value_lower for token in AI_GENERATOR_TOKENS) @@ -310,7 +307,7 @@ def _scan_head_impl(image_path: Path, size: int) -> bytes: with open(image_path, "rb") as f: head = f.read(size) # Lazy import: isobmff imports this module's constants at top level. - from remove_ai_watermarks.noai import isobmff + from remove_ai_watermarks._internal import isobmff if isobmff.is_isobmff(head): region = isobmff.scan_c2pa_region(image_path) @@ -348,7 +345,7 @@ def has_ai_metadata(image_path: Path) -> bool: # Check C2PA — via the official c2pa-python reader first (spec-tracking, every # container it supports), then a binary scan that also catches AVIF/HEIF/JPEG-XL # containers and synthetic/partial blobs the validator rejects. - from remove_ai_watermarks.noai.c2pa import read_manifest_store_json + from remove_ai_watermarks._internal.c2pa import read_manifest_store_json if read_manifest_store_json(image_path) is not None: return True @@ -466,7 +463,7 @@ def aigc_label(image_path: Path) -> dict[str, str] | None: # in ``moov.udta.meta.keys`` and points to a raw JSON value in ``ilst``. # Read it through the bounded box walker so a tail ``moov`` after a large # ``mdat`` is found without loading or scanning the media payload. - from remove_ai_watermarks.noai.isobmff import tc260_aigc_payloads + from remove_ai_watermarks._internal.isobmff import tc260_aigc_payloads isobmff_candidates = tuple(payload.decode("utf-8", "replace") for payload in tc260_aigc_payloads(image_path)) if result := aigc_label_from_metadata(b"", isobmff_candidates): @@ -475,7 +472,7 @@ def aigc_label(image_path: Path) -> dict[str, str] | None: # Native MKV/WebM TC260 metadata: ``Segment.Tags.Tag.SimpleTag`` carries # ``TagName=AIGC`` and the raw JSON in ``TagString``. The EBML walker seeks # over clusters and reads only bounded metadata values. - from remove_ai_watermarks.noai.ebml import tc260_aigc_payloads as ebml_tc260_aigc_payloads + from remove_ai_watermarks._internal.ebml import tc260_aigc_payloads as ebml_tc260_aigc_payloads ebml_candidates = tuple(payload.decode("utf-8", "replace") for payload in ebml_tc260_aigc_payloads(image_path)) if result := aigc_label_from_metadata(b"", ebml_candidates): @@ -486,11 +483,11 @@ def aigc_label(image_path: Path) -> dict[str, str] | None: # that could collide inside compressed video. legacy_payloads: tuple[bytes, ...] = () if image_path.suffix.lower() == ".avi": - from remove_ai_watermarks.noai.riff import tc260_aigc_payloads as riff_tc260_aigc_payloads + from remove_ai_watermarks._internal.riff import tc260_aigc_payloads as riff_tc260_aigc_payloads legacy_payloads = riff_tc260_aigc_payloads(image_path) elif image_path.suffix.lower() == ".flv": - from remove_ai_watermarks.noai.flv import tc260_aigc_payloads as flv_tc260_aigc_payloads + from remove_ai_watermarks._internal.flv import tc260_aigc_payloads as flv_tc260_aigc_payloads legacy_payloads = flv_tc260_aigc_payloads(image_path) legacy_candidates = tuple(payload.decode("utf-8", "replace") for payload in legacy_payloads) @@ -672,7 +669,7 @@ def synthid_source(image_path: Path) -> str | None: Returns: Comma-joined vendor name(s) (e.g. ``"OpenAI"``) or None. """ - from remove_ai_watermarks.noai.c2pa import extract_c2pa_info, synthid_vendors_in + from remove_ai_watermarks._internal.c2pa import extract_c2pa_info, synthid_vendors_in # PNG: the caBX chunk parser gives a clean, structured issuer. vendors = extract_c2pa_info(image_path).get("synthid_vendors") @@ -694,7 +691,7 @@ def synthid_source(image_path: Path) -> str | None: def generator_from_metadata(candidates: Iterable[str], scan: bytes = b"") -> str | None: """Return a known AI generator from collected EXIF, PNG, or XMP values.""" - from remove_ai_watermarks.noai.constants import AI_GENERATOR_TOKENS + from remove_ai_watermarks._internal.constants import AI_GENERATOR_TOKENS creator_tools = ( match.group(1).decode("latin1", "replace") @@ -844,7 +841,7 @@ def _ai_exif_targets(loaded: dict[str, Any]) -> list[tuple[str, int, bytes, str] """ import piexif - from remove_ai_watermarks.noai.constants import AI_GENERATOR_TOKENS + from remove_ai_watermarks._internal.constants import AI_GENERATOR_TOKENS ifd0: dict[int, Any] = loaded.get("0th") or {} ifde: dict[int, Any] = loaded.get("Exif") or {} @@ -903,7 +900,7 @@ def get_ai_metadata(image_path: Path) -> dict[str, str]: """ from PIL import Image - from remove_ai_watermarks.noai.c2pa import extract_c2pa_info, soft_binding_vendors_in, synthid_verdict + from remove_ai_watermarks._internal.c2pa import extract_c2pa_info, soft_binding_vendors_in, synthid_verdict result: dict[str, str] = {} @@ -923,7 +920,7 @@ def get_ai_metadata(image_path: Path) -> dict[str, str]: except Exception as exc: logger.debug("PIL could not open %s for AI-metadata scan: %s", image_path, exc) - # C2PA manifest fields from the single canonical parser (noai/c2pa.py). + # C2PA manifest fields from the single canonical parser (_internal/c2pa.py). c2pa = extract_c2pa_info(image_path) for key in ( "c2pa_manifest", @@ -1215,7 +1212,7 @@ def remove_ai_metadata( # offset-preserving streaming path; images retain the in-memory item scrub # needed for XMP/EXIF inside mdat/idat. Route the remaining formats by suffix # OR by an ``ftyp`` content sniff. - from remove_ai_watermarks.noai.isobmff import ( + from remove_ai_watermarks._internal.isobmff import ( blank_ai_exif_tokens, blank_ai_xmp_packets, blank_tc260_aigc_tags, diff --git a/src/remove_ai_watermarks/noai/c2pa.py b/src/remove_ai_watermarks/noai/c2pa.py deleted file mode 100644 index 81da1de..0000000 --- a/src/remove_ai_watermarks/noai/c2pa.py +++ /dev/null @@ -1,562 +0,0 @@ -"""C2PA (Coalition for Content Provenance and Authenticity) metadata handling. - -Reading goes through the official c2pa-python ``Reader`` first (any container it -supports), via ``extract_c2pa_info`` / ``read_manifest_store_json``. The -hand-rolled PNG ``caBX`` JUMBF-chunk tools below (``has_c2pa_metadata`` / -``extract_c2pa_chunk`` / ``inject_c2pa_chunk`` and the ``_extract_c2pa_info_png`` -fallback) cover raw-chunk extraction, re-injection, and the cases the validator -rejects (synthetic/partial blobs, a broken/absent wheel). Known issuers: - -- Google Imagen -- Adobe Firefly -- Microsoft Designer -- OpenAI (ChatGPT, GPT-4o, Sora, DALL-E) -- Truepic (signing authority) - -The fallback parser uses byte-level scanning — it does not validate JUMBF/CBOR -structure but reliably identifies known signatures, issuers, tools, and actions. -The vendor / source-type / SynthID / soft-binding registry scan -(``_populate_registry_fields``) is shared by both the reader and fallback paths. -""" - -from __future__ import annotations - -import contextlib -import functools -import json -import logging -import re -import struct -from pathlib import Path -from typing import Any, cast - -from remove_ai_watermarks.noai.constants import ( - C2PA_ACTIONS, - C2PA_AI_TOOLS, - C2PA_CHUNK_TYPE, - C2PA_ISSUERS, - C2PA_SIGNATURES, - C2PA_SOFT_BINDINGS, - PNG_SIGNATURE, - SYNTHID_C2PA_ISSUERS, -) - -logger = logging.getLogger(__name__) - -# Official C2PA reader (c2pa-python, a default dependency). It is the primary, -# spec-tracking manifest parser; the hand-rolled caBX/CBOR scanner below stays as -# a fallback for synthetic/partial blobs the validator rejects. The import is -# guarded so a partially-broken install degrades to the byte-scan rather than -# crashing the dependency-light identify path. -_C2paReader: Any = None -with contextlib.suppress(Exception): # broken/absent wheel -> byte-scan fallback - from c2pa import Reader as _C2paReader # pyright: ignore[reportMissingTypeStubs] -_C2PA_READER_AVAILABLE = _C2paReader is not None - - -def reader_available() -> bool: - """True when the official c2pa-python Reader imported successfully.""" - return _C2PA_READER_AVAILABLE - - -def read_manifest_store_json(image_path: Path) -> str | None: - """Return the full C2PA manifest-store JSON for ``image_path``, or None. - - Uses the official c2pa-python ``Reader`` (any container it supports: PNG, - JPEG, WebP, AVIF/HEIF, MP4, ...). Returns None when the reader is unavailable, - the file carries no parseable manifest, or parsing fails. The JSON is the - WHOLE store (every manifest plus ingredient manifests), matching the - whole-chunk semantics of the legacy byte scan -- an AI-source marker in a - parent/ingredient manifest (e.g. a ChatGPT edit of a Sora generation) is - still seen. - - Memoized per (path, mtime): one identify/get_ai_metadata call invokes the - structured parser ~3 times on the same file, so the cache turns the repeated - crypto-validating reads into one. - """ - if not _C2PA_READER_AVAILABLE: - return None - try: - mtime = image_path.stat().st_mtime_ns - except OSError: - return _read_manifest_store_impl(str(image_path)) - return _read_manifest_store_cached(str(image_path), mtime) - - -@functools.lru_cache(maxsize=8) -def _read_manifest_store_cached(path_str: str, _mtime_ns: int) -> str | None: - """Cache shim: ``_mtime_ns`` is part of the key only (invalidates on change).""" - return _read_manifest_store_impl(path_str) - - -def _read_manifest_store_impl(path_str: str) -> str | None: - # try_create returns None when there is no manifest; a default Reader does no - # trust enforcement, so an untrusted signer still yields the manifest content - # (we report what is in the file, we do not gate on certificate trust). - try: - reader = _C2paReader.try_create(path_str) - except Exception as exc: # malformed manifest, unsupported container, etc. - logger.debug("c2pa Reader could not parse %s: %s", path_str, exc) - return None - if reader is None: - return None - try: - with reader: - return reader.json() - except Exception as exc: # pragma: no cover - reader opened but json() failed - logger.debug("c2pa Reader.json() failed on %s: %s", path_str, exc) - return None - - -def has_c2pa_metadata(image_path: Path) -> bool: - """ - Check if an image contains C2PA metadata. - - Args: - image_path: Path to the image file. - - Returns: - True if C2PA metadata is detected, False otherwise. - """ - image_path = Path(image_path) - - if image_path.suffix.lower() != ".png": - return False - - try: - with open(image_path, "rb") as f: - signature = f.read(8) - if signature != PNG_SIGNATURE: - return False - - file_size = f.seek(0, 2) - f.seek(8) - - while True: - chunk_header = f.read(8) - if len(chunk_header) < 8: - break - - length = struct.unpack(">I", chunk_header[:4])[0] - chunk_type = chunk_header[4:8] - # Clamp the attacker-controlled 32-bit length to the bytes that - # actually remain, so a malformed huge length can't allocate GBs. - safe_length = max(0, min(length, file_size - f.tell())) - - if chunk_type == C2PA_CHUNK_TYPE: - chunk_data = f.read(safe_length) - # Check for any C2PA signature - for sig in C2PA_SIGNATURES: - if sig in chunk_data: - return True - # Also check if chunk_data itself contains C2PA-like patterns - if b"jumb" in chunk_data.lower() or b"c2pa" in chunk_data.lower(): - return True - f.read(4) - else: - f.seek(safe_length + 4, 1) - - if chunk_type == b"IEND": - break - except Exception: - pass - - return False - - -def _claim_generator_from_store(store: dict[str, Any]) -> str | None: - """Structured claim-generator name from the active manifest of a store dict. - - Prefers the top-level ``claim_generator`` string (Firefly: "Adobe_Firefly"), - falling back to the first ``claim_generator_info[].name`` (ChatGPT keys it - only there). isprintable() guards against odd binary-ish values. - """ - active = _active_manifest(store) - generator: Any = active.get("claim_generator") - if not (isinstance(generator, str) and generator): - info_list: list[Any] = active.get("claim_generator_info") or [] - if info_list and isinstance(first := info_list[0], dict): - generator = cast("dict[str, Any]", first).get("name") - return generator if isinstance(generator, str) and generator and generator.isprintable() else None - - -def _active_manifest(store: dict[str, Any]) -> dict[str, Any]: - """The active manifest dict from a manifest-store dict, or {} when absent.""" - manifests: Any = store.get("manifests") - if not isinstance(manifests, dict): - return {} - active = cast("dict[str, Any]", manifests).get(store.get("active_manifest", "")) - return cast("dict[str, Any]", active) if isinstance(active, dict) else {} - - -def _info_from_store(store: dict[str, Any], store_bytes: bytes) -> dict[str, Any]: - """Build normalized C2PA info from one parsed manifest store.""" - c2pa_info: dict[str, Any] = { - "has_c2pa": True, - "type": "C2PA (Coalition for Content Provenance and Authenticity)", - "c2pa_manifest": f"C2PA manifest store ({len(store_bytes)} bytes)", - } - # The whole-store JSON carries every vendor / source-type / SynthID / - # soft-binding signature (across active + ingredient manifests), so the same - # registry scan that runs on the raw caBX chunk applies unchanged here. - _populate_registry_fields(store_bytes, c2pa_info) - - if generator := _claim_generator_from_store(store): - c2pa_info["claim_generator"] = generator - sig: Any = _active_manifest(store).get("signature_info") - if isinstance(sig, dict) and (time := cast("dict[str, Any]", sig).get("time")): - c2pa_info["timestamp"] = str(time) - return c2pa_info - - -def _info_from_store_json(store_json: str) -> dict[str, Any]: - """Build the C2PA info dict from a c2pa-python manifest-store JSON string.""" - store_bytes = store_json.encode("utf-8") - try: - parsed: Any = json.loads(store_json) - except (ValueError, TypeError): - parsed = {} - store = cast("dict[str, Any]", parsed) if isinstance(parsed, dict) else {} - return _info_from_store(store, store_bytes) - - -def c2pa_info_from_manifest_store(store: str | dict[str, Any]) -> dict[str, Any]: - """Build normalized C2PA evidence from an externally collected manifest store. - - ``store`` may be the JSON string returned by ``c2pa.Reader.json()`` or its - decoded dictionary form. This is the non-file-backed counterpart to - :func:`extract_c2pa_info`. - """ - if isinstance(store, dict): - parsed = store - try: - store_json = json.dumps(store, ensure_ascii=False) - except (TypeError, ValueError): - return {} - else: - store_json = store - try: - decoded: Any = json.loads(store_json) - except (TypeError, ValueError): - return {} - if not isinstance(decoded, dict): - return {} - parsed = cast("dict[str, Any]", decoded) - if not store_json or not parsed or parsed.get("error"): - return {} - return _info_from_store(parsed, store_json.encode("utf-8")) - - -def extract_c2pa_info(image_path: Path) -> dict[str, Any]: - """ - Extract C2PA metadata information from an image. - - Uses the official c2pa-python reader first (any supported container), falling - back to the hand-rolled PNG caBX parser when the reader is unavailable or the - file carries no parseable manifest (synthetic/partial blobs). - - Args: - image_path: Path to the image file. - - Returns: - Dictionary containing C2PA metadata info, or {} when none is found. - """ - image_path = Path(image_path) - - if (store_json := read_manifest_store_json(image_path)) is not None: - return _info_from_store_json(store_json) - - return _extract_c2pa_info_png(image_path) - - -def _extract_c2pa_info_png(image_path: Path) -> dict[str, Any]: - """Fallback PNG caBX parser, used when the c2pa-python reader finds nothing.""" - c2pa_info: dict[str, Any] = {} - - if not has_c2pa_metadata(image_path): - return c2pa_info - - c2pa_info["has_c2pa"] = True - c2pa_info["type"] = "C2PA (Coalition for Content Provenance and Authenticity)" - - try: - with open(image_path, "rb") as f: - signature = f.read(8) - if signature != PNG_SIGNATURE: - return c2pa_info - - file_size = f.seek(0, 2) - f.seek(8) - - while True: - chunk_header = f.read(8) - if len(chunk_header) < 8: - break - - length = struct.unpack(">I", chunk_header[:4])[0] - chunk_type = chunk_header[4:8] - # Clamp the attacker-controlled 32-bit length to the bytes that - # actually remain, so a malformed huge length can't allocate GBs. - safe_length = max(0, min(length, file_size - f.tell())) - - if chunk_type == C2PA_CHUNK_TYPE: - chunk_data = f.read(safe_length) - _parse_c2pa_chunk(chunk_data, c2pa_info) - f.read(4) - else: - f.seek(safe_length + 4, 1) - - if chunk_type == b"IEND": - break - except Exception: - pass - - return c2pa_info - - -def cbor_text_after(payload: bytes, key: bytes) -> str | None: - """Return the CBOR text-string immediately following ``key`` in ``payload``. - - Handles CBOR major-type 3 length prefixes: direct (0x60-0x77), 1-byte - (0x78 NN), and 2-byte (0x79 NN NN). This reads the actual encoded value, so - it avoids the byte-grabbing artifacts a loose regex produces (e.g. the - leading length byte showing up as ``fGPT-4o``). - """ - idx = payload.find(key) - if idx < 0: - return None - p = idx + len(key) - if p >= len(payload): - return None - head = payload[p] - if 0x60 <= head <= 0x77: - length, start = head - 0x60, p + 1 - elif head == 0x78 and p + 1 < len(payload): - length, start = payload[p + 1], p + 2 - elif head == 0x79 and p + 2 < len(payload): - length, start = (payload[p + 1] << 8) | payload[p + 2], p + 3 - else: - return None - raw_str = payload[start : start + length] - try: - return raw_str.decode("utf-8") - except UnicodeDecodeError: - return raw_str.decode("latin1", errors="replace") - - -def synthid_verdict(vendors: str) -> str: - """Human-readable SynthID-source verdict, shared by all callers.""" - return f"likely present ({vendors} embeds SynthID with C2PA)" - - -def synthid_vendors_in(buffer: bytes) -> list[str]: - """Return SynthID-using C2PA issuer names whose signature appears in ``buffer``. - - Shared by the PNG caBX parser and the format-agnostic binary scan so both - apply the same SYNTHID_C2PA_ISSUERS rule against their respective bytes. - """ - return sorted({name for sig, name in C2PA_ISSUERS.items() if sig in buffer and sig in SYNTHID_C2PA_ISSUERS}) - - -def soft_binding_vendors_in(buffer: bytes) -> list[str]: - """Return forensic-watermark vendor names whose C2PA soft-binding ``alg`` - identifier appears in ``buffer``. - - A ``c2pa.soft-binding`` assertion names the watermark scheme that stamped the - pixels (Adobe TrustMark, Digimarc, Imatag, Steg.AI, ...). Shared by the PNG - caBX parser and the format-agnostic binary scan so both apply the same - C2PA_SOFT_BINDINGS rule against their respective bytes. - """ - return sorted({name for sig, name in C2PA_SOFT_BINDINGS.items() if sig in buffer}) - - -def _populate_registry_fields(buf: bytes, c2pa_info: dict[str, Any]) -> bool: - """Populate the registry-driven C2PA fields by scanning ``buf``. - - Shared by the legacy caBX-chunk parser and the c2pa-python store-JSON path so - both produce an identical dict shape. ``buf`` is the raw manifest bytes for - the former and the manifest-store JSON (UTF-8) for the latter; the vendor / - tool / action / source-type / SynthID / soft-binding signatures appear in - both. Sets ``issuer``, ``ai_tool``, ``actions``, ``source_type``, - ``synthid_vendors`` / ``synthid_watermark``, ``soft_binding_vendors`` / - ``soft_binding`` when present and returns whether the source type is AI. - """ - if issuers := [name for sig, name in C2PA_ISSUERS.items() if sig in buf]: - c2pa_info["issuer"] = ", ".join(dict.fromkeys(issuers)) - - if ai_tools := [name for sig, name in C2PA_AI_TOOLS.items() if sig in buf]: - c2pa_info["ai_tool"] = ", ".join(dict.fromkeys(ai_tools)) - - if actions := [name for sig, name in C2PA_ACTIONS.items() if sig in buf]: - c2pa_info["actions"] = ", ".join(actions) - - # Digital source type (matched anywhere in the store, including ingredient - # manifests -- a ChatGPT edit of a Sora generation carries the AI marker on - # the parent, not the active manifest). - # ``ai_source_kind`` is the structured generated-vs-enhanced split the caller - # branches on (full-frame scrub vs region-targeted clean); ``source_type`` is the - # human-readable form. The two byte strings are unambiguous: - # "compositeWithTrainedAlgorithmicMedia" capitalizes the inner "Trained", so a - # lowercase "trainedAlgorithmicMedia" match is standalone full generation, which - # wins when both appear (an edit chain). - ai_source = False - if b"trainedAlgorithmicMedia" in buf: - c2pa_info["source_type"] = "trainedAlgorithmicMedia (AI-generated)" - c2pa_info["ai_source_kind"] = "generated" - ai_source = True - elif b"compositeWithTrainedAlgorithmicMedia" in buf: - # Checked BEFORE bare ``algorithmicMedia``: a manifest can carry both tokens - # (an AI-enhanced composite with a procedural ingredient), and the bare-token - # branch would otherwise fire first and misclassify the AI composite as non-AI. - c2pa_info["source_type"] = "compositeWithTrainedAlgorithmicMedia (AI-enhanced)" - c2pa_info["ai_source_kind"] = "enhanced" - ai_source = True - elif b"algorithmicMedia" in buf: - c2pa_info["source_type"] = "algorithmicMedia" - - # SynthID pixel-watermark proxy: a C2PA manifest from a SynthID-using - # vendor (Google/OpenAI) on AI-generated content implies an invisible - # SynthID watermark in the pixels (see SYNTHID_C2PA_ISSUERS). - synthid_vendors = synthid_vendors_in(buf) - if synthid_vendors and ai_source: - c2pa_info["synthid_vendors"] = synthid_vendors - c2pa_info["synthid_watermark"] = synthid_verdict(", ".join(synthid_vendors)) - - # Soft-binding: a forensic/third-party watermark vendor named in the - # manifest (Adobe TrustMark, Digimarc, ...), independent of the issuer. - soft_binding_vendors = soft_binding_vendors_in(buf) - if soft_binding_vendors: - c2pa_info["soft_binding_vendors"] = soft_binding_vendors - c2pa_info["soft_binding"] = ", ".join(soft_binding_vendors) - - return ai_source - - -def _parse_c2pa_chunk(chunk_data: bytes, c2pa_info: dict[str, Any]) -> None: - """Parse a raw caBX chunk payload and populate the info dictionary. - - The fallback path, used when the official c2pa-python reader is unavailable - or rejects the file (synthetic/partial blobs, broken installs). - """ - c2pa_info["c2pa_manifest"] = f"C2PA manifest ({len(chunk_data)} bytes)" - - _populate_registry_fields(chunk_data, c2pa_info) - - # Claim generator and spec version: read the CBOR text-string values - # directly (regex byte-grabbing produced artifacts like ``fGPT-4o``). - # Guard with isprintable(): on some manifests (e.g. Microsoft Designer) the - # first ``name`` key precedes a binary field (a hash), not the generator - # string, which would otherwise surface as control-char garbage. - if (generator := cbor_text_after(chunk_data, b"name")) and generator.isprintable(): - c2pa_info["claim_generator"] = generator - if (spec := cbor_text_after(chunk_data, b"specVersion")) and spec.isprintable(): - c2pa_info["c2pa_spec"] = spec - - # Find timestamps - timestamp_matches = re.findall(rb"(\d{14}Z)", chunk_data) - if timestamp_matches: - c2pa_info["timestamp"] = timestamp_matches[0].decode("utf-8") - if len(timestamp_matches) > 1: - c2pa_info["timestamps"] = [t.decode("utf-8") for t in timestamp_matches[:3]] - - -def extract_c2pa_chunk(image_path: Path) -> bytes | None: - """ - Extract the raw C2PA JUMBF chunk from a PNG file. - - Args: - image_path: Path to the source PNG file. - - Returns: - Raw bytes of the C2PA chunk or None. - """ - if image_path.suffix.lower() != ".png": - return None - - try: - with open(image_path, "rb") as f: - signature = f.read(8) - if signature != PNG_SIGNATURE: - return None - - file_size = f.seek(0, 2) - f.seek(8) - - while True: - chunk_header = f.read(8) - if len(chunk_header) < 8: - break - - length = struct.unpack(">I", chunk_header[:4])[0] - chunk_type = chunk_header[4:8] - # Clamp the attacker-controlled 32-bit length to the bytes that - # actually remain, so a malformed huge length can't allocate GBs. - safe_length = max(0, min(length, file_size - f.tell())) - - if chunk_type == C2PA_CHUNK_TYPE: - chunk_data = f.read(safe_length) - crc = f.read(4) - - # Check for any C2PA signature - for sig in C2PA_SIGNATURES: - if sig in chunk_data: - return chunk_header + chunk_data + crc - - # Also check lowercase variants - if b"jumb" in chunk_data.lower() or b"c2pa" in chunk_data.lower(): - return chunk_header + chunk_data + crc - else: - f.seek(safe_length + 4, 1) - - if chunk_type == b"IEND": - break - except Exception: - pass - - return None - - -def inject_c2pa_chunk(target_path: Path, output_path: Path, c2pa_chunk: bytes) -> None: - """ - Inject a C2PA JUMBF chunk into a PNG file. - - Args: - target_path: Path to the target PNG file. - output_path: Path where the output file will be saved. - c2pa_chunk: Raw bytes of the C2PA chunk to inject. - - Raises: - ValueError: If not PNG files. - """ - if target_path.suffix.lower() != ".png" or output_path.suffix.lower() != ".png": - raise ValueError("C2PA chunk injection is only supported for PNG files") - - output_path.parent.mkdir(parents=True, exist_ok=True) - - with open(target_path, "rb") as f_in, open(output_path, "wb") as f_out: - f_out.write(f_in.read(8)) - - c2pa_injected = False - while True: - chunk_header = f_in.read(8) - if len(chunk_header) < 8: - break - - length = struct.unpack(">I", chunk_header[:4])[0] - chunk_type = chunk_header[4:8] - chunk_data = f_in.read(length) - crc = f_in.read(4) - - if chunk_type == b"IDAT" and not c2pa_injected: - f_out.write(c2pa_chunk) - c2pa_injected = True - - if chunk_type == C2PA_CHUNK_TYPE: - continue - - f_out.write(chunk_header) - f_out.write(chunk_data) - f_out.write(crc) - - if chunk_type == b"IEND": - break diff --git a/src/remove_ai_watermarks/noai/constants.py b/src/remove_ai_watermarks/noai/constants.py deleted file mode 100644 index 541a77f..0000000 --- a/src/remove_ai_watermarks/noai/constants.py +++ /dev/null @@ -1,342 +0,0 @@ -"""Shared constants for AI metadata detection, C2PA parsing, and format support. - -All modules reference these constants rather than hard-coding values, -so adding a new AI tool or metadata key requires updating only this file. -""" - -from typing import NamedTuple - -# Supported image formats for the pixel/removal path (CLI input validation + batch -# discovery). PNG/JPEG/WebP decode+encode via cv2; HEIC/HEIF/AVIF via the optional -# pillow-heif dep (image_io.imread Pillow fallback + imwrite _pil_write), so batch -# now picks them up and the CLI no longer warns on an iPhone HEIC. JPEG-XL is left -# out on purpose -- it is metadata/strip-only (no pixel decoder without pillow-jxl). -SUPPORTED_FORMATS = {".png", ".jpg", ".jpeg", ".webp", ".heic", ".heif", ".avif"} - -# AI-generated image metadata keys (Stable Diffusion, ComfyUI, Midjourney, etc.) -AI_METADATA_KEYS = [ - "parameters", # Stable Diffusion WebUI (AUTOMATIC1111, Vladmandic) - "postprocessing", # SD WebUI post-processing info - "extras", # SD WebUI extras - "workflow", # ComfyUI workflow JSON - "prompt", # Some AI tools - "Dream", # DreamStudio - "SD:mode", # Stability AI - "StableDiffusionVersion", # SD version info - "generation_time", # Generation time info - "Model", # Model name - "Model hash", # Model hash - "Seed", # Seed value -] - -# Standard PNG metadata keys -PNG_METADATA_KEYS = [ - "Author", - "Title", - "Description", - "Copyright", - "Creation Time", - "Software", - "Disclaimer", - "Warning", - "Source", - "Comment", -] - -# AI-related keywords for detection -AI_KEYWORDS = [ - "prompt", - "negative_prompt", - "sampler", - "cfg_scale", - "lora", - "diffusion", - "comfy", - "midjourney", - "dall-e", - "dalle", - "imagen", - "firefly", - "c2pa", - "chatgpt", - "gpt-4", - "sora", - "openai", - "truepic", - "stable_diffusion", - "invokeai", -] - -# C2PA (Coalition for Content Provenance and Authenticity) constants -# Used by Google Imagen, Adobe Firefly, Microsoft Designer, OpenAI, etc. -C2PA_CHUNK_TYPE = b"caBX" # JUMBF container chunk type for C2PA -C2PA_SIGNATURES = [ - b"c2pa", - b"C2PA", - b"jumb", - b"jumd", - b"JUMBF", - b"jumbf", - b"cbor", - b"contentcreds", - b"digid", - b"assertions", - b"manifest", -] - - -# Single source of truth for every C2PA-signing vendor. The three per-vendor -# facts that used to live in separate tables -- the issuer byte signature -# (C2PA_ISSUERS), the SynthID pairing (SYNTHID_C2PA_ISSUERS), and the human -# platform label (identify._ISSUER_PLATFORM) -- are all fields here, so adding a -# new C2PA vendor is a single append below; the views derive automatically. -class C2paAiVendor(NamedTuple): - issuer: bytes # distinctive byte signature scanned in the manifest (cert org / signer) - org: str # resolved issuer/cert-org display name (the old C2PA_ISSUERS value) - # Human platform label for identify; None marks a signing authority / non-generator - # (e.g. Truepic), which never names an AI platform on its own. - platform: str | None - # Substring matched against the joined issuer-org names for platform attribution - # (usually a shorter form of org, e.g. "Google" for "Google LLC"); None when platform is. - needle: str | None - synthid: bool = False # vendor pairs an invisible SynthID pixel watermark with its C2PA manifest - # The vendor's mere presence in the manifest asserts AI generation even without - # a digitalSourceType (``trainedAlgorithmicMedia``) assertion. Set ONLY for a - # pure-generator brand whose issuer/generator byte string is unambiguous (e.g. - # "Dreamina"). Do NOT set for common-word issuers (Adobe/Google/OpenAI/Microsoft): - # those appear incidentally in unrelated XMP/trust-chain bytes, so they stay - # source-type-gated in identify._attribute_platform. - asserts_ai: bool = False - - -# C2PA known vendors, ORDERED for first-match-wins platform attribution: when a -# manifest names several issuers (Microsoft Designer signs as "OpenAI, Microsoft"), -# the earlier entry wins so the product, not the backend engine, is named. -# Used by Google Imagen, Adobe Firefly, Microsoft Designer, OpenAI, etc. -C2PA_AI_VENDORS: tuple[C2paAiVendor, ...] = ( - # Microsoft signs both Designer and Bing Image Creator; Bing now runs its own - # MAI-Image model (not DALL-E), so the label stays model-neutral. - C2paAiVendor(b"Microsoft", "Microsoft", "Microsoft (Bing Image Creator / Designer)", "Microsoft"), - C2paAiVendor(b"Adobe", "Adobe", "Adobe Firefly", "Adobe"), - C2paAiVendor(b"OpenAI", "OpenAI", "OpenAI (ChatGPT / gpt-image / DALL-E / Sora)", "OpenAI", synthid=True), - C2paAiVendor(b"Google", "Google LLC", "Google (Gemini / Imagen)", "Google", synthid=True), - # Stability AI signs C2PA as "Stability AI" (cert org "Stability AI Ltd"). - # Verified on a live Brand Studio (DreamStudio successor) output, 2026-05-24. - C2paAiVendor(b"Stability AI", "Stability AI", "Stability AI (Stable Image / DreamStudio)", "Stability AI"), - # Black Forest Labs (FLUX) API output: claim_generator_info "Black Forest - # Labs API" + a c2pa.ai_generated_content assertion + trainedAlgorithmicMedia. - # Verified on a real signed FLUX JPEG, 2026-05-29. - C2paAiVendor(b"Black Forest Labs", "Black Forest Labs", "Black Forest Labs (FLUX)", "Black Forest Labs"), - # ByteDance's Volcano Engine (Volcengine) signs its AI image output with a - # cert from certificate_center@volcengine.com -- the platform behind Doubao / - # Jimeng. Verified on two real signed JPEGs, 2026-05-29. - C2paAiVendor( - b"volcengine", "ByteDance (Volcano Engine)", "ByteDance (Doubao / Jimeng / Volcano Engine)", "ByteDance" - ), - # Some Volcano Engine certs name the signer with the Chinese legal entity - # "北京火山引擎科技有限公司" (Beijing Volcano Engine Technology Co., Ltd.) rather - # than the latin "volcengine" -- the latin needle misses it entirely. The issuer is the - # UTF-8 of the Chinese name (it appears UTF-8-encoded in the manifest-store - # JSON and the raw caBX bytes alike); it normalizes to the same "ByteDance" - # needle and platform as the volcengine row, so the two collapse together for - # clash detection. Verified against compatible signed samples. - C2paAiVendor( - "北京火山引擎科技有限公司".encode(), - "ByteDance (Volcano Engine)", - "ByteDance (Doubao / Jimeng / Volcano Engine)", - "ByteDance", - ), - # ByteDance's international brand (BytePlus / Seedream / Seededit) signs its - # cert as "Byteplus Pte. Ltd." -- the bare ``volcengine`` needle misses it, so - # real BytePlus AI output was mis-attributed (an incidental "Adobe XMP" string - # in the file's XMP made it read "Adobe Firefly"). Adding the issuer means the - # clean manifest issuer matches "BytePlus (ByteDance)" directly. The platform - # string mirrors the volcengine row: both share the "ByteDance" needle, so the - # earlier row's label wins anyway -- they normalize together for clash - # detection. Verified on compatible signed samples. - C2paAiVendor(b"Byteplus", "BytePlus (ByteDance)", "ByteDance (Doubao / Jimeng / Volcano Engine)", "ByteDance"), - # Dreamina (ByteDance's international Jimeng brand) signs C2PA as "Bytedance - # Pte. Ltd." with a "Dreamina/x.y" claim generator and, unlike the Volcano - # Engine output, NO digitalSourceType assertion -- so the generator name is the - # only AI signal. It is registered by that generator token (which the caBX / - # store-JSON byte scan sees across active + ingredient manifests, where the - # active manifest is often a plain c2pa-tool transcode). ``asserts_ai`` lets the - # issuer alone flag AI without trainedAlgorithmicMedia; "Dreamina" is a - # distinctive brand string, so it does not risk the incidental-mention problem - # the common-word issuers have. Verified on compatible signed samples. - # Normalizes to the same "ByteDance" needle/platform as the - # volcengine row (they collapse together for clash detection). - C2paAiVendor( - b"Dreamina", - "ByteDance (Dreamina)", - "ByteDance (Doubao / Jimeng / Volcano Engine)", - "ByteDance", - asserts_ai=True, - ), - # Canva Magic Media signs AI-generated images as "Canva" with a generic - # c2pa-rs claim generator + trainedAlgorithmicMedia; without this entry the - # source read AI but no platform was attributed. Verified on compatible signed - # samples. Canva does not use SynthID. - C2paAiVendor(b"Canva", "Canva", "Canva (Magic Media)", "Canva"), - # ElevenLabs is a pure generative-AI company (AI voice / audio, and image / - # video via its API); it signs output as "Eleven Labs Inc.", so the C2PA - # manifest alone marks AI generation. Verified on compatible signed samples. - # ElevenLabs does not use SynthID. - C2paAiVendor(b"Eleven Labs", "ElevenLabs", "ElevenLabs", "ElevenLabs"), - # fal.ai (generative inference platform, issuer "fal - Features & Labels - # Inc." / common name "fal.ai", claim generators like "fal-ai/seedvr", - # "fal-ai/gpt-image-2"). The files carry trainedAlgorithmicMedia, so the - # verdict already fired, but the platform stayed unattributed. fal.ai is - # a pure generative platform, so ``asserts_ai`` also covers its output - # that omits the source-type. - C2paAiVendor(b"fal-ai", "fal.ai", "fal.ai", "fal.ai", asserts_ai=True), - # Bria AI (bria.ai, generative platform) signs as "Bria Artificial - # Intelligence" with a "Bria Ai" claim generator and source type - # ``empty`` (NOT trainedAlgorithmicMedia), so a real signed file was - # completely missed by identify. A pure-AI - # vendor with distinctive strings, so ``asserts_ai`` is safe here. - C2paAiVendor(b"Bria", "Bria Artificial Intelligence", "Bria AI", "Bria", asserts_ai=True), - # Truepic is a C2PA signing authority, not an AI generator: no platform label, - # never asserts is_ai (the verdict comes from the digital-source-type). - C2paAiVendor(b"Truepic", "Truepic", None, None), -) - -# Deliberately NOT registered as AI-generation vendors: -# - TikTok Inc.: signs C2PA as a content-provenance / AI-labeling authority on -# uploads, not as an image generator. The is_ai verdict keys off the -# digitalSourceType (trainedAlgorithmicMedia), which is already honored; a -# bare TikTok signer marks distribution provenance, not generation, so adding -# it as a generator needle would mis-label human uploads as AI. -# - PixelBin.io (issuer "Fynd"): an image transformation / optimization / CDN -# service. Its C2PA stamps a transform/upload step, not a generation event. -# Both are excluded to avoid false-positive AI attribution; re-evaluate only -# against a real signed file whose manifest carries a trainedAlgorithmicMedia -# digital-source type produced by the vendor itself. - -# Derived view -- add a vendor to C2PA_AI_VENDORS above, not here. -# C2PA issuer signature -> resolved org name, for the manifest byte-scan. -C2PA_ISSUERS: dict[bytes, str] = {v.issuer: v.org for v in C2PA_AI_VENDORS} - -# Resolved org names of the vendors whose presence asserts AI generation on its -# own (no digitalSourceType needed) -- see the ``asserts_ai`` field. identify uses -# this to lift the AI verdict for an identity-AI issuer (e.g. Dreamina) that ships -# no trainedAlgorithmicMedia. Derived from the flag -- set it on the vendor, not here. -C2PA_IDENTITY_AI_ORGS: frozenset[str] = frozenset(v.org for v in C2PA_AI_VENDORS if v.asserts_ai) - -# C2PA issuers whose signed outputs also carry an invisible SynthID pixel -# watermark -- a metadata proxy for "SynthID is in the pixels": -# - Google (Imagen/Gemini): embeds SynthID, long-standing (DeepMind docs). -# - OpenAI (ChatGPT/Codex/API): pairs SynthID with C2PA since ~2026-05-20. -# Confirmed by OpenAI's Help Center ("C2PA and SynthID in OpenAI-generated -# images", updated 2026-05-21): "Images generated with ChatGPT, Codex, and -# our API include both C2PA metadata and SynthID watermarks." OpenAI also -# notes a signal may be absent if "the image was created before these -# signals were available" -- so OpenAI images from before the rollout can -# carry C2PA without SynthID. For OpenAI the proxy is therefore "likely", -# not certain; the verdict string is hedged accordingly. OpenAI's own oracle -# is openai.com/verify (Google's is the Gemini app "Verify with SynthID"). -# The issuer byte ("OpenAI"/"Google") is verified locally against data/fixtures/provenance; -# the SynthID pairing is documented behavior (Google: DeepMind; OpenAI: above). -# Adobe Firefly and Microsoft Designer sign C2PA but do NOT use SynthID, so a -# C2PA manifest alone is not a SynthID signal -- the issuer is. The pixel -# watermark is not locally detectable (proprietary decoder); the C2PA companion -# is the proxy, and only while the manifest is intact. -# Derived from the `synthid` flag on C2PA_AI_VENDORS -- set it there, not here. -SYNTHID_C2PA_ISSUERS: frozenset[bytes] = frozenset(v.issuer for v in C2PA_AI_VENDORS if v.synthid) - -# C2PA known AI tools -C2PA_AI_TOOLS = { - b"GPT-4o": "GPT-4o", - b"ChatGPT": "ChatGPT", - b"Sora": "Sora", - b"DALL-E": "DALL-E", - b"DALL": "DALL-E", - b"Imagen": "Imagen", - b"Firefly": "Firefly", -} - -# C2PA ``c2pa.soft-binding`` algorithm identifiers -> the forensic-watermark -# vendor that stamped the pixels. The manifest's ``alg`` field names the -# watermark scheme even when the watermark itself cannot be decoded locally, so -# a byte-scan for these (keyed on a distinctive prefix to catch all variants) -# tells us a third-party forensic watermark is present and whose. Verified -# against the official C2PA registry (github.com/c2pa-org/softbinding-algorithm-list). -# Adobe TrustMark is additionally decodable locally (see ``trustmark_detector``); -# the rest (Digimarc, Imatag, Steg.AI, etc.) are proprietary oracle-only decoders. -C2PA_SOFT_BINDINGS = { - b"com.adobe.trustmark": "Adobe TrustMark", - b"com.adobe.icn": "Adobe (content fingerprint)", - b"com.digimarc": "Digimarc", - b"com.imatag.lamark": "Imatag (Lamark)", - b"ai.steg": "Steg.AI", - b"com.microsoft.invismark": "Microsoft InvisMark", - b"com.microsoft.wavmark": "Microsoft WavMark", - b"com.verimatrix": "Verimatrix", - b"com.nagra.nexguard": "NAGRA NexGuard", - b"com.aiwatermark": "AIWatermark (Meta PixelSeal)", - b"ai.trufo": "Trufo", - b"app.overlai": "Overlai", - b"com.markany": "MarkAny", - b"com.mentaport": "Mentaport", - b"es.lumatrace": "LumaTrace", - b"ai.verda": "VerdaAI", - b"ai.contentlens": "ContentLens", - b"io.iscc": "ISCC (content code)", -} - -# Lowercased substrings that mark an AI generator when found in an EXIF -# ``Software`` / XMP ``CreatorTool`` value. Conservative on purpose: plain -# editors like "Adobe Photoshop" or "GIMP" must NOT match (no AI token), so only -# generator names land here. Add new generators here, not inline. -AI_GENERATOR_TOKENS: frozenset[str] = frozenset( - { - "firefly", - "dall-e", - "dalle", - "midjourney", - "stable diffusion", - "stable-diffusion", - "stablediffusion", - "comfyui", - "automatic1111", - "invokeai", - "imagen", - "gpt-image", - "nightcafe", - "ideogram", - "leonardo", - "flux", - "dreamstudio", - # Generator stamps without C2PA: - # - NovelAI (anime SD): PNG tEXt Software="NovelAI", Source="NovelAI - # Diffusion V4.5 ", Title="NovelAI generated image". - # - Reve Image (reve.com): EXIF Software / XMP CreatorTool = "reve.com" - # (the bare token "reve" would false-positive on "forever"/"reverie"). - # - Aphrodite AI: EXIF Make / Software = "Aphrodite AI[ v1.0]". - "novelai", - "reve.com", - "aphrodite ai", - # Additional verified markers: - # - Apple Photos Clean Up (Apple Intelligence object removal): XMP - # photoshop:Credit / IPTC credit value; composite source-type - # covered detection, this token covers removal parity. - # - fal-ai: generative-platform generator string. - "apple photos clean up", - "fal-ai", - } -) - -# C2PA action types -C2PA_ACTIONS = { - b"c2pa.created": "created", - b"c2pa.converted": "converted", - b"c2pa.edited": "edited", - b"c2pa.filtered": "filtered", - b"c2pa.cropped": "cropped", - b"c2pa.resized": "resized", - b"c2pa.opened": "opened", - b"c2pa.placed": "placed", -} - -# PNG signature -PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" diff --git a/src/remove_ai_watermarks/noai/extractor.py b/src/remove_ai_watermarks/noai/extractor.py deleted file mode 100644 index 0c4909b..0000000 --- a/src/remove_ai_watermarks/noai/extractor.py +++ /dev/null @@ -1,155 +0,0 @@ -"""Read-only metadata extraction from PNG and JPEG images. - -Provides functions to pull all metadata, AI-only metadata, or a -human-readable summary without modifying the source file. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, cast - -if TYPE_CHECKING: - from pathlib import Path - -import piexif -from PIL import Image - -from remove_ai_watermarks.noai.c2pa import extract_c2pa_chunk, extract_c2pa_info, has_c2pa_metadata -from remove_ai_watermarks.noai.constants import AI_KEYWORDS, AI_METADATA_KEYS, PNG_METADATA_KEYS - - -def extract_metadata(source_path: Path) -> dict[str, Any]: - """ - Extract all metadata from a PNG or JPG file. - - Args: - source_path: Path to the source image file. - - Returns: - Dictionary containing all extracted metadata. - """ - metadata: dict[str, Any] = {} - - with Image.open(source_path) as img: - # Extract EXIF data - if "exif" in img.info: - try: - exif_dict = piexif.load(img.info["exif"]) - metadata["exif"] = exif_dict - except Exception: - metadata["exif_raw"] = img.info["exif"] - - # Extract standard PNG metadata - for key in PNG_METADATA_KEYS: - if key in img.info: - metadata[key] = img.info[key] - - # Extract all other metadata including AI-specific - for key, value in img.info.items(): - if not isinstance(key, str): - continue - if key not in metadata and key not in ["exif"]: - metadata[key] = value - - # Extract DPI and gamma if present - if "dpi" in img.info: - metadata["dpi"] = img.info["dpi"] - if "gamma" in img.info: - metadata["gamma"] = img.info["gamma"] - - # Check for C2PA metadata - if has_c2pa_metadata(source_path): - metadata["c2pa"] = extract_c2pa_info(source_path) - c2pa_chunk = extract_c2pa_chunk(source_path) - if c2pa_chunk: - metadata["c2pa_chunk"] = c2pa_chunk - - return metadata - - -def extract_ai_metadata(source_path: Path) -> dict[str, Any]: - """ - Extract only AI-generated metadata from a PNG or JPG file. - - Args: - source_path: Path to the source image file. - - Returns: - Dictionary containing only AI-related metadata. - """ - ai_metadata: dict[str, Any] = {} - - with Image.open(source_path) as img: - for key in AI_METADATA_KEYS: - if key in img.info: - ai_metadata[key] = img.info[key] - - for key, value in img.info.items(): - if not isinstance(key, str): - continue - key_lower = key.lower() - if key not in ai_metadata and any(kw in key_lower for kw in AI_KEYWORDS): - ai_metadata[key] = value - - # Check for C2PA metadata - if has_c2pa_metadata(source_path): - ai_metadata["c2pa"] = extract_c2pa_info(source_path) - c2pa_chunk = extract_c2pa_chunk(source_path) - if c2pa_chunk: - ai_metadata["c2pa_chunk"] = c2pa_chunk - - return ai_metadata - - -def has_ai_metadata(image_path: Path) -> bool: - """ - Check if an image contains AI-generated metadata. - - Args: - image_path: Path to the image file. - - Returns: - True if AI metadata is detected, False otherwise. - """ - with Image.open(image_path) as img: - for key in AI_METADATA_KEYS: - if key in img.info: - return True - - return bool(has_c2pa_metadata(image_path)) - - -def get_ai_metadata_summary(source_path: Path) -> str: - """ - Get a human-readable summary of AI metadata. - - Args: - source_path: Path to the source image file. - - Returns: - Formatted string with AI metadata summary. - """ - ai_meta = extract_ai_metadata(source_path) - - if not ai_meta: - return "No AI metadata found." - - lines = ["AI Image Metadata:"] - lines.append("-" * 40) - - for key, value in ai_meta.items(): - if key == "c2pa_chunk": - continue - if key == "c2pa" and isinstance(value, dict): - lines.append("C2PA Metadata:") - for ck, cv in cast("dict[str, Any]", value).items(): - lines.append(f" {ck}: {cv}") - elif isinstance(value, str) and len(value) > 100: - value = value[:100] + "..." - lines.append(f"{key}: {value}") - elif isinstance(value, bytes): - lines.append(f"{key}: ") - else: - lines.append(f"{key}: {value}") - - return "\n".join(lines) diff --git a/src/remove_ai_watermarks/noai/img2img_runner.py b/src/remove_ai_watermarks/noai/img2img_runner.py deleted file mode 100644 index a5aed90..0000000 --- a/src/remove_ai_watermarks/noai/img2img_runner.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Img2img pipeline execution with progress monitoring and MPS fallback. - -Extracted from ``watermark_remover.py`` to keep the ``WatermarkRemover`` -class focused on orchestration. -""" - -from __future__ import annotations - -import contextlib -import logging -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Callable - - from PIL import Image - -from remove_ai_watermarks.noai.progress import is_mps_error, make_pipeline_progress - -logger = logging.getLogger(__name__) - - -def run_img2img( - pipeline: Any, - image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - device: str, - set_progress: Callable[[str], None], - extra_kwargs: dict[str, Any] | None = None, -) -> Image.Image: - """Execute img2img with live progress and return the generated image. - - ``extra_kwargs`` overlays additional pipeline arguments (e.g. the ControlNet - ``control_image`` / ``controlnet_conditioning_scale`` and a non-empty prompt), - so a ControlNet img2img pass reuses the same progress + fallback machinery. - """ - effective_steps = max(1, int(num_inference_steps * strength)) - - step_cb, first_step, done_ev, start_updater = make_pipeline_progress( - effective_steps, - device, - set_progress, - ) - start_updater() - - try: - result = _call_pipeline( - pipeline, image, strength, num_inference_steps, guidance_scale, generator, step_cb, extra_kwargs - ) - done_ev.set() - return result.images[0] - except TypeError as exc: - # The only TypeError we retry is the deprecated-callback case: `_call_pipeline` - # passes the legacy `callback`/`callback_steps` kwargs, and a diffusers version - # that removed them raises TypeError("... unexpected keyword argument - # 'callback'"). We then re-run once WITHOUT the progress callback. Any OTHER - # TypeError (e.g. a bad control_image/dtype in the forward pass) is a real error - # -- re-raise it instead of silently re-running the whole diffusion pass and - # masking the cause. - if "callback" not in str(exc): - raise - first_step.set() - result = _call_pipeline( - pipeline, image, strength, num_inference_steps, guidance_scale, generator, None, extra_kwargs - ) - done_ev.set() - return result.images[0] - finally: - first_step.set() - done_ev.set() - - -def run_img2img_with_mps_fallback( - load_pipeline: Callable[[], Any], - image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - device: str, - set_progress: Callable[[str], None], - *, - reload_on_cpu: Callable[[], Any], - extra_kwargs: dict[str, Any] | None = None, -) -> tuple[Image.Image, str]: - """Run img2img; on MPS error, fall back to CPU. - - ``extra_kwargs`` overlays extra pipeline arguments (used by the ControlNet - path). Returns ``(result_image, final_device)`` — device may change to - ``"cpu"`` on fallback. - """ - pipeline = load_pipeline() - - try: - img = run_img2img( - pipeline, - image, - strength, - num_inference_steps, - guidance_scale, - generator, - device, - set_progress, - extra_kwargs, - ) - return img, device - except RuntimeError as error: - if device == "mps" and is_mps_error(error): - logger.warning("MPS error detected: %s. Falling back to CPU.", error) - set_progress("MPS error! Clearing cache and retrying on CPU...") - try_empty_device_cache("mps") - pipeline = reload_on_cpu() - img = run_img2img( - pipeline, image, strength, num_inference_steps, guidance_scale, None, "cpu", set_progress, extra_kwargs - ) - return img, "cpu" - raise - - -def _call_pipeline( - pipeline: Any, - image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - step_callback: Any, - extra_kwargs: dict[str, Any] | None = None, -) -> Any: - kwargs: dict[str, Any] = { - "prompt": "", - "image": image, - "strength": strength, - "num_inference_steps": num_inference_steps, - "guidance_scale": guidance_scale, - "generator": generator, - } - if extra_kwargs: - kwargs.update(extra_kwargs) - if step_callback is not None: - kwargs["callback"] = step_callback - kwargs["callback_steps"] = 1 - return pipeline(**kwargs) - - -def try_empty_device_cache(device: str) -> None: - """Best-effort free of cached GPU/MPS/XPU memory for ``device``. - - ``torch..empty_cache()`` exists for cuda/mps/xpu but not cpu (the - hasattr guard skips the cpu no-op). Never raises -- callers use it as cleanup - (the MPS->CPU fallback here, and the batch loop in watermark_remover). - """ - with contextlib.suppress(Exception): - import torch - - backend = getattr(torch, device, None) - if backend is not None and hasattr(backend, "empty_cache"): - backend.empty_cache() # type: ignore[attr-defined] diff --git a/src/remove_ai_watermarks/noai/progress.py b/src/remove_ai_watermarks/noai/progress.py deleted file mode 100644 index 262a306..0000000 --- a/src/remove_ai_watermarks/noai/progress.py +++ /dev/null @@ -1,332 +0,0 @@ -"""Terminal progress animation and library output suppression. - -This module provides two main capabilities for the CLI: - -1. ``run_with_progress`` — a styled two-line terminal animation that - displays a bouncing highlight bar, a braille spinner, elapsed time, - and a live operation message while a background task executes. - -2. ``silence_library_output`` — a wrapper that suppresses noisy log - output produced by third-party ML libraries (transformers, diffusers, - huggingface_hub, tqdm) so the user only sees our own progress messages. -""" - -from __future__ import annotations - -import contextlib -import io -import os -import sys -import threading -import time -import warnings -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Callable - -# ── ANSI color constants ──────────────────────────────────────────── -_CYAN = "\033[36m" -_YELLOW = "\033[33m" -_GREEN = "\033[32m" -_DIM = "\033[2m" -_BOLD = "\033[1m" -_RESET = "\033[0m" - -# Bar geometry -_BAR_WIDTH = 32 -_HIGHLIGHT_WIDTH = 5 - - -def _no_color() -> bool: - """Respect the NO_COLOR convention (https://no-color.org/).""" - return bool(os.environ.get("NO_COLOR")) - - -def _truncate(text: str, max_len: int = 72) -> str: - """Shorten a string with an ellipsis if it exceeds *max_len*.""" - return text if len(text) <= max_len else text[: max_len - 1] + "…" - - -def _build_bar(step: int) -> str: - """Build a flowing highlight bar that bounces across the width. - - The highlight segment (5 chars wide) travels left→right→left - continuously, giving the user a visual "working" signal. - """ - cycle = _BAR_WIDTH * 2 - 2 - pos = step % cycle - if pos >= _BAR_WIDTH: - pos = cycle - pos - - hl_start = max(0, pos - _HIGHLIGHT_WIDTH // 2) - hl_end = min(_BAR_WIDTH, pos + _HIGHLIGHT_WIDTH // 2 + 1) - before = "━" * hl_start - highlight = "━" * (hl_end - hl_start) - after = "━" * (_BAR_WIDTH - hl_end) - - if _no_color(): - return before + highlight + after - return f"{_DIM}{before}{_RESET}{_BOLD}{_YELLOW}{highlight}{_RESET}{_DIM}{after}{_RESET}" - - -def run_with_progress( - task: Callable[[], Any], - progress_state: dict[str, str] | None = None, -) -> Any: - """Execute *task* in a background thread while showing a progress animation. - - The animation renders two lines to ``sys.__stderr__``: - - - **Line 1**: braille spinner + bouncing bar + elapsed seconds - - **Line 2**: current operation message from *progress_state* - - When the task finishes, a green "Completed" line replaces the animation. - - Args: - task: A zero-argument callable to run in the background. - progress_state: Mutable dict whose ``"message"`` key is read - by the animation loop to display the current operation. - - Returns: - Whatever *task* returns. - - Raises: - Any exception raised by *task* is re-raised after the animation - is cleaned up. - """ - done = threading.Event() - output_holder: dict[str, Any] = {"result": None, "error": None} - - def worker() -> None: - try: - output_holder["result"] = task() - except Exception as error: # pragma: no cover - passthrough - output_holder["error"] = error - finally: - done.set() - - thread = threading.Thread(target=worker, daemon=True) - thread.start() - - spinner_frames = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏" - idx = 0 - start_time = time.time() - no_color = _no_color() - - def _get_operation() -> str: - if isinstance(progress_state, dict): - return progress_state.get("message", "Processing...") - return "Processing..." - - # ── Animation loop ────────────────────────────────────────────── - while not done.is_set(): - spinner = spinner_frames[idx % len(spinner_frames)] - elapsed = int(time.time() - start_time) - bar_str = _build_bar(idx) - operation = _truncate(_get_operation()) - - if no_color: - line1 = f" {spinner} Processing {bar_str} {elapsed:>3}s" - line2 = f" ╰─ {operation}" - else: - line1 = f" {_CYAN}{spinner}{_RESET} Processing {bar_str} {_BOLD}{_YELLOW}{elapsed:>3}s{_RESET}" - line2 = f" {_DIM}╰─ {operation}{_RESET}" - - print( - f"\r\033[2K{line1}\n\033[2K{line2}\033[1A\r", - end="", - flush=True, - file=sys.__stderr__, - ) - time.sleep(0.08) - idx += 1 - - # ── Final "done" frame ────────────────────────────────────────── - thread.join() - total = int(time.time() - start_time) - final_operation = _truncate(_get_operation()) - done_bar = "━" * _BAR_WIDTH - - if no_color: - final_line1 = f" ✓ Completed {done_bar} {total:>3}s" - final_line2 = f" ╰─ {final_operation}" - else: - final_line1 = ( - f" {_GREEN}{_BOLD}✓{_RESET} {_GREEN}Completed{_RESET} " - f"{_GREEN}{done_bar}{_RESET} {_BOLD}{_GREEN}{total:>3}s{_RESET}" - ) - final_line2 = f" {_DIM}╰─ {final_operation}{_RESET}" - - print( - f"\r\033[2K{final_line1}\n\033[2K{final_line2}", - file=sys.__stderr__, - ) - - if output_holder["error"] is not None: - raise output_holder["error"] - - return output_holder["result"] - - -def silence_library_output( - run_func: Callable[[], Any], - set_progress: Callable[[str], None] | None = None, -) -> Callable[[], Any]: - """Return a wrapper that silences noisy ML library output. - - The wrapper: - - 1. Disables HuggingFace Hub progress bars via env var. - 2. Sets ``transformers``, ``diffusers``, and ``huggingface_hub`` - loggers to *error* level. - 3. Redirects ``stdout`` and ``stderr`` to ``io.StringIO`` sinks so - that stray ``tqdm`` bars and model-loading chatter are invisible. - 4. Suppresses all Python warnings during the call. - - Args: - run_func: The callable to execute silently. - set_progress: Optional callback to report phase changes. - - Returns: - A zero-argument callable that, when invoked, runs *run_func* - inside the silent context. - """ - - def wrapped() -> Any: - if set_progress: - set_progress("Configuring runtime and suppressing noisy logs...") - - os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") - - for _silence in ( - lambda: __import__("transformers").logging.set_verbosity_error(), - lambda: _silence_diffusers(), - lambda: __import__("huggingface_hub").logging.set_verbosity_error(), - ): - with contextlib.suppress(Exception): - _silence() - - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): - if set_progress: - set_progress("Executing watermark removal pipeline...") - return run_func() - - return wrapped - - -def _silence_diffusers() -> None: - """Silence diffusers logging and progress bars.""" - from diffusers.utils import logging as diffusers_logging - - diffusers_logging.set_verbosity_error() - if hasattr(diffusers_logging, "disable_progress_bar"): - diffusers_logging.disable_progress_bar() - - -# ── Shared pipeline progress helpers ───────────────────────────────── - -_DEFAULT_PRE_PHASES: list[tuple[int, str]] = [ - (0, "Encoding image with VAE encoder"), - (3, "Mapping pixel data → latent space"), - (7, "Injecting noise into latent representation"), - (12, "Building denoiser schedule"), - (18, "Starting reverse diffusion sampler"), - (30, "Running first denoising iteration"), - (50, "Still processing — this can take a while"), - (90, "Pipeline running — may take a few minutes"), -] - -_DEFAULT_POST_PHASES: list[tuple[int, str]] = [ - (0, "Denoising complete · Running VAE decoder"), - (2, "Decoding latent channels → RGB color space"), - (5, "Reconstructing pixel grid from latents"), - (10, "Applying color space conversion and normalization"), - (18, "Finalizing pixel output"), - (30, "Still decoding — large images take longer"), - (60, "Almost done — large images take longer to decode"), -] - - -def make_pipeline_progress( - effective_steps: int, - device: str, - set_progress: Callable[[str], None], - *, - bar_len: int = 20, - label: str = "Denoising", - pre_phases: list[tuple[int, str]] | None = None, - post_phases: list[tuple[int, str]] | None = None, -) -> tuple[Callable[..., None], threading.Event, threading.Event, Callable[[], threading.Thread]]: - """Create step callback and background updater for a diffusion pipeline. - - Returns: - (step_callback, first_step_event, pipeline_done_event, start_updater) - where ``start_updater()`` launches and returns the background thread. - """ - pre = pre_phases or [(s, f"{m} on {device}") for s, m in _DEFAULT_PRE_PHASES] - post = post_phases or [(s, f"{m} on {device}") for s, m in _DEFAULT_POST_PHASES] - - t0_holder: list[float] = [time.monotonic()] - first_step = threading.Event() - pipeline_done = threading.Event() - last_cb_time: list[float] = [t0_holder[0]] - - def _background_updater() -> None: - idx = 0 - while not first_step.is_set(): - elapsed = time.monotonic() - t0_holder[0] - while idx < len(pre) - 1 and elapsed >= pre[idx + 1][0]: - idx += 1 - set_progress(pre[idx][1]) - first_step.wait(timeout=0.4) - - idx = 0 - post_start: float | None = None - while not pipeline_done.is_set(): - since_cb = time.monotonic() - last_cb_time[0] - if since_cb >= 1.5: - if post_start is None: - post_start = time.monotonic() - elapsed = time.monotonic() - post_start - while idx < len(post) - 1 and elapsed >= post[idx + 1][0]: - idx += 1 - set_progress(post[idx][1]) - else: - post_start = None - idx = 0 - pipeline_done.wait(timeout=0.4) - - def step_callback(step: int, timestep: int, latents: Any) -> None: - first_step.set() - last_cb_time[0] = time.monotonic() - elapsed = time.monotonic() - t0_holder[0] - current = step + 1 - per_step = elapsed / max(1, current) - remaining = per_step * max(0, effective_steps - current) - filled = int(bar_len * current / max(1, effective_steps)) - bar = "█" * filled + "░" * (bar_len - filled) - set_progress( - f"{label} [{bar}] {current}/{effective_steps} | {elapsed:.0f}s elapsed, ~{remaining:.0f}s left | {device}" - ) - - def start_updater() -> threading.Thread: - t0_holder[0] = time.monotonic() - last_cb_time[0] = t0_holder[0] - first_step.clear() - pipeline_done.clear() - t = threading.Thread(target=_background_updater, daemon=True) - t.start() - return t - - return step_callback, first_step, pipeline_done, start_updater - - -# ── MPS fallback helper ────────────────────────────────────────────── - - -def is_mps_error(error: Exception) -> bool: - """Check whether an exception is an MPS-related runtime error.""" - return "mps" in str(error).lower() diff --git a/src/remove_ai_watermarks/noai/utils.py b/src/remove_ai_watermarks/noai/utils.py deleted file mode 100644 index fdce354..0000000 --- a/src/remove_ai_watermarks/noai/utils.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Low-level utility helpers used across the metadata pipeline. - -Kept deliberately small — only format detection lives here so that -higher-level modules can import without circular dependencies. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from pathlib import Path - -from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS - - -def is_supported_format(file_path: Path) -> bool: - """ - Check if the file format is supported. - - Args: - file_path: Path to the image file. - - Returns: - True if the format is supported, False otherwise. - """ - return file_path.suffix.lower() in SUPPORTED_FORMATS - - -def get_image_format(file_path: Path) -> str: - """ - Get the image format from file path. - - Args: - file_path: Path to the image file. - - Returns: - Format string (PNG, JPEG, etc.). - """ - suffix = file_path.suffix.lower() - if suffix in {".jpg", ".jpeg"}: - return "JPEG" - return "PNG" diff --git a/src/remove_ai_watermarks/noai/watermark_profiles.py b/src/remove_ai_watermarks/noai/watermark_profiles.py deleted file mode 100644 index f540c84..0000000 --- a/src/remove_ai_watermarks/noai/watermark_profiles.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Watermark removal model profiles and the default strength. - -Pure configuration and lookup functions with no ML dependencies. -""" - -from __future__ import annotations - -import math -from typing import TYPE_CHECKING, Literal - -if TYPE_CHECKING: - from pathlib import Path - -DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0" - -# Qwen-Image (20B MMDiT, Apache-2.0 code AND weights) base for the ``qwen`` pipeline: -# an img2img alternative to SDXL with native text rendering (incl. CJK). Loaded only -# when ``--pipeline qwen`` is selected; CUDA/cloud-class (does not fit MPS). CERTIFIED -# oracle floors (2026-06-20): OpenAI **0.10** (seed-robust -- clean on seeds 0-4) and -# Google/Gemini **0.25** (seed 0 verified on 2 images; pin a seed in prod, the Gemini -# oracle rate-limits volume seed-repeat). The Gemini floor (0.25) is HIGHER than the -# certified controlnet Gemini floor (0.15); ``resolve_strength(..., pipeline="qwen")`` -# now carries this via ``_QWEN_VENDOR_STRENGTH`` (below), so ``--pipeline qwen`` gets the -# right floor automatically -- the old manual "pass --strength 0.25 for Gemini on qwen" -# workaround is retired. -# (Dispatch uses the bare "qwen" literal, matching the sdxl/controlnet sites, so there -# is no QWEN_PROFILE constant -- only the model id is referenced from code.) -QWEN_MODEL_ID = "Qwen/Qwen-Image" - -# Canonical pipeline-profile names + the back-compat alias. The plain SDXL img2img -# profile is ``sdxl``; ``default`` is kept as an accepted alias (it was the profile's -# name before ``controlnet`` became the default-selected pipeline, 2026-06-09). -SDXL_PROFILE = "sdxl" -QWEN_ZIMAGE_PROFILE = "qwen-zimage" -_PROFILE_ALIASES = { - "default": SDXL_PROFILE, - "qwen_zimage": QWEN_ZIMAGE_PROFILE, -} - - -def normalize_profile(profile: str) -> str: - """Canonicalize a pipeline-profile name, resolving the ``default`` -> ``sdxl`` alias.""" - normalized = profile.strip().lower() - return _PROFILE_ALIASES.get(normalized, normalized) - - -def resolve_steps(num_inference_steps: int | None, pipeline: str) -> int: - """Resolve a profile-specific step default while preserving explicit values. - - The Lightning LoRA in ``qwen-zimage`` is distilled for four steps. Existing - SDXL and Qwen profiles keep the long-standing 50-step CLI default. - """ - if num_inference_steps is not None: - return num_inference_steps - return 4 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else 50 - - -def resolve_seed(seed: int | None, pipeline: str) -> int | None: - """Keep the oracle-verified qwen-zimage profile deterministic by default.""" - if seed is not None: - return seed - return 0 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else None - - -# The SDXL-native canny ControlNet used by the ``controlnet`` pipeline. The -# ControlNet is an add-on to the SDXL base checkpoint (DEFAULT_MODEL_ID), not a -# separate base model, so both the ``sdxl`` and ``controlnet`` profiles load the -# same base weights and share the same vendor-adaptive strength ladder (see below). -CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0" - -# Vendor-adaptive default denoising strength for the SDXL img2img scrub, overridable -# from the CLI (`--strength`). The right strength depends on which vendor's SynthID is -# present (detected from the C2PA issuer, metadata.synthid_source). The SAME ladder -# applies to BOTH pipelines (`sdxl` plain img2img and `controlnet`) -- see "why one -# ladder" below. -# -# Data basis (see docs/synthid.md sections 2.2 / 5.5): ORACLE-CERTIFIED controlnet floors. -# Oracle re-testing -# LOWERED the ladder back to OpenAI 0.10 / Google 0.15: each output verified on its own -# oracle (openai.com/verify for OpenAI, the Google Gemini app for Google), all clean -> -# - OpenAI 0.10: 2 photoreal images (1402 / 1448 px), SynthID not found on either. -# - Google 0.15: 2 NATIVE-resolution images (both 2816x1536), SynthID not found on -# either -- this directly retires the earlier "native ~2816 likely needs ~0.35+" -# guess, which was speculative and never oracle-checked at that resolution. -# This supersedes the 2026-06-04 cert (OpenAI 0.20 / Google 0.30), whose higher floor a -# pixel-fidelity sweep showed was ~2x the removal floor and over-regenerated for no -# efficacy gain (Google MAE -20% at 0.15 vs 0.30, no SynthID returning). Unknown vendor -# tracks the Google (more robust watermark) value -> 0.15, still safe-by-default and the -# floor that real (no-vendor) photos hit, so it also minimizes damage when there is in -# fact nothing to remove. CAVEAT: the re-test is n=2 per vendor on photoreal / landscape -# content; FLAT-GRAPHIC hard cases (the historical `sdxl` weak spot) were NOT in the -# sample, so if an oracle still reads SynthID on a flat output, raise `--strength`. -# -# Why ONE ladder for both pipelines (2026-06-09): the certification was run on -# controlnet, and it does NOT transfer to `sdxl` by symmetry -- the two pipelines have -# OPPOSITE hard cases (controlnet leaves SynthID on photoreal, `sdxl` leaves it on flat -# graphics; the content-x-pipeline table in docs/synthid.md §5.1). BUT on its OWN hard -# case (flat fills) `sdxl` is the WEAKER remover -- plain img2img at low strength barely -# perturbs a flat region -- so it needs AT LEAST as much strength as controlnet, not -# less. Hence the certified controlnet floor is the right floor for `sdxl` too. The -# higher strength costs little quality where it matters. `controlnet` is now the default -# pipeline and `sdxl` is reached only through an explicit `--pipeline sdxl`. NOTE: -# this is a MARGIN argument for `sdxl`, not a fresh certification -- there is no local -# SynthID detector, so if an oracle still reads SynthID on a flat `sdxl` output, raise -# `--strength`. -OPENAI_STRENGTH = 0.10 -GEMINI_STRENGTH = 0.15 -UNKNOWN_STRENGTH = 0.15 -# Backwards-compatible alias: the vendor-unknown value (what a caller gets without a -# detected vendor). Kept as DEFAULT_STRENGTH for existing references. -DEFAULT_STRENGTH = UNKNOWN_STRENGTH - -# Detected-vendor -> default strength. Vendor strings come from `vendor_for_strength`. -_VENDOR_STRENGTH = {"openai": OPENAI_STRENGTH, "google": GEMINI_STRENGTH} - -# Qwen has its OWN certified floors (Modal A100-80GB, 2026-06-20), DIFFERENT from the -# SDXL ladder above: OpenAI 0.10 (seed-robust), Gemini 0.25 (HIGHER than controlnet's -# 0.15 -- the 20B MMDiT perturbs less per denoising step, so it needs more strength to -# clear Gemini SynthID). Unknown vendor tracks the higher (Gemini) value, safe-by-default. -# `resolve_strength(..., pipeline="qwen")` uses this table so `--pipeline qwen` carries the -# right floor automatically -- retiring the old manual "pass --strength 0.25 for Gemini on -# qwen" workaround. -QWEN_OPENAI_STRENGTH = 0.10 -QWEN_GEMINI_STRENGTH = 0.25 -QWEN_UNKNOWN_STRENGTH = 0.25 -_QWEN_VENDOR_STRENGTH = {"openai": QWEN_OPENAI_STRENGTH, "google": QWEN_GEMINI_STRENGTH} - - -def strength_default_help() -> str: - """One-line description of the vendor-adaptive default, derived from the constants. - - Single source of truth for the CLI ``--strength`` help so the numbers can never - drift from the actual ladder (they did once when the per-pipeline split was unified). - """ - return ( - f"vendor-adaptive (OpenAI {OPENAI_STRENGTH} / Google {GEMINI_STRENGTH} / " - f"unknown {UNKNOWN_STRENGTH}, from the C2PA issuer; qwen-zimage instead uses " - "resolution-adaptive denoise)" - ) - - -def resolve_strength(strength: float | None, vendor: str | None = None, pipeline: str | None = None) -> float: - """Resolve the denoising strength, applying the vendor default when unset. - - ``None`` means "the user did not pass ``--strength``", which resolves - **vendor-adaptively**: ``vendor`` (``"openai"`` / ``"google"`` / None, from - ``vendor_for_strength``) selects the per-vendor floor. The ``sdxl`` and ``controlnet`` - pipelines share ONE ladder (``OPENAI_STRENGTH`` / ``GEMINI_STRENGTH`` / - ``UNKNOWN_STRENGTH`` -- see the module comment for why); ``qwen`` has its OWN higher - ladder (``_QWEN_VENDOR_STRENGTH``, Gemini 0.25 vs controlnet 0.15), selected when - ``pipeline`` normalizes to ``"qwen"``. An explicit value always wins (including - ``0.0`` -- the check is ``is None``, not falsiness). Shared by the CLI (for display) - and the engine (for execution) so the two never disagree -- both must pass the SAME - ``vendor`` and ``pipeline``. - """ - if strength is not None: - return strength - if pipeline is not None and normalize_profile(pipeline) == "qwen": - return _QWEN_VENDOR_STRENGTH.get(vendor or "", QWEN_UNKNOWN_STRENGTH) - return _VENDOR_STRENGTH.get(vendor or "", UNKNOWN_STRENGTH) - - -def viable_steps(num_inference_steps: int, strength: float) -> int: - """The smallest step count >= ``num_inference_steps`` that actually denoises. - - diffusers derives its img2img timesteps as ``int(steps * strength)``. When that - rounds to ZERO the pipeline builds an empty latent and dies deep inside attention - with ``cannot reshape tensor of 0 elements into shape [0, -1, 1, 512]`` -- an opaque - torch error for what is really "these two options cannot work together". - - The combination is reachable with entirely valid CLI arguments: at the default - strength 0.15 every ``--steps`` below 7 crashed, and nothing told the user that - ``--steps`` and ``--strength`` interact (found by the release smoke matrix, - 2026-07-19). Raising the count to the minimum that denoises keeps the caller's intent - -- they asked for "few steps", not "zero" -- and the engine logs the adjustment. - - A non-positive ``strength`` cannot denoise at any step count; return the caller's - value unchanged rather than dividing by zero. - """ - if strength <= 0: - return num_inference_steps - if int(num_inference_steps * strength) >= 1: - return num_inference_steps - return math.ceil(1 / strength) - - -def vendor_for_strength(image_path: Path) -> Literal["openai", "google"] | None: - """Detect the SynthID vendor for strength selection: ``"openai"`` / ``"google"`` / None. - - Reads the C2PA SynthID proxy (``metadata.synthid_source``) on the ORIGINAL input, - so it must run before any pass that strips metadata. When both issuers appear (a - rare multi-sign anomaly) Google wins -- the more-robust watermark -> safer (higher) - strength. Returns None when metadata is stripped or the issuer is neither vendor, - which maps to ``UNKNOWN_STRENGTH``. Lazy-imports ``metadata`` to keep this module - dependency-light. - """ - try: - from remove_ai_watermarks.metadata import synthid_source - - src = (synthid_source(image_path) or "").lower() - except Exception: # metadata unreadable -> treat as unknown vendor - return None - if "google" in src: - return "google" - if "openai" in src: - return "openai" - return None diff --git a/src/remove_ai_watermarks/noai/watermark_remover.py b/src/remove_ai_watermarks/noai/watermark_remover.py deleted file mode 100644 index b5c3031..0000000 --- a/src/remove_ai_watermarks/noai/watermark_remover.py +++ /dev/null @@ -1,1117 +0,0 @@ -"""Watermark removal using diffusion model regeneration attack. - -Four pipelines (selected by the explicit ``pipeline`` ctor arg): - -0. ``qwen-zimage`` -- Qwen-Image-2512 Lightning + DiffSynth Canny regenerates the - frame, then SAM-masked Z-Image Turbo regenerates original face crops and feathers - them into the global result. CUDA-only and installed through its own optional extra. -1. ``qwen`` -- Qwen-Image (20B MMDiT, Apache-2.0) img2img. The scrub still comes from - the img2img ``strength``; Qwen preserves text (incl. CJK) and structure markedly - better than SDXL at the scrub floor, so it over-regenerates real photos far less. - CUDA/cloud-class (does not fit MPS). See ``watermark_profiles`` for the certified - oracle floors. (Near-threshold scrub is seed-non-deterministic, but the prod path - pins ONE fixed seed, so a certified floor reproduces run-to-run -- pinning the seed - is the release gate, not sweeping seeds.) - -Two SDXL pipelines: -2. ``controlnet`` (DEFAULT) -- SDXL img2img with a canny ControlNet. The watermark - REMOVAL still comes from the img2img regeneration (``strength``); the ControlNet - only PRESERVES structure (text/faces) by conditioning on the edge map. No original - pixels are ever copied or frozen. Because the edge map keeps the regeneration - closer to the original, it needs a higher ``strength`` floor than ``default`` to - destroy SynthID (the certified controlnet ladder; see ``watermark_profiles``). - ``controlnet_conditioning_scale`` is the preservation knob. -3. ``default`` -- plain SDXL img2img. Partial-noise regeneration scrubs the - invisible watermark; ``strength`` controls how much is regenerated. Lighter (no - ControlNet weights), but at the low default strength it leaves SynthID on - flat-graphic content -- use it for inputs without text/faces. -""" - -# torch/diffusers/cv2 boundary: these libs ship no usable types for the tensor and -# array ops below; relax the unknown-type rules for this file only. -# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false -from __future__ import annotations - -import contextlib -import logging -import os -import sys -import time -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Callable - from pathlib import Path - -from PIL import Image - -from remove_ai_watermarks.noai.watermark_profiles import ( - CONTROLNET_CANNY_MODEL, - DEFAULT_MODEL_ID, - DEFAULT_STRENGTH, - QWEN_MODEL_ID, - QWEN_ZIMAGE_PROFILE, - normalize_profile, - resolve_seed, - resolve_steps, - resolve_strength, - viable_steps, -) - -logger = logging.getLogger(__name__) - -# Check for optional dependencies -_HAS_TORCH = False -_HAS_DIFFUSERS = False - -try: - import torch - - _HAS_TORCH = True -except ImportError: - torch = None # type: ignore - -try: - from diffusers import AutoPipelineForImage2Image as AutoImg2ImgPipeline - - _HAS_DIFFUSERS = True -except ImportError: - AutoImg2ImgPipeline = None # type: ignore - - -def is_watermark_removal_available() -> bool: - """Check if watermark removal dependencies are installed.""" - return _HAS_TORCH and _HAS_DIFFUSERS - - -# Drop-in fp16-safe replacement for the SDXL VAE. The stock SDXL VAE overflows -# to NaN in fp16 and decodes to an all-black image (issue #29, reproduced on a -# CUDA fp16 backend). This community VAE is numerically rescaled to -# stay in fp16 range. SDXL-architecture only. -_SDXL_FP16_VAE_ID = "madebyollin/sdxl-vae-fp16-fix" - - -def _needs_fp16_vae_fix(model_id: str, default_model_id: str, is_fp16: bool) -> bool: - """Whether the plain img2img pipeline must swap in the fp16-fixed SDXL VAE. - - Gated to the default SDXL checkpoint running in fp16: cpu/mps run fp32 (the - stock VAE is fine there) and the differential pipeline upcasts the VAE on its - own, so only this path on a fp16 GPU (CUDA/XPU) hits the NaN/black decode. - A custom non-SDXL ``model_id`` keeps its own VAE (the fix is SDXL-specific). - """ - return is_fp16 and model_id == default_model_id - - -# An fp16 VAE/UNet overflow decodes to NaN, which diffusers' postprocess casts to 0 -# -> a uniform all-black frame (issues #29, #41). The VAE swap above prevents it for -# the default checkpoint, but a custom model_id, a stale install, or a fal/custom -# loader can still bypass it. Detecting a degenerate output and retrying in fp32 (the -# path verified clean) is the model-agnostic safety net: never hand back a black image. -# One threshold serves both guards: a NaN->0 collapse drives mean and variance to ~0. -_DEGENERATE_THRESHOLD = 1.0 - - -def _is_degenerate_image(image: Image.Image) -> bool: - """True if a generated image collapsed to an all-black/NaN frame (#29/#41). - - A NaN fp16 decode casts to 0, so the output is a uniform near-zero image: an - extremely low mean AND near-zero variance. The variance guard keeps a - legitimately dark-but-textured photo (low mean, real detail) from being flagged. - """ - import numpy as np - - arr = np.asarray(image.convert("RGB"), dtype=np.float32) - return float(arr.mean()) < _DEGENERATE_THRESHOLD and float(arr.std()) < _DEGENERATE_THRESHOLD - - -_CUDA_FIX_ENV_KEY = "NOAI_CUDA_FIXED" - - -def _auto_install(packages: list[str], index_url: str | None = None) -> bool: - """Attempt to install missing packages via pip. Returns True on success.""" - import subprocess - - cmd = [sys.executable, "-m", "pip", "install", "-q", *packages] - if index_url: - cmd.extend(["--index-url", index_url]) - try: - subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) - return True - except (subprocess.CalledProcessError, FileNotFoundError): - return False - - -def _has_nvidia_gpu() -> bool: - """Check if an NVIDIA GPU is present via nvidia-smi.""" - import subprocess - - try: - subprocess.check_call( - ["nvidia-smi"], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - return True - except (subprocess.CalledProcessError, FileNotFoundError): - return False - - -def _detect_cuda_index_url() -> str: - """Detect the appropriate PyTorch CUDA index URL from nvidia-smi output.""" - import subprocess - - try: - out = subprocess.check_output( - ["nvidia-smi"], - stderr=subprocess.DEVNULL, - text=True, - ) - for line in out.splitlines(): - if "CUDA Version" in line: - version_str = line.split("CUDA Version:")[-1].strip().rstrip("|").strip() - major, minor = version_str.split(".")[:2] - cuda_tag = f"cu{major}{minor}" - return f"https://download.pytorch.org/whl/{cuda_tag}" - except Exception: # noqa: S110 - pass - return "https://download.pytorch.org/whl/cu121" - - -def _reinstall_torch_cuda_and_restart() -> None: - """Reinstall torch with CUDA support showing live progress, then restart.""" - import re - import subprocess - - from remove_ai_watermarks.noai.progress import run_with_progress - - index_url = _detect_cuda_index_url() - progress_state: dict[str, str] = {"message": "NVIDIA GPU detected — installing CUDA-enabled PyTorch..."} - - pct_re = re.compile(r"(\d+)%") - pkg_re = re.compile(r"(?:Collecting|Downloading|Installing)\s+(\S+)") - - def _run_pip() -> bool: - cmd = [ - sys.executable, - "-m", - "pip", - "install", - "--force-reinstall", - "torch", - "--index-url", - index_url, - ] - proc = subprocess.Popen( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - for line in iter(proc.stdout.readline, ""): # type: ignore[union-attr] - stripped = line.strip() - if not stripped: - continue - pkg_m = pkg_re.search(stripped) - pct_m = pct_re.search(stripped) - if pct_m and pkg_m: - progress_state["message"] = f"Downloading {pkg_m.group(1)} ({pct_m.group(1)}%)" - elif pct_m: - progress_state["message"] = f"Downloading CUDA packages ({pct_m.group(1)}%)" - elif pkg_m: - action = "Installing" if stripped.startswith("Installing") else "Downloading" - progress_state["message"] = f"{action} {pkg_m.group(1)}" - elif "Successfully installed" in stripped: - progress_state["message"] = "CUDA-enabled PyTorch installed successfully" - proc.wait() - return proc.returncode == 0 - - try: - success = run_with_progress(_run_pip, progress_state) - except Exception: - success = False - - if not success: - print( - f"\n Failed to install CUDA-enabled PyTorch.\n" - f" Install manually:\n" - f" pip install torch --index-url {index_url}\n", - file=sys.stderr, - ) - return - - os.environ[_CUDA_FIX_ENV_KEY] = "1" - # Re-exec via ``-m`` rather than building a ``-c`` string from repr(sys.argv). - # ``-m`` makes Python set argv[0] to the module path, so forward only the - # actual args (sys.argv[1:]); passing the full argv would re-inject the - # program name as a spurious first argument to Click. - os.execv(sys.executable, [sys.executable, "-m", "remove_ai_watermarks.cli", *sys.argv[1:]]) - - -def _ensure_watermark_deps() -> None: - """Auto-install and re-import missing watermark removal dependencies.""" - global _HAS_TORCH, _HAS_DIFFUSERS, torch, AutoImg2ImgPipeline - missing_pkgs: list[str] = [] - if not _HAS_TORCH: - missing_pkgs.append("torch") - if not _HAS_DIFFUSERS: - missing_pkgs.extend(["diffusers", "transformers", "accelerate"]) - logger.info("Auto-installing missing dependencies: %s", missing_pkgs) - if not _auto_install(missing_pkgs): - raise ImportError( - f"Failed to auto-install missing dependencies: {', '.join(missing_pkgs)}. " - "Try manually: pip install --force-reinstall noai-watermark" - ) - import torch as _torch - - torch = _torch - _HAS_TORCH = True - from diffusers import AutoPipelineForImage2Image - - AutoImg2ImgPipeline = AutoPipelineForImage2Image - _HAS_DIFFUSERS = True - - -def get_device() -> str: - """Get the best available device for inference.""" - if not _HAS_TORCH: - return "cpu" - if torch.cuda.is_available(): # type: ignore - try: - t = torch.tensor([1.0], device="cuda") - _ = t + t - del t - return "cuda" - except (AssertionError, RuntimeError): - pass - # Intel GPU (Arc / Data Center) via the torch XPU backend. The torch.xpu - # namespace exists in stock wheels, but is_available() is only True on an - # XPU-enabled build (download.pytorch.org/whl/xpu), so this is inert on the - # default CPU/CUDA install. Checked before the nvidia-smi path so an Intel - # box never triggers the CUDA reinstaller. - if hasattr(torch, "xpu") and torch.xpu.is_available(): # type: ignore - try: - t = torch.tensor([1.0], device="xpu") - _ = t + t - del t - return "xpu" - except (AssertionError, RuntimeError): - pass - if _has_nvidia_gpu() and not os.environ.get(_CUDA_FIX_ENV_KEY): - _reinstall_torch_cuda_and_restart() - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return "mps" - return "cpu" - - -def _make_seed_generator(device: str, seed: int) -> Any: - """Build a seeded ``torch.Generator``, falling back to a CPU generator. - - Some backends have no device-side RNG (notably certain torch-xpu builds), - so ``torch.Generator(device="xpu")`` can raise. A CPU generator is - backend-agnostic and still seeds the pipeline reproducibly, so fall back to - it rather than failing the run when ``--seed`` is used on such a device. - """ - try: - return torch.Generator(device=device).manual_seed(seed) # type: ignore - except (RuntimeError, TypeError): - return torch.Generator().manual_seed(seed) # type: ignore - - -# Canny edge thresholds for the ControlNet control image (xinsir canny recipe: -# cv2.Canny(gray, 100, 200) -> a 3-channel edge map). -_CANNY_LOW = 100 -_CANNY_HIGH = 200 - -# A neutral quality prompt: the goal is faithful regeneration, not creative edits. -_CONTROLNET_PROMPT = "best quality, high quality, sharp, detailed, photographic" -_CONTROLNET_NEGATIVE = "blurry, lowres, deformed, distorted text, garbled text, watermark, jpeg artifacts" - -# Neutral prompts for the Qwen-Image img2img pass (faithful regeneration, not an edit). -_QWEN_PROMPT = "high quality, sharp, detailed, faithful to the original" -_QWEN_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts" - - -def _qwen_target_size(width: int, height: int) -> tuple[int, int]: - """Floor (width, height) to a multiple of 16 for Qwen's VAE/patchifier (>= 16). - - Pure; unit-tested. Without explicit dims the img2img pipeline defaults to a 1024x1024 - SQUARE and silently distorts any non-square input. - """ - return max(16, (width // 16) * 16), max(16, (height // 16) * 16) - - -def _build_qwen_kwargs( - image: Image.Image, strength: float, num_inference_steps: int, true_cfg_scale: float, generator: Any -) -> dict[str, Any]: - """Build the QwenImageImg2ImgPipeline call kwargs (pure; unit-tested without torch). - - Qwen-Image uses ``true_cfg_scale`` (not SDXL's ``guidance_scale``) and takes an - explicit ``negative_prompt``; the scrub still comes from the img2img ``strength``. - Passes an explicit ``height``/``width`` derived from the input (floored to /16): the - pipeline otherwise defaults to a 1024x1024 SQUARE, squishing any non-square input - (the abba mixed-seam test: a 2816x1536 poster came back 1024x1024, distorting the - scene and garbling text). So qwen regenerates at the input's own geometry. - """ - qw, qh = _qwen_target_size(image.width, image.height) - return { - "prompt": _QWEN_PROMPT, - "negative_prompt": _QWEN_NEGATIVE, - "image": image, - "strength": strength, - "num_inference_steps": num_inference_steps, - "true_cfg_scale": true_cfg_scale, - "generator": generator, - "height": qh, - "width": qw, - } - - -class WatermarkRemover: - """Remove watermarks from images using diffusion model regeneration. - - Attributes: - model_id: HuggingFace model ID for the diffusion model. - device: Device to run inference on (cuda, xpu, mps, or cpu). - """ - - DEFAULT_MODEL_ID = DEFAULT_MODEL_ID - DEFAULT_STRENGTH = DEFAULT_STRENGTH - CONTROLNET_CANNY_MODEL = CONTROLNET_CANNY_MODEL - - def __init__( - self, - model_id: str | None = None, - device: str | None = None, - torch_dtype: Any = None, - progress_callback: Callable[[str], None] | None = None, - hf_token: str | None = None, - pipeline: str = "controlnet", - controlnet_conditioning_scale: float = 1.0, - cpu_offload: bool = False, - ) -> None: - self.model_id = model_id or self.DEFAULT_MODEL_ID - # Diffusers offloads whole model components between CUDA calls. The custom - # qwen-zimage runtime uses the same flag to keep its face stack off VRAM. - self.cpu_offload = cpu_offload - # The pipeline profile is threaded explicitly (not inferred from model_id): - # both "sdxl" and "controlnet" use the same SDXL base checkpoint. Normalize so - # the legacy "default" alias resolves to "sdxl". - self.model_profile = normalize_profile(pipeline) - self.controlnet_conditioning_scale = controlnet_conditioning_scale - if self.model_profile == QWEN_ZIMAGE_PROFILE and self.model_id != self.DEFAULT_MODEL_ID: - raise ValueError( - "The qwen-zimage pipeline uses a fixed Qwen-Image-2512 + Z-Image model stack; " - "--model is not supported for this profile." - ) - if self.model_profile == QWEN_ZIMAGE_PROFILE: - self.model_id = "Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo" - - if not is_watermark_removal_available(): - _ensure_watermark_deps() - self.device = (device or get_device()).lower() - if self.device == "auto": - self.device = get_device() - if self.device not in {"cpu", "mps", "cuda", "xpu"}: - raise ValueError(f"Unsupported device '{device}'. Use one of: auto, cpu, mps, cuda, xpu.") - if torch_dtype is None: - if self.device == "cpu" or self.device == "mps": - self.torch_dtype = torch.float32 # type: ignore - elif self.model_profile in {"qwen", QWEN_ZIMAGE_PROFILE}: - # Qwen-Image and Z-Image are published in bf16; fp16 risks overflow. - # cuda/xpu-only by construction: the cpu/mps guard above already forced - # fp32, and the 20B model does not fit MPS anyway. - self.torch_dtype = torch.bfloat16 # type: ignore - else: - self.torch_dtype = torch.float16 # type: ignore - else: - self.torch_dtype = torch_dtype - - self._pipeline: AutoImg2ImgPipeline | None = None - self._controlnet_pipeline: Any = None - self._qwen_pipeline: Any = None - self._qwen_zimage_pipeline: Any = None - self._progress_callback = progress_callback - self.hf_token: str | None = hf_token or os.environ.get("HF_TOKEN") - - def _set_progress(self, message: str) -> None: - """Send a progress update through callback when available.""" - if self._progress_callback is None: - return - with contextlib.suppress(Exception): - self._progress_callback(message) - - # ── Preload ────────────────────────────────────────────────────── - - def preload(self, *, global_only: bool = False) -> None: - """Eagerly load the pipeline so download progress bars are visible. - - ``global_only`` applies to qwen-zimage, whose face stage is optional. - """ - if self.model_profile == QWEN_ZIMAGE_PROFILE: - self._load_qwen_zimage_pipeline().preload(global_only=global_only) - elif self.model_profile == "qwen": - self._load_qwen_pipeline() - elif self.model_profile == "controlnet": - self._load_controlnet_pipeline() - else: - self._load_pipeline() - - # ── Pipeline loading ───────────────────────────────────────────── - - def _maybe_add_fp16_vae(self, load_kwargs: dict[str, Any]) -> None: - """Swap in the fp16-fixed SDXL VAE for the default checkpoint on a fp16 GPU. - - The stock SDXL VAE overflows to NaN in fp16 and decodes to an all-black - image (issue #29). Shared by both pipeline loaders; a no-op on fp32 (cpu/mps) - or a non-SDXL checkpoint. - """ - if _needs_fp16_vae_fix(self.model_id, self.DEFAULT_MODEL_ID, self.torch_dtype == torch.float16): - from diffusers import AutoencoderKL - - self._set_progress("Loading fp16-fixed SDXL VAE (avoids black output)...") - load_kwargs["vae"] = AutoencoderKL.from_pretrained(_SDXL_FP16_VAE_ID, torch_dtype=torch.float16) - - @staticmethod - def _disable_sdxl_watermarker(load_kwargs: dict[str, Any]) -> None: - """Turn off the diffusers default invisible watermarker on an SDXL pipeline. - - diffusers embeds an open "Stable Diffusion XL" DWT-DCT invisible watermark on - EVERY SDXL output whenever ``invisible-watermark`` is installed (kept as a - development parity dependency). A watermark REMOVER must not re-stamp a - detectable AI watermark, or the - cleaned output re-reads as AI (``identify`` -> "Open invisible watermark: Stable - Diffusion XL"). Shared by both SDXL loaders; the ``ControlNetModel`` sub-model - and the Qwen loader never call it (only the pipeline accepts the kwarg). - """ - load_kwargs["add_watermarker"] = False - - def _move_to_device_and_optimize(self, pipeline: Any) -> Any: - """Move a freshly-loaded pipeline to ``self.device`` + enable memory opts. - - Shared by both loaders. On a CUDA move failure (missing CUDA torch build), - trigger the torch-CUDA reinstall+restart. Returns the moved pipeline. - """ - self._set_progress(f"Moving model to device: {self.device}") - if self.cpu_offload and self.device == "cuda": - enable_cpu_offload = getattr(pipeline, "enable_model_cpu_offload", None) - if not callable(enable_cpu_offload): - raise RuntimeError( - "CPU offload was requested, but this pipeline does not support enable_model_cpu_offload()." - ) - self._set_progress("Enabling CUDA model CPU offload (low-VRAM mode)...") - try: - enable_cpu_offload(device=self.device) - except (RuntimeError, AssertionError) as exc: - raise RuntimeError(f"Failed to enable model CPU offload ({exc}).") from exc - else: - try: - pipeline = pipeline.to(self.device) - except (RuntimeError, AssertionError) as exc: - if self.device == "cuda" and not os.environ.get(_CUDA_FIX_ENV_KEY): - self._set_progress("CUDA failed. Reinstalling torch with CUDA support...") - _reinstall_torch_cuda_and_restart() - raise RuntimeError( - f"Failed to move model to {self.device} ({exc}). " - "Install CUDA-enabled PyTorch manually:\n" - f" pip install torch --index-url {_detect_cuda_index_url()}" - ) from exc - - if hasattr(pipeline, "enable_xformers_memory_efficient_attention"): - with contextlib.suppress(Exception): - self._set_progress("Enabling memory optimizations...") - pipeline.enable_xformers_memory_efficient_attention() - - # Mac Float32 memory slicing - if self.device == "mps" and hasattr(pipeline, "enable_attention_slicing"): - with contextlib.suppress(Exception): - pipeline.enable_attention_slicing("max") - - return pipeline - - def _base_load_kwargs(self) -> dict[str, Any]: - """The ``from_pretrained`` kwargs shared by all three loaders (dtype + token). - - Each loader adds its own extras (SDXL safety_checker + fp16 VAE, the ControlNet - model, etc.). Centralizing the dtype/token pair avoids the drift trap of three - copies (a token forgotten on one loader silently breaks gated downloads there). - """ - load_kwargs: dict[str, Any] = {"torch_dtype": self.torch_dtype} - if self.hf_token: - load_kwargs["token"] = self.hf_token - return load_kwargs - - def _load_from_pretrained(self, cls: Any, model_id: str, **load_kwargs: Any) -> Any: - """Call ``cls.from_pretrained`` reading the fp16 weight VARIANT when on fp16. - - When ``torch_dtype`` is float16 (the CUDA/XPU SDXL default), pass - ``variant="fp16"`` so diffusers fetches/reads the half-precision weight files - (~half the bytes of the fp32 defaults) instead of reading the full fp32 files - and downcasting in memory. On a warm weights cache this roughly halves the - cold-start weight read + host->device transfer, which a phase-timed Modal run - measured as ~half of the ~25s cold start. Not every checkpoint publishes an - fp16 variant (a custom ``--model``, or the canny ControlNet if it ships only - default files), so fall back to the default weights if the variant is missing - - worst case is the prior behavior (read fp32, downcast). fp32 (cpu/mps) and bf16 - (qwen) never request the variant. - """ - if self.torch_dtype == torch.float16: - try: - return cls.from_pretrained(model_id, variant="fp16", **load_kwargs) - except Exception as exc: - logger.info("No fp16 weight variant for %s (%s); loading default weights", model_id, exc) - return cls.from_pretrained(model_id, **load_kwargs) - - def _load_pipeline(self) -> AutoImg2ImgPipeline: - """Load the plain SDXL img2img pipeline lazily.""" - if self._pipeline is None: - logger.info("Loading model %s on %s...", self.model_id, self.device) - self._set_progress(f"Loading model weights: {self.model_id}") - - load_kwargs = self._base_load_kwargs() - load_kwargs["safety_checker"] = None - load_kwargs["requires_safety_checker"] = False - self._disable_sdxl_watermarker(load_kwargs) - self._maybe_add_fp16_vae(load_kwargs) - - pipeline = self._load_from_pretrained(AutoImg2ImgPipeline, self.model_id, **load_kwargs) # type: ignore - self._pipeline = self._move_to_device_and_optimize(pipeline) - - logger.info("Model loaded successfully") - self._set_progress("Model initialized. Preparing input image...") - - return self._pipeline # type: ignore - - def _load_controlnet_pipeline(self) -> Any: - """Load the SDXL + canny-ControlNet img2img pipeline lazily. - - Mirrors ``_load_pipeline`` (same fp16-fix VAE, device move, attention - slicing via the shared helpers) but loads the canny ControlNet on top of - the SDXL base. The ControlNet only preserves structure via the edge map; - removal still comes from the img2img regeneration (``strength``). - """ - if self._controlnet_pipeline is None: - from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline - - logger.info("Loading SDXL + ControlNet (%s) on %s...", CONTROLNET_CANNY_MODEL, self.device) - self._set_progress(f"Loading ControlNet: {CONTROLNET_CANNY_MODEL}") - controlnet = self._load_from_pretrained( - ControlNetModel, CONTROLNET_CANNY_MODEL, torch_dtype=self.torch_dtype - ) - - load_kwargs = self._base_load_kwargs() - load_kwargs["controlnet"] = controlnet - self._disable_sdxl_watermarker(load_kwargs) - self._maybe_add_fp16_vae(load_kwargs) - - self._set_progress(f"Loading model weights: {self.model_id}") - pipeline = self._load_from_pretrained( - StableDiffusionXLControlNetImg2ImgPipeline, self.model_id, **load_kwargs - ) - pipeline = self._move_to_device_and_optimize(pipeline) - with contextlib.suppress(Exception): - pipeline.set_progress_bar_config(disable=True) - - logger.info("ControlNet model loaded successfully") - self._controlnet_pipeline = pipeline - - return self._controlnet_pipeline - - def _load_qwen_pipeline(self) -> Any: - """Load the Qwen-Image img2img pipeline lazily. - - Qwen-Image is its OWN base model (not an SDXL add-on), so it loads - ``QWEN_MODEL_ID`` unless the caller passed a custom ``--model``. Needs a - diffusers build that ships ``QwenImageImg2ImgPipeline``; raises a clear error - otherwise. CUDA/cloud-class (the 20B MMDiT does not fit MPS). - """ - if self._qwen_pipeline is None: - try: - from diffusers import QwenImageImg2ImgPipeline - except ImportError as exc: - raise ImportError( - "The 'qwen' pipeline needs a diffusers version that ships " - "QwenImageImg2ImgPipeline. Upgrade: pip install -U diffusers" - ) from exc - - # Use the Qwen base unless the user explicitly overrode --model. - model = self.model_id if self.model_id != self.DEFAULT_MODEL_ID else QWEN_MODEL_ID - logger.info("Loading Qwen-Image (%s) on %s...", model, self.device) - self._set_progress(f"Loading model weights: {model}") - pipeline = QwenImageImg2ImgPipeline.from_pretrained(model, **self._base_load_kwargs()) - pipeline = self._move_to_device_and_optimize(pipeline) - with contextlib.suppress(Exception): - pipeline.set_progress_bar_config(disable=True) - - logger.info("Qwen-Image model loaded successfully") - self._qwen_pipeline = pipeline - - return self._qwen_pipeline - - def _load_qwen_zimage_pipeline(self) -> Any: - """Load the two-stage Qwen-Image-2512 + Z-Image runtime lazily.""" - if self._qwen_zimage_pipeline is None: - from remove_ai_watermarks.noai.qwen_zimage_pipeline import QwenZImagePipeline - - self._qwen_zimage_pipeline = QwenZImagePipeline( - device=self.device, - torch_dtype=self.torch_dtype, - hf_token=self.hf_token, - progress_callback=self._progress_callback, - controlnet_conditioning_scale=self.controlnet_conditioning_scale, - keep_face_models_on_device=False if self.cpu_offload else None, - ) - return self._qwen_zimage_pipeline - - # ── Core removal ───────────────────────────────────────────────── - - def remove_watermark( - self, - image_path: Path, - output_path: Path | None = None, - strength: float | None = None, - num_inference_steps: int | None = None, - guidance_scale: float | None = None, - seed: int | None = None, - vendor: str | None = None, - tile: bool = False, - tile_size: int = 1024, - tile_overlap: int = 128, - region: tuple[int, int, int, int] | None = None, - region_feather: int = 64, - ) -> Path: - """Remove watermark from an image using regeneration attack. - - Args: - image_path: Path to the watermarked image. - output_path: Path for the cleaned image. If None, modifies in place. - strength: Denoising strength (0.0-1.0). None -> the vendor-adaptive - default (see ``vendor``). - num_inference_steps: Number of denoising steps. - guidance_scale: Classifier-free guidance scale. - seed: Random seed for reproducibility. None resolves to 0 for - qwen-zimage and stays random for the other profiles. - vendor: SynthID vendor (``"openai"`` / ``"google"`` / None) used to pick the - default strength when ``strength`` is None. Detect it from the ORIGINAL - input with ``watermark_profiles.vendor_for_strength`` before processing - strips the metadata; the caller passes it down so display and execution - agree. - tile: Process the image in overlapping tiles instead of one forward pass. - This keeps the input's native dimensions instead of applying a - ``--max-resolution`` downscale, but every tile is still regenerated. - Only engages when the long side exceeds ``tile_size``; smaller images - run a single pass unchanged. - tile_size: Tile dimension in px (default 1024). - tile_overlap: Overlap between adjacent tiles in px (default 128), feather- - blended so there is no visible seam. - region: Restrict the regeneration to the AI-composited box ``(x, y, w, h)`` - and feather-composite it back over the ORIGINAL pixels everywhere else. - For AI-ENHANCED composites (digitalSourceType - ``compositeWithTrainedAlgorithmicMedia``, surfaced as - ``identify.ProvenanceReport.ai_source_kind == "enhanced"``): the real - photo outside the box is preserved exactly, only the AI region is - scrubbed. The box is supplied by the caller (a C2PA composite manifest - does not carry a reliable machine-readable region). None -> whole frame. - region_feather: Seam taper in px for ``region`` compositing (default 64). - - Returns: - Path to the cleaned image. - - Raises: - FileNotFoundError: If input image doesn't exist. - ValueError: If strength is not in valid range. - """ - if not image_path.exists(): - raise FileNotFoundError(f"Image not found: {image_path}") - - if output_path is None: - output_path = image_path - - self._set_progress("Loading and preprocessing input image...") - init_image = Image.open(image_path).convert("RGB") - w, h = init_image.size - self._set_progress(f"Image loaded: {w}x{h}px | Model: {self.model_id}") - - if self.model_profile == QWEN_ZIMAGE_PROFILE: - from remove_ai_watermarks.noai.qwen_zimage_pipeline import resolution_adaptive_denoise - - strength = strength if strength is not None else resolution_adaptive_denoise(w, h) - else: - strength = resolve_strength(strength, vendor, self.model_profile) - seed = resolve_seed(seed, self.model_profile) - if not 0.0 <= strength <= 1.0: - raise ValueError(f"Strength must be between 0.0 and 1.0, got {strength}") - - num_inference_steps = resolve_steps(num_inference_steps, self.model_profile) - if guidance_scale is None: - guidance_scale = 1.0 if self.model_profile == QWEN_ZIMAGE_PROFILE else 7.5 - elif self.model_profile == QWEN_ZIMAGE_PROFILE and guidance_scale != 1.0: - raise ValueError("The qwen-zimage profile fixes both diffusion stages at CFG 1.0.") - - generator = None - if seed is not None and _HAS_TORCH: - self._set_progress(f"Setting reproducible seed: {seed}") - generator = _make_seed_generator(self.device, seed) - - # A step count whose product with strength rounds to zero kills the pipeline - # inside attention with an opaque reshape error, so raise it to the minimum that - # denoises. Must be applied to the value HANDED TO THE PIPELINE -- the old - # max(1, ...) below only clamped the number in the log line. - if self.model_profile != QWEN_ZIMAGE_PROFILE: - adjusted = viable_steps(num_inference_steps, strength) - if adjusted != num_inference_steps: - logger.warning( - "steps=%s at strength=%s denoises 0 steps and would crash; using steps=%s (1 effective)", - num_inference_steps, - strength, - adjusted, - ) - num_inference_steps = adjusted - elif num_inference_steps != 4: - raise ValueError("The qwen-zimage profile uses the 4-step Lightning LoRA, so --steps must be 4.") - - # DiffSynth keeps all timesteps and compresses their sigma range for a low - # denoise value. Diffusers instead truncates the schedule by strength. - effective_steps = ( - num_inference_steps - if self.model_profile == QWEN_ZIMAGE_PROFILE - else max(1, int(num_inference_steps * strength)) - ) - self._set_progress( - f"Config: strength={strength}, steps={num_inference_steps} " - f"(~{effective_steps} effective), guidance={guidance_scale}, device={self.device}" - ) - - _total_start = time.monotonic() - - def _generate_one(img: Image.Image) -> Image.Image: - if self.model_profile == "qwen": - return self._run_qwen(img, strength, num_inference_steps, guidance_scale, generator) - if self.model_profile == "controlnet": - return self._run_controlnet(img, strength, num_inference_steps, guidance_scale, generator) - return self._run_img2img(img, strength, num_inference_steps, guidance_scale, generator) - - def _generate() -> Image.Image: - # qwen-zimage owns its global-only tiling because its face stage must run - # once after the tiles are blended. Other profiles tile their whole pass. - if self.model_profile == QWEN_ZIMAGE_PROFILE: - return self._run_qwen_zimage( - init_image, - strength, - seed, - tile=tile, - tile_size=tile_size, - tile_overlap=tile_overlap, - ) - # Tile only when asked AND the image is larger than one tile; otherwise a - # single full-image pass (tiling a sub-tile image is pure overhead). - if tile and max(init_image.size) > tile_size: - from remove_ai_watermarks.noai.tiling import run_tiled - - return run_tiled(_generate_one, init_image, tile_size, tile_overlap, self._set_progress) - return _generate_one(init_image) - - cleaned_image = _generate() - - # Safety net for the fp16 all-black/NaN decode (#29/#41): if an fp16 run - # produced a degenerate (uniform black) frame -- the VAE swap did not engage - # for this model/version -- retry once in fp32 on the same device (verified - # clean) so the user never gets a black image. Skipped when an MPS->CPU - # fallback already moved us to fp32. - if self.torch_dtype == torch.float16 and _is_degenerate_image(cleaned_image): - logger.warning("fp16 output was degenerate (all-black/NaN, #29/#41); retrying in fp32 on %s.", self.device) - self._set_progress("Output was black (fp16 overflow); retrying in fp32...") - self.torch_dtype = torch.float32 - self._pipeline = None - self._controlnet_pipeline = None - self._qwen_zimage_pipeline = None - cleaned_image = _generate() - - # Region-targeted regeneration for AI-enhanced composites: keep the real photo - # outside the AI box pixel-exact, blend only the regenerated AI region back in. - if region is not None: - import numpy as np - - from remove_ai_watermarks.noai.tiling import feather_region_composite - - gen = cleaned_image.convert("RGB") - if gen.size != init_image.size: # a downscaled/tiled pass can resize - gen = gen.resize(init_image.size) - cleaned_image = gen - base_rgb = np.asarray(init_image) # original RGB, untouched outside the box - merged = feather_region_composite(base_rgb, np.asarray(gen), region, feather=region_feather) - cleaned_image = Image.fromarray(merged) - self._set_progress(f"Region-targeted regeneration: AI box {region}, real photo preserved") - - self._set_progress(f"Regeneration complete · Output: {w}x{h}px {cleaned_image.mode}") - - output_path.parent.mkdir(parents=True, exist_ok=True) - fmt = output_path.suffix.lower() - if fmt in (".jpg", ".jpeg"): - self._set_progress(f"Encoding as JPEG → {output_path.name}...") - else: - self._set_progress(f"Encoding as PNG → {output_path.name}...") - # Encode through image_io so the regenerated image gets the same quality- - # preserving write as the visible path (JPEG q100 / 4:4:4, HEIC/AVIF via Pillow) - # instead of PIL's default JPEG quality 75, which would crush the diffusion - # output further for no reason. - import numpy as np - - from remove_ai_watermarks import image_io - - rgb = np.asarray(cleaned_image.convert("RGB"))[:, :, ::-1] # PIL RGB -> cv2 BGR - if not image_io.imwrite(str(output_path), np.ascontiguousarray(rgb)): - cleaned_image.save(output_path) # fallback for a container image_io/cv2 cannot encode - - if output_path.exists(): - self._set_progress("Stripping AI metadata from output...") - try: - # The single, robust stripper (byte-level: lossless for JPEG, handles - # every container) -- not the legacy PIL-re-encoding noai.cleaner one. - from remove_ai_watermarks.metadata import remove_ai_metadata - - remove_ai_metadata(output_path, output_path, keep_standard=True) - except Exception: - logger.debug("AI metadata stripping skipped", exc_info=True) - - total_time = time.monotonic() - _total_start - - size_str = "" - try: - file_size = output_path.stat().st_size - if file_size < 1024 * 1024: - size_str = f" ({file_size / 1024:.0f}KB)" - else: - size_str = f" ({file_size / (1024 * 1024):.1f}MB)" - except OSError: - pass - - logger.info("Cleaned image saved to %s", output_path) - self._set_progress(f"✓ Saved {output_path.name}{size_str} · {w}x{h}px · {total_time:.0f}s total") - - return output_path - - # ── Img2img runner ─────────────────────────────────────────────── - - def _run_img2img( - self, - init_image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - ) -> Image.Image: - """Execute the img2img pipeline with progress and MPS fallback.""" - from remove_ai_watermarks.noai.img2img_runner import run_img2img_with_mps_fallback - - result_image, final_device = run_img2img_with_mps_fallback( - load_pipeline=self._load_pipeline, - image=init_image, - strength=strength, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - device=self.device, - set_progress=self._set_progress, - reload_on_cpu=self._reload_pipeline_on_cpu, - ) - - if final_device != self.device: - self.device = final_device - self.torch_dtype = torch.float32 # type: ignore[assignment] - - return result_image - - def _reload_pipeline_on_cpu(self) -> Any: - """Reload pipeline on CPU after MPS failure.""" - self.device = "cpu" - self.torch_dtype = torch.float32 # type: ignore[assignment] - self._pipeline = None - return self._load_pipeline() - - # ── ControlNet runner ──────────────────────────────────────────── - - def _build_canny_control_image(self, init_image: Image.Image) -> Image.Image: - """Build the canny ControlNet conditioning image (xinsir recipe). - - cv2.Canny on the RGB->gray array, stacked to 3 channels, wrapped as a PIL - image. The edge map only PRESERVES structure; it never copies pixels. - ``init_image`` is already RGB (``remove_watermark`` converts on load). - """ - import cv2 - import numpy as np - - rgb = np.array(init_image) - gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) - edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH) - edges_rgb = np.stack([edges, edges, edges], axis=-1) - return Image.fromarray(edges_rgb) - - def _run_controlnet( - self, - init_image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - ) -> Image.Image: - """Run the SDXL + canny-ControlNet img2img pass. - - Removal still comes from the img2img regeneration (``strength``); the canny - ControlNet only PRESERVES text and face STRUCTURE via the edge map. No - original pixels are copied/frozen, so SynthID does not survive (canny holds - structure, not face identity). ``controlnet_conditioning_scale`` is the - structure-preservation knob. Shares the img2img runner (live progress + - MPS->CPU fallback) with ``_run_img2img``; the only delta is the extra - ControlNet kwargs (canny control image + conditioning scale + a non-empty - prompt) overlaid via ``extra_kwargs``. - """ - from remove_ai_watermarks.noai.img2img_runner import run_img2img_with_mps_fallback - - extra_kwargs = { - "prompt": _CONTROLNET_PROMPT, - "negative_prompt": _CONTROLNET_NEGATIVE, - "control_image": self._build_canny_control_image(init_image), - "controlnet_conditioning_scale": float(self.controlnet_conditioning_scale), - } - result_image, final_device = run_img2img_with_mps_fallback( - load_pipeline=self._load_controlnet_pipeline, - image=init_image, - strength=strength, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - device=self.device, - set_progress=self._set_progress, - reload_on_cpu=self._reload_controlnet_on_cpu, - extra_kwargs=extra_kwargs, - ) - - if final_device != self.device: - self.device = final_device - self.torch_dtype = torch.float32 # type: ignore[assignment] - - return result_image - - def _reload_controlnet_on_cpu(self) -> Any: - """Reload the controlnet pipeline on CPU after an MPS failure.""" - self.device = "cpu" - self.torch_dtype = torch.float32 # type: ignore[assignment] - self._controlnet_pipeline = None - return self._load_controlnet_pipeline() - - # ── Qwen runner ────────────────────────────────────────────────── - - def _run_qwen( - self, - init_image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - ) -> Image.Image: - """Run the Qwen-Image img2img pass. - - Removal comes from the img2img ``strength`` (same lever as the SDXL paths); - Qwen-Image preserves text/structure markedly better at the scrub floor. The - CLI ``guidance_scale`` maps to Qwen's ``true_cfg_scale`` (~4.0 is typical; - the SDXL default of 7.5 is high for Qwen). No MPS->CPU fallback: the 20B MMDiT - is CUDA/cloud-class and does not run on MPS, so an error here propagates. - """ - pipeline = self._load_qwen_pipeline() - self._set_progress(f"Running Qwen-Image img2img (strength={strength}, true_cfg={guidance_scale})...") - kwargs = _build_qwen_kwargs(init_image, strength, num_inference_steps, guidance_scale, generator) - result = pipeline(**kwargs) - return result.images[0] - - def _run_qwen_zimage( - self, - init_image: Image.Image, - strength: float, - seed: int | None, - *, - tile: bool = False, - tile_size: int = 1024, - tile_overlap: int = 128, - ) -> Image.Image: - """Run the Qwen 2512 Canny pass and masked Z-Image face repair.""" - pipeline = self._load_qwen_zimage_pipeline() - return pipeline.run( - init_image, - strength=strength, - seed=seed, - tile=tile, - tile_size=tile_size, - tile_overlap=tile_overlap, - ) - - # ── Batch ──────────────────────────────────────────────────────── - - def remove_watermark_batch( - self, - input_dir: Path, - output_dir: Path, - strength: float | None = None, - num_inference_steps: int | None = None, - extensions: tuple[str, ...] = (".png", ".jpg", ".jpeg", ".webp"), - ) -> list[Path]: - """Remove watermarks from all images in a directory.""" - if not input_dir.exists(): - raise FileNotFoundError(f"Input directory not found: {input_dir}") - - output_dir.mkdir(parents=True, exist_ok=True) - cleaned_paths: list[Path] = [] - - # Lazy import keeps this module torch-optional; frees device cache per image. - from remove_ai_watermarks.noai.img2img_runner import try_empty_device_cache - - for ext in extensions: - for image_path in input_dir.glob(f"*{ext}"): - output_path = output_dir / image_path.name - try: - result_path = self.remove_watermark( - image_path=image_path, - output_path=output_path, - strength=strength, - num_inference_steps=num_inference_steps, - ) - cleaned_paths.append(result_path) - except Exception as e: - logger.error("Failed to process %s: %s", image_path, e) - try_empty_device_cache(self.device) - - return cleaned_paths - - -# ── Convenience function ───────────────────────────────────────────── - - -def remove_watermark( - image_path: Path, - output_path: Path | None = None, - strength: float | None = None, - model_id: str | None = None, - device: str | None = None, - hf_token: str | None = None, - region: tuple[int, int, int, int] | None = None, -) -> Path: - """Convenience function to remove watermark from an image. - - ``strength=None`` lets the profile pick its vendor-adaptive default - (0.10 OpenAI / 0.15 Google / 0.15 unknown, from the C2PA SynthID proxy on the - input; same ladder for the controlnet and sdxl pipelines -- the single source of - truth is ``watermark_profiles.py``). Pass a value to override. - - ``region=(x, y, w, h)`` restricts the regeneration to that box and preserves the - real photo elsewhere -- for AI-enhanced composites (see - ``WatermarkRemover.remove_watermark``). - """ - from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength - - remover = WatermarkRemover(model_id=model_id, device=device, hf_token=hf_token) - return remover.remove_watermark( - image_path=image_path, - output_path=output_path, - strength=strength, - vendor=vendor_for_strength(image_path), - region=region, - ) diff --git a/src/remove_ai_watermarks/video.py b/src/remove_ai_watermarks/video.py index 59fff6b..6f5bd35 100644 --- a/src/remove_ai_watermarks/video.py +++ b/src/remove_ai_watermarks/video.py @@ -318,7 +318,7 @@ def _select_stable_visible_mark( def _platform_from_video_metadata(markers: dict[str, str]) -> str | None: """Map supported C2PA-derived marker text to its generating platform.""" - from remove_ai_watermarks.noai.constants import C2PA_AI_VENDORS + from remove_ai_watermarks._internal.constants import C2PA_AI_VENDORS marker_text = "\n".join(markers.values()).casefold() if not marker_text: diff --git a/tests/test_cpu_offload.py b/tests/test_cpu_offload.py index f2c6ffd..8a50b49 100644 --- a/tests/test_cpu_offload.py +++ b/tests/test_cpu_offload.py @@ -12,7 +12,7 @@ from unittest.mock import Mock import pytest -from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover +from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover def _remover(device: str, cpu_offload: bool) -> WatermarkRemover: diff --git a/tests/test_identify.py b/tests/test_identify.py index c32494d..92ce88d 100644 --- a/tests/test_identify.py +++ b/tests/test_identify.py @@ -718,7 +718,7 @@ class TestSparkleDetectRemoveAlignment: class TestIdentifyImportIsLight: - """`import identify` must stay torch-free (lazy noai/__init__): the package + """`import identify` must stay torch-free (lazy _internal/__init__): the package is deployed on a 512 MB host where eagerly pulling torch/diffusers OOMs.""" def test_import_identify_does_not_pull_torch(self): diff --git a/tests/test_img2img_runner.py b/tests/test_img2img_runner.py index 1f964b7..7adeea2 100644 --- a/tests/test_img2img_runner.py +++ b/tests/test_img2img_runner.py @@ -13,8 +13,8 @@ from unittest.mock import Mock import pytest -from remove_ai_watermarks.noai import img2img_runner -from remove_ai_watermarks.noai.img2img_runner import ( +from remove_ai_watermarks._internal import img2img_runner +from remove_ai_watermarks._internal.img2img_runner import ( run_img2img, run_img2img_with_mps_fallback, ) diff --git a/tests/test_invisible_engine.py b/tests/test_invisible_engine.py index b44f8ce..077bcbe 100644 --- a/tests/test_invisible_engine.py +++ b/tests/test_invisible_engine.py @@ -215,7 +215,7 @@ class TestCannyControlImage: pytest.skip("diffusion extra (torch/diffusers) not installed") import numpy as np - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover rng = np.random.default_rng(0) img = Image.fromarray(rng.integers(0, 256, (64, 80, 3), dtype=np.uint8)) @@ -224,4 +224,10 @@ class TestCannyControlImage: arr = np.array(out) assert out.mode == "RGB" assert arr.shape == (64, 80, 3) - assert arr.max() <= 255 + import cv2 + + gray = cv2.cvtColor(np.asarray(img.convert("RGB")), cv2.COLOR_RGB2GRAY) + expected = cv2.Canny(gray, 100, 200) + assert np.array_equal(arr[:, :, 0], expected) + assert np.array_equal(arr[:, :, 1], expected) + assert np.array_equal(arr[:, :, 2], expected) diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 2f6fccc..4c863c6 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -110,8 +110,8 @@ class TestHasAiMetadata: def test_strip_c2pa_boxes_removes_uuid_box(self, tmp_path: Path): """ISOBMFF strip should drop the C2PA uuid box and keep everything else.""" + from remove_ai_watermarks._internal.isobmff import strip_c2pa_boxes from remove_ai_watermarks.metadata import C2PA_UUID - from remove_ai_watermarks.noai.isobmff import strip_c2pa_boxes ftyp = b"\x00\x00\x00\x18ftypavif\x00\x00\x00\x00avifmif1" # uuid box: size(4) + 'uuid' + 16-byte UUID + minimal payload (8 bytes -> total 32) @@ -123,7 +123,7 @@ class TestHasAiMetadata: def test_strip_c2pa_boxes_passthrough_for_non_isobmff(self): """Non-ISOBMFF input must be returned unchanged.""" - from remove_ai_watermarks.noai.isobmff import strip_c2pa_boxes + from remove_ai_watermarks._internal.isobmff import strip_c2pa_boxes data = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR" + b"\x00" * 100 cleaned, stripped = strip_c2pa_boxes(data) @@ -1375,7 +1375,7 @@ class TestSoftBinding: """C2PA soft-binding alg identifier -> forensic-watermark vendor name.""" def test_vendors_in_recognizes_known_algs(self): - from remove_ai_watermarks.noai.c2pa import soft_binding_vendors_in + from remove_ai_watermarks._internal.c2pa import soft_binding_vendors_in assert soft_binding_vendors_in(b"...alg...com.adobe.trustmark.P...") == ["Adobe TrustMark"] assert soft_binding_vendors_in(b"com.digimarc.validate.1") == ["Digimarc"] @@ -1385,7 +1385,7 @@ class TestSoftBinding: assert soft_binding_vendors_in(b"io.iscc.v0") == ["ISCC (content code)"] def test_vendors_in_empty_when_absent(self): - from remove_ai_watermarks.noai.c2pa import soft_binding_vendors_in + from remove_ai_watermarks._internal.c2pa import soft_binding_vendors_in assert soft_binding_vendors_in(b"no soft binding here") == [] @@ -1472,8 +1472,8 @@ class TestLateProvenanceBox: return p def test_scan_c2pa_region_finds_late_box(self, tmp_path: Path): + from remove_ai_watermarks._internal.isobmff import scan_c2pa_region from remove_ai_watermarks.metadata import C2PA_UUID - from remove_ai_watermarks.noai.isobmff import scan_c2pa_region region = scan_c2pa_region(self._mp4_late_c2pa(tmp_path)) assert C2PA_UUID in region @@ -1495,7 +1495,7 @@ class TestLateProvenanceBox: assert has_ai_metadata(self._mp4_late_c2pa(tmp_path)) is True def test_scan_c2pa_region_non_isobmff_is_empty(self, tmp_path: Path): - from remove_ai_watermarks.noai.isobmff import scan_c2pa_region + from remove_ai_watermarks._internal.isobmff import scan_c2pa_region p = tmp_path / "not.bin" p.write_bytes(b"\x89PNG\r\n\x1a\n not an isobmff file") @@ -1505,7 +1505,7 @@ class TestLateProvenanceBox: """A 64-bit largesize (size32 == 1) uuid box must be walked and collected.""" import struct - from remove_ai_watermarks.noai.isobmff import scan_c2pa_region + from remove_ai_watermarks._internal.isobmff import scan_c2pa_region payload = b"LARGESIZE-C2PA-MANIFEST" total = 16 + len(payload) # 4 (size32=1) + 4 (type) + 8 (largesize) + payload @@ -1516,7 +1516,7 @@ class TestLateProvenanceBox: def test_scan_c2pa_region_caps_at_max_total(self, tmp_path: Path): """The collected payload is bounded by ``max_total`` (never unbounded).""" - from remove_ai_watermarks.noai.isobmff import scan_c2pa_region + from remove_ai_watermarks._internal.isobmff import scan_c2pa_region p = tmp_path / "big.mp4" p.write_bytes(_MP4_FTYP + _box(b"uuid", b"A" * 5000)) @@ -1550,7 +1550,7 @@ class TestMetaBoxXmpBlanking: in place (same length -> iloc offsets and image data stay intact).""" def test_blanks_ai_packet_only(self): - from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets + from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets before, after = b"IMG_BEFORE" * 4, b"IMG_AFTER" * 4 data = before + _AI_XMP + after + _PLAIN_XMP @@ -1563,13 +1563,13 @@ class TestMetaBoxXmpBlanking: assert b"dc:rights" in out # plain XMP left alone def test_no_packet_is_noop(self): - from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets + from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets data = b"just some mdat bytes, no xmp here" assert blank_ai_xmp_packets(data) == (data, 0) def test_plain_xmp_untouched(self): - from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets + from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets out, n = blank_ai_xmp_packets(_PLAIN_XMP) assert n == 0 @@ -1600,7 +1600,7 @@ class TestIsobmffMetadataRemoval: def test_strips_ai_xmp_uuid_box(self): # A uuid box carrying a TC260 AIGC label is dropped by content match, # regardless of the (non-C2PA) XMP UUID's byte order. - from remove_ai_watermarks.noai.isobmff import strip_c2pa_boxes + from remove_ai_watermarks._internal.isobmff import strip_c2pa_boxes xmp_uuid = bytes(range(16)) # arbitrary, not the C2PA UUID payload = b'{"Label":"1"}' @@ -1611,7 +1611,7 @@ class TestIsobmffMetadataRemoval: def test_keeps_plain_non_ai_xmp(self): # A uuid box with ordinary (non-AI) XMP must be preserved. - from remove_ai_watermarks.noai.isobmff import strip_c2pa_boxes + from remove_ai_watermarks._internal.isobmff import strip_c2pa_boxes xmp_uuid = bytes(range(16)) payload = b"(c) me" diff --git a/tests/test_noai.py b/tests/test_metadata_internals.py similarity index 93% rename from tests/test_noai.py rename to tests/test_metadata_internals.py index a99a93c..c1900fa 100644 --- a/tests/test_noai.py +++ b/tests/test_metadata_internals.py @@ -1,5 +1,5 @@ -"""Tests for vendored noai submodules: constants, extractor, c2pa, plus the -consolidated metadata strip (formerly noai.cleaner).""" +"""Tests for metadata compatibility submodules: constants, extractor, C2PA, plus the +consolidated metadata strip (formerly legacy metadata helper).""" from __future__ import annotations @@ -8,10 +8,7 @@ from pathlib import Path import pytest -from remove_ai_watermarks.metadata import ( - remove_ai_metadata as noai_remove_ai_metadata, -) -from remove_ai_watermarks.noai.c2pa import ( +from remove_ai_watermarks._internal.c2pa import ( _parse_c2pa_chunk, cbor_text_after, extract_c2pa_chunk, @@ -20,24 +17,27 @@ from remove_ai_watermarks.noai.c2pa import ( inject_c2pa_chunk, synthid_verdict, ) -from remove_ai_watermarks.noai.constants import ( +from remove_ai_watermarks._internal.constants import ( AI_KEYWORDS, AI_METADATA_KEYS, C2PA_CHUNK_TYPE, PNG_SIGNATURE, SUPPORTED_FORMATS, ) -from remove_ai_watermarks.noai.extractor import ( +from remove_ai_watermarks._internal.extractor import ( extract_ai_metadata, extract_metadata, get_ai_metadata_summary, has_ai_metadata, ) -from remove_ai_watermarks.noai.isobmff import ( +from remove_ai_watermarks._internal.isobmff import ( blank_ai_exif_tokens, is_isobmff, strip_c2pa_boxes, ) +from remove_ai_watermarks.metadata import ( + remove_ai_metadata as remove_metadata, +) # ── Constants ─────────────────────────────────────────────────────── @@ -77,7 +77,7 @@ class TestConstants: class TestExtractor: - """Tests for noai.extractor functions.""" + """Tests for internal metadata extraction helpers.""" def test_extract_metadata_returns_dict(self, tmp_clean_png): meta = extract_metadata(tmp_clean_png) @@ -115,11 +115,11 @@ class TestExtractor: class TestCleaner: """Metadata stripping via the single, consolidated ``metadata.remove_ai_metadata`` - (the legacy ``noai.cleaner`` duplicate was retired).""" + (the legacy ``legacy metadata helper`` duplicate was retired).""" def test_remove_ai_metadata(self, tmp_png_with_ai_metadata, tmp_path): output = tmp_path / "cleaned.png" - noai_remove_ai_metadata(tmp_png_with_ai_metadata, output) + remove_metadata(tmp_png_with_ai_metadata, output) assert output.exists() # Verify AI metadata removed meta = extract_ai_metadata(output) @@ -201,7 +201,7 @@ class TestC2PARealSamples: def test_extract_info_uses_reader_store(self): """The c2pa-python reader path: structured (not heuristic) extraction.""" - from remove_ai_watermarks.noai import c2pa + from remove_ai_watermarks._internal import c2pa assert c2pa.reader_available() info = extract_c2pa_info(SAMPLES_DIR / "chatgpt-1.png") @@ -213,7 +213,7 @@ class TestC2PARealSamples: def test_fallback_to_png_parser_when_reader_unavailable(self, monkeypatch): """With the reader disabled, the hand-rolled PNG parser still works.""" - from remove_ai_watermarks.noai import c2pa + from remove_ai_watermarks._internal import c2pa monkeypatch.setattr(c2pa, "_C2PA_READER_AVAILABLE", False) info = extract_c2pa_info(SAMPLES_DIR / "chatgpt-1.png") @@ -326,7 +326,7 @@ class TestC2PADigitalSourceType: (AI-enhanced) and a bare procedural ``algorithmicMedia`` token must classify as AI-enhanced. Before the reorder the bare-token elif fired first and returned non-AI, dropping the composite AI signal (a false negative).""" - from remove_ai_watermarks.noai.c2pa import _populate_registry_fields + from remove_ai_watermarks._internal.c2pa import _populate_registry_fields info: dict = {} _populate_registry_fields(b"x compositeWithTrainedAlgorithmicMedia x algorithmicMedia x", info) @@ -453,7 +453,7 @@ class TestISOBMFF: assert out == FTYP + b"\x00\x00\x00\x0cmdat" + b"pixels!!" def test_streaming_malformed_walk_copies_input_unchanged(self, tmp_path: Path): - from remove_ai_watermarks.noai.isobmff import strip_isobmff_media_file + from remove_ai_watermarks._internal.isobmff import strip_isobmff_media_file source = tmp_path / "malformed.mp4" output = tmp_path / "clean.mp4" @@ -471,7 +471,7 @@ class TestISOBMFF: monkeypatch: pytest.MonkeyPatch, ): from remove_ai_watermarks import metadata - from remove_ai_watermarks.noai import isobmff + from remove_ai_watermarks._internal import isobmff source = tmp_path / "source.mp4" output = tmp_path / "clean.mp4" @@ -495,7 +495,7 @@ class TestIterTopLevelBoxes: """The box walker's three size encodings and its underflow/overflow guards.""" def test_64bit_largesize(self): - from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes + from remove_ai_watermarks._internal.isobmff import _iter_top_level_boxes # size32 == 1 -> a 64-bit largesize follows the type; total box length = 24. box = struct.pack(">I", 1) + b"uuid" + struct.pack(">Q", 24) + b"payload!" @@ -505,7 +505,7 @@ class TestIterTopLevelBoxes: assert (start, end, btype, payload_off) == (0, 24, b"uuid", 16) def test_size0_runs_to_eof(self): - from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes + from remove_ai_watermarks._internal.isobmff import _iter_top_level_boxes box = struct.pack(">I", 0) + b"mdat" + b"tail-to-eof" boxes = list(_iter_top_level_boxes(box)) @@ -514,13 +514,13 @@ class TestIterTopLevelBoxes: assert (start, end, btype, payload_off) == (0, len(box), b"mdat", 8) def test_underflow_size_stops_safely(self): - from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes + from remove_ai_watermarks._internal.isobmff import _iter_top_level_boxes # size (4) < the 8-byte header -> the guard returns without yielding a box. assert list(_iter_top_level_boxes(struct.pack(">I", 4) + b"ftyp" + b"more")) == [] def test_overflow_size_stops_safely(self): - from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes + from remove_ai_watermarks._internal.isobmff import _iter_top_level_boxes # size claims 999 but the buffer is far shorter -> guard returns, no partial box. assert list(_iter_top_level_boxes(struct.pack(">I", 999) + b"uuid" + b"x")) == [] @@ -533,7 +533,7 @@ class TestBlankAiXmpPackets: AIMARK = b"trainedAlgorithmicMedia" def test_ai_packet_blanked_same_length(self): - from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets + from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets packet = b'' + self.AIMARK + b'' data = b"boxhdr" + packet + b"tail" @@ -545,7 +545,7 @@ class TestBlankAiXmpPackets: assert b"tail" in out def test_clean_packet_left_intact(self): - from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets + from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets packet = b'plain copyright' out, n = blank_ai_xmp_packets(packet) @@ -553,7 +553,7 @@ class TestBlankAiXmpPackets: assert out == packet def test_missing_end_delimiter_not_blanked(self): - from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets + from remove_ai_watermarks._internal.isobmff import blank_ai_xmp_packets # No -> the packet regex cannot match, so it is left unchanged. data = b'' + self.AIMARK + b"" @@ -568,7 +568,7 @@ class TestC2paBufferScans: as vendors are added.""" def test_soft_binding_vendors_in(self): - from remove_ai_watermarks.noai.c2pa import C2PA_SOFT_BINDINGS, soft_binding_vendors_in + from remove_ai_watermarks._internal.c2pa import C2PA_SOFT_BINDINGS, soft_binding_vendors_in sig, name = next(iter(C2PA_SOFT_BINDINGS.items())) assert name in soft_binding_vendors_in(b"...manifest..." + sig + b"...tail...") @@ -576,7 +576,7 @@ class TestC2paBufferScans: assert soft_binding_vendors_in(b"no soft-binding assertion here") == [] def test_synthid_vendors_in_requires_synthid_issuer(self): - from remove_ai_watermarks.noai.c2pa import C2PA_ISSUERS, SYNTHID_C2PA_ISSUERS, synthid_vendors_in + from remove_ai_watermarks._internal.c2pa import C2PA_ISSUERS, SYNTHID_C2PA_ISSUERS, synthid_vendors_in syn_sig = next(s for s in C2PA_ISSUERS if s in SYNTHID_C2PA_ISSUERS) non_sig = next(s for s in C2PA_ISSUERS if s not in SYNTHID_C2PA_ISSUERS) @@ -585,7 +585,7 @@ class TestC2paBufferScans: assert C2PA_ISSUERS[non_sig] not in synthid_vendors_in(b"x" + non_sig + b"x") def test_synthid_verdict_format(self): - from remove_ai_watermarks.noai.c2pa import synthid_verdict + from remove_ai_watermarks._internal.c2pa import synthid_verdict assert synthid_verdict("Google LLC") == "likely present (Google LLC embeds SynthID with C2PA)" diff --git a/tests/test_platform.py b/tests/test_platform.py index 2464649..694f03b 100644 --- a/tests/test_platform.py +++ b/tests/test_platform.py @@ -13,9 +13,9 @@ import numpy as np import pytest from PIL import Image -from remove_ai_watermarks.noai.progress import is_mps_error -from remove_ai_watermarks.noai.utils import get_image_format, is_supported_format -from remove_ai_watermarks.noai.watermark_profiles import ( +from remove_ai_watermarks._internal.progress import is_mps_error +from remove_ai_watermarks._internal.utils import get_image_format, is_supported_format +from remove_ai_watermarks._internal.watermark_profiles import ( DEFAULT_STRENGTH, GEMINI_STRENGTH, OPENAI_STRENGTH, @@ -24,7 +24,7 @@ from remove_ai_watermarks.noai.watermark_profiles import ( resolve_strength, strength_default_help, ) -from remove_ai_watermarks.noai.watermark_remover import get_device, is_watermark_removal_available +from remove_ai_watermarks._internal.watermark_remover import get_device, is_watermark_removal_available # ── Device detection ──────────────────────────────────────────────── @@ -42,7 +42,7 @@ class TestDeviceDetection: # Just verify it doesn't crash and returns a valid string assert isinstance(device, str) - @patch("remove_ai_watermarks.noai.watermark_remover._HAS_TORCH", False) + @patch("remove_ai_watermarks._internal.watermark_remover._HAS_TORCH", False) def test_no_torch_returns_cpu(self): assert get_device() == "cpu" @@ -55,7 +55,7 @@ class TestDeviceDetection: fake_torch = MagicMock() fake_torch.cuda.is_available.return_value = False fake_torch.xpu.is_available.return_value = True - with patch("remove_ai_watermarks.noai.watermark_remover.torch", fake_torch): + with patch("remove_ai_watermarks._internal.watermark_remover.torch", fake_torch): assert get_device() == "xpu" fake_torch.tensor.assert_called_with([1.0], device="xpu") @@ -65,7 +65,7 @@ class TestDeviceDetection: pytest.skip("torch/diffusers not installed") import torch - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover remover = WatermarkRemover(device="xpu") assert remover.device == "xpu" @@ -74,7 +74,7 @@ class TestDeviceDetection: def test_seed_generator_falls_back_to_cpu_when_device_rng_unsupported(self): """A device with no RNG backend (e.g. some torch-xpu builds) falls back to a CPU generator instead of raising when --seed is used.""" - from remove_ai_watermarks.noai import watermark_remover as wr + from remove_ai_watermarks._internal import watermark_remover as wr def fake_generator(device="cpu"): if device == "xpu": @@ -138,7 +138,7 @@ class TestFp16WeightVariant: def _remover(self, dtype: object): if not is_watermark_removal_available(): pytest.skip("torch/diffusers not installed") - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover # device="cpu" alone would force fp32; the explicit torch_dtype override lets us # exercise the fp16 path with no GPU (construction loads no weights). @@ -192,12 +192,12 @@ class TestNoReembeddedWatermark: def _remover(self, profile: str): if not is_watermark_removal_available(): pytest.skip("torch/diffusers not installed") - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover return WatermarkRemover(device="cpu", pipeline=profile) def _capture(self, monkeypatch, remover): - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover calls: list[tuple[str, dict]] = [] @@ -242,7 +242,7 @@ class TestQwenKwargs: """ def test_uses_true_cfg_not_guidance_scale(self): - from remove_ai_watermarks.noai.watermark_remover import _build_qwen_kwargs + from remove_ai_watermarks._internal.watermark_remover import _build_qwen_kwargs gen = object() img = _StubImage(2816, 1536) @@ -254,14 +254,13 @@ class TestQwenKwargs: assert kwargs["strength"] == 0.3 assert kwargs["image"] is img assert kwargs["generator"] is gen - # Faithful-regeneration prompt + an explicit negative prompt. - assert kwargs["prompt"] - assert kwargs["negative_prompt"] + assert kwargs["prompt"] == "high quality, sharp, detailed, faithful to the original" + assert kwargs["negative_prompt"] == "blurry, lowres, distorted text, garbled text, artifacts" def test_passes_explicit_aspect_preserving_size(self): # Without height/width the pipeline defaults to 1024x1024 and squishes non-square # input (the abba mixed-seam regression). Both already multiples of 16 -> unchanged. - from remove_ai_watermarks.noai.watermark_remover import _build_qwen_kwargs + from remove_ai_watermarks._internal.watermark_remover import _build_qwen_kwargs kwargs = _build_qwen_kwargs( _StubImage(2816, 1536), strength=0.25, num_inference_steps=40, true_cfg_scale=4.0, generator=None @@ -270,14 +269,14 @@ class TestQwenKwargs: assert kwargs["height"] == 1536 def test_qwen_target_size_floors_to_multiple_of_16(self): - from remove_ai_watermarks.noai.watermark_remover import _qwen_target_size + from remove_ai_watermarks._internal.watermark_remover import _qwen_target_size assert _qwen_target_size(2816, 1536) == (2816, 1536) # already /16 assert _qwen_target_size(1122, 1402) == (1120, 1392) # floored assert _qwen_target_size(10, 10) == (16, 16) # min clamp, never 0 def test_qwen_model_id_is_qwen_image(self): - from remove_ai_watermarks.noai.watermark_profiles import QWEN_MODEL_ID + from remove_ai_watermarks._internal.watermark_profiles import QWEN_MODEL_ID assert QWEN_MODEL_ID == "Qwen/Qwen-Image" @@ -303,7 +302,7 @@ class TestResolveStrength: # Qwen's certified Gemini floor (0.25) is HIGHER than controlnet's (0.15); OpenAI # matches (0.10). Unknown vendor on qwen tracks the higher Gemini value. This retires # the old manual "pass --strength 0.25 for Gemini on qwen" workaround. - from remove_ai_watermarks.noai.watermark_profiles import QWEN_GEMINI_STRENGTH, QWEN_OPENAI_STRENGTH + from remove_ai_watermarks._internal.watermark_profiles import QWEN_GEMINI_STRENGTH, QWEN_OPENAI_STRENGTH assert QWEN_GEMINI_STRENGTH == 0.25 assert QWEN_OPENAI_STRENGTH == 0.10 @@ -354,32 +353,32 @@ class TestVendorForStrength: return patch("remove_ai_watermarks.metadata.synthid_source", return_value=value) def test_openai(self): - from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength + from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength with self._patch("OpenAI"): assert vendor_for_strength(Path("x.png")) == "openai" def test_google(self): - from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength + from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength with self._patch("Google"): assert vendor_for_strength(Path("x.png")) == "google" def test_both_issuers_google_wins(self): # The more-robust watermark wins -> safer (higher) strength. - from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength + from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength with self._patch("OpenAI, Google"): assert vendor_for_strength(Path("x.png")) == "google" def test_none_when_no_synthid_source(self): - from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength + from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength with self._patch(None): assert vendor_for_strength(Path("x.png")) is None def test_unreadable_metadata_is_none(self): - from remove_ai_watermarks.noai.watermark_profiles import vendor_for_strength + from remove_ai_watermarks._internal.watermark_profiles import vendor_for_strength with patch("remove_ai_watermarks.metadata.synthid_source", side_effect=OSError): assert vendor_for_strength(Path("x.png")) is None @@ -478,19 +477,19 @@ class TestFp16VaeFix: DEFAULT = "stabilityai/stable-diffusion-xl-base-1.0" def test_default_sdxl_on_fp16_needs_fix(self): - from remove_ai_watermarks.noai.watermark_remover import _needs_fp16_vae_fix + from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix assert _needs_fp16_vae_fix(self.DEFAULT, self.DEFAULT, is_fp16=True) is True def test_fp32_does_not_need_fix(self): """cpu/mps run fp32, where the stock SDXL VAE is fine.""" - from remove_ai_watermarks.noai.watermark_remover import _needs_fp16_vae_fix + from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix assert _needs_fp16_vae_fix(self.DEFAULT, self.DEFAULT, is_fp16=False) is False def test_non_default_model_keeps_own_vae(self): """A custom (non-SDXL) checkpoint must not get the SDXL-specific VAE.""" - from remove_ai_watermarks.noai.watermark_remover import _needs_fp16_vae_fix + from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix assert _needs_fp16_vae_fix("runwayml/stable-diffusion-v1-5", self.DEFAULT, is_fp16=True) is False @@ -500,13 +499,13 @@ class TestDegenerateOutputGuard: ``remove_watermark`` can retry in fp32. Pure image statistics, no model needed.""" def test_all_black_is_degenerate(self): - from remove_ai_watermarks.noai.watermark_remover import _is_degenerate_image + from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image black = Image.fromarray(np.zeros((64, 64, 3), np.uint8)) assert _is_degenerate_image(black) is True def test_normal_image_is_not_degenerate(self): - from remove_ai_watermarks.noai.watermark_remover import _is_degenerate_image + from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image rng = np.random.default_rng(0) normal = Image.fromarray(rng.integers(0, 256, (64, 64, 3), dtype=np.uint8)) @@ -514,7 +513,7 @@ class TestDegenerateOutputGuard: def test_dark_but_textured_image_is_not_degenerate(self): """A legitimately dark photo with real detail must NOT be flagged (variance guard).""" - from remove_ai_watermarks.noai.watermark_remover import _is_degenerate_image + from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image rng = np.random.default_rng(1) dark = Image.fromarray(rng.integers(0, 40, (64, 64, 3), dtype=np.uint8)) diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index cb3d520..2cf3a18 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/tests/test_qwen_zimage_pipeline.py @@ -13,7 +13,7 @@ from PIL import Image def _mock_watermark_runtime_deps(monkeypatch): """Bypass optional GPU imports while testing Qwen Z-Image routing.""" - from remove_ai_watermarks.noai import watermark_remover + from remove_ai_watermarks._internal import watermark_remover fake_torch = MagicMock() fake_torch.float16 = object() @@ -23,26 +23,24 @@ def _mock_watermark_runtime_deps(monkeypatch): monkeypatch.setattr(watermark_remover, "is_watermark_removal_available", lambda: True) -def test_resolution_adaptive_denoise_matches_reference_formula(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import resolution_adaptive_denoise +def test_resolution_adaptive_denoise_preserves_calibrated_values(): + from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise - # The reference node maps 0.30 MP to the lower bound and 3.70 MP to the - # upper bound. Level 6 adds one fifth of the configured upward spread. assert resolution_adaptive_denoise(600, 500, adaptive_level=6) == pytest.approx(0.084) assert resolution_adaptive_denoise(2000, 1850, adaptive_level=6) == pytest.approx(0.154) -def test_largest_face_denoise_matches_reference_formula(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import largest_face_denoise +def test_largest_face_denoise_preserves_calibrated_values(): + from remove_ai_watermarks._internal.qwen_zimage_pipeline import largest_face_denoise image_size = (1000, 1000) - assert largest_face_denoise([(0, 0, 300, 100)], image_size) == 0.10 - assert largest_face_denoise([(0, 0, 150, 100)], image_size) == 0.05 - assert largest_face_denoise([(0, 0, 900, 900)], image_size) == 0.28 + assert largest_face_denoise([(0, 0, 300, 100)], image_size) == pytest.approx(0.10) + assert largest_face_denoise([(0, 0, 150, 100)], image_size) == pytest.approx(0.05) + assert largest_face_denoise([(0, 0, 900, 900)], image_size) == pytest.approx(0.28) def test_global_kwargs_use_lightning_and_diffsynth_controlnet_shape(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import build_global_kwargs + from remove_ai_watermarks._internal.qwen_zimage_pipeline import build_global_kwargs image = Image.new("RGB", (1122, 1402)) kwargs = build_global_kwargs(image, strength=0.11, seed=7, controlnet_input="CONTROL") @@ -56,10 +54,12 @@ def test_global_kwargs_use_lightning_and_diffsynth_controlnet_shape(): assert kwargs["width"] == 1120 assert kwargs["height"] == 1392 assert kwargs["exponential_shift_mu"] == pytest.approx(math.log(3.0)) + assert kwargs["prompt"] == "ultra clear and smoothe skin, spotless skin" + assert kwargs["negative_prompt"] == "moles, freckes, high detail skin" -def test_face_kwargs_use_zimage_reference_settings(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import build_face_kwargs +def test_face_kwargs_use_project_zimage_settings(): + from remove_ai_watermarks._internal.qwen_zimage_pipeline import build_face_kwargs crop = Image.new("RGB", (713, 941)) kwargs = build_face_kwargs(crop, strength=0.17, seed=9) @@ -71,10 +71,12 @@ def test_face_kwargs_use_zimage_reference_settings(): assert kwargs["seed"] == 9 assert kwargs["width"] == 704 assert kwargs["height"] == 928 + assert kwargs["prompt"] == "" + assert kwargs["negative_prompt"] == "blurry, ugly, bad quality," -def test_canny_control_image_matches_reference_thresholds(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import build_canny_control_image +def test_canny_control_image_is_three_channel_and_detects_an_edge(): + from remove_ai_watermarks._internal.qwen_zimage_pipeline import build_canny_control_image source = np.zeros((64, 80, 3), dtype=np.uint8) source[:, 40:] = 255 @@ -83,11 +85,21 @@ def test_canny_control_image_matches_reference_thresholds(): assert result.shape == (64, 80, 3) assert np.array_equal(result[:, :, 0], result[:, :, 1]) assert np.array_equal(result[:, :, 1], result[:, :, 2]) - assert result.max() == 255 + import cv2 + + expected = cv2.Canny(cv2.cvtColor(source, cv2.COLOR_RGB2GRAY), 13, 64) + assert np.array_equal(result[:, :, 0], expected) + + +def test_face_crop_geometry_preserves_calibrated_values(): + from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline, _expanded_box + + assert _expanded_box((100, 100, 200, 200), (500, 500)) == (25, 25, 275, 275) + assert QwenZImagePipeline._detail_size((500, 400), (100, 80)) == (1024, 816) def test_yunet_download_targets_verified_lfs_artifact(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import ( + from remove_ai_watermarks._internal.qwen_zimage_pipeline import ( YUNET_MODEL_SHA256, YUNET_MODEL_URL, YUNET_SCORE_THRESHOLD, @@ -95,14 +107,11 @@ def test_yunet_download_targets_verified_lfs_artifact(): assert YUNET_MODEL_URL.startswith("https://media.githubusercontent.com/media/opencv/opencv_zoo/") assert YUNET_MODEL_SHA256 == "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4" - # YuNet scores are not calibrated like the upstream YOLO detector's scores. - # A 0.2 YuNet threshold admitted background and decorative false positives, - # multiplying the serial Z-Image face-stage cost on crowded scenes. assert pytest.approx(0.5) == YUNET_SCORE_THRESHOLD def test_resident_face_models_disable_vram_offload(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import ( + from remove_ai_watermarks._internal.qwen_zimage_pipeline import ( QwenZImagePipeline, _pin_vram_managed_models, resolve_face_model_residency, @@ -154,7 +163,7 @@ def test_resident_face_models_disable_vram_offload(): def test_static_prompt_cache_reuses_embeddings_without_caching_image_edits(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import _cache_static_prompt_embeddings + from remove_ai_watermarks._internal.qwen_zimage_pipeline import _cache_static_prompt_embeddings class PromptUnit: output_params = ("prompt_embeds",) @@ -186,7 +195,7 @@ def test_static_prompt_cache_reuses_embeddings_without_caching_image_edits(): def test_sam_pixels_match_model_dtype_without_casting_boxes(): import torch - from remove_ai_watermarks.noai.qwen_zimage_pipeline import _prepare_sam_inputs + from remove_ai_watermarks._internal.qwen_zimage_pipeline import _prepare_sam_inputs class Inputs(dict[str, torch.Tensor]): def to(self, device: str): @@ -206,7 +215,7 @@ def test_sam_pixels_match_model_dtype_without_casting_boxes(): def test_sam_prompts_match_impact_center_and_clip_masks_to_boxes(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import ( + from remove_ai_watermarks._internal.qwen_zimage_pipeline import ( _clip_sam_masks_to_boxes, _sam_point_prompts, ) @@ -226,7 +235,7 @@ def test_sam_prompts_match_impact_center_and_clip_masks_to_boxes(): def test_sam_proposal_selection_matches_impact_sub_threshold(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import _select_sam_masks + from remove_ai_watermarks._internal.qwen_zimage_pipeline import _select_sam_masks masks = np.zeros((2, 3, 8, 8), dtype=np.float32) masks[0, 0, 1:3, 1:3] = 1.0 @@ -256,7 +265,7 @@ def test_sam_proposal_selection_matches_impact_sub_threshold(): def test_sam_bfloat16_outputs_convert_to_numpy_float32(): import torch - from remove_ai_watermarks.noai.qwen_zimage_pipeline import _sam_outputs_to_numpy + from remove_ai_watermarks._internal.qwen_zimage_pipeline import _sam_outputs_to_numpy masks = torch.ones((1, 2, 3, 4, 4), dtype=torch.bfloat16) scores = torch.tensor([[[0.95, 0.75, 0.50], [0.99, 0.80, 0.60]]], dtype=torch.bfloat16) @@ -269,7 +278,7 @@ def test_sam_bfloat16_outputs_convert_to_numpy_float32(): def test_face_composite_preserves_every_pixel_outside_mask(): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import composite_face + from remove_ai_watermarks._internal.qwen_zimage_pipeline import composite_face base = np.full((32, 32, 3), 10, dtype=np.uint8) detail = np.full((32, 32, 3), 240, dtype=np.uint8) @@ -284,7 +293,7 @@ def test_face_composite_preserves_every_pixel_outside_mask(): def test_profile_defaults_to_four_global_steps(): - from remove_ai_watermarks.noai.watermark_profiles import ( + from remove_ai_watermarks._internal.watermark_profiles import ( normalize_profile, resolve_seed, resolve_steps, @@ -305,7 +314,7 @@ def test_cli_exposes_qwen_zimage_profile(): assert "qwen-zimage" in _PIPELINE_CHOICES -def test_cli_qwen_zimage_keeps_upstream_postprocess_default(tmp_image_path, monkeypatch): +def test_cli_qwen_zimage_keeps_profile_postprocess_default(tmp_image_path, monkeypatch): from remove_ai_watermarks import cli mock_engine = MagicMock() @@ -338,7 +347,7 @@ def test_cli_qwen_zimage_keeps_upstream_postprocess_default(tmp_image_path, monk def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch): - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover _mock_watermark_runtime_deps(monkeypatch) source = tmp_path / "source.png" @@ -364,7 +373,7 @@ def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch): def test_watermark_remover_dispatches_qwen_tiling_to_full_pipeline(tmp_path, monkeypatch): - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover _mock_watermark_runtime_deps(monkeypatch) source = tmp_path / "source.png" @@ -395,11 +404,11 @@ def test_watermark_remover_dispatches_qwen_tiling_to_full_pipeline(tmp_path, mon def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatch): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import ( + from remove_ai_watermarks._internal.qwen_zimage_pipeline import ( QwenZImagePipeline, resolution_adaptive_denoise, ) - from remove_ai_watermarks.noai.tiling import plan_tiles + from remove_ai_watermarks._internal.tiling import plan_tiles image = Image.new("RGB", (1500, 1500), (20, 30, 40)) runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16") @@ -415,7 +424,7 @@ def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatc monkeypatch.setattr(runtime, "_run_global", fake_global) monkeypatch.setattr(runtime, "_run_faces", face_stage) monkeypatch.setattr( - "remove_ai_watermarks.noai.qwen_zimage_pipeline.detect_faces", + "remove_ai_watermarks._internal.qwen_zimage_pipeline.detect_faces", lambda _image: [(100, 100, 300, 300)], ) monkeypatch.setattr(runtime, "_sam_masks", lambda _image, _boxes: [np.ones((1500, 1500), dtype=np.uint8)]) @@ -435,9 +444,6 @@ def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatc assert all(strength == pytest.approx(resolution_adaptive_denoise(1500, 1500)) for _, strength, _ in global_calls) assert all(seed == 0 for _, _, seed in global_calls) face_stage.assert_called_once() - # Keep this literal independent from the runtime helper: the port deliberately - # uses half the upstream face denoise because it lacks the reference latent - # noise-mask feather and uses a different sampler/runtime. assert face_stage.call_args.kwargs["strength"] == pytest.approx(0.0296296296) assert face_stage.call_args.args[0] is image assert face_stage.call_args.args[1].size == image.size @@ -445,7 +451,7 @@ def test_qwen_tiling_runs_global_tiles_then_one_full_frame_face_stage(monkeypatc def test_global_only_preload_skips_face_models(monkeypatch): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import QwenZImagePipeline + from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16") qwen = MagicMock() @@ -456,7 +462,7 @@ def test_global_only_preload_skips_face_models(monkeypatch): monkeypatch.setattr(runtime, "_load_zimage", zimage) monkeypatch.setattr(runtime, "_load_sam", sam) monkeypatch.setattr( - "remove_ai_watermarks.noai.qwen_zimage_pipeline._yunet_model_path", + "remove_ai_watermarks._internal.qwen_zimage_pipeline._yunet_model_path", yunet, ) @@ -469,7 +475,7 @@ def test_global_only_preload_skips_face_models(monkeypatch): def test_full_preload_still_loads_face_models(monkeypatch): - from remove_ai_watermarks.noai.qwen_zimage_pipeline import QwenZImagePipeline + from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline runtime = QwenZImagePipeline(device="cuda", torch_dtype="bf16") qwen = MagicMock() @@ -480,7 +486,7 @@ def test_full_preload_still_loads_face_models(monkeypatch): monkeypatch.setattr(runtime, "_load_zimage", zimage) monkeypatch.setattr(runtime, "_load_sam", sam) monkeypatch.setattr( - "remove_ai_watermarks.noai.qwen_zimage_pipeline._yunet_model_path", + "remove_ai_watermarks._internal.qwen_zimage_pipeline._yunet_model_path", yunet, ) @@ -493,7 +499,7 @@ def test_full_preload_still_loads_face_models(monkeypatch): def test_watermark_remover_forwards_global_only_preload(monkeypatch): - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover runtime = MagicMock() remover = WatermarkRemover.__new__(WatermarkRemover) @@ -506,7 +512,7 @@ def test_watermark_remover_forwards_global_only_preload(monkeypatch): def test_qwen_zimage_rejects_runtime_knobs_that_change_fixed_graph(tmp_path, monkeypatch): - from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover _mock_watermark_runtime_deps(monkeypatch) with pytest.raises(ValueError, match="fixed Qwen-Image-2512"): @@ -517,7 +523,7 @@ def test_qwen_zimage_rejects_runtime_knobs_that_change_fixed_graph(tmp_path, mon remover = WatermarkRemover(device="cpu", pipeline="qwen-zimage") with pytest.raises(ValueError, match=r"CFG 1\.0"): remover.remove_watermark(source, guidance_scale=2.0) - with pytest.raises(ValueError, match="4-step Lightning"): + with pytest.raises(ValueError, match="requires 4 steps"): remover.remove_watermark(source, num_inference_steps=8) diff --git a/tests/test_security_clamp.py b/tests/test_security_clamp.py index 9e3c194..ffdaae5 100644 --- a/tests/test_security_clamp.py +++ b/tests/test_security_clamp.py @@ -20,7 +20,7 @@ import struct import tracemalloc from remove_ai_watermarks import metadata -from remove_ai_watermarks.noai import c2pa, isobmff +from remove_ai_watermarks._internal import c2pa, isobmff PNG_SIG = b"\x89PNG\r\n\x1a\n" _HUGE = 0x7FFFFFFF # ~2 GiB declared length on a tiny file diff --git a/tests/test_tiling.py b/tests/test_tiling.py index ab6df75..fd2cccd 100644 --- a/tests/test_tiling.py +++ b/tests/test_tiling.py @@ -12,7 +12,7 @@ import numpy as np import pytest from PIL import Image -from remove_ai_watermarks.noai.tiling import ( +from remove_ai_watermarks._internal.tiling import ( Tile, _axis_positions, feather_region_composite, diff --git a/tests/test_watermark_profiles.py b/tests/test_watermark_profiles.py index 3459b58..85964f1 100644 --- a/tests/test_watermark_profiles.py +++ b/tests/test_watermark_profiles.py @@ -4,7 +4,7 @@ from __future__ import annotations import pytest -from remove_ai_watermarks.noai.watermark_profiles import resolve_strength, viable_steps +from remove_ai_watermarks._internal.watermark_profiles import resolve_strength, viable_steps class TestViableSteps: