object-remove 0.1.0__py3-none-any.whl

object_remove/mask/brush_rasterizer.py

@@ -0,0 +1,97 @@
+"""Brush stroke to coverage mask.
+
+Converts a polyline of stamp centres (a user's drag) plus a radius into a
+coverage mask by stamping capsule (round cap) segments. Also emits an
+approximating polygon so the result can feed :class:`PolygonSelection` when the
+caller wants the exact polygon representation.
+"""
+
+from __future__ import annotations
+
+from typing import List, Sequence, Tuple
+
+import numpy as np
+
+Point = Tuple[float, float]
+
+
+class BrushRasterizer:
+    def __init__(self, width: int, height: int):
+        self.width = width
+        self.height = height
+
+    def stamp_stroke(self, points: Sequence[Point], radius: float) -> np.ndarray:
+        """Rasterise a brush drag as a coverage mask of round capped segments."""
+        mask = np.zeros((self.height, self.width), dtype=bool)
+        pts = [tuple(map(float, p)) for p in points]
+        if not pts:
+            return mask
+        if len(pts) == 1:
+            self._stamp_disk(mask, pts[0], radius)
+            return mask
+        for start, end in zip(pts[:-1], pts[1:]):
+            self._stamp_capsule(mask, start, end, radius)
+        return mask
+
+    def _stamp_disk(self, mask: np.ndarray, center: Point, r: float) -> None:
+        cx, cy = center
+        x0 = max(0, int(np.floor(cx - r)))
+        x1 = min(self.width - 1, int(np.ceil(cx + r)))
+        y0 = max(0, int(np.floor(cy - r)))
+        y1 = min(self.height - 1, int(np.ceil(cy + r)))
+        if x1 < x0 or y1 < y0:
+            return
+        yy, xx = np.mgrid[y0:y1 + 1, x0:x1 + 1]
+        inside = (xx - cx) ** 2 + (yy - cy) ** 2 <= r * r
+        mask[y0:y1 + 1, x0:x1 + 1] |= inside
+
+    def _stamp_capsule(self, mask: np.ndarray, start: Point, end: Point, r: float) -> None:
+        ax, ay = start
+        bx, by = end
+        x0 = max(0, int(np.floor(min(ax, bx) - r)))
+        x1 = min(self.width - 1, int(np.ceil(max(ax, bx) + r)))
+        y0 = max(0, int(np.floor(min(ay, by) - r)))
+        y1 = min(self.height - 1, int(np.ceil(max(ay, by) + r)))
+        if x1 < x0 or y1 < y0:
+            return
+        yy, xx = np.mgrid[y0:y1 + 1, x0:x1 + 1]
+        dx, dy = bx - ax, by - ay
+        seg_len2 = dx * dx + dy * dy
+        if seg_len2 == 0:
+            self._stamp_disk(mask, start, r)
+            return
+        t = ((xx - ax) * dx + (yy - ay) * dy) / seg_len2
+        t = np.clip(t, 0.0, 1.0)
+        px = ax + t * dx
+        py = ay + t * dy
+        dist2 = (xx - px) ** 2 + (yy - py) ** 2
+        mask[y0:y1 + 1, x0:x1 + 1] |= dist2 <= r * r
+
+    @staticmethod
+    def stroke_to_polygon(points: Sequence[Point], radius: float, segments: int = 8) -> List[Point]:
+        """Approximate a capsule chain as a single closed polygon outline.
+
+        Useful when the caller wants the exact polygon representation to hand to
+        :class:`PolygonSelection` rather than a raster. Builds left/right offset
+        curves plus rounded end caps.
+        """
+        pts = np.asarray(points, dtype=np.float64)
+        if len(pts) == 1:
+            c = pts[0]
+            ang = np.linspace(0, 2 * np.pi, max(8, segments * 4), endpoint=False)
+            return [(float(c[0] + radius * np.cos(a)), float(c[1] + radius * np.sin(a))) for a in ang]
+        left: List[Point] = []
+        right: List[Point] = []
+        for i in range(len(pts) - 1):
+            start, end = pts[i], pts[i + 1]
+            seg_vec = end - start
+            normal = np.array([-seg_vec[1], seg_vec[0]])
+            norm = np.hypot(*normal)
+            if norm == 0:
+                continue
+            normal = normal / norm * radius
+            left.append((start[0] + normal[0], start[1] + normal[1]))
+            left.append((end[0] + normal[0], end[1] + normal[1]))
+            right.append((start[0] - normal[0], start[1] - normal[1]))
+            right.append((end[0] - normal[0], end[1] - normal[1]))
+        return left + right[::-1]

