tuned pipeline to 60 fps
This commit is contained in:
@@ -166,6 +166,13 @@ class nHHD(object):
|
||||
# compute potentials
|
||||
if mask is None and self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE):
|
||||
self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=self.debug)
|
||||
elif mask is not None and self.masked_poisson_solver == 'sparse':
|
||||
self.nD, self.nRu = self.psolver.solve_masked_pair(
|
||||
self.div,
|
||||
self.curlw,
|
||||
mask,
|
||||
verbose > 0,
|
||||
)
|
||||
else:
|
||||
self.nD = self._solve_poisson_backend(self.div, verbose > 0, valid_mask=mask)
|
||||
self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0, valid_mask=mask)
|
||||
|
||||
@@ -57,6 +57,11 @@ class MaskedDomain:
|
||||
"num_active",
|
||||
"pinned",
|
||||
"active_coords",
|
||||
"active_linear_indices",
|
||||
"_rhs_buffer",
|
||||
"_rhs_pair_buffer",
|
||||
"_scatter_buffer",
|
||||
"_scatter_pair_buffer",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
@@ -69,6 +74,7 @@ class MaskedDomain:
|
||||
num_active,
|
||||
pinned,
|
||||
active_coords,
|
||||
active_linear_indices,
|
||||
):
|
||||
self.shape = shape
|
||||
self.mask = mask
|
||||
@@ -78,6 +84,11 @@ class MaskedDomain:
|
||||
self.num_active = num_active
|
||||
self.pinned = pinned
|
||||
self.active_coords = active_coords
|
||||
self.active_linear_indices = active_linear_indices
|
||||
self._rhs_buffer = None
|
||||
self._rhs_pair_buffer = None
|
||||
self._scatter_buffer = None
|
||||
self._scatter_pair_buffer = None
|
||||
|
||||
@classmethod
|
||||
def build(cls, mask, spacings):
|
||||
@@ -101,10 +112,12 @@ class MaskedDomain:
|
||||
num_active=0,
|
||||
pinned=0,
|
||||
active_coords=active_coords,
|
||||
active_linear_indices=numpy.flatnonzero(active),
|
||||
)
|
||||
|
||||
index_map = numpy.full(active.shape, -1, dtype=numpy.int32)
|
||||
index_map[active] = numpy.arange(num_active, dtype=numpy.int32)
|
||||
active_linear_indices = numpy.flatnonzero(active)
|
||||
|
||||
dy, dx = float(spacings[0]), float(spacings[1])
|
||||
coeff_y = 1.0 / (dy * dy)
|
||||
@@ -153,6 +166,7 @@ class MaskedDomain:
|
||||
num_active=num_active,
|
||||
pinned=pinned,
|
||||
active_coords=active_coords,
|
||||
active_linear_indices=active_linear_indices,
|
||||
)
|
||||
|
||||
def matches(self, mask):
|
||||
@@ -166,13 +180,34 @@ class MaskedDomain:
|
||||
f: (H, W) scalar source field
|
||||
returns: (N,) RHS on active pixels; empty if no active pixels
|
||||
"""
|
||||
rhs = numpy.asarray(f, dtype=numpy.float64)[self.mask]
|
||||
if rhs.size == 0:
|
||||
return rhs
|
||||
if self.num_active == 0:
|
||||
return numpy.zeros(0, dtype=numpy.float64)
|
||||
if self._rhs_buffer is None or self._rhs_buffer.shape != (self.num_active,):
|
||||
self._rhs_buffer = numpy.empty(self.num_active, dtype=numpy.float64)
|
||||
numpy.take(
|
||||
numpy.asarray(f, dtype=numpy.float64).ravel(),
|
||||
self.active_linear_indices,
|
||||
out=self._rhs_buffer,
|
||||
)
|
||||
rhs = self._rhs_buffer
|
||||
rhs -= rhs.mean()
|
||||
rhs[self.pinned] = 0.0
|
||||
return rhs
|
||||
|
||||
def prepare_rhs_pair(self, f0, f1):
|
||||
"""Extract two gauge-fixed RHS vectors into one reusable (N, 2) buffer."""
|
||||
if self.num_active == 0:
|
||||
return numpy.zeros((0, 2), dtype=numpy.float64)
|
||||
if self._rhs_pair_buffer is None or self._rhs_pair_buffer.shape != (self.num_active, 2):
|
||||
self._rhs_pair_buffer = numpy.empty((self.num_active, 2), dtype=numpy.float64)
|
||||
flat0 = numpy.asarray(f0, dtype=numpy.float64).ravel()
|
||||
flat1 = numpy.asarray(f1, dtype=numpy.float64).ravel()
|
||||
numpy.take(flat0, self.active_linear_indices, out=self._rhs_pair_buffer[:, 0])
|
||||
numpy.take(flat1, self.active_linear_indices, out=self._rhs_pair_buffer[:, 1])
|
||||
self._rhs_pair_buffer -= self._rhs_pair_buffer.mean(axis=0, keepdims=True)
|
||||
self._rhs_pair_buffer[self.pinned, :] = 0.0
|
||||
return self._rhs_pair_buffer
|
||||
|
||||
def scatter(self, sol_active, template):
|
||||
"""Map reduced solution back to full grid.
|
||||
|
||||
@@ -180,10 +215,25 @@ class MaskedDomain:
|
||||
template: (H, W) array used for output shape/dtype
|
||||
returns: (H, W) potential; zero outside mask
|
||||
"""
|
||||
out = numpy.zeros_like(template, dtype=numpy.float32)
|
||||
if self._scatter_buffer is None or self._scatter_buffer.shape != template.shape:
|
||||
self._scatter_buffer = numpy.zeros(template.shape, dtype=numpy.float32)
|
||||
else:
|
||||
self._scatter_buffer.fill(0)
|
||||
if sol_active.size:
|
||||
out[self.mask] = numpy.asarray(sol_active, dtype=numpy.float32)
|
||||
return out
|
||||
self._scatter_buffer.ravel()[self.active_linear_indices] = numpy.asarray(sol_active, dtype=numpy.float32)
|
||||
return self._scatter_buffer.copy()
|
||||
|
||||
def scatter_pair(self, sol_active_pair, template):
|
||||
"""Map two reduced solutions back to a pair of full-grid outputs."""
|
||||
if self._scatter_pair_buffer is None or self._scatter_pair_buffer.shape != (2,) + template.shape:
|
||||
self._scatter_pair_buffer = numpy.zeros((2,) + template.shape, dtype=numpy.float32)
|
||||
else:
|
||||
self._scatter_pair_buffer.fill(0)
|
||||
if sol_active_pair.size:
|
||||
values = numpy.asarray(sol_active_pair, dtype=numpy.float32)
|
||||
self._scatter_pair_buffer[0].ravel()[self.active_linear_indices] = values[:, 0]
|
||||
self._scatter_pair_buffer[1].ravel()[self.active_linear_indices] = values[:, 1]
|
||||
return self._scatter_pair_buffer[0].copy(), self._scatter_pair_buffer[1].copy()
|
||||
|
||||
def residual(self, field, rhs_active):
|
||||
"""Compute masked Laplacian residual L u - rhs on active pixels.
|
||||
|
||||
@@ -237,6 +237,22 @@ class PoissonSolver(object):
|
||||
return self._cap_solver.solve(f, domain)
|
||||
return self._sparse_solver.solve(f, domain)
|
||||
|
||||
def solve_masked_pair(self, f0, f1, mask, verbose=False):
|
||||
"""Solve two masked RHS fields, using a batched sparse path when available."""
|
||||
mask_arr = self._normalize_mask(mask, f0.shape)
|
||||
if mask_arr is None:
|
||||
return None, None
|
||||
if f1.shape != f0.shape:
|
||||
raise ValueError("Both masked RHS fields must have the same shape")
|
||||
if not numpy.any(mask_arr):
|
||||
empty = numpy.zeros_like(f0, dtype=numpy.float32)
|
||||
return empty, empty.copy()
|
||||
|
||||
domain = self._get_or_build_domain(mask_arr)
|
||||
if self.masked_config.method == "capfft":
|
||||
return self._cap_solver.solve(f0, domain), self._cap_solver.solve(f1, domain)
|
||||
return self._sparse_solver.solve_pair(f0, f1, domain)
|
||||
|
||||
def solve(self, f, verbose=False, mask=None):
|
||||
"""Solve Poisson equation on the prepared grid.
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import numpy
|
||||
from scipy import linalg as dense_linalg
|
||||
|
||||
from ...backend_availability import CUPY_AVAILABLE, cp
|
||||
from ..config import MaskedPoissonConfig
|
||||
from ..domain import MaskedDomain, compute_boundary_mask, sample_boundary_coords
|
||||
|
||||
@@ -14,6 +15,25 @@ class CapacitanceFFTSolver:
|
||||
self._capfft_cache = None
|
||||
self._capfft_domain = None
|
||||
|
||||
def _use_cupy_backend(self):
|
||||
return (
|
||||
getattr(self.bulk_solver, "poisson_backend", "cpu") == "cupy"
|
||||
and CUPY_AVAILABLE
|
||||
)
|
||||
|
||||
def _bulk_solve(self, rhs):
|
||||
"""Solve one full-grid Poisson problem, using CuPy when configured."""
|
||||
if not self._use_cupy_backend():
|
||||
return self.bulk_solver.solve(rhs)
|
||||
|
||||
if not getattr(self.bulk_solver, "ready", False):
|
||||
self.bulk_solver.prepare(False)
|
||||
|
||||
rhs_gpu = cp.asarray(rhs, dtype=cp.float32)
|
||||
field_gpu = self.bulk_solver._solve_gpu_internal(rhs_gpu)
|
||||
cp.cuda.get_current_stream().synchronize()
|
||||
return cp.asnumpy(field_gpu)
|
||||
|
||||
def prepare(self, domain: MaskedDomain):
|
||||
"""Precompute capacitance matrix for boundary correction.
|
||||
|
||||
@@ -40,7 +60,7 @@ class CapacitanceFFTSolver:
|
||||
for col, (row, col_idx) in enumerate(boundary_coords):
|
||||
charge = numpy.zeros(active.shape, dtype=numpy.float32)
|
||||
charge[row, col_idx] = 1.0
|
||||
response = numpy.asarray(self.bulk_solver.solve(charge), dtype=numpy.float64)
|
||||
response = numpy.asarray(self._bulk_solve(charge), dtype=numpy.float64)
|
||||
residual = domain.residual(response, zero_rhs)
|
||||
cap_matrix[:, col] = residual[boundary_active_indices]
|
||||
|
||||
@@ -82,7 +102,7 @@ class CapacitanceFFTSolver:
|
||||
rhs_full = numpy.zeros_like(f, dtype=numpy.float32)
|
||||
rhs_full[domain.mask] = rhs_active.astype(numpy.float32)
|
||||
|
||||
field = numpy.asarray(self.bulk_solver.solve(rhs_full), dtype=numpy.float64)
|
||||
field = numpy.asarray(self._bulk_solve(rhs_full), dtype=numpy.float64)
|
||||
cap_cache = self.get_cache(domain)
|
||||
if not cap_cache or cap_cache["boundary_count"] == 0:
|
||||
field_active = field[domain.mask]
|
||||
@@ -103,7 +123,7 @@ class CapacitanceFFTSolver:
|
||||
charges = numpy.zeros_like(f, dtype=numpy.float32)
|
||||
coords = cap_cache["boundary_coords"]
|
||||
charges[coords[:, 0], coords[:, 1]] = alpha.astype(numpy.float32)
|
||||
correction = numpy.asarray(self.bulk_solver.solve(charges), dtype=numpy.float64)
|
||||
correction = numpy.asarray(self._bulk_solve(charges), dtype=numpy.float64)
|
||||
field -= correction
|
||||
|
||||
field_active = field[domain.mask]
|
||||
|
||||
@@ -19,3 +19,18 @@ class MaskedSparseSolver:
|
||||
|
||||
sol_active = domain.factor(rhs) # (N,) reduced-system potential
|
||||
return domain.scatter(sol_active, f)
|
||||
|
||||
def solve_pair(self, f0, f1, domain: MaskedDomain):
|
||||
"""Solve two masked Poisson RHS with one reused sparse factorization call."""
|
||||
rhs_pair = domain.prepare_rhs_pair(f0, f1)
|
||||
if rhs_pair.shape[0] == 0:
|
||||
empty = numpy.zeros_like(f0, dtype=numpy.float32)
|
||||
return empty, empty.copy()
|
||||
|
||||
try:
|
||||
sol_pair = domain.factor(rhs_pair) # (N, 2) reduced-system potentials
|
||||
except Exception:
|
||||
sol_pair = numpy.column_stack(
|
||||
[domain.factor(rhs_pair[:, 0]), domain.factor(rhs_pair[:, 1])]
|
||||
)
|
||||
return domain.scatter_pair(sol_pair, f0)
|
||||
|
||||
@@ -15,6 +15,7 @@ from .precompute import (
|
||||
load_exported_data,
|
||||
recommend_edge_expand,
|
||||
)
|
||||
from .geometry import resolve_uv_runtime_sizes, scale_camera_geometry, scaled_map_cache_path
|
||||
from .remap import UVRemap, flatten_frame, remap_valid_ratio
|
||||
from .surface import UVSurfaceMap
|
||||
|
||||
@@ -33,5 +34,8 @@ __all__ = [
|
||||
"load_exported_data",
|
||||
"recommend_edge_expand",
|
||||
"remap_valid_ratio",
|
||||
"resolve_uv_runtime_sizes",
|
||||
"scale_camera_geometry",
|
||||
"scaled_map_cache_path",
|
||||
"uv_triangle_to_image",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
"""Camera geometry scaling for faster UV remap at reduced source resolution."""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def scale_camera_geometry(source_K, source_width, source_height, scale):
|
||||
"""Scale pinhole intrinsics and image size for a downsampled source frame.
|
||||
|
||||
When the live frame is resized by ``scale`` before undistort/remap, ``fx``,
|
||||
``fy``, ``cx``, and ``cy`` must be scaled by the same factor. OpenCV
|
||||
distortion coefficients are unchanged.
|
||||
|
||||
source_K: (3, 3) camera matrix
|
||||
source_width, source_height: native calibration / capture size
|
||||
scale: (0, 1] resize factor applied to width and height
|
||||
returns: (K_scaled, width_scaled, height_scaled)
|
||||
"""
|
||||
scale = float(scale)
|
||||
if scale <= 0.0 or scale > 1.0:
|
||||
raise ValueError("source scale must be in (0, 1]")
|
||||
if np.isclose(scale, 1.0):
|
||||
return np.asarray(source_K, dtype=np.float64), int(source_width), int(source_height)
|
||||
|
||||
k = np.asarray(source_K, dtype=np.float64).copy()
|
||||
k[0, 0] *= scale
|
||||
k[1, 1] *= scale
|
||||
k[0, 2] *= scale
|
||||
k[1, 2] *= scale
|
||||
|
||||
width_scaled = max(1, int(round(int(source_width) * scale)))
|
||||
height_scaled = max(1, int(round(int(source_height) * scale)))
|
||||
return k, width_scaled, height_scaled
|
||||
|
||||
|
||||
def resolve_uv_runtime_sizes(camera_config, out_width=None, out_height=None, source_scale=None):
|
||||
"""Resolve UV output size and source scale from config and CLI overrides.
|
||||
|
||||
Prefers explicit CLI values, then ``camera.uv_output`` or
|
||||
``camera.resolution_stitch``, then legacy ``camera.work4``.
|
||||
Source scale prefers CLI, then ``camera.source_scale``, then ``work4[3]``.
|
||||
"""
|
||||
camera_config = camera_config or {}
|
||||
|
||||
stitch = camera_config.get("resolution_stitch")
|
||||
uv_output = camera_config.get("uv_output")
|
||||
if out_width is None or out_height is None:
|
||||
if uv_output and len(uv_output) >= 2:
|
||||
out_width = out_width if out_width is not None else int(uv_output[0])
|
||||
out_height = out_height if out_height is not None else int(uv_output[1])
|
||||
elif stitch and len(stitch) >= 2:
|
||||
out_width = out_width if out_width is not None else int(stitch[0])
|
||||
out_height = out_height if out_height is not None else int(stitch[1])
|
||||
|
||||
if source_scale is None:
|
||||
if "source_scale" in camera_config:
|
||||
source_scale = float(camera_config["source_scale"])
|
||||
else:
|
||||
work4 = camera_config.get("work4")
|
||||
if work4 and len(work4) >= 4:
|
||||
source_scale = float(work4[3])
|
||||
|
||||
if out_width is None or out_height is None:
|
||||
raise ValueError("UV output width/height must be set via CLI or config resolution_stitch")
|
||||
if source_scale is None:
|
||||
source_scale = 1.0
|
||||
|
||||
source_scale = float(source_scale)
|
||||
if source_scale <= 0.0 or source_scale > 1.0:
|
||||
raise ValueError("source_scale must be in (0, 1]")
|
||||
|
||||
return int(out_width), int(out_height), source_scale
|
||||
|
||||
|
||||
def scaled_map_cache_path(path, source_scale):
|
||||
"""Append a scale suffix so caches for different source scales do not collide."""
|
||||
if path is None or np.isclose(source_scale, 1.0):
|
||||
return path
|
||||
scale_tag = f"{source_scale:.3f}".rstrip("0").rstrip(".")
|
||||
base, ext = path.rsplit(".", 1) if "." in path else (path, "npz")
|
||||
return f"{base}_s{scale_tag}.{ext}"
|
||||
@@ -181,6 +181,7 @@ def build_uv_domain(
|
||||
edge_expand=1.0,
|
||||
include_surface=True,
|
||||
verbose=True,
|
||||
source_scale=1.0,
|
||||
):
|
||||
"""Build image-to-UV remap and optional UV-to-world surface map."""
|
||||
faces = data["faces"]
|
||||
@@ -270,6 +271,7 @@ def build_uv_domain(
|
||||
"edge_band": np.array(edge_band, dtype=np.float64),
|
||||
"edge_expand": np.array(edge_expand, dtype=np.float64),
|
||||
"projection_K": source_K.astype(np.float64),
|
||||
"source_scale": np.array(float(source_scale), dtype=np.float64),
|
||||
"method": np.array("uv"),
|
||||
},
|
||||
)
|
||||
@@ -292,6 +294,7 @@ def build_or_load_domain(
|
||||
edge_expand=0.0,
|
||||
include_surface=True,
|
||||
verbose=True,
|
||||
source_scale=1.0,
|
||||
):
|
||||
"""Build or load full UV domain (2D remap + optional 3D surface).
|
||||
|
||||
@@ -317,6 +320,7 @@ def build_or_load_domain(
|
||||
flip_v=flip_v,
|
||||
edge_band=edge_band,
|
||||
edge_expand=resolved_edge_expand,
|
||||
source_scale=source_scale,
|
||||
)
|
||||
surface_ok = (not include_surface) or (cached.surface is not None)
|
||||
if cache_ok and surface_ok:
|
||||
@@ -338,6 +342,7 @@ def build_or_load_domain(
|
||||
edge_expand=resolved_edge_expand,
|
||||
include_surface=include_surface,
|
||||
verbose=verbose,
|
||||
source_scale=source_scale,
|
||||
)
|
||||
valid_ratio = remap_valid_ratio(domain.remap)
|
||||
if verbose:
|
||||
|
||||
@@ -37,6 +37,7 @@ class UVRemap:
|
||||
meta_keys = {
|
||||
"source_width",
|
||||
"source_height",
|
||||
"source_scale",
|
||||
"out_width",
|
||||
"out_height",
|
||||
"flip_v",
|
||||
@@ -121,6 +122,7 @@ def _cache_matches(
|
||||
flip_v,
|
||||
edge_band,
|
||||
edge_expand,
|
||||
source_scale=1.0,
|
||||
) -> bool:
|
||||
"""True if cached remap matches the requested build parameters."""
|
||||
if isinstance(remap, UVRemap):
|
||||
@@ -145,6 +147,9 @@ def _cache_matches(
|
||||
return False
|
||||
if "projection_K" in rd and not np.allclose(rd["projection_K"], source_K, rtol=1e-6, atol=1e-6):
|
||||
return False
|
||||
cached_scale = float(rd["source_scale"]) if "source_scale" in rd else 1.0
|
||||
if not np.isclose(cached_scale, float(source_scale), rtol=1e-6, atol=1e-6):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user