CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-10 16:10:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: complete product video watermark pipeline

Browse Source
This commit is contained in:
Victor Kuznetsov
2026-07-31 10:42:39 -07:00
parent 8975916b11
commit 6315b58628
33 changed files with 3793 additions and 515 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/remove_ai_watermarks/__init__.py
+19 -2
View File
@@ -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)
src/remove_ai_watermarks/cli.py
+256 -52
View File
@@ -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: <source>_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: <source>_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: <source>_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: <source>_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: <directory>_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))
src/remove_ai_watermarks/metadata.py
+15 -2
View File
@@ -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
src/remove_ai_watermarks/noai/isobmff.py
+150 -4
View File
@@ -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.
src/remove_ai_watermarks/video.py
+523 -120
View File
@@ -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 ``<source_stem>_clean<source_suffix>``. 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,
src/remove_ai_watermarks/video_encoding.py
+264 -16
View File
@@ -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
src/remove_ai_watermarks/video_invisible.py
+83 -106
View File
@@ -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:
src/remove_ai_watermarks/video_synthid.py
+3 -2
View File
@@ -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."
)
src/remove_ai_watermarks/video_temporal.py
+189
View File
@@ -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
src/remove_ai_watermarks/video_visible.py
+171 -15
View File
@@ -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