object-remove 0.1.0__py3-none-any.whl

object_remove/__init__.py

@@ -0,0 +1,37 @@
+"""object_remove — a from scratch, CPU implementation of an on device object
+removal pipeline.
+
+See ``docs/ARCHITECTURE.md`` for the full design: a polygon/scanline mask
+layer, three fill engines of increasing sophistication, two interchangeable
+blend backends, and a locked/queued session persistence layer.
+
+Public entry point::
+
+    from object_remove import RemoveObjectJob, EngineTier
+    job = RemoveObjectJob(image_rgb)
+    result = job.run(mask, tier=EngineTier.AUTO)
+"""
+
+import importlib
+from typing import Any
+
+__version__ = "0.1.0"
+
+# Lazy attribute map: name maps to (submodule, attribute). Loaded on first access so
+# importing any subpackage never eagerly drags in the whole tree.
+_LAZY = {
+    "RemoveObjectJob": ("object_remove.pipeline.remove_object_job", "RemoveObjectJob"),
+    "EngineTier": ("object_remove.pipeline.remove_object_job", "EngineTier"),
+    "JobResult": ("object_remove.pipeline.remove_object_job", "JobResult"),
+    "PolygonSelection": ("object_remove.mask.polygon_selection", "PolygonSelection"),
+    "BrushRasterizer": ("object_remove.mask.brush_rasterizer", "BrushRasterizer"),
+}
+
+__all__ = list(_LAZY)
+
+
+def __getattr__(name: str) -> Any:
+    if name in _LAZY:
+        module_name, attr = _LAZY[name]
+        return getattr(importlib.import_module(module_name), attr)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

object_remove/fillengines/__init__.py

@@ -0,0 +1,20 @@
+"""Stage C — the three fill engine tiers.
+
+* Tier 1 :class:`SegmentAwarePatchFillEngine` — legacy CPU segment aware SSD.
+* Tier 2 :class:`SelfSimilarShiftEngine` — ORB+BRIEF rigid shift + DP seam carve
+  (the fast default path for repetitive texture).
+* Tier 3 :class:`PatchMatchEngine` — multi scale, multi group PatchMatch (quality).
+
+Every engine exposes the same ``fill(image, mask, progress=None, cancel=None)``
+signature so the pipeline can pick a tier without special casing.
+"""
+
+from .patch_comparator import SegmentAwarePatchFillEngine
+from .self_similar_shift import SelfSimilarShiftEngine
+from .patchmatch import PatchMatchEngine
+
+__all__ = [
+    "SegmentAwarePatchFillEngine",
+    "SelfSimilarShiftEngine",
+    "PatchMatchEngine",
+]

object_remove/fillengines/patch_comparator.py

