mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-09-04 03:20:36 +02:00
Add project files, tests, and documentation for GitHub release
- CLI with visible, invisible, all, metadata, and batch commands - Gemini watermark removal via reverse alpha blending - Invisible watermark removal via diffusion regeneration (SynthID, TreeRing) - AI metadata stripping (EXIF, PNG text, C2PA) - Face protection (YOLO/Haar) and analog humanizer - 137 tests covering all CLI modes and core engines - Ruff and Pyright clean
This commit is contained in:
test-user
@@ -0,0 +1,598 @@
|
||||
"""Unified CLI for remove-ai-watermarks.
|
||||
|
||||
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
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
|
||||
from rich.table import Table
|
||||
|
||||
from remove_ai_watermarks import __version__
|
||||
|
||||
console = Console()
|
||||
|
||||
SUPPORTED_FORMATS = {".png", ".jpg", ".jpeg", ".webp"}
|
||||
|
||||
|
||||
def _setup_logging(verbose: bool) -> None:
|
||||
level = logging.DEBUG if verbose else logging.WARNING
|
||||
logging.basicConfig(
|
||||
level=level,
|
||||
format="%(name)s | %(message)s",
|
||||
handlers=[logging.StreamHandler()],
|
||||
)
|
||||
|
||||
|
||||
def _banner() -> None:
|
||||
console.print(
|
||||
Panel(
|
||||
f"[bold cyan]Remove-AI-Watermarks[/] [dim]v{__version__}[/]\n[dim]Visible & invisible watermark removal[/]",
|
||||
border_style="cyan",
|
||||
padding=(0, 2),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _validate_image(path: Path) -> Path:
|
||||
if not path.exists():
|
||||
console.print(f"[red]Error:[/] File not found: {path}")
|
||||
raise SystemExit(1)
|
||||
if path.suffix.lower() not in SUPPORTED_FORMATS:
|
||||
console.print(
|
||||
f"[yellow]Warning:[/] {path.suffix} may not be supported (expected: {', '.join(SUPPORTED_FORMATS)})"
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
# ── Main group ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@click.group(invoke_without_command=True)
|
||||
@click.version_option(__version__, prog_name="remove-ai-watermarks")
|
||||
@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."""
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv() # Load .env (e.g. HF_TOKEN)
|
||||
|
||||
ctx.ensure_object(dict)
|
||||
ctx.obj["verbose"] = verbose
|
||||
_setup_logging(verbose)
|
||||
|
||||
if ctx.invoked_subcommand is None:
|
||||
_banner()
|
||||
click.echo(ctx.get_help())
|
||||
|
||||
|
||||
# ── Visible (Gemini) watermark removal ───────────────────────────────
|
||||
|
||||
|
||||
@main.command("visible")
|
||||
@click.argument("source", type=click.Path(exists=True, path_type=Path))
|
||||
@click.option(
|
||||
"-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)."
|
||||
)
|
||||
@click.option("--inpaint/--no-inpaint", default=True, help="Apply inpainting cleanup after removal.")
|
||||
@click.option(
|
||||
"--inpaint-method", type=click.Choice(["ns", "telea", "gaussian"]), default="ns", help="Inpainting method."
|
||||
)
|
||||
@click.option("--inpaint-strength", type=float, default=0.85, help="Inpainting blend strength (0.0–1.0).")
|
||||
@click.option("--detect/--no-detect", default=True, help="Detect watermark before removal.")
|
||||
@click.option("--detect-threshold", type=float, default=0.25, help="Detection confidence threshold.")
|
||||
@click.option("--strip-metadata/--keep-metadata", default=True, help="Strip AI metadata from output.")
|
||||
@click.pass_context
|
||||
def cmd_visible(
|
||||
ctx: click.Context,
|
||||
source: Path,
|
||||
output: Path | None,
|
||||
inpaint: bool,
|
||||
inpaint_method: str,
|
||||
inpaint_strength: float,
|
||||
detect: bool,
|
||||
detect_threshold: float,
|
||||
strip_metadata: bool,
|
||||
) -> None:
|
||||
"""Remove visible Gemini watermark (sparkle logo) from an image.
|
||||
|
||||
Uses reverse alpha blending — fast, deterministic, offline.
|
||||
"""
|
||||
import cv2
|
||||
|
||||
from remove_ai_watermarks.gemini_engine import GeminiEngine
|
||||
|
||||
_banner()
|
||||
source = _validate_image(source)
|
||||
|
||||
if output is None:
|
||||
output = source.with_stem(source.stem + "_clean")
|
||||
|
||||
engine = GeminiEngine()
|
||||
|
||||
# Load image
|
||||
image = cv2.imread(str(source), cv2.IMREAD_COLOR)
|
||||
if image is None:
|
||||
console.print(f"[red]Error:[/] Failed to read image: {source}")
|
||||
raise SystemExit(1)
|
||||
|
||||
h, w = image.shape[:2]
|
||||
console.print(f" [dim]Input:[/] {source.name} ({w}×{h})")
|
||||
|
||||
# Detection (we always detect softly, to find dynamic region for inpainting)
|
||||
with console.status("[cyan]Detecting watermark…[/]"):
|
||||
det = engine.detect_watermark(image)
|
||||
|
||||
if detect:
|
||||
if det.detected:
|
||||
console.print(
|
||||
f" [green]✓[/] Watermark detected "
|
||||
f"[dim](confidence: {det.confidence:.1%}, "
|
||||
f"spatial: {det.spatial_score:.3f}, "
|
||||
f"gradient: {det.gradient_score:.3f})[/]"
|
||||
)
|
||||
else:
|
||||
console.print(f" [yellow]⚠[/] Watermark not detected [dim](confidence: {det.confidence:.1%})[/]")
|
||||
if det.confidence < detect_threshold:
|
||||
console.print(" [dim]Skipping. Use --no-detect to force removal.[/]")
|
||||
raise SystemExit(0)
|
||||
|
||||
# Removal
|
||||
t0 = time.monotonic()
|
||||
with console.status("[cyan]Removing watermark…[/]"):
|
||||
result = engine.remove_watermark(image)
|
||||
|
||||
if inpaint:
|
||||
if det.confidence > 0.15:
|
||||
region = det.region
|
||||
else:
|
||||
from remove_ai_watermarks.gemini_engine import get_watermark_config
|
||||
|
||||
config = get_watermark_config(w, h)
|
||||
pos = config.get_position(w, h)
|
||||
region = (pos[0], pos[1], config.logo_size, config.logo_size)
|
||||
result = engine.inpaint_residual(
|
||||
result,
|
||||
region,
|
||||
strength=inpaint_strength,
|
||||
method=inpaint_method,
|
||||
)
|
||||
|
||||
elapsed = time.monotonic() - t0
|
||||
|
||||
# Save
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
cv2.imwrite(str(output), result)
|
||||
|
||||
# Strip metadata
|
||||
if strip_metadata:
|
||||
try:
|
||||
from remove_ai_watermarks.metadata import remove_ai_metadata
|
||||
|
||||
remove_ai_metadata(output, output)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
size_kb = output.stat().st_size / 1024
|
||||
console.print(f" [green]✓[/] Saved: {output} [dim]({size_kb:.0f} KB, {elapsed:.2f}s)[/]")
|
||||
|
||||
|
||||
# ── Invisible watermark removal ─────────────────────────────────────
|
||||
|
||||
|
||||
@main.command("invisible")
|
||||
@click.argument("source", type=click.Path(exists=True, path_type=Path))
|
||||
@click.option(
|
||||
"-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)."
|
||||
)
|
||||
@click.option("--strength", type=float, default=0.02, help="Denoising strength (0.0–1.0). Default: 0.02.")
|
||||
@click.option("--steps", type=int, default=100, help="Number of denoising steps. Default: 100.")
|
||||
@click.option("--pipeline", type=click.Choice(["default", "ctrlregen"]), default="default", help="Pipeline profile.")
|
||||
@click.option("--device", type=click.Choice(["auto", "cpu", "mps", "cuda"]), default="auto", help="Inference device.")
|
||||
@click.option("--seed", type=int, default=None, help="Random seed for reproducibility.")
|
||||
@click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
|
||||
@click.option("--humanize", type=float, default=0.0, help="Humanization strength (0.0–1.0) for invisible removal.")
|
||||
@click.pass_context
|
||||
def cmd_invisible(
|
||||
ctx: click.Context,
|
||||
source: Path,
|
||||
output: Path | None,
|
||||
strength: float,
|
||||
steps: int,
|
||||
pipeline: str,
|
||||
device: str,
|
||||
seed: int | None,
|
||||
hf_token: str | None,
|
||||
humanize: float,
|
||||
) -> None:
|
||||
"""Remove invisible AI watermarks (SynthID, StableSignature, TreeRing).
|
||||
|
||||
Uses diffusion-based regeneration. Requires GPU for reasonable speed.
|
||||
"""
|
||||
from remove_ai_watermarks.invisible_engine import InvisibleEngine
|
||||
|
||||
source = _validate_image(source)
|
||||
if output is None:
|
||||
output = source.with_stem(source.stem + "_clean")
|
||||
|
||||
device_str = None if device == "auto" else device
|
||||
|
||||
def progress_cb(msg: str) -> None:
|
||||
console.print(f" [dim]{msg}[/]")
|
||||
|
||||
engine = InvisibleEngine(
|
||||
device=device_str,
|
||||
pipeline=pipeline,
|
||||
hf_token=hf_token,
|
||||
progress_callback=progress_cb,
|
||||
)
|
||||
|
||||
console.print(f" [dim]Input:[/] {source.name}")
|
||||
console.print(f" [dim]Pipeline:[/] {pipeline}")
|
||||
console.print(f" [dim]Strength:[/] {strength} Steps: {steps}")
|
||||
|
||||
t0 = time.monotonic()
|
||||
result_path = engine.remove_watermark(
|
||||
image_path=source,
|
||||
output_path=output,
|
||||
strength=strength,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=None,
|
||||
seed=seed,
|
||||
humanize=humanize,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
|
||||
size_kb = result_path.stat().st_size / 1024
|
||||
console.print(f"\n [green]✓[/] Saved: {result_path} [dim]({size_kb:.0f} KB, {elapsed:.1f}s)[/]")
|
||||
|
||||
|
||||
# ── Metadata operations ─────────────────────────────────────────────
|
||||
|
||||
|
||||
@main.command("metadata")
|
||||
@click.argument("source", type=click.Path(exists=True, 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: overwrite source)."
|
||||
)
|
||||
@click.option("--keep-standard/--remove-all", default=True, help="Keep standard metadata (Author, Title, etc.).")
|
||||
@click.pass_context
|
||||
def cmd_metadata(
|
||||
ctx: click.Context,
|
||||
source: Path,
|
||||
check: bool,
|
||||
remove: bool,
|
||||
output: Path | None,
|
||||
keep_standard: bool,
|
||||
) -> None:
|
||||
"""Check or remove AI-generation metadata from images.
|
||||
|
||||
Strips EXIF AI tags, PNG text chunks, and C2PA provenance manifests.
|
||||
"""
|
||||
from remove_ai_watermarks.metadata import get_ai_metadata, has_ai_metadata, remove_ai_metadata
|
||||
|
||||
_banner()
|
||||
source = _validate_image(source)
|
||||
|
||||
if check or (not remove):
|
||||
has_ai = has_ai_metadata(source)
|
||||
if has_ai:
|
||||
console.print(f" [yellow]⚠[/] AI metadata detected in {source.name}:")
|
||||
meta = get_ai_metadata(source)
|
||||
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" [green]✓[/] No AI metadata found in {source.name}")
|
||||
|
||||
if not remove:
|
||||
return
|
||||
|
||||
# Remove
|
||||
out = remove_ai_metadata(source, output, keep_standard=keep_standard)
|
||||
console.print(f" [green]✓[/] AI metadata stripped → {out}")
|
||||
|
||||
|
||||
# ── Combined "all" mode ──────────────────────────────────────────────
|
||||
|
||||
|
||||
@main.command("all")
|
||||
@click.argument("source", type=click.Path(exists=True, path_type=Path))
|
||||
@click.option(
|
||||
"-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)."
|
||||
)
|
||||
@click.option("--inpaint/--no-inpaint", default=True, help="Apply inpainting cleanup after visible removal.")
|
||||
@click.option(
|
||||
"--inpaint-method", type=click.Choice(["ns", "telea", "gaussian"]), default="ns", help="Inpainting method."
|
||||
)
|
||||
@click.option("--strength", type=float, default=0.04, help="Invisible watermark denoising strength (0.0–1.0).")
|
||||
@click.option("--steps", type=int, default=50, help="Number of denoising steps for invisible removal.")
|
||||
@click.option(
|
||||
"--pipeline",
|
||||
type=click.Choice(["default", "ctrlregen"]),
|
||||
default="default",
|
||||
help="Pipeline profile for invisible removal.",
|
||||
)
|
||||
@click.option("--model", type=str, default=None, help="HuggingFace model ID for invisible removal.")
|
||||
@click.option("--device", type=click.Choice(["auto", "cpu", "mps", "cuda"]), default="auto", help="Inference device.")
|
||||
@click.option("--seed", type=int, default=None, help="Random seed for reproducibility.")
|
||||
@click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
|
||||
@click.option("--humanize", type=float, default=0.0, help="Humanization strength (0.0–1.0) for invisible removal.")
|
||||
@click.pass_context
|
||||
def cmd_all(
|
||||
ctx: click.Context,
|
||||
source: Path,
|
||||
output: Path | None,
|
||||
inpaint: bool,
|
||||
inpaint_method: str,
|
||||
strength: float,
|
||||
steps: int,
|
||||
pipeline: str,
|
||||
model: str | None,
|
||||
device: str,
|
||||
seed: int | None,
|
||||
hf_token: str | None,
|
||||
humanize: float,
|
||||
) -> None:
|
||||
"""Remove ALL watermarks: visible + invisible + metadata.
|
||||
|
||||
Runs the full pipeline in order:
|
||||
1. Visible watermark removal (Gemini sparkle, reverse alpha blending)
|
||||
2. Invisible watermark removal (SynthID etc., diffusion regeneration)
|
||||
3. AI metadata stripping (EXIF, PNG text, C2PA)
|
||||
|
||||
If invisible watermark deps are not installed, skips step 2 with a warning.
|
||||
"""
|
||||
import cv2
|
||||
|
||||
from remove_ai_watermarks.gemini_engine import GeminiEngine, get_watermark_config
|
||||
|
||||
_banner()
|
||||
source = _validate_image(source)
|
||||
|
||||
if output is None:
|
||||
output = source.with_stem(source.stem + "_clean")
|
||||
|
||||
t0 = time.monotonic()
|
||||
|
||||
# ── Step 1: Visible watermark ────────────────────────────────
|
||||
console.print("\n [bold cyan]① Visible watermark removal[/]")
|
||||
engine = GeminiEngine()
|
||||
image = cv2.imread(str(source), cv2.IMREAD_COLOR)
|
||||
if image is None:
|
||||
console.print(f"[red]Error:[/] Failed to read image: {source}")
|
||||
raise SystemExit(1)
|
||||
|
||||
h, w = image.shape[:2]
|
||||
console.print(f" [dim]Input:[/] {source.name} ({w}×{h})")
|
||||
|
||||
with console.status("[cyan]Removing visible watermark…[/]"):
|
||||
det = engine.detect_watermark(image)
|
||||
if det.detected:
|
||||
result = engine.remove_watermark(image)
|
||||
if inpaint:
|
||||
if det.confidence > 0.15:
|
||||
region = det.region
|
||||
else:
|
||||
config = get_watermark_config(w, h)
|
||||
pos = config.get_position(w, h)
|
||||
region = (pos[0], pos[1], config.logo_size, config.logo_size)
|
||||
result = engine.inpaint_residual(result, region, method=inpaint_method)
|
||||
console.print(" [green]✓[/] Visible watermark removed")
|
||||
else:
|
||||
result = image.copy()
|
||||
console.print(" [dim]Skipped (no visible watermark detected)[/]")
|
||||
|
||||
# Save intermediate
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
cv2.imwrite(str(output), result)
|
||||
|
||||
# ── Step 2: Invisible watermark ──────────────────────────────
|
||||
console.print("\n [bold cyan]② Invisible watermark removal[/]")
|
||||
from remove_ai_watermarks.invisible_engine import InvisibleEngine
|
||||
|
||||
device_str = None if device == "auto" else device
|
||||
|
||||
def progress_cb(msg: str) -> None:
|
||||
console.print(f" [dim]{msg}[/]")
|
||||
|
||||
inv_engine = InvisibleEngine(
|
||||
model_id=model,
|
||||
device=device_str,
|
||||
pipeline=pipeline,
|
||||
hf_token=hf_token,
|
||||
progress_callback=progress_cb,
|
||||
)
|
||||
|
||||
console.print(f" [dim]Strength:[/] {strength} Steps: {steps}")
|
||||
inv_engine.remove_watermark(
|
||||
image_path=output,
|
||||
output_path=output,
|
||||
strength=strength,
|
||||
num_inference_steps=steps,
|
||||
seed=seed,
|
||||
humanize=humanize,
|
||||
)
|
||||
console.print(" [green]✓[/] Invisible watermark removed")
|
||||
|
||||
# ── Step 3: Metadata ─────────────────────────────────────────
|
||||
console.print("\n [bold cyan]③ AI metadata stripping[/]")
|
||||
try:
|
||||
from remove_ai_watermarks.metadata import remove_ai_metadata
|
||||
|
||||
remove_ai_metadata(output, output)
|
||||
console.print(" [green]✓[/] AI metadata stripped")
|
||||
except Exception as e:
|
||||
console.print(f" [yellow]⚠[/] Metadata strip failed: {e}")
|
||||
|
||||
# ── Done ─────────────────────────────────────────────────────
|
||||
elapsed = time.monotonic() - t0
|
||||
size_kb = output.stat().st_size / 1024
|
||||
console.print(f"\n [bold green]✓ Done:[/] {output} [dim]({size_kb:.0f} KB, {elapsed:.1f}s total)[/]")
|
||||
|
||||
|
||||
# ── Batch command ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@main.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: <dir>_clean/).",
|
||||
)
|
||||
@click.option(
|
||||
"--mode", type=click.Choice(["visible", "invisible", "metadata", "all"]), default="visible", help="Processing mode."
|
||||
)
|
||||
@click.option("--strength", type=float, default=None, help="Denoising strength (invisible mode).")
|
||||
@click.option("--steps", type=int, default=50, help="Number of denoising steps (invisible mode).")
|
||||
@click.option("--inpaint/--no-inpaint", default=True, help="Apply inpainting (visible mode).")
|
||||
@click.option("--humanize", type=float, default=0.0, help="Humanization strength (0.0–1.0) for invisible removal.")
|
||||
@click.option("--pipeline", type=click.Choice(["default", "ctrlregen"]), default="default", help="Pipeline profile.")
|
||||
@click.option("--device", type=click.Choice(["auto", "cpu", "mps", "cuda"]), default="auto", help="Inference device.")
|
||||
@click.option("--seed", type=int, default=None, help="Random seed for reproducibility.")
|
||||
@click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
|
||||
@click.pass_context
|
||||
def cmd_batch(
|
||||
ctx: click.Context,
|
||||
directory: Path,
|
||||
mode: str,
|
||||
output_dir: Path | None,
|
||||
strength: float | None,
|
||||
steps: int,
|
||||
pipeline: str,
|
||||
device: str,
|
||||
seed: int | None,
|
||||
hf_token: str | None,
|
||||
inpaint: bool,
|
||||
humanize: float,
|
||||
) -> None:
|
||||
"""Process all images in a directory."""
|
||||
_banner()
|
||||
|
||||
if output_dir is None:
|
||||
output_dir = directory.parent / (directory.name + "_clean")
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Assuming supported_formats is defined elsewhere
|
||||
supported_formats = [".jpg", ".jpeg", ".png", ".webp"]
|
||||
|
||||
images = sorted(p for p in directory.iterdir() if p.suffix.lower() in supported_formats)
|
||||
|
||||
if not images:
|
||||
console.print(f"[yellow]No supported images found in {directory}[/]")
|
||||
return
|
||||
|
||||
console.print(f" Found [bold]{len(images)}[/] images in {directory}")
|
||||
console.print(f" Output → {output_dir}")
|
||||
console.print(f" Mode: [cyan]{mode}[/]")
|
||||
|
||||
processed = 0
|
||||
errors = 0
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
task = progress.add_task("Processing…", total=len(images))
|
||||
|
||||
for img_path in images:
|
||||
out_path = output_dir / img_path.name
|
||||
progress.update(task, description=f"[cyan]{img_path.name}[/]")
|
||||
|
||||
try:
|
||||
if mode in ("visible", "all"):
|
||||
import cv2
|
||||
|
||||
from remove_ai_watermarks.gemini_engine import (
|
||||
GeminiEngine,
|
||||
get_watermark_config,
|
||||
)
|
||||
|
||||
if "_vis_engine" not in ctx.obj:
|
||||
ctx.obj["_vis_engine"] = GeminiEngine()
|
||||
engine = ctx.obj["_vis_engine"]
|
||||
read_path = img_path
|
||||
if mode == "all" and out_path.exists():
|
||||
read_path = out_path
|
||||
image = cv2.imread(str(read_path), cv2.IMREAD_COLOR)
|
||||
if image is None:
|
||||
raise ValueError("Failed to read image")
|
||||
|
||||
det = engine.detect_watermark(image)
|
||||
if det.detected:
|
||||
result = engine.remove_watermark(image)
|
||||
if inpaint:
|
||||
if det.confidence > 0.15:
|
||||
region = det.region
|
||||
else:
|
||||
h, w = image.shape[:2]
|
||||
config = get_watermark_config(w, h)
|
||||
pos = config.get_position(w, h)
|
||||
region = (pos[0], pos[1], config.logo_size, config.logo_size)
|
||||
|
||||
result = engine.inpaint_residual(result, region)
|
||||
else:
|
||||
result = image.copy()
|
||||
|
||||
cv2.imwrite(str(out_path), result)
|
||||
|
||||
if mode in ("invisible", "all"):
|
||||
from remove_ai_watermarks.invisible_engine import (
|
||||
is_available as invisible_available,
|
||||
)
|
||||
|
||||
if invisible_available():
|
||||
from remove_ai_watermarks.invisible_engine import InvisibleEngine
|
||||
|
||||
engine_inv = InvisibleEngine()
|
||||
engine_inv.remove_watermark(
|
||||
img_path if mode == "invisible" else out_path,
|
||||
out_path,
|
||||
strength=strength,
|
||||
num_inference_steps=steps,
|
||||
seed=seed,
|
||||
humanize=humanize,
|
||||
)
|
||||
|
||||
if mode in ("metadata", "all"):
|
||||
from remove_ai_watermarks.metadata import remove_ai_metadata
|
||||
|
||||
remove_ai_metadata(img_path if mode == "metadata" else out_path, out_path)
|
||||
|
||||
processed += 1
|
||||
|
||||
except Exception as e:
|
||||
errors += 1
|
||||
if ctx.obj.get("verbose"):
|
||||
console.print(f" [red]✗[/] {img_path.name}: {e}")
|
||||
|
||||
progress.advance(task)
|
||||
|
||||
console.print(f"\n [green]✓[/] {processed} processed" + (f" [red]✗[/] {errors} errors" if errors else ""))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user