from __future__ import annotations import logging import time from dataclasses import dataclass import cv2 import numpy as np from PIL import Image from config import Config logger = logging.getLogger("image_processor") class WatermarkRemovalError(RuntimeError): pass @dataclass class WatermarkResult: image: Image.Image elapsed_seconds: float region: tuple[int, int, int, int] # x, y, w, h in pixels (search area) mask_pixels: int # number of pixels actually inpainted class WatermarkRemover: """Removes a thin text watermark from a fixed bottom-left region. Masking — color-agnostic text detection: 1. Crop a bottom-left search region (fractions of width/height). 2. Grayscale, then morphological top-hat (bright text on any background) + black-hat (dark text on any background) with a kernel slightly larger than the text stroke height. 3. Threshold, morphologically join neighbouring glyphs into text bands, keep bands up to a fraction of the ROI (rejects product/background bleed and tiny noise), then mask only the ink strokes inside them. Inpainting — two engines (``config.watermark_engine``): * ``"classical"`` — ``cv2.inpaint`` (Telea). Fast, no extra deps, but smears on textured backgrounds and can't recover detail. On a heavily textured strip the stroke detector responds almost everywhere, the band is rejected as "too big", and the image is returned unchanged. * ``"lama"`` (default) — the big-lama neural inpainter. Reconstructs texture (brick, fabric, gradients) instead of smearing, so it handles both smooth and busy backgrounds. When stroke detection finds nothing (a textured strip), we fall back to masking the whole search region and let LaMa rebuild it — something cv2.inpaint must NOT do (it would smear), which is why the fallback is engine-specific. Detection is color-agnostic (white, black, or gray text) and does not over-mask uniform background areas. """ def __init__(self, config: Config) -> None: self.config = config self.engine = str(getattr(config, "watermark_engine", "lama")).lower() self._lama = None if self.engine == "lama": try: from lama_inpainter import LamaInpainter self._lama = LamaInpainter(device=str(getattr(config, "watermark_device", "cpu"))) except Exception as exc: # missing torch/model, load failure, ... logger.warning( "LaMa engine unavailable (%s); falling back to classical inpaint.", exc ) self.engine = "classical" def remove(self, image: Image.Image) -> WatermarkResult: start = time.perf_counter() width, height = image.size x, y, w, h = self._resolve_region(width, height) if w <= 0 or h <= 0: return WatermarkResult(image=image, elapsed_seconds=0.0, region=(x, y, w, h), mask_pixels=0) rgb = image.convert("RGB") if image.mode != "RGB" else image rgb_np = np.array(rgb) roi = rgb_np[y : y + h, x : x + w] text_mask_roi = self._detect_text(roi) full_mask = np.zeros((height, width), dtype=np.uint8) full_mask[y : y + h, x : x + w] = text_mask_roi mask_pixels = int(np.count_nonzero(full_mask)) if mask_pixels == 0: # Detection found nothing. For LaMa this is usually a textured strip # that drowns the stroke detector — mask the whole search region and # let the model rebuild it. cv2.inpaint would only smear such a box, # so the classical engine instead returns the image untouched. if self._lama is not None: full_mask[y : y + h, x : x + w] = 255 mask_pixels = int(np.count_nonzero(full_mask)) else: return WatermarkResult(image=image, elapsed_seconds=time.perf_counter() - start, region=(x, y, w, h), mask_pixels=0) feather = max(0, int(self.config.watermark_mask_feather)) if feather > 0: kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (feather * 2 + 1, feather * 2 + 1)) full_mask = cv2.dilate(full_mask, kernel) if self._lama is not None: result_rgb = self._inpaint_lama(rgb, rgb_np, full_mask) else: result_rgb = self._inpaint_classical(rgb_np, full_mask) return WatermarkResult( image=Image.fromarray(result_rgb), elapsed_seconds=time.perf_counter() - start, region=(x, y, w, h), mask_pixels=mask_pixels, ) def _inpaint_classical(self, rgb_np: np.ndarray, full_mask: np.ndarray) -> np.ndarray: bgr = cv2.cvtColor(rgb_np, cv2.COLOR_RGB2BGR) try: inpainted = cv2.inpaint( bgr, full_mask, self.config.watermark_inpaint_radius, cv2.INPAINT_TELEA ) except cv2.error as exc: # pragma: no cover - defensive raise WatermarkRemovalError(f"cv2.inpaint failed: {exc}") from exc return cv2.cvtColor(inpainted, cv2.COLOR_BGR2RGB) def _inpaint_lama(self, rgb: Image.Image, rgb_np: np.ndarray, full_mask: np.ndarray) -> np.ndarray: try: lama_rgb = self._lama.inpaint(rgb, full_mask) except Exception as exc: raise WatermarkRemovalError(f"LaMa inpaint failed: {exc}") from exc # Composite: only the masked pixels come from LaMa; everything else is the # untouched original, so the product can never be altered outside the mask. keep = (full_mask > 0)[:, :, None] return np.where(keep, lama_rgb, rgb_np).astype(np.uint8) def _detect_text(self, roi_rgb: np.ndarray) -> np.ndarray: cfg = self.config gray = cv2.cvtColor(roi_rgb, cv2.COLOR_RGB2GRAY) roi_h, roi_w = gray.shape[:2] # Kernel sized to text stroke height. Slightly larger than stroke so # tophat/blackhat respond to the strokes but not big uniform areas. k = max(3, int(cfg.watermark_text_kernel) | 1) # force odd kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (k, k)) tophat = cv2.morphologyEx(gray, cv2.MORPH_TOPHAT, kernel) # bright on dark blackhat = cv2.morphologyEx(gray, cv2.MORPH_BLACKHAT, kernel) # dark on bright contrast = cv2.max(tophat, blackhat) thresh = max(1, min(255, int(cfg.watermark_text_threshold))) _, binary = cv2.threshold(contrast, thresh, 255, cv2.THRESH_BINARY) # Bridge the gaps between neighbouring glyphs so the SKU string becomes # one (or a few) solid text band(s). Without this the characters are # detected individually and the whole string also tends to merge into a # single blob whose size is unpredictable — making any per-component # area window fragile. Joining first lets us reason about the band, not # the accident of which glyphs happened to touch. bridge = cv2.getStructuringElement(cv2.MORPH_RECT, (k, max(3, (k // 2) | 1))) joined = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, bridge) # Keep text-sized bands. The upper bound is a fraction of the ROI so it # scales with image size and only rejects a blob that fills most of the # search area (real background / product bleed, never a thin watermark). # Note: tophat/blackhat already give ~no response on uniform regions, so # this is a safety net rather than the primary background rejection. num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(joined, connectivity=8) min_area = max(1, int(cfg.watermark_min_component_area)) max_area = max(min_area + 1, int(roi_h * roi_w * _clamp01(cfg.watermark_max_area_frac))) text_bands = np.zeros_like(binary) for label_id in range(1, num_labels): area = stats[label_id, cv2.CC_STAT_AREA] if min_area <= area <= max_area: text_bands[labels == label_id] = 255 # Inpaint only the actual ink strokes that fall inside an accepted band, # NOT the whole (closed) band rectangle. Masking the solid band makes # cv2.inpaint invent a large patch that rarely matches the surrounding # tone — the visible "blur". Restricting to the strokes keeps the real # background between glyphs as reference, so the fill blends in. The # caller dilates this slightly (mask_feather) to swallow anti-aliased # edges around each stroke. return cv2.bitwise_and(binary, text_bands) def _resolve_region(self, width: int, height: int) -> tuple[int, int, int, int]: cfg = self.config x = int(round(width * _clamp01(cfg.watermark_x_frac))) y = int(round(height * _clamp01(cfg.watermark_y_frac))) w = int(round(width * _clamp01(cfg.watermark_w_frac))) h = int(round(height * _clamp01(cfg.watermark_h_frac))) w = max(0, min(w, width - x)) h = max(0, min(h, height - y)) return x, y, w, h def _clamp01(value: float) -> float: if value < 0.0: return 0.0 if value > 1.0: return 1.0 return value