@@ -0,0 +1,182 @@
+"""Tier 1 — legacy CPU segment aware patch comparator.
+
+Patch distance is SSD style but computed *within a segmentation zone
+context*: candidates are penalised when they don't belong to the same segment
+as the query, and validity is checked by scanning a per pixel byte mask rather
+than a statistical flatness heuristic.
+
+The engine on top is a greedy onion peel exemplar fill: cheap, correct on small
+holes, and the natural CPU fallback tier for low end devices.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import numpy as np
+
+from ..imageutil import as_float_rgb, rgb_to_gray
+
+
+class SegmentValidityMap:
+    """Validity + segment zone bookkeeping alongside the color image."""
+
+    def __init__(self, image: np.ndarray, valid: np.ndarray,
+                 segments: Optional[np.ndarray] = None):
+        self.image = image
+        self.valid = valid  # bool (H,W): True = usable source pixel
+        if segments is None:
+            segments = self._quantize_segments(image)
+        self.segments = segments
+
+    @staticmethod
+    def _quantize_segments(image: np.ndarray, n_zones: int = 6) -> np.ndarray:
+        """Cheap segmentation zones: coarse luma quantisation.
+
+        Stands in for a real segmentation map; good enough to make the
+        same segment gate meaningful (grass vs sky vs wall).
+        """
+        luma = rgb_to_gray(image)
+        return np.clip((luma * n_zones).astype(np.int32), 0, n_zones - 1)
+
+    def is_rect_opaque_safe(self, x: int, y: int, w: int, h: int) -> bool:
+        """Bounds check then scan the byte mask; false on the first invalid pixel."""
+        H, W = self.valid.shape
+        if x < 0 or y < 0 or x + w > W or y + h > H:
+            return False
+        return bool(self.valid[y:y + h, x:x + w].all())
+
+    def is_same_segment(self, y0: int, x0: int, y1: int, x1: int) -> bool:
+        return self.segments[y0, x0] == self.segments[y1, x1]
+
+
+@dataclass
+class PatchSegmentationComparator:
+    patch_radius: int = 3
+    segment_penalty: float = 0.25  # added to distance when segments differ
+
+    def patch_sqr_distance(self, seg: SegmentValidityMap,
+                           target_y: int, target_x: int,
+                           source_y: int, source_x: int,
+                           target_valid: np.ndarray) -> float:
+        """Segment aware masked SSD between a target and a source patch center.
+
+        Only pixels valid in *both* the target neighbourhood (``target_valid``)
+        and the source are compared, matching the byte mask validity discipline.
+        """
+        r = self.patch_radius
+        H, W = seg.valid.shape
+        img = seg.image
+        acc = 0.0
+        valid_pair_count = 0
+        for dy in range(-r, r + 1):
+            ty, sy = target_y + dy, source_y + dy
+            if not (0 <= ty < H and 0 <= sy < H):
+                continue
+            for dx in range(-r, r + 1):
+                tx, sx = target_x + dx, source_x + dx
+                if not (0 <= tx < W and 0 <= sx < W):
+                    continue
+                if not target_valid[ty, tx] or not seg.valid[sy, sx]:
+                    continue
+                diff = img[ty, tx] - img[sy, sx]
+                acc += float(diff.dot(diff))
+                valid_pair_count += 1
+        if valid_pair_count == 0:
+            return float("inf")
+        dist = acc / valid_pair_count
+        if not seg.is_same_segment(target_y, target_x, source_y, source_x):
+            dist += self.segment_penalty
+        return dist
+
+
+class SegmentAwarePatchFillEngine:
+    """Greedy onion peel exemplar fill using the comparator (tier 1)."""
+
+    def __init__(self, patch_radius: int = 3, search_stride: int = 2,
+                 max_sources: int = 500, seed: int = 7):
+        self.comparator = PatchSegmentationComparator(patch_radius=patch_radius)
+        self.patch_radius = patch_radius
+        self.search_stride = search_stride
+        self.max_sources = max_sources
+        self.seed = seed
+
+    def fill(self, image: np.ndarray, hole: np.ndarray,
+             progress=None, cancel=None) -> np.ndarray:
+        img = as_float_rgb(image).astype(np.float64).copy()
+        hole = np.ascontiguousarray(hole > 0)
+        H, W = hole.shape
+        valid = ~hole
+        seg = SegmentValidityMap(img, valid.copy())
+
+        src_ys, src_xs = np.nonzero(valid)
+        src_pts = np.stack([src_ys, src_xs], axis=1)
+        src_pts = src_pts[::max(1, self.search_stride)]
+        if len(src_pts) > self.max_sources:
+            rng = np.random.default_rng(self.seed)
+            sel = rng.choice(len(src_pts), self.max_sources, replace=False)
+            src_pts = src_pts[sel]
+
+        remaining = int(hole.sum())
+        total = max(1, remaining)
+        filled_valid = valid.copy()  # grows as we fill
+
+        while remaining > 0:
+            if cancel is not None and cancel():
+                break
+            border = self._boundary_pixels(hole)
+            if border.size == 0:
+                break
+            for (target_y, target_x) in border:
+                if not hole[target_y, target_x]:
+                    continue
+                best = self._best_source(seg, filled_valid, target_y, target_x, src_pts)
+                if best is None:
+                    img[target_y, target_x] = self._local_mean(img, filled_valid, target_y, target_x)
+                else:
+                    source_y, source_x = best
+                    img[target_y, target_x] = img[source_y, source_x]
+                hole[target_y, target_x] = False
+                filled_valid[target_y, target_x] = True
+                seg.valid[target_y, target_x] = True
+                remaining -= 1
+                if progress is not None and remaining % 64 == 0:
+                    progress((total - remaining) / total)
+        if progress is not None:
+            progress(1.0)
+        return img
+
+    def _best_source(self, seg, target_valid, target_y, target_x, src_pts):
+        best_d = float("inf")
+        best = None
+        for (source_y, source_x) in src_pts:
+            d = self.comparator.patch_sqr_distance(
+                seg, target_y, target_x, source_y, source_x, target_valid)
+            if d < best_d:
+                best_d = d
+                best = (int(source_y), int(source_x))
+        return best
+
+    @staticmethod
+    def _boundary_pixels(hole: np.ndarray) -> np.ndarray:
+        """Hole pixels adjacent to a filled pixel (the fill front)."""
+        h = hole
+        front = np.zeros_like(h)
+        front[1:, :] |= h[1:, :] & ~h[:-1, :]
+        front[:-1, :] |= h[:-1, :] & ~h[1:, :]
+        front[:, 1:] |= h[:, 1:] & ~h[:, :-1]
+        front[:, :-1] |= h[:, :-1] & ~h[:, 1:]
+        ys, xs = np.nonzero(front)
+        return np.stack([ys, xs], axis=1)
+
+    @staticmethod
+    def _local_mean(img, valid, target_y, target_x, r=2):
+        H, W = valid.shape
+        y0, y1 = max(0, target_y - r), min(H, target_y + r + 1)
+        x0, x1 = max(0, target_x - r), min(W, target_x + r + 1)
+        block = img[y0:y1, x0:x1]
+        vmask = valid[y0:y1, x0:x1]
+        if vmask.any():
+            return block[vmask].mean(axis=0)
+        return img[target_y, target_x]

object_remove/fillengines/patchmatch/__init__.py

@@ -0,0 +1,41 @@
+"""Tier 3 — CPU PatchMatch fill engine.
+
+Get classic CPU PatchMatch correct first; GPU resident nearest neighbour field
+storage and learned feature preprocessing are natural follow on upgrades once
+the CPU algorithm is solid. This package is that CPU stage.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .solver import PatchMatchSolver
+from .group_solver import GroupSolver, multiscale_fill
+from .patch import Patch, NearestNeighborField
+from .patch_scorer import PatchScorer
+from .patch_search_index import PatchSearchIndex
+
+__all__ = [
+    "PatchMatchEngine", "PatchMatchSolver", "GroupSolver", "multiscale_fill",
+    "Patch", "NearestNeighborField", "PatchScorer", "PatchSearchIndex",
+]
+
+
+class PatchMatchEngine:
+    """Public tier 3 engine: multi scale, multi group PatchMatch inpainting."""
+
+    def __init__(self, patch_radius: int = 3, pm_iters: int = 4,
+                 em_iters: int = 4, levels: int = 4, seed: int = 1234,
+                 isolate_groups: bool = True):
+        self.solver = PatchMatchSolver(radius=patch_radius, pm_iters=pm_iters,
+                                       em_iters=em_iters, seed=seed)
+        self.levels = levels
+        self.isolate_groups = isolate_groups
+
+    def fill(self, image: np.ndarray, mask: np.ndarray,
+             progress=None, cancel=None) -> np.ndarray:
+        if self.isolate_groups:
+            return GroupSolver(self.solver, self.levels).fill(
+                image, mask, progress=progress, cancel=cancel)
+        return multiscale_fill(image, mask > 0, mask > 0, self.solver,
+                               self.levels, progress=progress, cancel=cancel)

