texture3d 0.2.0__py3-none-any.whl

texture3d/mapping/scoring.py

@@ -0,0 +1,229 @@
+"""Stage 2: which views see which faces, and how well.
+
+Per view the mesh is rasterised into a face map with a depth test, so a pixel
+belongs to the nearest face. A face's quality in that view is the blurred
+Sobel gradient magnitude summed over its pixels, scaled by the squared cosine
+between the view direction and the smoothed face normal: large, sharp,
+head-on projections score high. The mean colour of those pixels (YCbCr over
+BGR) feeds the outlier test that drops views disagreeing with the majority,
+which is what removes an object passing in front of the subject in one frame.
+"""
+
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Tuple
+
+import numpy as np
+
+from texture3d.mapping import imageops, raster
+from texture3d.mapping.topology import Topology
+from texture3d.mapping.types import NO_ID, FaceViews, Mesh, View, rgb2ycbcr
+
+
+@dataclass
+class ViewStats:
+    quality: np.ndarray  # (F,) float32
+    area: np.ndarray  # (F,) int64
+    color: np.ndarray  # (F, 3) float32, YCbCr of the mean BGR
+
+
+def face_map(mesh: Mesh, view: View, max_candidates: int = 1 << 24) -> Tuple[np.ndarray, np.ndarray]:
+    """The view's face map and depth map: per pixel the nearest face whose
+    three vertices project in front of the camera and at least three pixels
+    inside the image, or NO_ID."""
+    width, height = view.width, view.height
+    cam = view.camera.to_camera(mesh.positions)
+    z = cam[:, 2]
+    in_front = z > 0
+    safe = cam.copy()
+    safe[~in_front, 2] = 1.0
+    px = view.camera.to_image(safe).astype(np.float32)
+    inside_vertex = (
+        in_front & (px[:, 0] >= 3) & (px[:, 1] >= 3) & (px[:, 0] <= width - 4) & (px[:, 1] <= height - 4)
+    )
+    candidates = np.flatnonzero(inside_vertex[mesh.faces].all(axis=1))
+    faces = mesh.faces[candidates]
+    tri = px[faces]
+    zf = z[faces].astype(np.float32)
+
+    fmap = np.full(width * height, NO_ID, dtype=np.int64)
+    dmap = np.zeros(width * height, dtype=np.float32)
+    for ti, x, y, bary in raster.triangles_bary(tri, width, height, cull=True, max_candidates=max_candidates):
+        zt = zf[ti]
+        pbary = raster.perspective_correct(bary, zt)
+        depth = (pbary * zt).sum(axis=1).astype(np.float32)
+        pix = y * width + x
+        # The nearest face per pixel within the chunk; on equal depth the
+        # earlier face wins.
+        order = np.lexsort((ti, depth, pix))
+        pix_sorted = pix[order]
+        first = np.ones(len(order), dtype=bool)
+        first[1:] = pix_sorted[1:] != pix_sorted[:-1]
+        sel = order[first]
+        p = pix[sel]
+        d = depth[sel]
+        current = dmap[p]
+        replace = (current == 0) | (current > d)
+        dmap[p[replace]] = d[replace]
+        fmap[p[replace]] = candidates[ti[sel[replace]]]
+    return fmap.reshape(height, width), dmap.reshape(height, width)
+
+
+def score_view(mesh: Mesh, topology: Topology, view: View, max_candidates: int = 1 << 24) -> ViewStats:
+    face_count = mesh.face_count
+    gradient = imageops.gradient_magnitude_blurred(imageops.gray_normalized(view)).ravel()
+    mask = imageops.invalid_region_mask(view)
+    fmap, _ = face_map(mesh, view, max_candidates)
+    fmap = fmap.ravel()
+    if mask is not None:
+        # A face that touched a masked pixel is dropped so none of its
+        # remaining pixels contribute.
+        bad = np.unique(fmap[(mask.ravel() == 0) & (fmap != NO_ID)])
+        if len(bad):
+            fmap[np.isin(fmap, bad)] = NO_ID
+    valid = fmap != NO_ID
+    f = fmap[valid]
+    area = np.bincount(f, minlength=face_count)
+    quality = np.bincount(f, weights=gradient[valid], minlength=face_count).astype(np.float32)
+    bgr = view.bgr.reshape(-1, 3)[valid].astype(np.float64)
+    color_sum = np.stack([np.bincount(f, weights=bgr[:, c], minlength=face_count) for c in range(3)], axis=-1)
+
+    seen = area > 0
+    faces = mesh.faces[seen]
+    p = mesh.positions
+    centre = (p[faces[:, 0]] + p[faces[:, 1]] + p[faces[:, 2]]) / np.float32(3)
+    cam_dir = view.camera.C.astype(np.float32) - centre
+    n = topology.face_normals[seen]
+    dot = (cam_dir * n).sum(axis=1)
+    norms = np.sqrt((cam_dir * cam_dir).sum(axis=1) * (n * n).sum(axis=1))
+    # A zero-area face has no normal; give it the grazing minimum rather than
+    # a NaN that would poison the labelling costs.
+    cosine = np.clip(np.where(norms > 0, dot / np.where(norms > 0, norms, 1), 0), -1, 1)
+    cosine = np.maximum(np.float32(0.001), np.minimum(np.float32(1), cosine))
+    quality[seen] *= cosine * cosine
+    color = np.zeros((face_count, 3), dtype=np.float32)
+    color[seen] = rgb2ycbcr((color_sum[seen] / area[seen][:, None]).astype(np.float32))
+    return ViewStats(quality=quality, area=area, color=color)
+
+
+def score(
+    mesh: Mesh,
+    topology: Topology,
+    views: List[View],
+    outlier_threshold: float = 0.06,
+    workers: int = 4,
+    log: Optional[Callable[[str], None]] = None,
+) -> FaceViews:
+    """Every view scored, merged per face in view order, then filtered."""
+    face_count = mesh.face_count
+    stats: List[Optional[ViewStats]] = [None] * len(views)
+
+    def run(v: int) -> int:
+        stats[v] = score_view(mesh, topology, views[v])
+        return v
+
+    done = 0
+    with ThreadPoolExecutor(max_workers=max(1, workers)) as pool:
+        for _ in pool.map(run, range(len(views))):
+            done += 1
+            if log and (done % 10 == 0 or done == len(views)):
+                log(f"  scored {done}/{len(views)} views")
+
+    face_ids, view_ids, qualities, colors = [], [], [], []
+    for v, s in enumerate(stats):
+        seen = np.flatnonzero(s.area > 0)
+        face_ids.append(seen)
+        view_ids.append(np.full(len(seen), v, dtype=np.int64))
+        qualities.append(s.quality[seen])
+        colors.append(s.color[seen])
+    face_ids = np.concatenate(face_ids)
+    view_ids = np.concatenate(view_ids)
+    order = np.lexsort((view_ids, face_ids))
+    offsets = np.zeros(face_count + 1, dtype=np.int64)
+    np.cumsum(np.bincount(face_ids, minlength=face_count), out=offsets[1:])
+    face_views = FaceViews(
+        offsets=offsets,
+        view=view_ids[order],
+        quality=np.concatenate(qualities)[order],
+        color=np.concatenate(colors)[order],
+    )
+    if outlier_threshold > 0:
+        face_views = remove_outliers(face_views, outlier_threshold)
+    return face_views
+
+
+def remove_outliers(face_views: FaceViews, threshold: float) -> FaceViews:
+    """Drop views whose mean face colour is an outlier of the others.
+
+    Fit a Gaussian to the views' mean colours, drop views whose likelihood
+    is below the threshold, refit, up to ten times. A face whose test
+    abstains (too few views, or a degenerate covariance) keeps every view.
+    """
+    th = 6e-2 if threshold <= 0 else float(threshold)
+    min_covariance = 1e-3
+    max_iterations = 10
+    min_inliers = 4
+    eps = np.finfo(np.float64).eps
+
+    counts = face_views.counts
+    keep = np.ones(len(face_views.view), dtype=bool)
+    for n in np.unique(counts[counts > min_inliers]):
+        n = int(n)
+        faces_n = np.flatnonzero(counts == n)
+        slots = face_views.offsets[faces_n][:, None] + np.arange(n)[None, :]
+        colors = face_views.color[slots].astype(np.float64)  # (Fn, n, 3)
+        inliers = np.ones((len(faces_n), n), dtype=bool)
+        active = np.ones(len(faces_n), dtype=bool)
+        remove = np.zeros(len(faces_n), dtype=bool)
+        for _ in range(max_iterations):
+            a = np.flatnonzero(active)
+            if len(a) == 0:
+                break
+            c = colors[a]
+            w = inliers[a][..., None].astype(np.float64)
+            num = inliers[a].sum(axis=1)
+            mean = (c * w).sum(axis=1) / num[:, None]
+            d = (c - mean[:, None, :]) * w
+            cov = np.einsum("fni,fnj->fij", d, d) / (num - 1)[:, None, None]
+            degenerate = np.abs(cov).reshape(len(a), -1).max(axis=1) < min_covariance
+            active[a[degenerate]] = False
+            remove[a[degenerate]] = True
+            rest = ~degenerate
+            a, cov, mean = a[rest], cov[rest], mean[rest]
+            if len(a) == 0:
+                continue
+            # Treat the covariance as singular when the smallest eigenvalue
+            # falls below 3 eps of the largest.
+            eig = np.linalg.eigvalsh(cov)
+            singular = eig[:, 0] <= 3 * eps * eig[:, 2]
+            active[a[singular]] = False
+            good = ~singular
+            a, cov, mean = a[good], cov[good], mean[good]
+            if len(a) == 0:
+                continue
+            inv = np.linalg.inv(cov)
+            dd = colors[a] - mean[:, None, :]
+            q = np.einsum("fni,fij,fnj->fn", dd, inv, dd)
+            new_inliers = np.exp(-0.5 * q) > th
+            num_new = new_inliers.sum(axis=1)
+            changed = (new_inliers != inliers[a]).any(axis=1)
+            inliers[a] = new_inliers
+            all_in = num_new == n
+            too_few = ~all_in & (num_new < min_inliers)
+            stable = ~all_in & ~too_few & ~changed
+            active[a[all_in | too_few | stable]] = False
+            remove[a[stable]] = True
+        remove[active] = True
+        drop = remove[:, None] & ~inliers
+        keep[slots[drop]] = False
+
+    if keep.all():
+        return face_views
+    face_of_slot = np.repeat(np.arange(len(counts)), counts)[keep]
+    offsets = np.zeros(len(counts) + 1, dtype=np.int64)
+    np.cumsum(np.bincount(face_of_slot, minlength=len(counts)), out=offsets[1:])
+    return FaceViews(
+        offsets=offsets, view=face_views.view[keep], quality=face_views.quality[keep], color=face_views.color[keep]
+    )

