mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-05 21:48:36 +02:00
fix: harden metadata parsers and engines; sync docs (full-repo review)
Apply fixes from a full-repo review (code, tests, docs). Security / correctness: - Clamp attacker-controlled PNG/caBX chunk lengths to the remaining file size in metadata.py and noai/c2pa.py (a malformed length no longer drives a multi-GB read); skipped chunks seek instead of read. - noai/isobmff.strip_c2pa_boxes is now fail-safe on a malformed box: return the original bytes with a warning instead of silently truncating the tail, so metadata --remove can no longer emit a corrupt file. - doubao_engine._fixed_alpha_map clamps the glyph box to the image (no crash on degenerate width-vs-height). - watermark_remover._run_region_hires gates the phaseCorrelate offset on response and magnitude (a spurious shift no longer garbles text) and drops the generator after a CPU fallback (no MPS/CPU device mismatch). Robustness: - gemini_engine, doubao_engine, region_eraser normalize grayscale and RGBA inputs to BGR at the engine entry points. - image_io.imwrite returns False on an unwritable path (matches cv2). - invisible_engine guards a None imread result before use. - trustmark_detector._decoder uses a double-checked threading lock. - ctrlregen.tiling.tile_positions raises on overlap >= tile. - humanizer chromatic shift no longer wraps opposite-edge pixels. - identify OpenAI caveat keyed on the normalized vendor, not a substring. - Remove the dead "visible --detect-threshold" CLI option. - publish.yml verifies the release tag matches the package version. Docs: - README strength 0.05 to 0.10; .env.example HF_TOKEN marked optional; doubao_capture README updated to reverse-alpha-only; CLAUDE.md synced with the new behaviors and the batch command. Tests: new test_security_clamp.py for the read clamp and isobmff fail-safe; erase CLI coverage; integrity-clash rule 2 end-to-end; multi-tag EXIF survival and cross-format strip guards; channel/size, tiling, humanizer, and imwrite regressions. Full suite 493 passed, 2 skipped; ruff and pyright src/ clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
5298dcc6a3
commit
5d0e6c3a65
@@ -160,7 +160,6 @@ def main(ctx: click.Context, verbose: bool) -> None:
|
||||
)
|
||||
@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(
|
||||
"--mark",
|
||||
type=click.Choice(["auto", *watermark_registry.mark_keys()]),
|
||||
@@ -178,7 +177,6 @@ def cmd_visible(
|
||||
inpaint_method: Literal["ns", "telea", "gaussian"],
|
||||
inpaint_strength: float,
|
||||
detect: bool,
|
||||
detect_threshold: float,
|
||||
mark: str,
|
||||
strip_metadata: bool,
|
||||
) -> None:
|
||||
|
||||
@@ -222,7 +222,14 @@ class DoubaoEngine:
|
||||
"""
|
||||
h, w = image.shape[:2]
|
||||
x, y, bw, bh = loc.bbox
|
||||
roi = image[y : y + bh, x : x + bw].astype(np.float32)
|
||||
# Normalize the ROI to 3-channel BGR: a 2D grayscale or 4-channel BGRA
|
||||
# input would otherwise break the axis=2 channel reductions below.
|
||||
roi = image[y : y + bh, x : x + bw]
|
||||
if roi.ndim == 2:
|
||||
roi = cv2.cvtColor(roi, cv2.COLOR_GRAY2BGR)
|
||||
elif roi.shape[2] == 4:
|
||||
roi = cv2.cvtColor(roi, cv2.COLOR_BGRA2BGR)
|
||||
roi = roi.astype(np.float32)
|
||||
|
||||
luma = roi.mean(axis=2)
|
||||
sat = roi.max(axis=2) - roi.min(axis=2)
|
||||
@@ -290,7 +297,12 @@ class DoubaoEngine:
|
||||
if at is None:
|
||||
return None
|
||||
h, w = image.shape[:2]
|
||||
gw, gh = max(1, int(_ALPHA_WIDTH_FRAC * w)), max(1, int(_ALPHA_HEIGHT_FRAC * w))
|
||||
# Glyph box scales with WIDTH; on a wide/short image the height-from-width
|
||||
# box can exceed the image height. Clamp both dims so the slice assignment
|
||||
# below cannot overflow (a degenerate 2048x1 input otherwise raised
|
||||
# ValueError on the broadcast). Normal images are unaffected.
|
||||
gw = min(w, max(1, int(_ALPHA_WIDTH_FRAC * w)))
|
||||
gh = min(h, max(1, int(_ALPHA_HEIGHT_FRAC * w)))
|
||||
ax = max(0, w - int(_ALPHA_MARGIN_RIGHT_FRAC * w) - gw)
|
||||
ay = max(0, h - int(_ALPHA_MARGIN_BOTTOM_FRAC * w) - gh)
|
||||
amap = np.zeros((h, w), np.float32)
|
||||
@@ -353,6 +365,12 @@ class DoubaoEngine:
|
||||
inpaint there costs nothing and reliably clears the mark).
|
||||
Call only when :meth:`reverse_alpha_available` and the mark is detected.
|
||||
"""
|
||||
# Normalize to 3-channel BGR so a 2D grayscale or 4-channel BGRA input
|
||||
# does not break the reverse-alpha math (which assumes a 3-channel logo).
|
||||
if image.ndim == 2:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
||||
elif image.shape[2] == 4:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
|
||||
at_native = abs(image.shape[1] / _ALPHA_NATIVE_WIDTH - 1.0) <= _ALPHA_NATIVE_BAND
|
||||
if at_native:
|
||||
amap = self._fixed_alpha_map(image)
|
||||
|
||||
@@ -329,16 +329,21 @@ class GeminiEngine:
|
||||
"""
|
||||
result = image.copy()
|
||||
|
||||
# Handle alpha channel
|
||||
if result.shape[2] == 4:
|
||||
# Normalize to 3-channel BGR up front: 2D grayscale (no channel axis) and
|
||||
# 4-channel BGRA both reach this public entry point and would otherwise
|
||||
# crash on the channel-count checks / downstream 3-channel math.
|
||||
if result.ndim == 2:
|
||||
result = cv2.cvtColor(result, cv2.COLOR_GRAY2BGR)
|
||||
elif result.shape[2] == 4:
|
||||
result = cv2.cvtColor(result, cv2.COLOR_BGRA2BGR)
|
||||
elif result.shape[2] == 1:
|
||||
result = cv2.cvtColor(result, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
size = force_size or get_watermark_size(result.shape[1], result.shape[0])
|
||||
|
||||
# Detect dynamic position & size
|
||||
detection = self.detect_watermark(image, force_size=size)
|
||||
# Detect dynamic position & size (on the normalized 3-channel image so a
|
||||
# grayscale/BGRA input does not crash the detector).
|
||||
detection = self.detect_watermark(result, force_size=size)
|
||||
|
||||
if not detection.detected:
|
||||
logger.debug(
|
||||
|
||||
@@ -36,10 +36,14 @@ def apply_analog_humanizer(image: NDArray, grain_intensity: float = 4.0, chromat
|
||||
b, g, r = cv2.split(image)
|
||||
|
||||
# 1. Chromatic Aberration
|
||||
# Shift R channel left, B channel right
|
||||
# Shift R channel left, B channel right. np.roll is circular, so it wraps
|
||||
# the opposite edge into a thin colored fringe at the L/R borders; replicate
|
||||
# the original edge columns there to keep the intended offset interior-only.
|
||||
if chromatic_shift > 0:
|
||||
r = np.roll(r, -chromatic_shift, axis=1)
|
||||
r[:, -chromatic_shift:] = r[:, -chromatic_shift - 1 : -chromatic_shift]
|
||||
b = np.roll(b, chromatic_shift, axis=1)
|
||||
b[:, :chromatic_shift] = b[:, chromatic_shift : chromatic_shift + 1]
|
||||
|
||||
merged = cv2.merge((b, g, r))
|
||||
|
||||
|
||||
@@ -431,7 +431,7 @@ def identify(image_path: Path, *, check_visible: bool = True, check_invisible: b
|
||||
if synthid:
|
||||
watermarks.append(f"SynthID pixel watermark ({synthid})")
|
||||
caveats.append(_SYNTHID_CAVEAT)
|
||||
if "OpenAI" in (" ".join(issuers) + synthid):
|
||||
if _vendor_of(synthid) == "OpenAI":
|
||||
caveats.append(_OPENAI_CAVEAT)
|
||||
if v := _vendor_of(synthid):
|
||||
ai_vendor_claims["synthid"] = v
|
||||
|
||||
@@ -54,7 +54,8 @@ def imwrite(path: str | Path, img: NDArray[Any]) -> bool:
|
||||
|
||||
The output format is taken from the path extension (e.g. ``.png``), exactly
|
||||
like ``cv2.imwrite``. Returns ``True`` on success, ``False`` if the codec
|
||||
rejects the image.
|
||||
rejects the image or the path cannot be written (matching ``cv2.imwrite``,
|
||||
which returns ``False`` rather than raising on an unwritable path).
|
||||
"""
|
||||
import cv2
|
||||
|
||||
@@ -62,5 +63,8 @@ def imwrite(path: str | Path, img: NDArray[Any]) -> bool:
|
||||
ok, buf = cv2.imencode(ext, img)
|
||||
if not ok:
|
||||
return False
|
||||
buf.tofile(str(path))
|
||||
try:
|
||||
buf.tofile(str(path))
|
||||
except OSError:
|
||||
return False
|
||||
return True
|
||||
|
||||
@@ -73,7 +73,7 @@ class InvisibleEngine:
|
||||
"""
|
||||
|
||||
# SDXL base is the default since May 2026: empirically defeats SynthID v2
|
||||
# at strength=0.05 / steps=50 / native ~1024px. See CLAUDE.md "Known
|
||||
# at strength=0.10 / steps=50 / native ~1024px. See CLAUDE.md "Known
|
||||
# limitations" for the regression evidence ruling out SD-1.5 pipelines.
|
||||
DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
CTRLREGEN_MODEL_ID = "yepengliu/ctrlregen"
|
||||
@@ -227,6 +227,8 @@ class InvisibleEngine:
|
||||
from remove_ai_watermarks import image_io
|
||||
|
||||
out_cv = image_io.imread(out_path, cv2.IMREAD_COLOR)
|
||||
if out_cv is None:
|
||||
return out_path
|
||||
|
||||
if protect_faces and original_faces:
|
||||
if self._progress_callback:
|
||||
|
||||
@@ -190,6 +190,8 @@ def _png_late_metadata(image_path: Path, window: int) -> bytes:
|
||||
with open(image_path, "rb") as f:
|
||||
if f.read(8) != b"\x89PNG\r\n\x1a\n":
|
||||
return b""
|
||||
f.seek(0, 2)
|
||||
file_size = f.tell()
|
||||
pos = 8
|
||||
while True:
|
||||
f.seek(pos)
|
||||
@@ -201,9 +203,12 @@ def _png_late_metadata(image_path: Path, window: int) -> bytes:
|
||||
if chunk_type == b"IEND":
|
||||
break
|
||||
data_start = pos + 8
|
||||
# Clamp the attacker-controlled 32-bit length to the bytes that
|
||||
# actually remain, so a malformed huge length can't allocate GBs.
|
||||
safe_length = max(0, min(length, file_size - data_start))
|
||||
if chunk_type in _PNG_META_CHUNKS and data_start >= window:
|
||||
f.seek(data_start)
|
||||
out += f.read(length)
|
||||
out += f.read(safe_length)
|
||||
pos = data_start + length + 4 # data + CRC
|
||||
except OSError as exc:
|
||||
logger.debug("PNG late-metadata scan failed on %s: %s", image_path, exc)
|
||||
|
||||
@@ -55,6 +55,9 @@ def has_c2pa_metadata(image_path: Path) -> bool:
|
||||
if signature != PNG_SIGNATURE:
|
||||
return False
|
||||
|
||||
file_size = f.seek(0, 2)
|
||||
f.seek(8)
|
||||
|
||||
while True:
|
||||
chunk_header = f.read(8)
|
||||
if len(chunk_header) < 8:
|
||||
@@ -62,9 +65,12 @@ def has_c2pa_metadata(image_path: Path) -> bool:
|
||||
|
||||
length = struct.unpack(">I", chunk_header[:4])[0]
|
||||
chunk_type = chunk_header[4:8]
|
||||
# Clamp the attacker-controlled 32-bit length to the bytes that
|
||||
# actually remain, so a malformed huge length can't allocate GBs.
|
||||
safe_length = max(0, min(length, file_size - f.tell()))
|
||||
|
||||
if chunk_type == C2PA_CHUNK_TYPE:
|
||||
chunk_data = f.read(length)
|
||||
chunk_data = f.read(safe_length)
|
||||
# Check for any C2PA signature
|
||||
for sig in C2PA_SIGNATURES:
|
||||
if sig in chunk_data:
|
||||
@@ -74,7 +80,7 @@ def has_c2pa_metadata(image_path: Path) -> bool:
|
||||
return True
|
||||
f.read(4)
|
||||
else:
|
||||
f.read(length + 4)
|
||||
f.seek(safe_length + 4, 1)
|
||||
|
||||
if chunk_type == b"IEND":
|
||||
break
|
||||
@@ -108,6 +114,9 @@ def extract_c2pa_info(image_path: Path) -> dict[str, Any]:
|
||||
if signature != PNG_SIGNATURE:
|
||||
return c2pa_info
|
||||
|
||||
file_size = f.seek(0, 2)
|
||||
f.seek(8)
|
||||
|
||||
while True:
|
||||
chunk_header = f.read(8)
|
||||
if len(chunk_header) < 8:
|
||||
@@ -115,13 +124,16 @@ def extract_c2pa_info(image_path: Path) -> dict[str, Any]:
|
||||
|
||||
length = struct.unpack(">I", chunk_header[:4])[0]
|
||||
chunk_type = chunk_header[4:8]
|
||||
# Clamp the attacker-controlled 32-bit length to the bytes that
|
||||
# actually remain, so a malformed huge length can't allocate GBs.
|
||||
safe_length = max(0, min(length, file_size - f.tell()))
|
||||
|
||||
if chunk_type == C2PA_CHUNK_TYPE:
|
||||
chunk_data = f.read(length)
|
||||
chunk_data = f.read(safe_length)
|
||||
_parse_c2pa_chunk(chunk_data, c2pa_info)
|
||||
f.read(4)
|
||||
else:
|
||||
f.read(length + 4)
|
||||
f.seek(safe_length + 4, 1)
|
||||
|
||||
if chunk_type == b"IEND":
|
||||
break
|
||||
@@ -278,6 +290,9 @@ def extract_c2pa_chunk(image_path: Path) -> bytes | None:
|
||||
if signature != PNG_SIGNATURE:
|
||||
return None
|
||||
|
||||
file_size = f.seek(0, 2)
|
||||
f.seek(8)
|
||||
|
||||
while True:
|
||||
chunk_header = f.read(8)
|
||||
if len(chunk_header) < 8:
|
||||
@@ -285,9 +300,12 @@ def extract_c2pa_chunk(image_path: Path) -> bytes | None:
|
||||
|
||||
length = struct.unpack(">I", chunk_header[:4])[0]
|
||||
chunk_type = chunk_header[4:8]
|
||||
# Clamp the attacker-controlled 32-bit length to the bytes that
|
||||
# actually remain, so a malformed huge length can't allocate GBs.
|
||||
safe_length = max(0, min(length, file_size - f.tell()))
|
||||
|
||||
if chunk_type == C2PA_CHUNK_TYPE:
|
||||
chunk_data = f.read(length)
|
||||
chunk_data = f.read(safe_length)
|
||||
crc = f.read(4)
|
||||
|
||||
# Check for any C2PA signature
|
||||
@@ -299,7 +317,7 @@ def extract_c2pa_chunk(image_path: Path) -> bytes | None:
|
||||
if b"jumb" in chunk_data.lower() or b"c2pa" in chunk_data.lower():
|
||||
return chunk_header + chunk_data + crc
|
||||
else:
|
||||
f.read(length + 4)
|
||||
f.seek(safe_length + 4, 1)
|
||||
|
||||
if chunk_type == b"IEND":
|
||||
break
|
||||
|
||||
@@ -20,6 +20,8 @@ from PIL import Image
|
||||
|
||||
def tile_positions(total: int, tile: int, overlap: int) -> list[int]:
|
||||
"""Compute evenly-spaced tile start positions covering *total* pixels."""
|
||||
if not (0 <= overlap < tile):
|
||||
raise ValueError(f"overlap must satisfy 0 <= overlap < tile (got overlap={overlap}, tile={tile})")
|
||||
if total <= tile:
|
||||
return [0]
|
||||
n = max(2, math.ceil((total - overlap) / (tile - overlap)))
|
||||
|
||||
@@ -17,6 +17,7 @@ Reference: ISO/IEC 14496-12 (ISOBMFF) and C2PA 2.1 spec §11.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
import struct
|
||||
from typing import TYPE_CHECKING
|
||||
@@ -32,6 +33,8 @@ from remove_ai_watermarks.metadata import (
|
||||
IPTC_AI_MARKERS,
|
||||
)
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Top-level box types that may carry AI provenance. ``uuid`` boxes are checked
|
||||
# against ``C2PA_UUID`` / AI-label markers before being stripped; ``jumb`` boxes
|
||||
# are always stripped (JPEG-XL uses them exclusively for JUMBF).
|
||||
@@ -126,6 +129,8 @@ def scan_c2pa_region(path: str | Path, *, max_total: int = 4 * 1024 * 1024) -> b
|
||||
else:
|
||||
size = size32
|
||||
if size < (payload_off - pos) or pos + size > file_size:
|
||||
# Detection-only: a malformed box halts the walk, so a manifest
|
||||
# placed after it is missed (best-effort scan; no resync).
|
||||
break
|
||||
if box_type in C2PA_BOX_TYPES:
|
||||
f.seek(payload_off)
|
||||
@@ -162,7 +167,9 @@ def strip_c2pa_boxes(data: bytes) -> tuple[bytes, int]:
|
||||
|
||||
out = bytearray()
|
||||
stripped = 0
|
||||
consumed = 0
|
||||
for start, end, box_type, payload_off in _iter_top_level_boxes(data):
|
||||
consumed = end
|
||||
if box_type == b"uuid":
|
||||
# uuid boxes carry the 16-byte UUID immediately after the type.
|
||||
is_c2pa = payload_off + 16 <= end and data[payload_off : payload_off + 16] == C2PA_UUID
|
||||
@@ -174,6 +181,20 @@ def strip_c2pa_boxes(data: bytes) -> tuple[bytes, int]:
|
||||
stripped += 1
|
||||
continue
|
||||
out.extend(data[start:end])
|
||||
|
||||
# Fail-safe: the walker returns early on a malformed box (bad size, or a box
|
||||
# that runs past EOF), so anything after it was never visited. Emitting `out`
|
||||
# would silently truncate the file from the bad box to EOF -- worse than not
|
||||
# stripping. If the walk did not consume the whole input, return it unchanged.
|
||||
if consumed != len(data):
|
||||
log.warning(
|
||||
"ISOBMFF box walk stopped at offset %d of %d (malformed box); "
|
||||
"returning input unchanged to avoid truncation",
|
||||
consumed,
|
||||
len(data),
|
||||
)
|
||||
return data, 0
|
||||
|
||||
return bytes(out), stripped
|
||||
|
||||
|
||||
|
||||
@@ -272,6 +272,12 @@ def _make_seed_generator(device: str, seed: int) -> Any:
|
||||
return torch.Generator().manual_seed(seed) # type: ignore
|
||||
|
||||
|
||||
def _generator_device(generator: Any) -> str:
|
||||
"""Best-effort device type of a ``torch.Generator`` (e.g. ``"cpu"``, ``"mps"``)."""
|
||||
device = getattr(generator, "device", None)
|
||||
return getattr(device, "type", str(device)) if device is not None else "cpu"
|
||||
|
||||
|
||||
# Keep legacy name available for backwards compatibility
|
||||
_detect_model_profile_from_id = detect_model_profile
|
||||
|
||||
@@ -677,6 +683,14 @@ class WatermarkRemover:
|
||||
|
||||
base = self._run_img2img(init_image, strength, num_inference_steps, guidance_scale, generator)
|
||||
|
||||
# The base pass may have fallen back from MPS to CPU (it flips
|
||||
# self.device). The generator was built for the original device, and
|
||||
# diffusers rejects a device-mismatched generator ("Expected a 'cpu'
|
||||
# device generator but found 'mps'"), so drop it for the per-region
|
||||
# passes -- they then seed from the global RNG, which is fine here.
|
||||
if generator is not None and self.device == "cpu" and _generator_device(generator) != "cpu":
|
||||
generator = None
|
||||
|
||||
bgr = cv2.cvtColor(np.array(init_image), cv2.COLOR_RGB2BGR)
|
||||
try:
|
||||
boxes = text_protector.TextProtector().detect_text_boxes(bgr)
|
||||
@@ -718,8 +732,13 @@ class WatermarkRemover:
|
||||
# the composite even though the text is crisp.
|
||||
cg = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY).astype(np.float32)
|
||||
dg = cv2.cvtColor(down, cv2.COLOR_BGR2GRAY).astype(np.float32)
|
||||
(sx, sy), _resp = cv2.phaseCorrelate(cg, dg)
|
||||
if abs(sx) > 0.1 or abs(sy) > 0.1:
|
||||
(sx, sy), resp = cv2.phaseCorrelate(cg, dg)
|
||||
# Only correct for the real 1-2px round-trip shift. On a near-flat /
|
||||
# low-contrast crop phaseCorrelate returns a spurious large offset at
|
||||
# a tiny response (e.g. (19,19) at resp ~0.005); warping by that
|
||||
# garbles the composite -- the exact failure this was meant to
|
||||
# prevent. Gate on both a confident response and a plausible offset.
|
||||
if resp > 0.3 and abs(sx) < 4 and abs(sy) < 4 and (abs(sx) > 0.1 or abs(sy) > 0.1):
|
||||
m = np.float32([[1, 0, -sx], [0, 1, -sy]])
|
||||
down = cv2.warpAffine(down, m, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
|
||||
out_bgr = text_protector.feather_paste(out_bgr, down, x, y)
|
||||
|
||||
@@ -67,8 +67,16 @@ def erase_cv2(
|
||||
method: Literal["telea", "ns"] = "telea",
|
||||
radius: int = 6,
|
||||
) -> NDArray[Any]:
|
||||
"""Inpaint ``mask`` with classical cv2 inpainting (CPU, no extra deps)."""
|
||||
"""Inpaint ``mask`` with classical cv2 inpainting (CPU, no extra deps).
|
||||
|
||||
Accepts 1-/3-channel BGR (passed straight to ``cv2.inpaint``) and 4-channel
|
||||
BGRA: ``cv2.inpaint`` rejects 4 channels, so the alpha plane is split off,
|
||||
the BGR is inpainted, and alpha is re-attached unchanged.
|
||||
"""
|
||||
flag = cv2.INPAINT_TELEA if method == "telea" else cv2.INPAINT_NS
|
||||
if image_bgr.ndim == 3 and image_bgr.shape[2] == 4:
|
||||
bgr = cv2.inpaint(image_bgr[:, :, :3], mask, radius, flag)
|
||||
return np.dstack([bgr, image_bgr[:, :, 3]])
|
||||
return cv2.inpaint(image_bgr, mask, radius, flag)
|
||||
|
||||
|
||||
|
||||
@@ -22,6 +22,7 @@ signal, not proof of AI origin.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -32,7 +33,9 @@ log = logging.getLogger(__name__)
|
||||
# Adobe ships Variant P in production (com.adobe.trustmark.P).
|
||||
_MODEL_TYPE = "P"
|
||||
# Lazily constructed singleton -- model load + first-use download is expensive.
|
||||
# Guarded by a lock so concurrent callers don't double-construct/double-download.
|
||||
_tm: Any = None
|
||||
_tm_lock = threading.Lock()
|
||||
|
||||
|
||||
def is_available() -> bool:
|
||||
@@ -45,9 +48,11 @@ def is_available() -> bool:
|
||||
def _decoder() -> Any:
|
||||
global _tm
|
||||
if _tm is None:
|
||||
from trustmark import TrustMark
|
||||
with _tm_lock:
|
||||
if _tm is None:
|
||||
from trustmark import TrustMark
|
||||
|
||||
_tm = TrustMark(verbose=False, model_type=_MODEL_TYPE)
|
||||
_tm = TrustMark(verbose=False, model_type=_MODEL_TYPE)
|
||||
return _tm
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user