object_remove/fillengines/patchmatch/group_solver.py

@@ -0,0 +1,127 @@
+"""Multi scale driver + multi group isolation.
+
+* Coarse to fine: solve a tiny problem on the coarsest pyramid level, upsample
+  the fill as the initialisation for the next finer level, repeat — this is what
+  lets PatchMatch propagate structure across a large hole.
+* Multi group: each connected component is solved with the *other* components
+  held as invalid (unusable as source) so one object's pixels never leak into
+  another's fill — exactly the ``formQuickSelections`` /
+  ``separateObjectsFromSelection`` setup from Stage A.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from ...imageutil import as_float_rgb
+from ...pyramid.image_pyramid import ImagePyramid, upsample, downsample
+from ..patch_comparator import SegmentValidityMap  # noqa: F401 (kept for symmetry)
+from .solver import PatchMatchSolver
+
+
+def _hole_pyramid(hole: np.ndarray, levels: int):
+    pyr = [hole.astype(bool)]
+    level_hole = hole.astype(np.float64)
+    for _ in range(levels - 1):
+        level_hole = downsample(level_hole)
+        pyr.append(level_hole > 0.4)
+        if min(level_hole.shape) < 2:
+            break
+    return pyr
+
+
+def _pad(array: np.ndarray, r: int, value=0.0):
+    if array.ndim == 3:
+        return np.pad(array, ((r, r), (r, r), (0, 0)), mode="edge")
+    return np.pad(array, ((r, r), (r, r)), mode="constant", constant_values=value)
+
+
+def multiscale_fill(image: np.ndarray, target_hole: np.ndarray,
+                    invalid_hole: np.ndarray, solver: PatchMatchSolver,
+                    levels: int, progress=None, cancel=None) -> np.ndarray:
+    """Fill ``target_hole`` coarse to fine; ``invalid_hole`` = pixels no source may use."""
+    r = solver.radius
+    img = as_float_rgb(image).astype(np.float64)
+    img_pyr = ImagePyramid(img, levels)
+    L = img_pyr.levels
+    target_hole_pyr = _hole_pyramid(target_hole, L)
+    invalid_pyr = _hole_pyramid(invalid_hole, L)
+
+    level_estimate = None
+    for li in range(L - 1, -1, -1):
+        if cancel is not None and cancel():
+            break
+        level_img = img_pyr.gaussian[li].copy()
+        target_hole_lvl = target_hole_pyr[li]
+        invalid_lvl = invalid_pyr[li]
+        if target_hole_lvl.shape != level_img.shape[:2]:
+            target_hole_lvl = _resize_mask(target_hole_lvl, level_img.shape[:2])
+            invalid_lvl = _resize_mask(invalid_lvl, level_img.shape[:2])
+
+        if level_estimate is None:
+            valid = ~invalid_lvl
+            seed = level_img[valid].mean(axis=0) if valid.any() else np.zeros(3)
+            level_img[invalid_lvl] = seed
+        else:
+            upsampled_estimate = upsample(level_estimate, level_img.shape[:2])
+            level_img[invalid_lvl] = upsampled_estimate[invalid_lvl]
+
+        padded_img = _pad(level_img, r)
+        padded_hole = _pad(target_hole_lvl, r, value=False).astype(bool)
+        padded_invalid = _pad(invalid_lvl, r, value=False).astype(bool)
+
+        def level_progress(p, li=li):
+            if progress is not None:
+                done = (L - 1 - li) + p
+                progress(min(1.0, done / L))
+
+        filled = solver.solve_level(padded_img, padded_hole, padded_invalid,
+                                    progress=level_progress, cancel=cancel)
+        level_estimate = filled[r:-r, r:-r] if r > 0 else filled
+
+    if progress is not None:
+        progress(1.0)
+    out = as_float_rgb(image).astype(np.float64).copy()
+    hole_mask = target_hole > 0
+    out[hole_mask] = level_estimate[hole_mask]
+    return out
+
+
+def _resize_mask(mask: np.ndarray, shape) -> np.ndarray:
+    h, w = shape
+    ys = (np.linspace(0, mask.shape[0] - 1, h)).astype(int)
+    xs = (np.linspace(0, mask.shape[1] - 1, w)).astype(int)
+    return mask[np.ix_(ys, xs)]
+
+
+class GroupSolver:
+    """Solve each connected component independently, others held invalid."""
+
+    def __init__(self, solver: PatchMatchSolver, levels: int = 4):
+        self.solver = solver
+        self.levels = levels
+
+    def fill(self, image: np.ndarray, mask: np.ndarray, components=None,
+             progress=None, cancel=None) -> np.ndarray:
+        from ...mask.connected_components import split_components
+
+        img = as_float_rgb(image).astype(np.float64).copy()
+        full_hole = mask > 0
+        if components is None:
+            components = split_components(full_hole)
+        if not components:
+            return img
+        n = len(components)
+        for i, comp in enumerate(components):
+            if cancel is not None and cancel():
+                break
+
+            def comp_progress(p, i=i):
+                if progress is not None:
+                    progress((i + p) / n)
+
+            img = multiscale_fill(img, comp, full_hole, self.solver,
+                                  self.levels, progress=comp_progress, cancel=cancel)
+            # once a component is filled it becomes valid source for later ones
+            full_hole = full_hole & ~comp
+        return img

object_remove/fillengines/patchmatch/patch.py

@@ -0,0 +1,48 @@
+"""Patch + nearest neighbour field data structures.
+
+CPU first, for clarity and to keep the dependency footprint small. A GPU
+version would store this same field in a render target texture instead of a
+numpy array; the algorithm is identical.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class Patch:
+    """A square patch defined by its center and radius (in padded coords)."""
+
+    y: int
+    x: int
+    radius: int
+
+    @property
+    def size(self) -> int:
+        return 2 * self.radius + 1
+
+
+class NearestNeighborField:
+    """For each target pixel, the source patch center it currently maps to.
+
+    ``coords[y, x] = (sy, sx)`` source center; ``dist[y, x]`` = patch distance.
+    All coordinates are in the *padded* frame used by the solver.
+    """
+
+    def __init__(self, height: int, width: int):
+        self.height = height
+        self.width = width
+        self.coords = np.zeros((height, width, 2), dtype=np.int32)
+        self.dist = np.full((height, width), np.inf, dtype=np.float64)
+
+    def set(self, y: int, x: int, sy: int, sx: int, d: float) -> None:
+        self.coords[y, x, 0] = sy
+        self.coords[y, x, 1] = sx
+        self.dist[y, x] = d
+
+    def source_of(self, y: int, x: int) -> tuple[int, int]:
+        c = self.coords[y, x]
+        return int(c[0]), int(c[1])