texture3d/mapping/topology.py

@@ -0,0 +1,99 @@
+"""Stage 1: incidence, adjacency and smoothed normals of the mesh."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from texture3d.mapping.types import NO_ID, Mesh
+
+
+@dataclass
+class Topology:
+    #: CSR: the faces incident to vertex ``v`` are
+    #: ``vertex_face_list[vertex_face_offsets[v]:vertex_face_offsets[v + 1]]``, ascending.
+    vertex_face_offsets: np.ndarray
+    vertex_face_list: np.ndarray
+    #: Per face, the face across edge ``(face[v], face[(v + 1) % 3])`` for ``v``
+    #: in 0, 1, 2, or NO_ID at a boundary edge.
+    face_faces: np.ndarray
+    #: Unit face normals after three smoothing passes.
+    face_normals: np.ndarray
+
+    def faces_of_vertex(self, v: int) -> np.ndarray:
+        return self.vertex_face_list[self.vertex_face_offsets[v] : self.vertex_face_offsets[v + 1]]
+
+
+def cv_normalize(v: np.ndarray) -> np.ndarray:
+    """Unit-length vectors: the norm in double, the zero vector left at
+    zero, each component scaled in double and narrowed to float32."""
+    d = np.asarray(v, dtype=np.float64)
+    n = np.sqrt((d * d).sum(axis=-1, keepdims=True))
+    scale = np.where(n != 0, 1.0 / np.where(n != 0, n, 1.0), 0.0)
+    return (d * scale).astype(np.float32)
+
+
+def build(mesh: Mesh) -> Topology:
+    faces = mesh.faces
+    face_count = len(faces)
+    vertex_count = len(mesh.positions)
+
+    # Incident faces, CSR, ascending per vertex.
+    flat = faces.ravel()
+    order = np.argsort(flat, kind="stable")
+    vertex_face_list = (order // 3).astype(np.int64)
+    counts = np.bincount(flat, minlength=vertex_count)
+    vertex_face_offsets = np.zeros(vertex_count + 1, dtype=np.int64)
+    np.cumsum(counts, out=vertex_face_offsets[1:])
+
+    # Face adjacency across each edge slot. On a manifold edge the neighbour
+    # is the one other face; on an edge shared by more faces it is the
+    # lowest-numbered other face.
+    a = faces
+    b = np.roll(faces, -1, axis=1)
+    lo = np.minimum(a, b).ravel()
+    hi = np.maximum(a, b).ravel()
+    key = lo * vertex_count + hi
+    slot_face = np.repeat(np.arange(face_count, dtype=np.int64), 3)
+    order = np.lexsort((slot_face, key))
+    key_sorted = key[order]
+    face_sorted = slot_face[order]
+    group_start = np.ones(len(order), dtype=bool)
+    group_start[1:] = key_sorted[1:] != key_sorted[:-1]
+    group_id = np.cumsum(group_start) - 1
+    starts = np.flatnonzero(group_start)
+    sizes = np.diff(np.append(starts, len(order)))
+    first = face_sorted[starts]
+    second = np.where(sizes >= 2, face_sorted[np.minimum(starts + 1, len(order) - 1)], NO_ID)
+    own = face_sorted
+    neighbour = np.where(own == first[group_id], second[group_id], first[group_id])
+    neighbour = np.where(sizes[group_id] >= 2, neighbour, NO_ID)
+    face_faces = np.full(face_count * 3, NO_ID, dtype=np.int64)
+    face_faces[order] = neighbour
+    face_faces = face_faces.reshape(face_count, 3)
+
+    # Face normals, then three passes of averaging each normal with its edge
+    # neighbours within 25 degrees, half original and half average.
+    p = mesh.positions
+    v0 = p[faces[:, 0]]
+    v1 = p[faces[:, 1]]
+    v2 = p[faces[:, 2]]
+    normals = cv_normalize(np.cross(v1 - v0, v2 - v0))
+    cos_max_gradient = np.float32(np.cos(np.deg2rad(25.0)))
+    valid = face_faces != NO_ID
+    safe = np.where(valid, face_faces, 0)
+    for _ in range(3):
+        neighbour_normals = normals[safe]  # (F, 3, 3)
+        cosine = np.clip((normals[:, None, :] * neighbour_normals).sum(axis=-1), -1, 1)
+        take = valid & (cosine >= cos_max_gradient)
+        summed = (neighbour_normals * take[..., None].astype(np.float32)).sum(axis=1)
+        average = cv_normalize(summed)
+        normals = cv_normalize(normals * np.float32(0.5) + average * np.float32(0.5))
+
+    return Topology(
+        vertex_face_offsets=vertex_face_offsets,
+        vertex_face_list=vertex_face_list,
+        face_faces=face_faces,
+        face_normals=normals,
+    )

texture3d/mapping/types.py

@@ -0,0 +1,199 @@
+"""Shared types and colour helpers of the texture mapping.
+
+One photograph per face, then seams levelled. Stages, each its own module:
+
+1. ``topology``  incident faces, face adjacency, smoothed face normals
+2. ``scoring``   per view, rasterise the mesh and score each visible face by the
+                 image gradient it covers; reject views whose colour disagrees
+                 with the majority
+3. ``labeling``  one view per face, an MRF solved by loopy belief propagation
+                 with a Potts smoothness term
+4. ``patches``   connected same-view faces become patches with a rectangle in
+                 their view; seam vertices between them
+5. ``leveling``  global (least-squares colour offsets per vertex and patch) and
+                 local (Poisson blend of a border strip)
+6. ``atlas``     pack the patch rectangles into one texture, write the
+                 per-corner texture coordinates, sharpen
+
+Conventions: pixels are BGR, integer pixel coordinates are pixel centres,
+the camera looks down +Z with +Y down, and internal arithmetic is float32.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List
+
+import numpy as np
+
+#: An absent face, view, patch or label.
+NO_ID = -1
+
+
+@dataclass
+class Mesh:
+    positions: np.ndarray  # (V, 3) float32
+    faces: np.ndarray  # (F, 3) int64, counter-clockwise seen from outside
+
+    def __post_init__(self):
+        self.positions = np.ascontiguousarray(self.positions, dtype=np.float32)
+        self.faces = np.ascontiguousarray(self.faces, dtype=np.int64)
+
+    @property
+    def face_count(self) -> int:
+        return len(self.faces)
+
+    @property
+    def vertex_count(self) -> int:
+        return len(self.positions)
+
+
+@dataclass(frozen=True)
+class Camera:
+    """A pinhole camera: ``K`` in pixels of the image it is paired with, ``R``
+    rotating world to camera, ``C`` the centre in world coordinates."""
+
+    fx: float
+    fy: float
+    cx: float
+    cy: float
+    R: np.ndarray  # (3, 3) float64
+    C: np.ndarray  # (3,) float64
+
+    def to_camera(self, world: np.ndarray) -> np.ndarray:
+        """(N, 3) world points to camera space, in double."""
+        return (np.asarray(world, dtype=np.float64) - self.C) @ self.R.T
+
+    def to_image(self, cam: np.ndarray) -> np.ndarray:
+        """(N, 3) camera-space points to (N, 2) pixel coordinates, in double."""
+        z = cam[:, 2]
+        return np.stack([self.cx + self.fx * (cam[:, 0] / z), self.cy + self.fy * (cam[:, 1] / z)], axis=-1)
+
+    def project(self, world: np.ndarray) -> np.ndarray:
+        """World points to float32 pixel coordinates."""
+        return self.to_image(self.to_camera(world)).astype(np.float32)
+
+
+@dataclass
+class View:
+    """One keyframe: BGR8 pixels and the camera that took it."""
+
+    bgr: np.ndarray  # (H, W, 3) uint8
+    camera: Camera
+
+    @property
+    def width(self) -> int:
+        return self.bgr.shape[1]
+
+    @property
+    def height(self) -> int:
+        return self.bgr.shape[0]
+
+
+@dataclass
+class FaceViews:
+    """Per face, the views that see it, in view order (CSR)."""
+
+    offsets: np.ndarray  # (F + 1,) int64
+    view: np.ndarray  # (N,) int64
+    quality: np.ndarray  # (N,) float32
+    color: np.ndarray  # (N, 3) float32, YCbCr taken over the BGR channel order
+
+    @property
+    def counts(self) -> np.ndarray:
+        return np.diff(self.offsets)
+
+    @property
+    def visible(self) -> np.ndarray:
+        return self.counts > 0
+
+
+@dataclass
+class Rect:
+    """An integer pixel rectangle (x, y, width, height)."""
+
+    x: int
+    y: int
+    width: int
+    height: int
+
+    @property
+    def area(self) -> int:
+        return self.width * self.height
+
+    def is_contained_in(self, other: Rect) -> bool:
+        return (
+            self.x >= other.x
+            and self.y >= other.y
+            and self.x + self.width <= other.x + other.width
+            and self.y + self.height <= other.y + other.height
+        )
+
+
+@dataclass
+class Result:
+    #: Per face: the view it is textured from, or NO_ID.
+    labels: np.ndarray
+    #: Per face: its patch index; ``patch_count`` is the untextured patch.
+    face_patch: np.ndarray
+    patch_count: int
+    #: Square atlas, BGR8.
+    atlas_bgr: np.ndarray
+    atlas_size: int
+    #: (F, 3, 2) atlas pixel coordinates (integer = pixel centre, row 0 top).
+    face_texcoords: np.ndarray
+    #: Where each textured patch landed in the atlas.
+    patch_rects: List[Rect] = field(default_factory=list)
+    timings: dict = field(default_factory=dict)
+
+
+class MappingError(Exception):
+    pass
+
+
+class EmptyMesh(MappingError):
+    pass
+
+
+class NoViews(MappingError):
+    pass
+
+
+class NoVisibleFace(MappingError):
+    pass
+
+
+class PatchTooLarge(MappingError):
+    pass
+
+
+# Colour helpers shared by the stages.
+
+_RGB2YCBCR = np.array(
+    [[0.299, 0.587, 0.114], [-0.168736, -0.331264, 0.5], [0.5, -0.418688, -0.081312]], dtype=np.float32
+)
+_YCBCR_OFFSET = np.array([0.0, 128.0, 128.0], dtype=np.float32)
+
+
+def rgb2ycbcr(v: np.ndarray) -> np.ndarray:
+    """RGB-to-YCbCr on the last axis. Views hand this BGR, so the channels
+    stay in BGR order through the conversion."""
+    v = np.asarray(v, dtype=np.float32)
+    return v @ _RGB2YCBCR.T + _YCBCR_OFFSET
+
+
+def ycbcr2rgb(v: np.ndarray) -> np.ndarray:
+    v = np.asarray(v, dtype=np.float32)
+    y = v[..., 0]
+    cb = v[..., 1] - 128.0
+    cr = v[..., 2] - 128.0
+    return np.stack([y + cr * 1.402, y + cb * -0.34414 + cr * -0.71414, y + cb * 1.772], axis=-1)
+
+
+def round2int(x) -> np.ndarray:
+    """Round half up: ``floor(x + 0.5)``."""
+    return np.floor(np.asarray(x, dtype=np.float64) + 0.5).astype(np.int64)
+
+
+def clamp_byte(x: np.ndarray) -> np.ndarray:
+    return np.clip(x, 0, 255).astype(np.uint8)