object_remove/mask/connected_components.py

@@ -0,0 +1,110 @@
+"""Connected component labelling.
+
+Labelling itself is a from scratch two pass union find (no
+``scipy.ndimage.label``), and the per region *extraction* step is fanned out
+across a bounded thread pool so splitting a selection into its components
+stays fast on large masks.
+"""
+
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Tuple
+
+import numpy as np
+
+
+class _UnionFind:
+    __slots__ = ("parent",)
+
+    def __init__(self, n: int):
+        self.parent = list(range(n))
+
+    def find(self, x: int) -> int:
+        root = x
+        while self.parent[root] != root:
+            root = self.parent[root]
+        while self.parent[x] != root:
+            self.parent[x], x = root, self.parent[x]
+        return root
+
+    def union(self, a: int, b: int) -> None:
+        ra, rb = self.find(a), self.find(b)
+        if ra != rb:
+            self.parent[max(ra, rb)] = min(ra, rb)
+
+
+def label_components(mask: np.ndarray, connectivity: int = 8) -> Tuple[np.ndarray, int]:
+    """Two pass union find labelling. Returns (labels, count).
+
+    ``labels`` is int32, 0 = background, 1..count = components.
+    """
+    mask = np.ascontiguousarray(mask > 0)
+    h, w = mask.shape
+    labels = np.zeros((h, w), dtype=np.int32)
+    uf = _UnionFind(1)
+    next_label = 1
+
+    for y in range(h):
+        row = mask[y]
+        for x in range(w):
+            if not row[x]:
+                continue
+            neighbours: List[int] = []
+            if x > 0 and mask[y, x - 1]:
+                neighbours.append(labels[y, x - 1])
+            if y > 0 and mask[y - 1, x]:
+                neighbours.append(labels[y - 1, x])
+            if connectivity == 8 and y > 0:
+                if x > 0 and mask[y - 1, x - 1]:
+                    neighbours.append(labels[y - 1, x - 1])
+                if x < w - 1 and mask[y - 1, x + 1]:
+                    neighbours.append(labels[y - 1, x + 1])
+            if not neighbours:
+                labels[y, x] = next_label
+                uf.parent.append(next_label)
+                next_label += 1
+            else:
+                min_label = min(neighbours)
+                labels[y, x] = min_label
+                for nb in neighbours:
+                    uf.union(min_label, nb)
+
+    # second pass: flatten to contiguous labels
+    remap: dict[int, int] = {}
+    count = 0
+    for y in range(h):
+        for x in range(w):
+            lab = labels[y, x]
+            if lab == 0:
+                continue
+            root = uf.find(lab)
+            if root not in remap:
+                count += 1
+                remap[root] = count
+            labels[y, x] = remap[root]
+    return labels, count
+
+
+def split_components(mask: np.ndarray, max_selections: int = 0,
+                     min_area: int = 1, workers: int = 4) -> List[np.ndarray]:
+    """Split a mask into per component boolean masks.
+
+    Extraction is fanned across a bounded worker pool. ``max_selections`` (>0)
+    keeps only the largest N components; ``min_area`` drops specks.
+    """
+    labels, n = label_components(mask)
+    if n == 0:
+        return []
+
+    def extract(lab: int) -> np.ndarray:
+        return labels == lab
+
+    with ThreadPoolExecutor(max_workers=max(1, workers)) as ex:
+        comps = list(ex.map(extract, range(1, n + 1)))
+
+    comps = [c for c in comps if int(c.sum()) >= min_area]
+    comps.sort(key=lambda c: int(c.sum()), reverse=True)
+    if max_selections and max_selections > 0:
+        comps = comps[:max_selections]
+    return comps

object_remove/mask/polygon_selection.py