object_remove/fillengines/patchmatch/patch_scorer.py

@@ -0,0 +1,38 @@
+"""Patch distance scoring.
+
+Classic SSD on the current image estimate. A higher quality version would
+replace this with a *learned* feature distance computed once up front: that
+is a drop in swap here, precompute a feature image and point
+:class:`PatchScorer` at it instead of the raw RGB. The search algorithm
+doesn't change, which is exactly why learned features are a natural follow on
+upgrade rather than a redesign.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+class PatchScorer:
+    """Sum of squared differences over a fixed square window in padded coords."""
+
+    def __init__(self, radius: int = 3):
+        self.radius = radius
+
+    def distance(self, feat: np.ndarray, ty: int, tx: int,
+                 sy: int, sx: int) -> float:
+        r = self.radius
+        a = feat[ty - r:ty + r + 1, tx - r:tx + r + 1]
+        b = feat[sy - r:sy + r + 1, sx - r:sx + r + 1]
+        d = a - b
+        return float(np.einsum("ijk,ijk->", d, d))
+
+    def distance_weighted(self, feat: np.ndarray, weight: np.ndarray,
+                          ty: int, tx: int, sy: int, sx: int) -> float:
+        """SSD weighted per pixel (used to down weight still unknown pixels)."""
+        r = self.radius
+        a = feat[ty - r:ty + r + 1, tx - r:tx + r + 1]
+        b = feat[sy - r:sy + r + 1, sx - r:sx + r + 1]
+        w = weight[ty - r:ty + r + 1, tx - r:tx + r + 1]
+        d = (a - b)
+        return float(np.einsum("ij,ijk,ijk->", w, d, d))

object_remove/fillengines/patchmatch/patch_search_index.py

@@ -0,0 +1,56 @@
+"""Valid source bookkeeping + sampling.
+
+A source patch is usable only if its whole window is valid (outside every hole),
+the byte mask scan discipline from tier 1, here precomputed for the entire image
+via a box sum so per candidate validity is an O(1) lookup. Also provides uniform
+random sampling over valid source centers for PatchMatch initialisation and
+random search.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+class PatchSearchIndex:
+    def __init__(self, invalid: np.ndarray, radius: int):
+        """``invalid`` (H, W) bool in *padded* coords: True = cannot be sourced from."""
+        self.radius = radius
+        self.height, self.width = invalid.shape
+        self.source_ok = self._compute_source_ok(invalid, radius)
+        ys, xs = np.nonzero(self.source_ok)
+        self._src_ys = ys
+        self._src_xs = xs
+        self.n_sources = len(ys)
+
+    @staticmethod
+    def _compute_source_ok(invalid: np.ndarray, r: int) -> np.ndarray:
+        """True where a full (2r+1) window contains zero invalid pixels."""
+        H, W = invalid.shape
+        integral = np.zeros((H + 1, W + 1), dtype=np.int64)
+        integral[1:, 1:] = np.cumsum(np.cumsum(invalid.astype(np.int64), 0), 1)
+
+        def box_sum(y0, y1, x0, x1):
+            return (integral[y1, x1] - integral[y0, x1]
+                    - integral[y1, x0] + integral[y0, x0])
+
+        ok = np.zeros((H, W), dtype=bool)
+        for y in range(r, H - r):
+            y0, y1 = y - r, y + r + 1
+            row_ok = np.zeros(W, dtype=bool)
+            for x in range(r, W - r):
+                if box_sum(y0, y1, x - r, x + r + 1) == 0:
+                    row_ok[x] = True
+            ok[y] = row_ok
+        return ok
+
+    def is_source_ok(self, sy: int, sx: int) -> bool:
+        if 0 <= sy < self.height and 0 <= sx < self.width:
+            return bool(self.source_ok[sy, sx])
+        return False
+
+    def sample(self, rng: np.random.Generator, n: int = 1):
+        if self.n_sources == 0:
+            return np.empty((0, 2), dtype=np.int32)
+        idx = rng.integers(0, self.n_sources, size=n)
+        return np.stack([self._src_ys[idx], self._src_xs[idx]], axis=1).astype(np.int32)