watermark.py

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