mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-23 05:57:11 +02:00
Victor Kuznetsov
ee865e95e4
pre-commit: 1 maintain.sh - not rerun full, ruff + 34 targeted tests passed; 2) /simplify - pad lives in silhouette crop, stable=False reuses probes; 3) docs sync - no README/cli refs to update; 4) CLAUDE.md - no change EOF )
368 lines
17 KiB
Python
368 lines
17 KiB
Python
"""Opt-in restoration of verified text from a Qwen VAE reconstruction."""
|
|
|
|
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
from pathlib import Path
|
|
|
|
from numpy.typing import NDArray
|
|
|
|
TEXT_MANIFEST_SCHEMA = 1
|
|
FIDELITY_BLEND_ALPHA = 0.15
|
|
GLYPH_FEATHER = 0.5
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VerifiedTextLine:
|
|
"""One operator-verified source line in source-pixel coordinates."""
|
|
|
|
box: tuple[int, int, int, int]
|
|
text: str
|
|
script: str
|
|
angle: float = 0.0
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VerifiedTextManifest:
|
|
"""Text annotations cryptographically bound to one decoded RGB source."""
|
|
|
|
source_pixel_sha256: str
|
|
width: int
|
|
height: int
|
|
lines: tuple[VerifiedTextLine, ...]
|
|
|
|
|
|
def source_pixel_sha256(image: Image.Image) -> str:
|
|
"""Hash decoded RGB geometry and bytes, independent of container metadata."""
|
|
rgb = image.convert("RGB")
|
|
digest = hashlib.sha256()
|
|
digest.update(rgb.width.to_bytes(8, "big"))
|
|
digest.update(rgb.height.to_bytes(8, "big"))
|
|
digest.update(rgb.tobytes())
|
|
return digest.hexdigest()
|
|
|
|
|
|
def load_verified_text_manifest(path: Path, source: Image.Image) -> VerifiedTextManifest:
|
|
"""Load and validate a manually verified manifest for exactly ``source``."""
|
|
try:
|
|
payload = json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
raise ValueError(f"Cannot read text manifest {path}: {exc}") from exc
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("Text manifest must be a JSON object")
|
|
if payload.get("schema_version") != TEXT_MANIFEST_SCHEMA:
|
|
raise ValueError(f"Text manifest schema_version must be {TEXT_MANIFEST_SCHEMA}")
|
|
if payload.get("verified") is not True:
|
|
raise ValueError("Text manifest must contain verified=true after manual review")
|
|
|
|
rgb = source.convert("RGB")
|
|
width = _manifest_integer(payload, "width")
|
|
height = _manifest_integer(payload, "height")
|
|
if (width, height) != rgb.size:
|
|
raise ValueError(f"Text manifest dimensions {width}x{height} do not match source {rgb.width}x{rgb.height}")
|
|
expected_hash = payload.get("source_pixel_sha256")
|
|
if not isinstance(expected_hash, str) or len(expected_hash) != 64:
|
|
raise ValueError("Text manifest source_pixel_sha256 must be a 64-character SHA-256")
|
|
actual_hash = source_pixel_sha256(rgb)
|
|
if expected_hash.casefold() != actual_hash:
|
|
raise ValueError("Text manifest source_pixel_sha256 does not match the decoded source pixels")
|
|
|
|
raw_lines = payload.get("lines")
|
|
if not isinstance(raw_lines, list) or not raw_lines:
|
|
raise ValueError("Text manifest lines must be a non-empty list")
|
|
lines = tuple(_load_line(item, width, height, index) for index, item in enumerate(raw_lines))
|
|
if list(lines) != sorted(lines, key=lambda line: (line.box[1], line.box[0])):
|
|
raise ValueError("Text manifest lines must be in top-to-bottom, left-to-right reading order")
|
|
return VerifiedTextManifest(actual_hash, width, height, lines)
|
|
|
|
|
|
def _manifest_integer(payload: dict[str, Any], key: str) -> int:
|
|
value = payload.get(key)
|
|
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
|
|
raise ValueError(f"Text manifest {key} must be a positive integer")
|
|
return value
|
|
|
|
|
|
def _load_line(item: Any, width: int, height: int, index: int) -> VerifiedTextLine:
|
|
if not isinstance(item, dict):
|
|
raise ValueError(f"Text manifest line {index} must be an object")
|
|
raw_box = item.get("box")
|
|
if (
|
|
not isinstance(raw_box, list)
|
|
or len(raw_box) != 4
|
|
or any(isinstance(value, bool) or not isinstance(value, int) for value in raw_box)
|
|
):
|
|
raise ValueError(f"Text manifest line {index} box must contain four integers")
|
|
box = tuple(raw_box)
|
|
x1, y1, x2, y2 = box
|
|
if not (0 <= x1 < x2 <= width and 0 <= y1 < y2 <= height):
|
|
raise ValueError(f"Text manifest line {index} box is outside the source dimensions")
|
|
text = item.get("text")
|
|
script = item.get("script")
|
|
if not isinstance(text, str) or not text.strip():
|
|
raise ValueError(f"Text manifest line {index} text must be non-empty")
|
|
if not isinstance(script, str) or not script.strip():
|
|
raise ValueError(f"Text manifest line {index} script must be non-empty")
|
|
angle_value = item.get("angle", 0.0)
|
|
if isinstance(angle_value, bool) or not isinstance(angle_value, int | float):
|
|
raise ValueError(f"Text manifest line {index} angle must be numeric")
|
|
angle = float(angle_value)
|
|
if not math.isfinite(angle) or abs(angle) > 30.0:
|
|
raise ValueError(f"Text manifest line {index} angle must be between -30 and 30 degrees")
|
|
return VerifiedTextLine(box, text, script, angle)
|
|
|
|
|
|
def blend_fidelity_anchor(clean: Image.Image, donor: Image.Image) -> Image.Image:
|
|
"""Blend 15% Qwen-VAE reconstruction into the oracle-clean pipeline output."""
|
|
clean_rgb = np.asarray(clean.convert("RGB"), dtype=np.float32)
|
|
donor_rgb = np.asarray(donor.convert("RGB"), dtype=np.float32)
|
|
if clean_rgb.shape != donor_rgb.shape:
|
|
raise ValueError("Clean result and Qwen-VAE donor dimensions must match")
|
|
blended = np.rint(clean_rgb * (1.0 - FIDELITY_BLEND_ALPHA) + donor_rgb * FIDELITY_BLEND_ALPHA)
|
|
return Image.fromarray(np.clip(blended, 0, 255).astype(np.uint8))
|
|
|
|
|
|
def restore_verified_text(
|
|
source: Image.Image,
|
|
candidate: Image.Image,
|
|
donor: Image.Image,
|
|
lines: tuple[VerifiedTextLine, ...],
|
|
) -> Image.Image:
|
|
"""Erase candidate glyphs, then composite verified Qwen-VAE glyph cores."""
|
|
from remove_ai_watermarks import region_eraser
|
|
|
|
if not region_eraser.lama_available():
|
|
raise RuntimeError(
|
|
"Verified text restoration requires LaMa. Install: pip install 'remove-ai-watermarks[text-restoration]'"
|
|
)
|
|
source_rgb = np.asarray(source.convert("RGB"))
|
|
candidate_rgb = np.asarray(candidate.convert("RGB"))
|
|
donor_rgb = np.asarray(donor.convert("RGB"))
|
|
if source_rgb.shape != candidate_rgb.shape or source_rgb.shape != donor_rgb.shape:
|
|
raise ValueError("Source, candidate, and Qwen-VAE donor dimensions must match")
|
|
|
|
source_masks = [source_silhouette_mask(source_rgb, line.box, line.angle) for line in lines]
|
|
for index, mask in enumerate(source_masks):
|
|
if not np.any(mask):
|
|
raise ValueError(f"Verified text line {index} produced no source glyph pixels")
|
|
candidate_masks = [source_silhouette_mask(candidate_rgb, line.box, line.angle) for line in lines]
|
|
erase_masks = []
|
|
for line, source_mask, candidate_mask in zip(lines, source_masks, candidate_masks, strict=True):
|
|
radius = 5 if line.box[3] - line.box[1] >= 48 else 3
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1,) * 2)
|
|
erase_masks.append(cv2.dilate(np.maximum(source_mask, candidate_mask), kernel))
|
|
del candidate_masks
|
|
|
|
groups = group_text_lines(lines)
|
|
background = cv2.cvtColor(candidate_rgb, cv2.COLOR_RGB2BGR)
|
|
for group in groups:
|
|
background = region_eraser.erase_lama(background, np.maximum.reduce([erase_masks[index] for index in group]))
|
|
background_rgb = cv2.cvtColor(background, cv2.COLOR_BGR2RGB)
|
|
residual_masks = [
|
|
residual_glyph_mask(background_rgb, mask, line.box) for line, mask in zip(lines, erase_masks, strict=True)
|
|
]
|
|
for group in groups:
|
|
residual = np.maximum.reduce([residual_masks[index] for index in group])
|
|
if np.any(residual):
|
|
background = region_eraser.erase_lama(background, residual)
|
|
del erase_masks, residual_masks
|
|
restored = cv2.cvtColor(background, cv2.COLOR_BGR2RGB)
|
|
restored = composite_fresh_text_edges(source_rgb, restored, lines, source_masks)
|
|
source_glyph_mask = np.maximum.reduce(source_masks)
|
|
restored = composite_reconstructed_glyphs(donor_rgb, restored, source_glyph_mask)
|
|
return Image.fromarray(restored)
|
|
|
|
|
|
def _glyph_crop_box(box: tuple[int, int, int, int], width: int, height: int) -> tuple[int, int, int, int]:
|
|
"""Widen the detector box so descenders stay inside the silhouette crop.
|
|
|
|
Paddle line boxes sit 2-5 px above the true ink bottom on the poster
|
|
fixtures. The crop is the box itself, so those pixels never reached the
|
|
donor composite. Expand only on Y: 8% up, 25% down, clamped to the frame.
|
|
"""
|
|
x1, y1, x2, y2 = box
|
|
line_h = max(1, y2 - y1)
|
|
pad_top = max(1, round(line_h * 0.08))
|
|
pad_bot = max(2, round(line_h * 0.25))
|
|
return max(0, x1), max(0, y1 - pad_top), min(width, x2), min(height, y2 + pad_bot)
|
|
|
|
|
|
def source_silhouette_mask(
|
|
source_rgb: NDArray[Any],
|
|
box: tuple[int, int, int, int],
|
|
angle: float = 0.0,
|
|
) -> NDArray[Any]:
|
|
"""Recover a thresholded glyph shape without retaining source amplitudes."""
|
|
height, width = source_rgb.shape[:2]
|
|
x1, y1, x2, y2 = _glyph_crop_box(box, width, height)
|
|
gray = cv2.cvtColor(source_rgb[y1:y2, x1:x2], cv2.COLOR_RGB2GRAY)
|
|
support = np.ones(gray.shape, dtype=np.uint8)
|
|
if angle:
|
|
box_width, box_height = x2 - x1, y2 - y1
|
|
theta = math.radians(abs(angle))
|
|
cosine, sine = math.cos(theta), math.sin(theta)
|
|
denominator = cosine * cosine - sine * sine
|
|
rect_width = (box_width * cosine - box_height * sine) / denominator
|
|
rect_height = (box_height * cosine - box_width * sine) / denominator
|
|
rotated = cv2.boxPoints(
|
|
((box_width / 2, box_height / 2), (max(1.0, rect_width * 0.92), max(1.0, rect_height * 0.62)), -angle)
|
|
)
|
|
support.fill(0)
|
|
cv2.fillConvexPoly(support, np.rint(rotated).astype(np.int32), 1)
|
|
values = gray[support > 0]
|
|
background_luma = float(np.median(values))
|
|
else:
|
|
ring_pad = max(6, min(20, (y2 - y1) // 4))
|
|
rx1, ry1, rx2, ry2 = _clip_box((x1, y1, x2, y2), width, height, pad=ring_pad)
|
|
context = cv2.cvtColor(source_rgb[ry1:ry2, rx1:rx2], cv2.COLOR_RGB2GRAY)
|
|
ring = np.ones(context.shape, dtype=bool)
|
|
ring[y1 - ry1 : y2 - ry1, x1 - rx1 : x2 - rx1] = False
|
|
background_luma = float(np.median(context[ring])) if ring.any() else float(np.median(gray))
|
|
values = gray.reshape(-1)
|
|
low, high = float(np.percentile(values, 2)), float(np.percentile(values, 98))
|
|
dark_contrast, light_contrast = background_luma - low, high - background_luma
|
|
threshold = max(16.0, min(56.0, max(light_contrast, dark_contrast) * 0.22))
|
|
if light_contrast > dark_contrast:
|
|
crop_mask = (gray.astype(np.float32) >= background_luma + threshold).astype(np.uint8) * 255
|
|
else:
|
|
crop_mask = (gray.astype(np.float32) <= background_luma - threshold).astype(np.uint8) * 255
|
|
crop_mask[support == 0] = 0
|
|
result = np.zeros((height, width), dtype=np.uint8)
|
|
result[y1:y2, x1:x2] = crop_mask
|
|
return result
|
|
|
|
|
|
def residual_glyph_mask(
|
|
background_rgb: NDArray[Any],
|
|
original_mask: NDArray[Any],
|
|
box: tuple[int, int, int, int],
|
|
) -> NDArray[Any]:
|
|
"""Find glyph-like contrast left after the first inpaint pass."""
|
|
residual = _foreground_mask(background_rgb, box)
|
|
residual = cv2.bitwise_and(residual, original_mask)
|
|
return cv2.dilate(residual, np.ones((5, 5), np.uint8), iterations=1)
|
|
|
|
|
|
def composite_fresh_text_edges(
|
|
source_rgb: NDArray[Any],
|
|
background_rgb: NDArray[Any],
|
|
lines: tuple[VerifiedTextLine, ...],
|
|
masks: list[NDArray[Any]],
|
|
) -> NDArray[Any]:
|
|
"""Render fresh antialiased edges for source-derived glyph masks."""
|
|
restored = background_rgb
|
|
for line, mask in zip(lines, masks, strict=True):
|
|
color = _sample_text_color(source_rgb, mask, line.box)
|
|
restored = composite_fresh_silhouette(restored, mask, color)
|
|
return restored
|
|
|
|
|
|
def composite_reconstructed_glyphs(
|
|
donor_rgb: NDArray[Any],
|
|
background_rgb: NDArray[Any],
|
|
glyph_mask: NDArray[Any],
|
|
*,
|
|
feather: float = GLYPH_FEATHER,
|
|
) -> NDArray[Any]:
|
|
"""Composite an exact reconstructed core with a narrow donor edge."""
|
|
if donor_rgb.shape != background_rgb.shape or donor_rgb.shape[:2] != glyph_mask.shape:
|
|
raise ValueError("donor, background, and glyph mask dimensions must match")
|
|
blurred = cv2.GaussianBlur(glyph_mask, (0, 0), feather) if feather > 0 else glyph_mask
|
|
alpha = np.maximum(glyph_mask, blurred).astype(np.float32) / 255.0
|
|
combined = donor_rgb.astype(np.float32) * alpha[..., None] + background_rgb.astype(np.float32) * (
|
|
1.0 - alpha[..., None]
|
|
)
|
|
return np.clip(np.rint(combined), 0, 255).astype(np.uint8)
|
|
|
|
|
|
def composite_fresh_silhouette(
|
|
background_rgb: NDArray[Any],
|
|
glyph_mask: NDArray[Any],
|
|
color: tuple[int, int, int],
|
|
*,
|
|
feather: float = 0.35,
|
|
) -> NDArray[Any]:
|
|
"""Render a binary source shape with fresh color and antialiasing."""
|
|
if background_rgb.shape[:2] != glyph_mask.shape:
|
|
raise ValueError("background and glyph mask dimensions must match")
|
|
antialiased = cv2.GaussianBlur(glyph_mask, (0, 0), feather) if feather > 0 else glyph_mask
|
|
alpha = antialiased.astype(np.float32)[..., None] / 255.0
|
|
foreground = np.empty_like(background_rgb)
|
|
foreground[:, :] = color
|
|
combined = foreground.astype(np.float32) * alpha + background_rgb.astype(np.float32) * (1.0 - alpha)
|
|
return np.clip(combined, 0, 255).astype(np.uint8)
|
|
|
|
|
|
def _clip_box(box: tuple[int, int, int, int], width: int, height: int, pad: int = 0) -> tuple[int, int, int, int]:
|
|
x1, y1, x2, y2 = box
|
|
return max(0, x1 - pad), max(0, y1 - pad), min(width, x2 + pad), min(height, y2 + pad)
|
|
|
|
|
|
def _foreground_mask(source_rgb: NDArray[Any], box: tuple[int, int, int, int]) -> NDArray[Any]:
|
|
height, width = source_rgb.shape[:2]
|
|
line_height = box[3] - box[1]
|
|
x1, y1, x2, y2 = _clip_box(box, width, height, pad=max(6, int(line_height * 0.12)))
|
|
gray = cv2.cvtColor(source_rgb[y1:y2, x1:x2], cv2.COLOR_RGB2GRAY)
|
|
ring_pad = max(8, min(24, (y2 - y1) // 5))
|
|
rx1, ry1, rx2, ry2 = _clip_box((x1, y1, x2, y2), width, height, pad=ring_pad)
|
|
context = cv2.cvtColor(source_rgb[ry1:ry2, rx1:rx2], cv2.COLOR_RGB2GRAY)
|
|
ring = np.ones(context.shape, dtype=bool)
|
|
ring[y1 - ry1 : y2 - ry1, x1 - rx1 : x2 - rx1] = False
|
|
background_luma = float(np.median(context[ring])) if ring.any() else float(np.median(gray))
|
|
low, high = float(np.percentile(gray, 4)), float(np.percentile(gray, 96))
|
|
dark_contrast, light_contrast = background_luma - low, high - background_luma
|
|
threshold = max(24.0, min(72.0, max(light_contrast, dark_contrast) * 0.32))
|
|
if light_contrast > dark_contrast:
|
|
mask = (gray.astype(np.float32) >= background_luma + threshold).astype(np.uint8) * 255
|
|
else:
|
|
mask = (gray.astype(np.float32) <= background_luma - threshold).astype(np.uint8) * 255
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((2, 2), np.uint8))
|
|
dilation = 5 if line_height >= 48 else 3
|
|
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * dilation + 1,) * 2))
|
|
result = np.zeros((height, width), dtype=np.uint8)
|
|
result[y1:y2, x1:x2] = mask
|
|
return result
|
|
|
|
|
|
def _sample_text_color(
|
|
source_rgb: NDArray[Any], mask: NDArray[Any], box: tuple[int, int, int, int]
|
|
) -> tuple[int, int, int]:
|
|
height, width = source_rgb.shape[:2]
|
|
x1, y1, x2, y2 = _clip_box(box, width, height, pad=2)
|
|
crop = source_rgb[y1:y2, x1:x2]
|
|
pixels = crop[mask[y1:y2, x1:x2] > 0]
|
|
luma = pixels.mean(axis=1)
|
|
background_luma = float(crop[[0, -1], :, :].reshape(-1, 3).mean(axis=1).mean())
|
|
selected = (
|
|
pixels[luma <= np.percentile(luma, 20)] if background_luma >= 128 else pixels[luma >= np.percentile(luma, 80)]
|
|
)
|
|
return tuple(int(value) for value in np.median(selected, axis=0))
|
|
|
|
|
|
def group_text_lines(lines: Sequence[VerifiedTextLine]) -> list[list[int]]:
|
|
"""Group nearby same-script lines for a shared LaMa erase pass."""
|
|
groups: list[list[int]] = []
|
|
for index, line in enumerate(lines):
|
|
if not groups:
|
|
groups.append([index])
|
|
continue
|
|
previous = lines[groups[-1][-1]]
|
|
gap = line.box[1] - previous.box[3]
|
|
if line.script != previous.script or gap > max(60, int((previous.box[3] - previous.box[1]) * 1.1)):
|
|
groups.append([index])
|
|
else:
|
|
groups[-1].append(index)
|
|
return groups
|