@@ -0,0 +1,219 @@
+"""Polygon/scanline selection model.
+
+The selection is built from signed polygons, and its boundary geometry can be
+queried as unordered boundary points or ordered boundary cycles. The raster
+mask is a cached derivative, not the primary store.
+
+This class keeps a list of *signed* polygons (positive = add, negative =
+subtract). Scanlines and the raster mask are generated on demand via an
+even odd with sign coverage test evaluated per pixel row (scanline fill).
+Boolean ops (union of brush strokes, subtract of corrections) are therefore
+exact and resolution independent.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Iterable, List, Sequence, Tuple
+
+import numpy as np
+
+Point = Tuple[float, float]
+
+
+@dataclass
+class _Poly:
+    points: np.ndarray  # (N, 2) float, (x, y)
+    add: bool           # True => add region, False => subtract region
+
+
+@dataclass
+class PolygonSelection:
+    """A selection as a set of signed polygons over a fixed canvas size."""
+
+    width: int
+    height: int
+    polys: List[_Poly] = field(default_factory=list)
+
+    def add_poly(self, points: Sequence[Point]) -> "PolygonSelection":
+        """Add a filled polygon to the selection."""
+        self.polys.append(_Poly(np.asarray(points, dtype=np.float64), add=True))
+        return self
+
+    def remove_poly(self, points: Sequence[Point]) -> "PolygonSelection":
+        """Subtract a polygon from the selection."""
+        self.polys.append(_Poly(np.asarray(points, dtype=np.float64), add=False))
+        return self
+
+    @classmethod
+    def from_mask(cls, mask: np.ndarray) -> "PolygonSelection":
+        """Build a selection from a raster mask by tracing its boundary.
+
+        We keep the raster as the source of truth here (caller already has a
+        raster, e.g. from a brush rasteriser) but wrap it so the polygon API is
+        uniform. Boundary tracing is used for :meth:`get_boundary_cycles`.
+        """
+        h, w = mask.shape
+        sel = cls(width=w, height=h)
+        sel._raster_override = (mask > 0)
+        return sel
+
+    _raster_override: np.ndarray | None = field(default=None, repr=False)
+
+    def generate_scanlines(self) -> List[Tuple[int, int, int]]:
+        """Return spans as ``(y, x_start, x_end_exclusive)`` covering the selection."""
+        mask = self.to_mask()
+        spans: List[Tuple[int, int, int]] = []
+        for y in range(self.height):
+            row = mask[y]
+            if not row.any():
+                continue
+            idx = np.flatnonzero(np.diff(np.concatenate(([0], row.view(np.int8), [0]))))
+            for a, b in zip(idx[0::2], idx[1::2]):
+                spans.append((y, int(a), int(b)))
+        return spans
+
+    def to_mask(self) -> np.ndarray:
+        """Rasterise the signed polygon set to a boolean mask (cached derivative)."""
+        if self._raster_override is not None:
+            return self._raster_override
+        mask = np.zeros((self.height, self.width), dtype=bool)
+        for poly in self.polys:
+            filled = _rasterize_polygon(poly.points, self.width, self.height)
+            if poly.add:
+                mask |= filled
+            else:
+                mask &= ~filled
+        return mask
+
+    def union(self, other: "PolygonSelection") -> "PolygonSelection":
+        self.polys.extend(other.polys)
+        return self
+
+    def subtract(self, other: "PolygonSelection") -> "PolygonSelection":
+        for p in other.polys:
+            self.polys.append(_Poly(p.points, add=not p.add))
+        return self
+
+    def get_total_boundary_points(self) -> int:
+        return sum(len(c) for c in self.get_boundary_cycles())
+
+    def get_unordered_boundary_points(self) -> np.ndarray:
+        """All boundary pixels as an (N, 2) array of (x, y), in no particular order."""
+        mask = self.to_mask()
+        eroded = _erode(mask)
+        ys, xs = np.nonzero(mask & ~eroded)
+        return np.stack([xs, ys], axis=1)
+
+    def get_boundary_cycles(self) -> List[np.ndarray]:
+        """Ordered boundary loops, one per connected component."""
+        from .connected_components import label_components
+
+        labels, n = label_components(self.to_mask())
+        cycles: List[np.ndarray] = []
+        for lab in range(1, n + 1):
+            comp = labels == lab
+            cycle = _trace_boundary(comp)
+            if len(cycle):
+                cycles.append(cycle)
+        return cycles
+
+    def is_empty(self) -> bool:
+        return not self.to_mask().any()
+
+    def copy(self) -> "PolygonSelection":
+        new = PolygonSelection(self.width, self.height,
+                               [_Poly(p.points.copy(), p.add) for p in self.polys])
+        if self._raster_override is not None:
+            new._raster_override = self._raster_override.copy()
+        return new
+
+
+def _rasterize_polygon(points: np.ndarray, width: int, height: int) -> np.ndarray:
+    """Even odd scanline polygon fill. ``points`` is (N, 2) as (x, y)."""
+    mask = np.zeros((height, width), dtype=bool)
+    n = len(points)
+    if n < 3:
+        return mask
+    xs = points[:, 0]
+    ys = points[:, 1]
+    y_min = max(0, int(np.floor(ys.min())))
+    y_max = min(height - 1, int(np.ceil(ys.max())))
+    for y in range(y_min, y_max + 1):
+        yc = y + 0.5
+        nodes: List[float] = []
+        j = n - 1
+        for i in range(n):
+            yi, yj = ys[i], ys[j]
+            if (yi < yc <= yj) or (yj < yc <= yi):
+                t = (yc - yi) / (yj - yi)
+                nodes.append(xs[i] + t * (xs[j] - xs[i]))
+            j = i
+        nodes.sort()
+        for a, b in zip(nodes[0::2], nodes[1::2]):
+            xa = max(0, int(np.ceil(a - 0.5)))
+            xb = min(width - 1, int(np.floor(b - 0.5)))
+            if xb >= xa:
+                mask[y, xa:xb + 1] = True
+    return mask
+
+
+def _erode(mask: np.ndarray) -> np.ndarray:
+    """4 connected erosion (a pixel survives only if all N/S/E/W neighbours are set)."""
+    m = mask
+    out = m.copy()
+    out[1:, :] &= m[:-1, :]
+    out[:-1, :] &= m[1:, :]
+    out[:, 1:] &= m[:, :-1]
+    out[:, :-1] &= m[:, 1:]
+    # border pixels of a set region are boundary by definition
+    out[0, :] = False
+    out[-1, :] = False
+    out[:, 0] = False
+    out[:, -1] = False
+    return out & m
+
+
+def _trace_boundary(comp: np.ndarray) -> np.ndarray:
+    """Moore neighbour boundary trace of a single connected component.
+
+    Returns an ordered (M, 2) array of (x, y) boundary pixels. Robust but simple:
+    if tracing fails to close, falls back to unordered boundary pixels.
+    """
+    ys, xs = np.nonzero(comp)
+    if len(xs) == 0:
+        return np.empty((0, 2), dtype=int)
+    # start pixel: topmost leftmost
+    start_idx = np.lexsort((xs, ys))[0]
+    start = (int(xs[start_idx]), int(ys[start_idx]))
+    h, w = comp.shape
+
+    def is_set(px, py):
+        return 0 <= px < w and 0 <= py < h and comp[py, px]
+
+    # 8 neighbour offsets clockwise starting from "west"
+    neigh = [(-1, 0), (-1, -1), (0, -1), (1, -1),
+             (1, 0), (1, 1), (0, 1), (-1, 1)]
+    boundary = [start]
+    current_pixel = start
+    backdir = 0
+    max_steps = comp.sum() * 8 + 16
+    for _ in range(int(max_steps)):
+        found = False
+        for k in range(8):
+            d = (backdir + k) % 8
+            nx, ny = current_pixel[0] + neigh[d][0], current_pixel[1] + neigh[d][1]
+            if is_set(nx, ny):
+                backdir = (d + 5) % 8  # step back and one over
+                current_pixel = (nx, ny)
+                if current_pixel == start and len(boundary) > 2:
+                    return np.asarray(boundary, dtype=int)
+                boundary.append(current_pixel)
+                found = True
+                break
+        if not found:
+            break
+    # fallback: unordered boundary
+    eroded = _erode(comp)
+    bys, bxs = np.nonzero(comp & ~eroded)
+    return np.stack([bxs, bys], axis=1)

object_remove/mask/selection_enhancer.py

@@ -0,0 +1,127 @@
+"""Selection refinement operations.
+
+* ``denoise_and_smooth`` — jaggy edge cleanup via morphological open/close,
+  iteration count knob only.
+* ``soften`` — feathering, ``radius`` + ``iterations`` knobs, returns a *soft*
+  (float) mask, a distinct operation from denoise.
+* ``form_quick_selections`` — threaded connected component split (delegates to
+  :mod:`connected_components`).
+* ``separate_objects`` — explicit object separation, a distinct entry point
+  from quick select.
+* ``set_forbidden_area`` / ``shrink_with_brush`` — prohibited area mask and
+  per sub selection erosion.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+import numpy as np
+
+from .connected_components import split_components
+
+
+def _binary_dilate(mask: np.ndarray, iterations: int = 1) -> np.ndarray:
+    m = mask.copy()
+    for _ in range(iterations):
+        out = m.copy()
+        out[1:, :] |= m[:-1, :]
+        out[:-1, :] |= m[1:, :]
+        out[:, 1:] |= m[:, :-1]
+        out[:, :-1] |= m[:, 1:]
+        m = out
+    return m
+
+
+def _binary_erode(mask: np.ndarray, iterations: int = 1) -> np.ndarray:
+    m = mask.copy()
+    for _ in range(iterations):
+        out = m.copy()
+        out[1:, :] &= m[:-1, :]
+        out[:-1, :] &= m[1:, :]
+        out[:, 1:] &= m[:, :-1]
+        out[:, :-1] &= m[:, 1:]
+        # treat outside as empty, so borders erode away
+        out[0, :] &= m[0, :]
+        m = out
+    return m
+
+
+def _box_blur(img: np.ndarray, radius: int) -> np.ndarray:
+    """Separable box blur via cumulative sums (float input/output)."""
+    if radius <= 0:
+        return img
+    pad = radius
+    win = 2 * radius + 1
+    padded = np.pad(img, pad, mode="edge")
+    # horizontal sliding window mean (returns width back to original W)
+    cs = np.cumsum(padded, axis=1)
+    cs = np.pad(cs, ((0, 0), (1, 0)), mode="constant")
+    padded = (cs[:, win:] - cs[:, :-win]) / win
+    # vertical sliding window mean (returns height back to original H)
+    cs = np.cumsum(padded, axis=0)
+    cs = np.pad(cs, ((1, 0), (0, 0)), mode="constant")
+    padded = (cs[win:, :] - cs[:-win, :]) / win
+    return padded
+
+
+class SelectionEnhancer:
+    """Refinement operations over boolean/soft masks."""
+
+    def denoise_and_smooth(self, mask: np.ndarray, iterations: int = 1) -> np.ndarray:
+        """Morphological open then close — removes specks and fills pinholes.
+
+        Iteration count knob only, per PLAN (distinct from :meth:`soften`).
+        """
+        m = mask > 0
+        for _ in range(max(1, iterations)):
+            m = _binary_dilate(_binary_erode(m))  # open
+            m = _binary_erode(_binary_dilate(m))  # close
+        return m
+
+    def soften(self, mask: np.ndarray, radius: float = 3.0,
+               iterations: int = 1) -> np.ndarray:
+        """Feather the mask edge; returns a soft float mask in [0, 1].
+
+        ``radius`` + ``iterations`` knobs, per PLAN — a different operation from
+        denoise. Repeated box blur approximates a Gaussian feather.
+        """
+        soft = (mask > 0).astype(np.float32)
+        r = max(1, int(round(radius)))
+        for _ in range(max(1, iterations)):
+            soft = _box_blur(soft, r).astype(np.float32)
+        return np.clip(soft, 0.0, 1.0)
+
+    def shrink_with_brush(self, mask: np.ndarray, radius: int = 1) -> np.ndarray:
+        """Erode the mask by a brush radius."""
+        return _binary_erode(mask > 0, iterations=max(1, radius))
+
+    def form_quick_selections(self, mask: np.ndarray, max_selections: int = 0,
+                              min_area: int = 1, workers: int = 4) -> List[np.ndarray]:
+        """Split into sub selections across a bounded worker pool."""
+        return split_components(mask, max_selections=max_selections,
+                                min_area=min_area, workers=workers)
+
+    def separate_objects(self, mask: np.ndarray, min_area: int = 1) -> List[np.ndarray]:
+        """Explicit object separation, a distinct entry point from quick select.
+
+        Separates touching but thin
+        connections with an erosion before labelling, then dilates each object
+        back to its original extent (a watershed lite).
+        """
+        core = _binary_erode(mask > 0, iterations=1)
+        comps = split_components(core, min_area=min_area)
+        result = []
+        for c in comps:
+            grown = _binary_dilate(c, iterations=1) & (mask > 0)
+            result.append(grown)
+        return result
+
+    def set_forbidden_area(self, mask: np.ndarray,
+                           forbidden: np.ndarray) -> np.ndarray:
+        """Remove the forbidden/prohibited area from a selection.
+
+        The prohibited mask marks regions that must never be edited or used as a
+        fill source.
+        """
+        return (mask > 0) & ~(forbidden > 0)

object_remove/pipeline/__init__.py

@@ -0,0 +1,7 @@
+"""Pipeline orchestration."""
+
+from .remove_object_job import (
+    RemoveObjectJob, EngineTier, JobResult, ProgressFn, CancelFn,
+)
+
+__all__ = ["RemoveObjectJob", "EngineTier", "JobResult", "ProgressFn", "CancelFn"]