object-remove 0.1.0__py3-none-any.whl

object_remove/pipeline/remove_object_job.py

@@ -0,0 +1,154 @@
+"""Pipeline orchestrator — ``RemoveObjectJob``.
+
+Wires Stage A (mask) to B (scheduling) to C (fill, engine tier selection) to
+E (blend) to F (session), with progress reporting and cooperative cancellation.
+
+Engine tier selection heuristic:
+  * tiny holes / low end path to tier 1 (segment aware CPU comparator)
+  * homogeneous repetitive texture to tier 2 (self similar rigid shift), the
+    cheap default fast path
+  * structured / complex holes to tier 3 (multi scale PatchMatch)
+AUTO tries tier 2 first (cheap) and falls back to tier 3 when the rigid shift
+continuity score across the hole boundary is poor.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Callable, List, Optional
+
+import numpy as np
+
+from ..imageutil import as_float_rgb, rgb_to_gray
+from ..mask.connected_components import split_components
+from ..mask.selection_enhancer import SelectionEnhancer
+from ..fillengines import (
+    SegmentAwarePatchFillEngine, SelfSimilarShiftEngine, PatchMatchEngine,
+)
+from ..render.compositor import PatchCompositor, BlendBackend
+from ..session.undo_session_manager import SessionManager
+
+
+class EngineTier(str, Enum):
+    AUTO = "auto"
+    TIER1_COMPARATOR = "tier1"
+    TIER2_SHIFT = "tier2"
+    TIER3_PATCHMATCH = "tier3"
+
+
+ProgressFn = Callable[[str, float], None]
+CancelFn = Callable[[], bool]
+
+
+@dataclass
+class JobResult:
+    image: np.ndarray
+    tier_used: EngineTier
+    cancelled: bool = False
+    stages: List[str] = field(default_factory=list)
+    session_id: Optional[str] = None
+
+
+class RemoveObjectJob:
+    """One object removal edit. Reusable across masks on the same image."""
+
+    LABEL = "Remove objects"
+
+    def __init__(self, image: np.ndarray,
+                 manager: Optional[SessionManager] = None,
+                 blend_backend: BlendBackend = BlendBackend.POISSON2,
+                 tiny_hole_area: int = 64,
+                 shift_continuity_threshold: float = 0.004):
+        self.image = as_float_rgb(image).astype(np.float64)
+        self.manager = manager
+        self.blend_backend = blend_backend
+        self.tiny_hole_area = tiny_hole_area
+        # tier 2 rigid shift is kept only if its border continuity SSD is below
+        # this; otherwise AUTO falls back to tier 3 PatchMatch.
+        self.shift_continuity_threshold = shift_continuity_threshold
+        self._enh = SelectionEnhancer()
+
+    def run(self, mask: np.ndarray, tier: EngineTier = EngineTier.AUTO,
+            refine_mask: bool = True, progress: Optional[ProgressFn] = None,
+            cancel: Optional[CancelFn] = None) -> JobResult:
+        stages: List[str] = []
+        last_frac = [0.0]  # progress must never decrease, even across a
+        #                    tier 2 to tier 3 fallback that restarts fill at 0.
+
+        def report(stage: str, frac: float):
+            stages.append(stage) if stage not in stages else None
+            if progress is not None:
+                clamped_frac = max(last_frac[0], float(np.clip(frac, 0.0, 1.0)))
+                last_frac[0] = clamped_frac
+                progress(stage, clamped_frac)
+
+        def cancelled() -> bool:
+            return cancel is not None and cancel()
+
+        # Stage A: mask refinement
+        report("mask", 0.0)
+        hole = np.ascontiguousarray(mask > 0)
+        if refine_mask:
+            # Smooth jagged edges, but never drop user selected pixels — eroding
+            # the mask would leave slivers of the object unremoved. Union back
+            # with the original selection so refinement can only grow coverage.
+            hole = self._enh.denoise_and_smooth(hole, iterations=1) | (mask > 0)
+        if not hole.any():
+            return JobResult(self.image.copy(), tier, stages=stages)
+
+        # Stage F (open): session
+        session = None
+        if self.manager is not None:
+            session = self.manager.new_session()
+            session.record(hole, note=self.LABEL)
+
+        if cancelled():
+            return JobResult(self.image.copy(), tier, cancelled=True, stages=stages)
+
+        # Stage C: pick tier & fill
+        report("fill", 0.05)
+        filled, chosen = self._run_engine(tier, hole,
+                                          lambda f: report("fill", 0.05 + 0.8 * f),
+                                          cancelled)
+        if cancelled():
+            return JobResult(self.image.copy(), chosen, cancelled=True, stages=stages)
+
+        # Stage E: blend
+        report("blend", 0.9)
+        backend = self._backend_for_tier(chosen)
+        compositor = PatchCompositor(backend, bands=3, feather_radius=2.0)
+        out = compositor.composite(filled, self.image, hole)
+
+        report("done", 1.0)
+        return JobResult(out, chosen, stages=stages,
+                         session_id=session.session_id if session else None)
+
+    def _run_engine(self, tier: EngineTier, hole: np.ndarray,
+                    progress, cancelled):
+        """Run a fill engine. For AUTO: truly tiny holes go to tier 1; otherwise
+        try the cheap tier 2 rigid shift first and fall back to tier 3
+        PatchMatch when its boundary continuity is poor."""
+        area = int(hole.sum())
+
+        if tier == EngineTier.TIER1_COMPARATOR or (
+                tier == EngineTier.AUTO and area <= self.tiny_hole_area):
+            tier1_engine = SegmentAwarePatchFillEngine(patch_radius=3)
+            return tier1_engine.fill(self.image, hole, progress=progress,
+                                     cancel=cancelled), EngineTier.TIER1_COMPARATOR
+
+        if tier in (EngineTier.TIER2_SHIFT, EngineTier.AUTO):
+            tier2_engine = SelfSimilarShiftEngine(margin=8)
+            out = tier2_engine.fill(self.image, hole, progress=progress, cancel=cancelled)
+            good = tier2_engine.last_border_score <= self.shift_continuity_threshold
+            if tier == EngineTier.TIER2_SHIFT or good:
+                return out, EngineTier.TIER2_SHIFT
+            # AUTO + poor continuity, fall back to the quality tier
+        tier3_engine = PatchMatchEngine(levels=4)
+        return tier3_engine.fill(self.image, hole, progress=progress,
+                                 cancel=cancelled), EngineTier.TIER3_PATCHMATCH
+
+    def _backend_for_tier(self, tier: EngineTier) -> BlendBackend:
+        return self.blend_backend

object_remove/pyramid/__init__.py

@@ -0,0 +1,15 @@
+"""Stage B/D — multi scale (pyramid) and tile/slice scheduling.
+
+Two independent scalability knobs: a coarse to fine pyramid for the
+search/solve step, and tile/slice rendering for the compositing step on very
+large images.
+"""
+
+from .image_pyramid import ImagePyramid, gaussian_blur, downsample, upsample
+from .scale_scheduler import ScaleScheduler
+from .tile_scheduler import TileScheduler, Tile
+
+__all__ = [
+    "ImagePyramid", "gaussian_blur", "downsample", "upsample",
+    "ScaleScheduler", "TileScheduler", "Tile",
+]

object_remove/pyramid/image_pyramid.py

@@ -0,0 +1,101 @@
+"""Gaussian/Laplacian image pyramids.
+
+Used by two independent consumers:
+
+* the coarse to fine *search/solve* (fill engines), and
+* the *blend* stage (Burt-Adelson Laplacian blending in the multi band blender).
+
+The downsample blur uses a fixed **4 tap separable kernel** shared with the
+multi band blender so the pyramid math is consistent everywhere it's used:
+weights ``[0.15439, 0.15133, 0.14252, 0.12896]`` applied symmetrically as a
+7 tap ``[w3, w2, w1, w0, w1, w2, w3]`` kernel.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import numpy as np
+
+# Half kernel shared with MultiBandBlender.gaussian_blur.
+_HALF = np.array([0.15439, 0.15133, 0.14252, 0.12896], dtype=np.float64)
+# Full symmetric 7 tap kernel (renormalised so it sums to 1).
+_KERNEL = np.concatenate([_HALF[:0:-1], _HALF])
+_KERNEL = (_KERNEL / _KERNEL.sum()).astype(np.float64)
+
+
+def _convolve_sep(img: np.ndarray, kernel: np.ndarray) -> np.ndarray:
+    """Separable convolution with edge padding. Works on 2D or 3D (H,W[,C])."""
+    k = kernel
+    r = len(k) // 2
+    img_f = img.astype(np.float64)
+    if img_f.ndim == 2:
+        img_f = img_f[:, :, None]
+        squeeze = True
+    else:
+        squeeze = False
+    # horizontal
+    padded = np.pad(img_f, ((0, 0), (r, r), (0, 0)), mode="edge")
+    horiz_out = np.zeros_like(img_f)
+    for i, w in enumerate(k):
+        horiz_out += w * padded[:, i:i + img_f.shape[1], :]
+    # vertical
+    padded = np.pad(horiz_out, ((r, r), (0, 0), (0, 0)), mode="edge")
+    result = np.zeros_like(img_f)
+    for i, w in enumerate(k):
+        result += w * padded[i:i + img_f.shape[0], :, :]
+    return result[:, :, 0] if squeeze else result
+
+
+def gaussian_blur(img: np.ndarray) -> np.ndarray:
+    """Blur with the exact binary kernel (no sigma parameter, by design)."""
+    return _convolve_sep(img, _KERNEL)
+
+
+def downsample(img: np.ndarray) -> np.ndarray:
+    """Blur then drop every other pixel (Burt-Adelson REDUCE)."""
+    return gaussian_blur(img)[::2, ::2]
+
+
+def upsample(img: np.ndarray, out_shape: tuple[int, int]) -> np.ndarray:
+    """Insert zeros then blur (Burt-Adelson EXPAND), resized to ``out_shape``."""
+    oh, ow = out_shape
+    if img.ndim == 3:
+        up = np.zeros((oh, ow, img.shape[2]), dtype=np.float64)
+    else:
+        up = np.zeros((oh, ow), dtype=np.float64)
+    up[::2, ::2] = img[: (oh + 1) // 2, : (ow + 1) // 2]
+    # weight *4 compensates the energy lost by zero insertion (2x per axis)
+    return _convolve_sep(up, _KERNEL) * 4.0
+
+
+class ImagePyramid:
+    """Gaussian pyramid + on demand Laplacian bands."""
+
+    def __init__(self, image: np.ndarray, levels: int):
+        self.levels = max(1, levels)
+        self.gaussian: List[np.ndarray] = [image.astype(np.float64)]
+        for _ in range(self.levels - 1):
+            nxt = downsample(self.gaussian[-1])
+            if min(nxt.shape[:2]) < 2:
+                break
+            self.gaussian.append(nxt)
+        self.levels = len(self.gaussian)
+
+    def laplacian(self) -> List[np.ndarray]:
+        """Laplacian bands: L[i] = G[i] - EXPAND(G[i+1]); top level is the residual."""
+        bands: List[np.ndarray] = []
+        for i in range(len(self.gaussian) - 1):
+            g = self.gaussian[i]
+            up = upsample(self.gaussian[i + 1], g.shape[:2])
+            bands.append(g - up)
+        bands.append(self.gaussian[-1])
+        return bands
+
+    @staticmethod
+    def collapse(bands: List[np.ndarray]) -> np.ndarray:
+        """Reconstruct an image from Laplacian bands (upscale + add back up)."""
+        img = bands[-1]
+        for i in range(len(bands) - 2, -1, -1):
+            img = bands[i] + upsample(img, bands[i].shape[:2])
+        return img

object_remove/pyramid/scale_scheduler.py

@@ -0,0 +1,56 @@
+"""Coarse to fine scheduling.
+
+Decides how many pyramid levels a solve should use from the hole size, and yields
+levels coarse to fine so a fill engine can solve a small problem first and
+upsample the result as an initialisation for the next finer level.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterator, List
+
+import numpy as np
+
+from .image_pyramid import ImagePyramid, upsample
+
+
+@dataclass
+class ScaleScheduler:
+    max_levels: int = 5
+    min_coarse_dim: int = 16  # stop coarsening when the smaller image side hits this
+
+    def choose_levels(self, image_shape: tuple[int, int],
+                      hole_bbox: tuple[int, int, int, int] | None = None) -> int:
+        """Pick a level count. Bigger holes / images warrant more levels."""
+        h, w = image_shape[:2]
+        smaller = min(h, w)
+        levels = 1
+        dim = smaller
+        while dim > self.min_coarse_dim and levels < self.max_levels:
+            dim //= 2
+            levels += 1
+        if hole_bbox is not None:
+            x0, y0, x1, y1 = hole_bbox
+            hole = max(1, max(x1 - x0, y1 - y0))
+            # ensure the coarsest level shrinks the hole to a handful of pixels
+            levels_needed = 1
+            hole_dim = hole
+            while hole_dim > 8 and levels_needed < self.max_levels:
+                hole_dim //= 2
+                levels_needed += 1
+            levels = min(self.max_levels, max(levels, levels_needed))
+        return max(1, levels)
+
+    def build(self, image: np.ndarray, levels: int) -> ImagePyramid:
+        return ImagePyramid(image, levels)
+
+    @staticmethod
+    def coarse_to_fine(pyramid: ImagePyramid) -> Iterator[int]:
+        """Yield level indices from coarsest to finest."""
+        return iter(range(pyramid.levels - 1, -1, -1))
+
+    @staticmethod
+    def upscale_result(coarse: np.ndarray, fine_shape: tuple[int, int]) -> np.ndarray:
+        """Upsample a coarse level solve to initialise the next finer level."""
+        return upsample(coarse, fine_shape[:2])

object_remove/pyramid/tile_scheduler.py

@@ -0,0 +1,80 @@
+"""Tile/slice rendering.
+
+This is an *independent* knob from the resolution pyramid: it tiles the
+*render* to cap peak memory on very large images, distinct from coarse to fine
+search. It partitions a region into overlapping tiles that a compositor can
+render one at a time and stitch, with the overlap feathered so tile seams don't
+show.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Iterator, List, Tuple
+
+import numpy as np
+
+
+@dataclass
+class Tile:
+    x0: int
+    y0: int
+    x1: int  # exclusive
+    y1: int  # exclusive
+    overlap: int  # halo width included in [x0,x1)/[y0,y1) on each interior side
+
+
+@dataclass
+class TileScheduler:
+    tile_size: int = 512
+    overlap: int = 32
+
+    def plan(self, width: int, height: int) -> List[Tile]:
+        """Partition (width, height) into overlapping tiles."""
+        tiles: List[Tile] = []
+        step = max(1, self.tile_size - self.overlap)
+        ys = list(range(0, height, step))
+        xs = list(range(0, width, step))
+        for y in ys:
+            for x in xs:
+                x0 = max(0, x - self.overlap)
+                y0 = max(0, y - self.overlap)
+                x1 = min(width, x + self.tile_size)
+                y1 = min(height, y + self.tile_size)
+                tiles.append(Tile(x0, y0, x1, y1, self.overlap))
+                if x1 >= width:
+                    break
+            if y1 >= height:
+                break
+        return tiles
+
+    def render_slices(self, width: int, height: int, channels: int,
+                      render_tile: Callable[[Tile], np.ndarray]) -> np.ndarray:
+        """Render each tile and blend overlaps with a linear feather (``renderSlices``).
+
+        ``render_tile(tile)`` must return an array shaped
+        ``(tile.y1-tile.y0, tile.x1-tile.x0, channels)``.
+        """
+        acc = np.zeros((height, width, channels), dtype=np.float64)
+        wacc = np.zeros((height, width, 1), dtype=np.float64)
+        for tile in self.plan(width, height):
+            patch = render_tile(tile)
+            w = self._feather_weights(tile)
+            acc[tile.y0:tile.y1, tile.x0:tile.x1] += patch * w[:, :, None]
+            wacc[tile.y0:tile.y1, tile.x0:tile.x1] += w[:, :, None]
+        wacc[wacc == 0] = 1.0
+        return acc / wacc
+
+    @staticmethod
+    def _feather_weights(tile: Tile) -> np.ndarray:
+        h = tile.y1 - tile.y0
+        w = tile.x1 - tile.x0
+        overlap_px = max(1, tile.overlap)
+        weights_y = np.ones(h)
+        weights_x = np.ones(w)
+        ramp = (np.arange(overlap_px) + 1) / (overlap_px + 1)
+        weights_y[:overlap_px] = np.minimum(weights_y[:overlap_px], ramp)
+        weights_y[-overlap_px:] = np.minimum(weights_y[-overlap_px:], ramp[::-1])
+        weights_x[:overlap_px] = np.minimum(weights_x[:overlap_px], ramp)
+        weights_x[-overlap_px:] = np.minimum(weights_x[-overlap_px:], ramp[::-1])
+        return np.outer(weights_y, weights_x)

object_remove/render/__init__.py

@@ -0,0 +1,16 @@
+"""Stage E — render & blend.
+
+Two interchangeable compositing backends:
+:class:`MultiBandBlender` (fixed kernel, per level mask feather) and
+:class:`PoissonBlend` / :class:`PoissonBlend2` (gradient domain). Selected
+via :class:`PatchCompositor`.
+"""
+
+from .multiband_blender import MultiBandBlender
+from .poisson_blender import PoissonBlend, PoissonBlend2
+from .compositor import PatchCompositor, BlendBackend
+
+__all__ = [
+    "MultiBandBlender", "PoissonBlend", "PoissonBlend2",
+    "PatchCompositor", "BlendBackend",
+]

object_remove/render/compositor.py

@@ -0,0 +1,62 @@
+"""Patch compositor — selects a blend backend and applies it.
+
+Multi band and Poisson blending are interchangeable finishing steps chosen
+per engine (tier 2 defaults to multi band, with an optional Poisson pass;
+tier 3 defaults to Poisson). This compositor exposes that choice.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+
+import numpy as np
+
+from ..imageutil import as_float_rgb
+from ..mask.selection_enhancer import SelectionEnhancer
+from .multiband_blender import MultiBandBlender
+from .poisson_blender import PoissonBlend, PoissonBlend2
+
+
+class BlendBackend(str, Enum):
+    MULTIBAND = "multiband"
+    POISSON = "poisson"
+    POISSON2 = "poisson2"
+
+
+class PatchCompositor:
+    def __init__(self, backend: BlendBackend = BlendBackend.MULTIBAND,
+                 bands: int = 3, feather_radius: float = 2.0,
+                 poisson_sigma: float = 0.0, poisson_multiplier: float = 1.0):
+        self.backend = backend
+        self.bands = bands
+        self.feather_radius = feather_radius
+        self.poisson_sigma = poisson_sigma
+        self.poisson_multiplier = poisson_multiplier
+        self.enhancer = SelectionEnhancer()
+
+    def composite(self, filled: np.ndarray, original: np.ndarray,
+                  hole: np.ndarray) -> np.ndarray:
+        """Blend the ``filled`` result into the ``original`` over the ``hole``."""
+        filled = as_float_rgb(filled).astype(np.float64)
+        original = as_float_rgb(original).astype(np.float64)
+        hole = np.ascontiguousarray(hole > 0)
+
+        # The pixels *under* the hole in ``original`` are the object being removed
+        # (or garbage) — no blend backend may ever read them. Repair the
+        # background with the fill first, so both the Poisson and multi band
+        # passes see a valid target everywhere.
+        safe_bg = original.copy()
+        safe_bg[hole] = filled[hole]
+
+        if self.backend == BlendBackend.MULTIBAND:
+            # feather strictly outward (hole stays fully filled)
+            soft = self.enhancer.soften(hole, radius=self.feather_radius, iterations=1)
+            soft = np.maximum(soft, hole.astype(np.float64))
+            return MultiBandBlender(self.bands).pyramid_blend(
+                filled, safe_bg, soft)
+        if self.backend == BlendBackend.POISSON2:
+            return PoissonBlend2(self.poisson_sigma,
+                                   self.poisson_multiplier).blend(
+                filled, safe_bg, hole)
+        return PoissonBlend(self.poisson_sigma, self.poisson_multiplier,
+                              algo_n=1).blend(filled, safe_bg, hole)

object_remove/render/multiband_blender.py

@@ -0,0 +1,75 @@
+"""Multi band (Laplacian pyramid) blending.
+
+* the per level blur is a **fixed 4 tap separable kernel**, weights
+  ``[0.15439, 0.15133, 0.14252, 0.12896]``;
+* Laplacian bands are formed by downscale to upscale to subtract (textbook
+  Burt-Adelson);
+* per band blending is a **masked lerp** before collapse;
+* the **mask itself is blurred per pyramid level** with the same kernel before
+  being used as blend weights — feathering happens per level, not once at full
+  res. That subtlety is what removes the seam without softening real detail.
+
+Band count is a runtime parameter, defaulted low since the tier 2 caller
+passes a small fixed level count.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from ..pyramid.image_pyramid import ImagePyramid, gaussian_blur, upsample, downsample
+
+
+class MultiBandBlender:
+    def __init__(self, bands: int = 3):
+        self.bands = max(1, bands)
+
+    def pyramid_blend(self, foreground: np.ndarray, background: np.ndarray,
+                      mask: np.ndarray) -> np.ndarray:
+        """Blend ``foreground`` into ``background`` where ``mask`` ~ 1.
+
+        ``mask`` may be soft (float in [0, 1]) or boolean. Returns float RGB.
+        """
+        fg = foreground.astype(np.float64)
+        bg = background.astype(np.float64)
+        m = mask.astype(np.float64)
+        if m.ndim == 2:
+            m = m[:, :, None]
+
+        levels = self.bands
+        fg_lap = ImagePyramid(fg, levels).laplacian()
+        bg_lap = ImagePyramid(bg, levels).laplacian()
+        # Gaussian pyramid of the mask (per level blur = per level feather)
+        mask_gauss = ImagePyramid(m, levels).gaussian
+
+        blended_bands = []
+        for fg_band, bg_band, mask_band in zip(fg_lap, bg_lap, mask_gauss):
+            mask_band_resized = (mask_band if mask_band.shape[:2] == fg_band.shape[:2]
+                                 else _resize(mask_band, fg_band.shape[:2]))
+            blended_bands.append(self._blend(fg_band, bg_band, mask_band_resized))
+        out = ImagePyramid.collapse(blended_bands)
+        return np.clip(out, 0.0, 1.0)
+
+    @staticmethod
+    def _blend(fg_band: np.ndarray, bg_band: np.ndarray, w: np.ndarray) -> np.ndarray:
+        """Masked lerp: w*fg + (1-w)*bg."""
+        return w * fg_band + (1.0 - w) * bg_band
+
+    # exposed as a method too, for direct testing convenience
+    @staticmethod
+    def gaussian_blur(img: np.ndarray) -> np.ndarray:
+        return gaussian_blur(img)
+
+    @staticmethod
+    def mix_frequencies(img: np.ndarray) -> np.ndarray:
+        """One Laplacian band: img - EXPAND(REDUCE(img))."""
+        down = downsample(img)
+        up = upsample(down, img.shape[:2])
+        return img - up
+
+
+def _resize(m: np.ndarray, shape) -> np.ndarray:
+    h, w = shape
+    ys = np.linspace(0, m.shape[0] - 1, h).astype(int)
+    xs = np.linspace(0, m.shape[1] - 1, w).astype(int)
+    return m[np.ix_(ys, xs)]