integrated finger process pipeline on remapped 2d image

This commit is contained in:
ykcheungab
2026-06-14 17:48:30 +08:00
parent 667bdf1961
commit 30bebcd439
36 changed files with 3593 additions and 2043 deletions
+2 -1
View File
@@ -34,4 +34,5 @@ __pycache__/
/config/compound_eye.png
# Legacy / scratch folders
/old_orisys/
/old_orisys/
/tools/calibration/
+4 -2
View File
@@ -5,7 +5,9 @@
"camera": {
"resolution_raw": [1280, 720],
"fps": 30,
"resolution_stitch": [672, 728]
"resolution_stitch": [504, 546],
"old":[672, 728],
"work": [504, 546]
},
"tracker": {
"adaptive": true,
@@ -16,7 +18,7 @@
"dis_preset": "medium",
"dis_mode": 3,
"algorithm_list": ["raft", "dis"],
"hhd_downsample": 1
"hhd_downsample": 3
},
"contact": {
"min_contour_area": 3000,
-631
View File
@@ -1,631 +0,0 @@
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import orisys
import numpy as np
import time
import argparse
from datetime import datetime
import pyqtgraph as pg
from PyQt5.QtCore import Qt, QSize, QEvent
from PyQt5.QtGui import QImage, QFont, QPixmap, QKeySequence, QCursor
from PyQt5.QtWidgets import QLabel, QVBoxLayout, QWidget, QSizePolicy, QShortcut, QApplication
from pyqtgraph.Qt import QtGui
# Window aspect: width : total height = 936 : (936 + plot strip)
ASPECT_W = 936
ASPECT_H_TOTAL = 936 + 624
ARROW_STRETCH = 936
PLOT_STRETCH = 624
npoints = 120 # 图表数据长度
dispaly_angle = -90 # 箭头图显示角度
touch_threshold = 5 # 触碰反应阈值
frame_to_recover = 10 # 触碰后等待回复帧数
# Match arrow panel / Qt style (not pure black)
PLOT_THEME_BG = "#0B1016"
PLOT_THEME_FG = "#A0A0A0"
# OS title bar cannot be tinted to #0B1016; frameless removes it (use Alt+drag to move).
USE_FRAMELESS_WINDOW = True
class CombinedTactileWindow(QWidget):
"""
Resizable window that keeps width:height = ASPECT_W : ASPECT_H_TOTAL.
Dragging the width or height edge updates the other dimension accordingly.
"""
RESIZE_MARGIN = 8
def __init__(self):
super().__init__()
self._aspect_hw = ASPECT_H_TOTAL / float(ASPECT_W)
self._resize_guard = False
self.user_closed = False
self._drag_anchor = None
self._frameless_resize_edge = None
self._frameless_resize_start_global = None
self._frameless_resize_start_geom = None
def _min_client_size(self):
w = max(320, self.minimumSizeHint().width())
h = int(round(w * self._aspect_hw))
return w, h
def _hit_test_resize(self, global_pos):
if not USE_FRAMELESS_WINDOW:
return None
p = self.mapFromGlobal(global_pos)
w, h = self.width(), self.height()
m = self.RESIZE_MARGIN
left = p.x() <= m
right = p.x() >= w - m
top = p.y() <= m
bottom = p.y() >= h - m
if top and left:
return "tl"
if top and right:
return "tr"
if bottom and left:
return "bl"
if bottom and right:
return "br"
if top:
return "t"
if bottom:
return "b"
if left:
return "l"
if right:
return "r"
return None
def _cursor_for_edge(self, edge):
return {
"tl": Qt.SizeFDiagCursor,
"br": Qt.SizeFDiagCursor,
"tr": Qt.SizeBDiagCursor,
"bl": Qt.SizeBDiagCursor,
"l": Qt.SizeHorCursor,
"r": Qt.SizeHorCursor,
"t": Qt.SizeVerCursor,
"b": Qt.SizeVerCursor,
}.get(edge, Qt.ArrowCursor)
def _apply_frameless_resize(self, dg):
edge = self._frameless_resize_edge
g0 = self._frameless_resize_start_geom
min_w, min_h = self._min_client_size()
dx, dy = dg.x(), dg.y()
x0, y0, w0, h0 = g0.x(), g0.y(), g0.width(), g0.height()
if edge == "br":
nw = max(min_w, w0 + dx)
nh = int(round(nw * self._aspect_hw))
self.setGeometry(x0, y0, nw, nh)
elif edge == "r":
nw = max(min_w, w0 + dx)
nh = int(round(nw * self._aspect_hw))
self.setGeometry(x0, y0, nw, nh)
elif edge == "b":
nh = max(min_h, h0 + dy)
nw = int(round(nh / self._aspect_hw))
self.setGeometry(x0, y0, nw, nh)
elif edge == "l":
nw = max(min_w, w0 - dx)
nh = int(round(nw * self._aspect_hw))
nx = g0.right() - nw
self.setGeometry(nx, y0, nw, nh)
elif edge == "t":
nh = max(min_h, h0 - dy)
nw = int(round(nh / self._aspect_hw))
ny = g0.bottom() - nh
self.setGeometry(x0, ny, nw, nh)
elif edge == "tr":
nw = max(min_w, w0 + dx)
nh = int(round(nw * self._aspect_hw))
ny = g0.bottom() - nh
self.setGeometry(x0, ny, nw, nh)
elif edge == "bl":
nw = max(min_w, w0 - dx)
nh = int(round(nw * self._aspect_hw))
nx = g0.right() - nw
self.setGeometry(nx, y0, nw, nh)
elif edge == "tl":
nw = max(min_w, w0 - dx)
nh = int(round(nw * self._aspect_hw))
nx = g0.right() - nw
ny = g0.bottom() - nh
self.setGeometry(nx, ny, nw, nh)
def install_drag_filters(self):
"""Alt+left-drag to move when frameless; grabs go to this widget."""
for w in self.findChildren(QWidget):
if w is not self:
w.installEventFilter(self)
def eventFilter(self, watched, event):
if not USE_FRAMELESS_WINDOW:
return super().eventFilter(watched, event)
et = event.type()
if et == QEvent.MouseButtonPress and event.button() == Qt.LeftButton:
edge = self._hit_test_resize(event.globalPos())
if edge is not None:
self._frameless_resize_edge = edge
self._frameless_resize_start_global = event.globalPos()
self._frameless_resize_start_geom = self.geometry()
self.grabMouse()
return True
if QApplication.keyboardModifiers() & Qt.AltModifier:
self._drag_anchor = event.globalPos() - self.frameGeometry().topLeft()
self.grabMouse()
return True
if et == QEvent.MouseMove and not event.buttons():
edge = self._hit_test_resize(event.globalPos())
watched.setCursor(QCursor(self._cursor_for_edge(edge)))
return super().eventFilter(watched, event)
def mouseMoveEvent(self, event):
if USE_FRAMELESS_WINDOW and self._frameless_resize_edge and event.buttons() & Qt.LeftButton:
dg = event.globalPos() - self._frameless_resize_start_global
self._apply_frameless_resize(dg)
event.accept()
return
if (
USE_FRAMELESS_WINDOW
and self._drag_anchor is not None
and event.buttons() & Qt.LeftButton
):
self.move(event.globalPos() - self._drag_anchor)
event.accept()
return
super().mouseMoveEvent(event)
def mouseReleaseEvent(self, event):
if USE_FRAMELESS_WINDOW and self._frameless_resize_edge and event.button() == Qt.LeftButton:
self.releaseMouse()
self._frameless_resize_edge = None
self._frameless_resize_start_global = None
self._frameless_resize_start_geom = None
event.accept()
return
if USE_FRAMELESS_WINDOW and self._drag_anchor is not None and event.button() == Qt.LeftButton:
self.releaseMouse()
self._drag_anchor = None
event.accept()
return
super().mouseReleaseEvent(event)
def closeEvent(self, event):
self.user_closed = True
event.accept()
def minimumSizeHint(self):
w = 320
return QSize(w, int(round(w * self._aspect_hw)))
def resizeEvent(self, event):
if self._resize_guard:
super().resizeEvent(event)
return
super().resizeEvent(event)
w, h = self.width(), self.height()
if w < 8 or h < 8:
return
old = event.oldSize()
self._resize_guard = True
try:
if old.isValid() and old.width() > 0 and old.height() > 0:
dw = abs(w - old.width())
dh = abs(h - old.height())
if dw >= dh:
new_w, new_h = w, int(round(w * self._aspect_hw))
else:
new_w, new_h = int(round(h / self._aspect_hw)), h
else:
new_w, new_h = w, int(round(w * self._aspect_hw))
if new_w != w or new_h != h:
self.resize(new_w, new_h)
finally:
self._resize_guard = False
def compute_arrow_canvas_side(win):
"""
Pixel size of the square arrow canvas from the window client area and the
fixed vertical split ARROW_STRETCH : PLOT_STRETCH (matches QVBoxLayout stretch).
"""
w = win.width()
h = win.height()
if w < 16 or h < 16:
return max(32, ASPECT_W)
h_arrow = int(round(h * ARROW_STRETCH / float(ARROW_STRETCH + PLOT_STRETCH)))
# With aspect ASPECT_W:ASPECT_H_TOTAL, h_arrow should equal w; min() absorbs rounding.
return max(32, min(w, h_arrow))
def draw_arrows_wuxi(
frame,
optical_flow,
fps=True,
threshold=5,
grid_spacing=40,
arrow_scale=1.0,
margin=0.8,
colormap=cv2.COLORMAP_PARULA,
below_threshold_color=(0, 0, 255),
edge_margin=15,
target_size=936,
):
"""
Draw colored optical flow arrows on the image by magnitude (JET colormap).
Resizes frame and flow to target_size x target_size; scales flow vectors and
grid_spacing / margins / threshold so arrow density matches the original resolution.
"""
image = frame.copy()
if len(image.shape) == 2 or image.shape[2] == 1:
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
h0, w0 = image.shape[:2]
if h0 != target_size or w0 != target_size:
sx = target_size / float(w0)
sy = target_size / float(h0)
image = cv2.resize(image, (target_size, target_size), interpolation=cv2.INTER_LINEAR)
u = cv2.resize(optical_flow[:, :, 0], (target_size, target_size), interpolation=cv2.INTER_LINEAR) * sx
v = cv2.resize(optical_flow[:, :, 1], (target_size, target_size), interpolation=cv2.INTER_LINEAR) * sy
optical_flow = np.stack([u, v], axis=2)
scale = max(sx, sy)
grid_spacing = max(4, int(round(grid_spacing * scale)))
edge_margin = max(1, int(round(edge_margin * scale)))
threshold = float(threshold) * scale
h, w = image.shape[:2]
# Arrow grid in pixel coordinates (assume frame and flow same size)
rr = np.arange(grid_spacing // 2, h, grid_spacing)
cc = np.arange(grid_spacing // 2, w, grid_spacing)
rr, cc = np.meshgrid(rr, cc, indexing="ij")
rr, cc = rr.reshape(-1), cc.reshape(-1)
# Displacement at each grid point (flow stored as half-displacement)
disp = 2.0 * optical_flow[rr, cc, :]
mag = np.linalg.norm(disp, axis=1)
dx = disp[:, 0] / arrow_scale
dy = disp[:, 1] / arrow_scale
length = np.sqrt(dx * dx + dy * dy)
# Cap arrow length at grid_spacing
scale = np.where(length > grid_spacing*margin, grid_spacing / np.maximum(length, 1e-9), 1.0)*margin
dx, dy = dx * scale, dy * scale
end_x = (cc + dx).astype(np.int32)
end_y = (rr + dy).astype(np.int32)
valid = (
(cc >= edge_margin) & (cc < w - edge_margin) & (rr >= edge_margin) & (rr < h - edge_margin)
& (end_x >= 0) & (end_x < w) & (end_y >= 0) & (end_y < h)
)
mag_range = mag.max() - mag.min()
norm_mag = ((mag - mag.min()) / mag_range * 255).astype(np.uint8) if mag_range > 0 else np.zeros_like(mag, dtype=np.uint8)
norm_mag[mag < threshold] = 0
colors = cv2.applyColorMap(norm_mag.reshape(-1, 1), colormap)
for i in np.where(valid)[0]:
pt1, pt2 = (int(cc[i]), int(rr[i])), (int(end_x[i]), int(end_y[i]))
color = tuple(int(c) for c in colors[i, 0])
if mag[i] < threshold:
cv2.circle(image, pt1, radius=1, color=below_threshold_color, thickness=-1)
else:
# thickness = max(1, int(2 * mag[i] / mag.max())) if mag.max() > 0 else 2
thickness = 1
cv2.arrowedLine(image, pt1, pt2, color, thickness=thickness, tipLength=0.2)
return image
def set_arrow_label_bgr(label, bgr_image):
"""Show a BGR OpenCV image in a QLabel (RGB for Qt), centered with equal margins."""
rgb = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
rgb = np.ascontiguousarray(rgb)
h, w, ch = rgb.shape
qimg = QImage(rgb.data, w, h, ch * w, QImage.Format_RGB888).copy()
pm = QPixmap.fromImage(qimg)
lw, lh = label.width(), label.height()
if lw > 0 and lh > 0 and (pm.width() != lw or pm.height() != lh):
# Keep square aspect; center in label (equal L/R or T/B bars if rect not square)
pm = pm.scaled(lw, lh, Qt.KeepAspectRatio, Qt.SmoothTransformation)
label.setPixmap(pm)
def make_gradient_brush(color_hex, max_alpha=150):
"""
创建一个从上到下渐变的笔刷 (兼容 PyQt5 和 PyQt6)
"""
c_top = QtGui.QColor(color_hex)
c_top.setAlpha(max_alpha) # 顶部半透明
c_bot = QtGui.QColor(color_hex)
c_bot.setAlpha(0) # 底部全透明
# 创建线性渐变:(0,0)到(0,1) 代表从控件顶部到底部
grad = QtGui.QLinearGradient(0, 0, 0, 1)
if hasattr(QtGui.QGradient, 'CoordinateMode'):
grad.setCoordinateMode(QtGui.QGradient.CoordinateMode.ObjectBoundingMode)
else:
grad.setCoordinateMode(QtGui.QGradient.ObjectBoundingMode)
# ==============================================
grad.setColorAt(0, c_top) # 0.0 (顶端) 颜色
grad.setColorAt(1, c_bot) # 1.0 (底端) 颜色
return QtGui.QBrush(grad)
def create_plot(plot_container):
pg.setConfigOption("background", PLOT_THEME_BG)
pg.setConfigOption("foreground", PLOT_THEME_FG)
pg.setConfigOptions(antialias=True) # 抗锯齿
theme_color = '#00E5FF' # '#00E5FF'
line_pen = pg.mkPen(color=theme_color, width=2)
grad_brush = make_gradient_brush(theme_color, max_alpha=120)
axis_font = QFont("Segoe UI", 11)
p1 = plot_container.addPlot()
p1.setLabel('left', 'N')
p1.getAxis('left').label.setFont(axis_font)
p1.getAxis('left').label.setRotation(dispaly_angle)
p1.showGrid(x=True, y=True, alpha=0.3)
plot_container.nextRow()
p2 = plot_container.addPlot()
p2.setLabel('left', 'X')
p2.getAxis('left').label.setFont(axis_font)
p2.getAxis('left').label.setRotation(dispaly_angle)
p2.showGrid(x=True, y=True, alpha=0.3)
plot_container.nextRow()
p3 = plot_container.addPlot()
p3.setLabel('left', 'Y')
# p3.setLabel('bottom', 'Time', units='s')
p3.getAxis('left').label.setFont(axis_font)
p3.getAxis('left').label.setRotation(dispaly_angle)
p3.getAxis('bottom').label.setFont(axis_font)
p3.showGrid(x=True, y=True, alpha=0.3)
_plot_bg = pg.mkColor(PLOT_THEME_BG)
for p in (p1, p2, p3):
p.getAxis("bottom").setStyle(showValues=False)
p.getAxis("left").setStyle(showValues=False)
# Keep title as "N" / "X" / "Y" — no PyQtGraph "(x0.001)" suffix on small ranges.
p.getAxis("left").enableAutoSIPrefix(False)
# PlotItem.setBackground exists only in newer pyqtgraph; ViewBox works on older installs.
p.vb.setBackgroundColor(_plot_bg)
plot_af = p1.plot(pen=line_pen, fillLevel=0, fillBrush=grad_brush)
plot_tfx = p2.plot(pen=line_pen, fillLevel=0, fillBrush=grad_brush)
plot_tfy = p3.plot(pen=line_pen, fillLevel=0, fillBrush=grad_brush)
return p1, p2, p3, plot_af, plot_tfx, plot_tfy
def _plot_y_range(y):
"""Avoid degenerate Y span; explicit x/y in setData so older PyQtGraph never plots time as y."""
mn = float(np.min(y))
mx = float(np.max(y))
if not (np.isfinite(mn) and np.isfinite(mx)):
return 0.0, 1.0
if mx - mn <= 1e-15:
mid = 0.5 * (mn + mx)
return mid - 0.5, mid + 0.5
pad = (mx - mn) * 0.05
return mn - pad, mx + pad
def update(plot_af, plot_tfx, plot_tfy, data, p1, p2, p3):
# global plot_af, plot_tfx, plot_tfy, data, p1, p2, p3
plot_af.setData(data[:, 0, 0], data[:, 1, 0])
p1.setYRange(*_plot_y_range(data[:, 1, 0]))
plot_tfx.setData(data[:, 0, 1], data[:, 1, 1])
p2.setYRange(*_plot_y_range(data[:, 1, 1]))
plot_tfy.setData(data[:, 0, 2], data[:, 1, 2])
p3.setYRange(*_plot_y_range(data[:, 1, 2]))
# 实时图表:单窗口,上为光流箭头,下为曲线
pg.setConfigOption("background", PLOT_THEME_BG)
pg.setConfigOption("foreground", PLOT_THEME_FG)
app = pg.mkQApp("Wuxi Demo")
combined_win = CombinedTactileWindow()
combined_win.setWindowTitle("Wuxi tactile — flow + forces")
main_layout = QVBoxLayout(combined_win)
main_layout.setContentsMargins(0, 0, 0, 0)
main_layout.setSpacing(0)
# Ignored: do not let pixmap/content drive QLabel size — layout stretch sets arrow vs plot heights.
_policy_layout = QSizePolicy(QSizePolicy.Ignored, QSizePolicy.Ignored)
arrow_label = QLabel()
arrow_label.setAlignment(Qt.AlignHCenter | Qt.AlignVCenter)
arrow_label.setStyleSheet(
"background-color: #0B1016; margin: 0; padding: 0; border: none;"
)
arrow_label.setSizePolicy(_policy_layout)
arrow_label.setMinimumSize(0, 0)
arrow_label.setScaledContents(False)
main_layout.addWidget(arrow_label, stretch=ARROW_STRETCH)
plot_widget = pg.GraphicsLayoutWidget()
plot_widget.setBackground(PLOT_THEME_BG)
plot_widget.setSizePolicy(_policy_layout)
plot_widget.setMinimumSize(0, 0)
main_layout.addWidget(plot_widget, stretch=PLOT_STRETCH)
if USE_FRAMELESS_WINDOW:
combined_win.setWindowFlags(Qt.Window | Qt.FramelessWindowHint)
# Client area matches theme (no bright border flash on frameless)
combined_win.setStyleSheet(f"background-color: {PLOT_THEME_BG};")
combined_win.resize(ASPECT_W, ASPECT_H_TOTAL)
combined_win.install_drag_filters()
combined_win.show()
p1, p2, p3, plot_af, plot_tfx, plot_tfy = create_plot(plot_widget)
data = np.zeros([npoints,2,3])
start_time = None # Will be set when sensor starts
fps = 30.0
for i in range(3):
data[:,0,i] = np.linspace(0.0, npoints / fps, npoints)
def main():
global start_time
parser = argparse.ArgumentParser()
parser.add_argument("--video","-v", type=int, default=0, help="视频文件路径")
args = parser.parse_args()
session_start = datetime.now()
print(
f"[session] START {session_start:%Y-%m-%d %H:%M:%S} (long-run / crash log marker)",
flush=True,
)
if USE_FRAMELESS_WINDOW:
print(
"[window] Frameless: resize = drag window edges/corners; move = Alt+left-drag; quit = Q.",
flush=True,
)
sensor1 = None
try:
sensor1 = orisys.Sensor(
args.video,
config_name="./config/ddjx01_wuxi.json",
cal_path="./config/dd01a_0306.npy",
isstitch=True,
cuda=False,
verbose=False,
)
def _shortcut_reset_tracker():
sensor1.reset()
print("传感器已重置 (快捷键 R)", flush=True)
def _shortcut_quit():
combined_win.close()
_sc_r = QShortcut(QKeySequence(Qt.Key_R), combined_win)
_sc_r.activated.connect(_shortcut_reset_tracker)
_sc_q = QShortcut(QKeySequence(Qt.Key_Q), combined_win)
_sc_q.activated.connect(_shortcut_quit)
start_time = time.time()
is_save = False
frame_count = 0
writer = None
while True:
sensor1.get_img() # 提取并拼接复眼图像
sensor1.compute_deformation(check_motion = True, tail_frames = frame_to_recover, threshold = touch_threshold) # 计算形变和分解
# 提取信息:
# 向量场: VRAW, VNORMAL, VSHEAR: 原始,法向,切向
# 帧率,时间戳: FPS, TIMESTAMP
# 变形量总和: FNORMAL , FSHEARX, FSHEARY : 法向,切向(X),切向(Y)
fps, fn, fx, fy, flow = sensor1.read_info(
sensor1.info.FPS,
sensor1.info.FNORMAL,
sensor1.info.FSHEARX,
sensor1.info.FSHEARY,
sensor1.info.VRAW,
)
# print(fps,fn,fx,fy)
# 显示向量场 — BGR background; img may be 2D (H,W) or 3D (H,W,C)
_bg_bgr = (22, 16, 11)
_img = sensor1.img
_h, _w = _img.shape[:2]
black_img = np.empty((_h, _w, 3), dtype=_img.dtype)
black_img[:, :, 0] = _bg_bgr[0]
black_img[:, :, 1] = _bg_bgr[1]
black_img[:, :, 2] = _bg_bgr[2]
# Square side = min(label w,h) so pixmap matches centered KeepAspectRatio scaling
app.processEvents()
lw, lh = arrow_label.width(), arrow_label.height()
if lw > 16 and lh > 16:
side = max(32, min(lw, lh))
else:
side = compute_arrow_canvas_side(combined_win)
arrows = draw_arrows_wuxi(
black_img,
flow,
threshold=2,
grid_spacing=20,
arrow_scale=0.5,
below_threshold_color=(160, 160, 160),
target_size=int(side),
)
set_arrow_label_bgr(arrow_label, arrows)
# 显示图表 - 更新时间轴
current_time = time.time() - start_time
for i in range(3):
data[:-1, 0, i] = data[1:, 0, i] # shift time data left
data[:-1, 1, i] = data[1:, 1, i] # shift force data left
data[-1, 0, :] = frame_count / 15.0 # set current time for all plots
data[-1, 1, 0] = fn
data[-1, 1, 1] = abs(fx)
data[-1, 1, 2] = abs(fy)
update(plot_af, plot_tfx, plot_tfy, data, p1, p2, p3) # Update plot
frame_count += 1
app.processEvents()
if combined_win.user_closed or not combined_win.isVisible():
print("\n窗口已关闭,退出程序")
if writer is not None:
writer.release()
writer = None
break
# Convert combined window to image (arrows + curves)
pixmap = combined_win.grab()
# Convert QPixmap to QImage
qimg = pixmap.toImage()
# Convert QImage to numpy array (RGB format)
qimg = qimg.convertToFormat(QImage.Format_RGB888)
width = qimg.width()
height = qimg.height()
# Convert QImage to numpy array
ptr = qimg.bits()
ptr.setsize(height * width * 3)
arr = np.array(ptr, dtype=np.uint8).reshape((height, width, 3))
key = cv2.waitKey(1)
if key & 0xFF == ord("q"):
if writer is not None:
writer.release()
writer = None
print("\n 停止程序")
break
if key & 0xFF == ord("r"):
# 重置光流追踪器(当追踪漂移时使用);焦点需在 OpenCV 窗口上
sensor1.reset()
print("传感器已重置")
except BaseException as exc:
session_abnormal = datetime.now()
print(
f"[session] ABNORMAL at {session_abnormal:%Y-%m-%d %H:%M:%S}: "
f"{type(exc).__name__}: {exc}",
flush=True,
)
raise
finally:
session_end = datetime.now()
elapsed = session_end - session_start
print(
f"[session] END {session_end:%Y-%m-%d %H:%M:%S} "
f"(elapsed {elapsed})",
flush=True,
)
if sensor1 is not None:
try:
sensor1.disconnect()
except Exception as disconnect_err:
print(f"[session] disconnect error: {disconnect_err}", flush=True)
if __name__ == '__main__':
main()
+1 -1
View File
@@ -25,7 +25,7 @@ pip install orisys-{version}-{python-version}-{python-version}-{os}_{machine}.wh
```
如
```bash
pip install orisys-0.3.2-cp310-cp310-win_amd64.whl
pip install orisys-0.4.0-cp310-cp310-win_amd64.whl
```
3. **安装 CuPy(根据您的 CUDA 版本选择):**
+5 -5
View File
@@ -1,5 +1,5 @@
叠动科技软件使用手册
版本: 0.3.2
版本: 0.4.0
更新日期: 2026年4月16日
# 简介
@@ -57,9 +57,9 @@ orisys SDK 以 .whl 文件形式提供,您需要根据操作系统、Python
## 下载正确的 .whl 文件
从官方渠道获取 .whl 文件,文件名格式通常为:
orisys-{版本}-{Python版本}-{Python版本}-{操作系统}_{硬件架构}.whl
例如:orisys-0.3.2-cp310-cp310-win_amd64.whl 表示:
例如:orisys-0.4.0-cp310-cp310-win_amd64.whl 表示:
版本:0.3.2
版本:0.4.0
Python 版本:3.10(cp310)
操作系统:Windows(win)
硬件架构:64位(amd64)
@@ -71,7 +71,7 @@ Python 版本:3.10(cp310)
``` bash
# 请将文件名替换为您实际下载的文件名
pip install orisys-0.3.2-cp310-cp310-win_amd64.whl
pip install orisys-0.4.0-cp310-cp310-win_amd64.whl
```
## 安装 CuPy(可选)
@@ -108,7 +108,7 @@ python
import orisys
print(orisys.__version__)
```
如果输出版本号(例如 0.3.2),则说明安装成功。
如果输出版本号(例如 0.4.0),则说明安装成功。
# 使用方法
## 运行示例程序
+195
View File
@@ -0,0 +1,195 @@
#!/usr/bin/env python3
"""
Minimum UV-unroll finger tactile pipeline.
Offline: load Blender export -> build_or_load_domain -> unroll_map_uv.npz
Runtime: camera frame -> UVUnrollFrameProcessor -> Sensor -> forces/contact
For the full demo (UV debug panels, dashboard, benchmarks) see:
tools/3d_patch_generation_blender/unroll_realtime.py
"""
import argparse
import os
import sys
os.environ.setdefault("OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS", "0")
import cv2
import orisys
from orisys.configs import load_config
from orisys.frame_processors import UVUnrollFrameProcessor
from orisys.unroll import (
build_or_load_domain,
compute_uv_distortion_metrics,
load_calib_json,
load_exported_data,
recommend_edge_expand,
)
def open_source(video_path, camera_id):
return cv2.VideoCapture(video_path) if video_path else cv2.VideoCapture(camera_id)
def resolve_export_paths(export_dir):
"""Return Blender export paths expected inside one folder."""
export_dir = os.path.abspath(export_dir)
paths = {
"mesh": os.path.join(export_dir, "mesh.npz"),
"intrinsics": os.path.join(export_dir, "camera_intrinsics.npz"),
"pose": os.path.join(export_dir, "world_to_cam.npy"),
}
missing = [name for name, path in paths.items() if not os.path.isfile(path)]
if missing:
raise FileNotFoundError(
f"Missing export file(s) in {export_dir}: {', '.join(missing)}"
)
return paths["mesh"], paths["intrinsics"], paths["pose"]
def main():
parser = argparse.ArgumentParser(description="Minimum UV-unroll finger tactile pipeline")
parser.add_argument(
"--export-dir",
required=True,
help="Blender export folder containing mesh.npz, camera_intrinsics.npz, world_to_cam.npy",
)
parser.add_argument("--calib", default=None, help="OpenCV calibration JSON for the live camera")
parser.add_argument("--map_cache", default="unroll_map_uv.npz", help="Cached UV domain .npz")
parser.add_argument("--camera_id", type=int, default=0)
parser.add_argument("--video", help="Video file path (overrides camera)")
parser.add_argument("--config", "-c", default="./config/dd02-ov.json", help="Orisys SDK config")
parser.add_argument("--no_flip_source_vertical", action="store_true", help="Skip vertical flip before UV remap")
parser.add_argument("--show", action="store_true", help="Show unrolled image window")
parser.add_argument("--verbose", action="store_true", help="SDK verbose logs")
args = parser.parse_args()
sdk_config = load_config(args.config)
out_width, out_height = sdk_config["camera"]["resolution_stitch"]
mesh_path, intrinsics_path, pose_path = resolve_export_paths(args.export_dir)
data = load_exported_data(mesh_path, intrinsics_path, pose_path)
distortion = compute_uv_distortion_metrics(data)
edge_expand = recommend_edge_expand(distortion)
if args.calib:
source_K, dist_coeffs, img_size = load_calib_json(args.calib)
source_width, source_height = img_size
else:
source_K = data["K"]
dist_coeffs = None
source_width = data["width"]
source_height = data["height"]
domain = build_or_load_domain(
mesh_path,
intrinsics_path,
pose_path,
args.map_cache,
source_K,
source_width,
source_height,
out_width,
out_height,
edge_expand=edge_expand,
include_surface=True,
verbose=args.verbose,
)
domain.remap.ensure_fast_maps()
cap = open_source(args.video, args.camera_id)
if not cap.isOpened():
print("Cannot open camera/video source", file=sys.stderr)
sys.exit(1)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, source_width)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, source_height)
sensor = orisys.Sensor("external", config_name=args.config, verbose=args.verbose, mirror=False)
processor = UVUnrollFrameProcessor(
domain.remap,
camera_matrix=source_K if dist_coeffs is not None else None,
dist_coeffs=dist_coeffs,
source_width=source_width,
source_height=source_height,
flip_source_vertical=not args.no_flip_source_vertical,
)
ok, frame = cap.read()
if ok:
if args.show:
prepared = processor.prepare_source(frame)
sensor.get_img(
frame=processor.flatten_prepared(prepared),
already_processed=True,
valid_mask=processor.valid_mask,
)
else:
prepared_gray = processor.prepare_source_gray(frame)
sensor.get_img(
frame=processor.flatten_prepared_gray(prepared_gray),
already_processed=True,
already_gray=True,
valid_mask=processor.valid_mask,
)
sensor.warmup_deformation()
print("Press 'q' to quit, 'r' to reset tracker.")
try:
while True:
ok, frame = cap.read()
if not ok:
break
if args.show:
prepared = processor.prepare_source(frame)
unrolled = processor.flatten_prepared(prepared)
sensor.get_img(
frame=unrolled,
already_processed=True,
valid_mask=processor.valid_mask,
)
else:
prepared_gray = processor.prepare_source_gray(frame)
unrolled = processor.flatten_prepared_gray(prepared_gray)
sensor.get_img(
frame=unrolled,
already_processed=True,
already_gray=True,
valid_mask=processor.valid_mask,
)
sensor.compute_deformation(check_motion=True, threshold=0)
sensor.compute_contact()
fps, fnormal, fshearx, fsheary = sensor.read_info(
sensor.info.FPS,
sensor.info.FNORMAL,
sensor.info.FSHEARX,
sensor.info.FSHEARY,
)
print(
f"FPS={fps:.2f} normal={fnormal:.2f} shear=({fshearx:.2f}, {fsheary:.2f}) "
f"mask={sensor.mask_stats['coverage']:.2f}",
end="\r",
)
if args.show:
cv2.imshow("UV Unrolled", unrolled)
key = cv2.waitKey(1) & 0xFF
if key == ord("q"):
print("\nExiting.")
break
if key == ord("r"):
sensor.reset()
print("\nTracker reset.")
finally:
cap.release()
sensor.disconnect()
cv2.destroyAllWindows()
if __name__ == "__main__":
main()
+1 -1
View File
@@ -10,7 +10,7 @@ if config_dir.exists():
setup(
name="orisys",
version="0.3.2",
version="0.4.0",
description="Orisys - Optical Tactile Sensor SDK",
author="Your Name",
package_dir={"": "src"},
+1 -1
View File
@@ -5,7 +5,7 @@ A Python library for processing and analyzing optical tactile sensor data,
including image stitching, optical flow tracking, and deformation decomposition.
"""
__version__ = "0.3.2"
__version__ = "0.4.0"
__author__ = "Orisys"
from .sensor import Sensor
+196 -1
View File
@@ -1,6 +1,8 @@
import cv2
import numpy as np
from .unroll import UVRemap
class CompoundEyeFrameProcessor:
"""Convert a raw compound-eye frame into one stitched BGR tactile image."""
@@ -25,7 +27,196 @@ class SingleCameraFrameProcessor:
return frame
def create_frame_processor(sensor_type, stitcher=None):
class UVUnrollFrameProcessor:
"""Remap a raw camera frame into a flattened UV-domain BGR image."""
def __init__(
self,
remap: UVRemap,
camera_matrix=None,
dist_coeffs=None,
source_width=None,
source_height=None,
flip_source_vertical=True,
border_value=(0, 0, 0),
):
self.remap = remap
self.camera_matrix = camera_matrix
self.dist_coeffs = dist_coeffs
meta = remap.metadata
self.source_width = source_width if source_width is not None else int(meta.get("source_width", 0) or 0)
self.source_height = source_height if source_height is not None else int(meta.get("source_height", 0) or 0)
self.flip_source_vertical = flip_source_vertical
self.border_value = border_value
self._flip_buffer = None
self._gray_source_buffer = None
self._remap_dst_bgr = None
self._remap_dst_gray = None
self.runtime_remap = remap
self._use_composed_source_map = False
self._build_runtime_remap()
@property
def valid_mask(self):
"""(H, W) bool mask in unrolled coordinates; pass to Sensor.get_img(valid_mask=...)."""
return self.runtime_remap.valid_mask
def _build_runtime_remap(self):
"""Compose flip/undistort into the UV remap when source geometry is known."""
can_compose = self.source_width > 0 and self.source_height > 0 and (
self.flip_source_vertical or (self.camera_matrix is not None and self.dist_coeffs is not None)
)
if not can_compose:
self.runtime_remap.ensure_fast_maps()
return
lookup_x = self.remap.map_x.astype(np.float32, copy=True)
lookup_y = self.remap.map_y.astype(np.float32, copy=True)
if self.flip_source_vertical:
lookup_y = (self.source_height - 1.0) - lookup_y
if self.camera_matrix is not None and self.dist_coeffs is not None:
undistort_x, undistort_y = cv2.initUndistortRectifyMap(
self.camera_matrix,
self.dist_coeffs,
None,
self.camera_matrix,
(self.source_width, self.source_height),
cv2.CV_32FC1,
)
composed_x = cv2.remap(
undistort_x,
lookup_x,
lookup_y,
interpolation=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=-1,
)
composed_y = cv2.remap(
undistort_y,
lookup_x,
lookup_y,
interpolation=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=-1,
)
else:
composed_x = lookup_x
composed_y = lookup_y
invalid = ~self.remap.valid_mask
composed_x[invalid] = -1.0
composed_y[invalid] = -1.0
self.runtime_remap = UVRemap(
map_x=composed_x,
map_y=composed_y,
valid_mask=self.remap.valid_mask,
metadata=dict(self.remap.metadata),
)
self.runtime_remap.ensure_fast_maps()
self._use_composed_source_map = True
def _flip_vertical(self, frame_bgr):
if not self.flip_source_vertical:
return frame_bgr
if (
self._flip_buffer is None
or self._flip_buffer.shape != frame_bgr.shape
or self._flip_buffer.dtype != frame_bgr.dtype
):
self._flip_buffer = np.empty_like(frame_bgr)
cv2.flip(frame_bgr, 0, dst=self._flip_buffer)
return self._flip_buffer
def prepare_source(self, frame_bgr):
"""Resize, undistort, and orient raw frame before UV remap.
frame_bgr: (H_src, W_src, 3) raw camera BGR
returns: (H_src, W_src, 3) prepared BGR (not unrolled)
"""
if frame_bgr is None:
return None
frame = frame_bgr
if self.source_width > 0 and self.source_height > 0:
if frame.shape[1] != self.source_width or frame.shape[0] != self.source_height:
frame = cv2.resize(
frame,
(self.source_width, self.source_height),
interpolation=cv2.INTER_LINEAR,
)
if self._use_composed_source_map:
return frame
if self.camera_matrix is not None and self.dist_coeffs is not None:
frame = cv2.undistort(frame, self.camera_matrix, self.dist_coeffs)
return self._flip_vertical(frame)
def prepare_source_gray(self, frame_bgr):
"""Resize/orient source frame and convert once to grayscale."""
frame = self.prepare_source(frame_bgr)
if frame is None:
return None
if frame.ndim == 2:
return frame
if frame.ndim == 3 and frame.shape[2] == 1:
return frame[:, :, 0]
gray_shape = frame.shape[:2]
if self._gray_source_buffer is None or self._gray_source_buffer.shape != gray_shape:
self._gray_source_buffer = np.empty(gray_shape, dtype=np.uint8)
cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY, dst=self._gray_source_buffer)
return self._gray_source_buffer
def _remap_dst(self, prepared):
if prepared.ndim == 2:
if (
self._remap_dst_gray is None
or self._remap_dst_gray.shape != self.runtime_remap.valid_mask.shape
or self._remap_dst_gray.dtype != prepared.dtype
):
self._remap_dst_gray = np.empty(self.runtime_remap.valid_mask.shape, dtype=prepared.dtype)
return self._remap_dst_gray
shape = self.runtime_remap.valid_mask.shape + (prepared.shape[2],)
if (
self._remap_dst_bgr is None
or self._remap_dst_bgr.shape != shape
or self._remap_dst_bgr.dtype != prepared.dtype
):
self._remap_dst_bgr = np.empty(shape, dtype=prepared.dtype)
return self._remap_dst_bgr
def flatten_prepared(self, prepared_bgr):
"""Remap an already prepared source frame into UV flatten coordinates."""
if prepared_bgr is None:
return None
return self.runtime_remap.flatten(
prepared_bgr,
border_value=self.border_value,
dst=self._remap_dst(prepared_bgr),
)
def flatten_prepared_gray(self, prepared_gray):
"""Remap an already prepared grayscale frame into UV coordinates."""
if prepared_gray is None:
return None
return self.runtime_remap.flatten(
prepared_gray,
border_value=0,
dst=self._remap_dst(prepared_gray),
)
def process(self, frame_bgr):
"""Remap raw camera frame into UV flatten coordinates.
frame_bgr: (H_src, W_src, 3) raw camera BGR
returns: (H_out, W_out, 3) unrolled BGR
"""
return self.flatten_prepared(self.prepare_source(frame_bgr))
def create_frame_processor(sensor_type, stitcher=None, uv_remap=None, **uv_kwargs):
sensor_type = (sensor_type or "compound_eye").lower()
if sensor_type in ("compound_eye", "compound-eye", "compound"):
if stitcher is None:
@@ -33,4 +224,8 @@ def create_frame_processor(sensor_type, stitcher=None):
return CompoundEyeFrameProcessor(stitcher)
if sensor_type in ("single_camera", "camera", "mono"):
return SingleCameraFrameProcessor()
if sensor_type in ("uv_unroll", "unroll"):
if uv_remap is None:
raise ValueError("UV unroll sensor requires a UVRemap instance.")
return UVUnrollFrameProcessor(uv_remap, **uv_kwargs)
raise ValueError(f"Unsupported sensor type: {sensor_type}")
+34 -11
View File
@@ -68,7 +68,15 @@ class Contact3D:
return result
def compute_foe(self, flow_field, return_all_contours=False):
def _normalize_domain_mask(self, valid_mask, shape):
if valid_mask is None:
return None
mask = np.asarray(valid_mask, dtype=np.uint8)
if mask.shape != shape:
mask = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST)
return mask.astype(bool)
def compute_foe(self, flow_field, return_all_contours=False, valid_mask=None):
"""
计算光流场的扩张焦点(Focus of Expansion)。
@@ -128,13 +136,20 @@ class Contact3D:
if self.denoise_sigma > 0:
divergence = cv2.GaussianBlur(divergence, (0, 0), self.denoise_sigma)
domain_mask = self._normalize_domain_mask(valid_mask, (H, W))
if domain_mask is not None:
magnitude = np.where(domain_mask, magnitude, 0)
divergence = np.where(domain_mask, divergence, 0)
# 优化:使用原地操作 - 只保留大于阈值的散度值,其他设为0
divergence = np.where(divergence > self.divergence_threshold, divergence, 0)
# 创建掩码(使用原地操作优化)
valid_mask = magnitude > self.flow_threshold
expansion_mask = (valid_mask & (divergence > self.expansion_threshold))
contraction_mask = (valid_mask & (divergence < -self.expansion_threshold))
signal_mask = magnitude > self.flow_threshold
if domain_mask is not None:
signal_mask &= domain_mask
expansion_mask = (signal_mask & (divergence > self.expansion_threshold))
contraction_mask = (signal_mask & (divergence < -self.expansion_threshold))
# 初始化结果字典
result = self._create_foe_result(expansion_mask, contraction_mask, divergence)
@@ -166,7 +181,7 @@ class Contact3D:
self.foe_results = result
return self.foe_results
def compute_foe_gpu(self, flow_field, return_all_contours=False):
def compute_foe_gpu(self, flow_field, return_all_contours=False, valid_mask=None):
"""
使用CUDA加速计算光流场的扩张焦点(Focus of Expansion)。
@@ -187,7 +202,7 @@ class Contact3D:
"""
# 检查CuPy是否可用
if not CUPY_AVAILABLE:
return self.compute_foe(flow_field, return_all_contours)
return self.compute_foe(flow_field, return_all_contours, valid_mask=valid_mask)
H, W = flow_field.shape[:2]
@@ -254,14 +269,22 @@ class Contact3D:
if self.denoise_sigma > 0:
divergence_gpu = cp_gaussian_filter(divergence_gpu, sigma=self.denoise_sigma)
domain_mask = self._normalize_domain_mask(valid_mask, (H, W))
if domain_mask is not None:
domain_mask_gpu = cp.asarray(domain_mask)
magnitude_gpu = cp.where(domain_mask_gpu, magnitude_gpu, 0)
divergence_gpu = cp.where(domain_mask_gpu, divergence_gpu, 0)
# 优化:在GPU上完成掩码操作,减少数据传输
# 只保留大于阈值的散度值,其他设为0
divergence_gpu = cp.where(divergence_gpu > self.divergence_threshold, divergence_gpu, 0)
# 在GPU上创建掩码
valid_mask_gpu = magnitude_gpu > self.flow_threshold
expansion_mask_gpu = (valid_mask_gpu & (divergence_gpu > self.expansion_threshold))
contraction_mask_gpu = (valid_mask_gpu & (divergence_gpu < -self.expansion_threshold))
signal_mask_gpu = magnitude_gpu > self.flow_threshold
if domain_mask is not None:
signal_mask_gpu = signal_mask_gpu & domain_mask_gpu
expansion_mask_gpu = (signal_mask_gpu & (divergence_gpu > self.expansion_threshold))
contraction_mask_gpu = (signal_mask_gpu & (divergence_gpu < -self.expansion_threshold))
# 优化:一次性传输所有需要的数据回CPU
magnitude = cp.asnumpy(magnitude_gpu)
@@ -299,7 +322,7 @@ class Contact3D:
self.foe_results = result
return self.foe_results
def compute_multiple_foe(self,flow_field):
def compute_multiple_foe(self, flow_field, valid_mask=None):
"""
计算光流场的多个扩张焦点(Foci of Expansion)。
@@ -321,7 +344,7 @@ class Contact3D:
- 'divergence_map': 散度图
"""
# 调用 compute_foe_gpu 获取基础数据和所有轮廓
foe_result = self.compute_foe_gpu(flow_field, return_all_contours=True)
foe_result = self.compute_foe_gpu(flow_field, return_all_contours=True, valid_mask=valid_mask)
# 构建多焦点结果,复用基础数据
result = self._create_foe_result(
+8 -5
View File
@@ -6,11 +6,14 @@ This package is intended to hold the refactored implementation currently living
is moved here gradually.
Suggested migration target:
- `api.py` -> public `nHHD` class
- `grid.py` -> `StructuredGrid`
- `poisson.py` -> `PoissonSolver`
- `timing.py` -> `Timer`
- `kernels.py` -> helper kernels / backend-specific math
- `api.py` -> public `nHHD` class
- `grid.py` -> `StructuredGrid`
- `config.py` -> `MaskedPoissonConfig`
- `domain.py` -> `MaskedDomain`
- `poisson.py` -> `PoissonSolver` facade
- `solvers/` -> rectangular, masked sparse, capacitance FFT solvers
- `timing.py` -> `Timer`
- `kernels.py` -> helper kernels / backend-specific math
"""
from .api import nHHD
from .grid import StructuredGrid
+99 -28
View File
@@ -32,6 +32,12 @@ class nHHD(object):
self.debug = kwargs.get('debug', False)
self.poisson_backend = backend
self.masked_poisson_solver = kwargs.get('masked_poisson_solver', 'sparse')
self.masked_poisson_max_boundary = kwargs.get('masked_poisson_max_boundary', 96)
self.masked_poisson_boundary_stride = kwargs.get('masked_poisson_boundary_stride', 1)
self.masked_poisson_max_iters = kwargs.get('masked_poisson_max_iters', 3)
self.masked_poisson_tol = kwargs.get('masked_poisson_tol', 1e-4)
self.masked_poisson_reg = kwargs.get('masked_poisson_reg', 1e-8)
# Initialize CUDA streams for cupy backend
if self.poisson_backend == 'cupy':
@@ -50,7 +56,17 @@ class nHHD(object):
if 'grid' in args:
self.mesh = StructuredGrid(grid=kwargs['grid'], spacings=kwargs['spacings'])
self.psolver = PoissonSolver(grid=self.mesh.dims, spacings=self.mesh.dx, poisson_backend=self.poisson_backend)
self.psolver = PoissonSolver(
grid=self.mesh.dims,
spacings=self.mesh.dx,
poisson_backend=self.poisson_backend,
masked_poisson_solver=self.masked_poisson_solver,
masked_poisson_max_boundary=self.masked_poisson_max_boundary,
masked_poisson_boundary_stride=self.masked_poisson_boundary_stride,
masked_poisson_max_iters=self.masked_poisson_max_iters,
masked_poisson_tol=self.masked_poisson_tol,
masked_poisson_reg=self.masked_poisson_reg,
)
self.dim = self.mesh.dim
if (self.dim != 2) and (self.dim != 3):
@@ -59,7 +75,24 @@ class nHHD(object):
self._force_cpu_cache = None
self.psolver.prepare(False)
def get_force(self, vfield, verbose=False):
def _normalize_mask(self, valid_mask, shape):
if valid_mask is None:
return None
mask = numpy.asarray(valid_mask, dtype=bool)
if mask.shape != shape:
raise ValueError("valid_mask must match the vector field grid shape")
return mask
def get_force(self, vfield, verbose=False, valid_mask=None):
"""HHD force metrics from a 2D vector field.
vfield: (H, W, 2) optical flow or displacement field
valid_mask: optional (H, W) bool; True = active UV/domain pixels
returns: (flag, output) where output is [Sn, Stx, Sty, Stor] float32
"""
mask = self._normalize_mask(valid_mask, vfield.shape[:2]) if self.dim == 2 else None
if mask is not None:
return self.get_force_cpu(vfield, verbose, valid_mask=mask)
if self.poisson_backend == 'cupy':
return self.get_force_cuda(vfield, verbose)
else:
@@ -83,10 +116,16 @@ class nHHD(object):
return self.d
return cp.asnumpy(self._d_gpu)
def _solve_poisson_backend(self, field, verbose=False):
"""
Dispatch Poisson solve using the selected backend.
def _solve_poisson_backend(self, field, verbose=False, valid_mask=None):
"""Dispatch Poisson solve using the selected backend.
field: (H, W) scalar source (divergence or curl weight)
valid_mask: optional (H, W) bool; forces CPU masked solver when set
returns: (H, W) scalar potential
"""
if valid_mask is not None:
return self.psolver.solve(field, verbose, mask=valid_mask)
if self.poisson_backend == 'opencl':
return self.psolver.solve_opencl(field, verbose, profile=self.debug)
@@ -96,20 +135,25 @@ class nHHD(object):
if self.poisson_backend == 'cpu':
return self.psolver.solve(field, verbose)
# create the 3-component decomposition
def decompose(self, vfield, verbose=0):
# decompose for cpu and opencl backend
def decompose(self, vfield, verbose=0, valid_mask=None):
"""Helmholtz-Hodge decomposition of a vector field.
vfield: (H, W, 2) velocity/displacement field
valid_mask: optional (H, W) bool; True = active domain pixels
returns: None; sets self.d, self.r, self.h and intermediate potentials
"""
if vfield.shape[-1] != self.dim:
raise ValueError("nHHD.decompose requires a valid-dimensional vector field")
mask = self._normalize_mask(valid_mask, vfield.shape[:2]) if (valid_mask is not None and self.dim == 2) else None
if verbose > 0:
print('\n ==> nhhd.decompose', vfield.shape)
if self.dim == 2:
# compute div curl
(self.div, self.curlw) = self.mesh.divcurl(vfield)
(self.div, self.curlw) = self.mesh.divcurl(vfield, mask=mask)
if self.poisson_backend == 'opencl':
self.div = numpy.ascontiguousarray(self.div, dtype=numpy.float32)
@@ -120,21 +164,21 @@ class nHHD(object):
print('curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# compute potentials
if self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE):
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)
else:
self.nD = self._solve_poisson_backend(self.div, verbose > 0)
self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0)
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)
if verbose > 0:
print('nD =', self.nD.shape, self.nD.min(), self.nD.max())
print('nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())
# compute fields as gradients of potentials
self.d = self.mesh.gradient(self.nD)
self.d = self.mesh.gradient(self.nD, mask=mask)
if verbose > 0:
print('d =', self.d.shape, self.d.min(), self.d.max())
self.r = self.mesh.rotated_gradient(self.nRu)
self.r = self.mesh.rotated_gradient(self.nRu, mask=mask)
if verbose > 0:
print('r =', self.r.shape, self.r.min(), self.r.max())
else:
@@ -185,11 +229,17 @@ class nHHD(object):
self.h = numpy.add(self.d, self.r)
numpy.subtract(vfield, self.h, self.h)
if mask is not None:
self.d[~mask] = 0
self.r[~mask] = 0
self.h[~mask] = 0
self.nD[~mask] = 0
self.nRu[~mask] = 0
if verbose > 0:
print('h =', self.h.shape, self.h.min(), self.h.max())
def decompose_cuda(self, vfield, verbose=0):
def decompose_cuda(self, vfield, verbose=0, valid_mask=None):
"""
GPU-optimized decomposition that keeps intermediate results on GPU.
Returns GPU arrays to avoid CPU-GPU transfers.
@@ -201,6 +251,10 @@ class nHHD(object):
Returns:
Tuple of GPU arrays: (d_gpu, h_gpu, r_gpu, nRu_gpu) for 2D fields
"""
if valid_mask is not None:
self.decompose(vfield, verbose, valid_mask=valid_mask)
return None
if not CUPY_AVAILABLE:
if verbose > 0:
print("CuPy not available, falling back to CPU decompose")
@@ -283,27 +337,38 @@ class nHHD(object):
self.decompose(vfield, verbose)
return None
def get_force_cpu(self, vfield, verbose=False):
def get_force_cpu(self, vfield, verbose=False, valid_mask=None):
shape = (vfield.shape[0], vfield.shape[1])
if (not hasattr(self, '_force_cpu_cache') or
self._force_cpu_cache is None or
self._force_cpu_cache['shape'] != shape):
self._init_force_cpu_cache(shape)
grid = self._force_cpu_cache['grid']
mask = self._normalize_mask(valid_mask, shape) if valid_mask is not None else None
self.decompose(vfield, verbose)
self.decompose(vfield, verbose, valid_mask=mask)
# Optimized: use abs instead of absolute (slightly faster)
self.Sn = numpy.sum(numpy.abs(self.d))
if mask is None:
self.Sn = numpy.sum(numpy.abs(self.d))
self.Stx = numpy.sum(self.h[:, :, 0])
self.Sty = numpy.sum(self.h[:, :, 1])
nRu_flat = self.nRu.ravel()
nR_max_idx = numpy.argmax(nRu_flat)
nR_min_idx = numpy.argmin(nRu_flat)
else:
if not numpy.any(mask):
output = numpy.zeros(4, dtype=numpy.float32)
self.Sn = self.Stx = self.Sty = self.Stor = 0.0
return False, output
self.Sn = numpy.sum(numpy.abs(self.d[mask]))
self.Stx = numpy.sum(self.h[:, :, 0][mask])
self.Sty = numpy.sum(self.h[:, :, 1][mask])
masked_hi = numpy.where(mask, self.nRu, -numpy.inf)
masked_lo = numpy.where(mask, self.nRu, numpy.inf)
nR_max_idx = numpy.argmax(masked_hi.ravel())
nR_min_idx = numpy.argmin(masked_lo.ravel())
# Optimized: direct indexing and sum (already efficient)
self.Stx = numpy.sum(self.h[:, :, 0])
self.Sty = numpy.sum(self.h[:, :, 1])
# Optimized: find both max and min indices in single pass through flattened array
nRu_flat = self.nRu.ravel()
nR_max_idx = numpy.argmax(nRu_flat)
nR_min_idx = numpy.argmin(nRu_flat)
# Convert flat indices to 2D coordinates (faster than unravel_index)
nR_max_loc = (nR_max_idx // shape[1], nR_max_idx % shape[1])
nR_min_loc = (nR_min_idx // shape[1], nR_min_idx % shape[1])
@@ -325,6 +390,9 @@ class nHHD(object):
min_y, min_x = nR_min_loc[0], nR_min_loc[1]
cross_n = (grid_x - min_x) * r_y - (grid_y - min_y) * r_x
if mask is not None:
cross_p = numpy.where(mask, cross_p, 0)
cross_n = numpy.where(mask, cross_n, 0)
self.Stor = numpy.sum(cross_p) + numpy.sum(cross_n)
# Optimized: pre-allocate output array with explicit dtype
@@ -371,7 +439,7 @@ class nHHD(object):
self._force_gpu_cache['mgrid_1'], shape
).copy()
def get_force_cuda(self, vfield, verbose=False):
def get_force_cuda(self, vfield, verbose=False, valid_mask=None):
"""
Further optimized GPU force computation using CuPy ElementwiseKernel.
Fuses all operations into a single kernel for maximum performance.
@@ -385,10 +453,13 @@ class nHHD(object):
- flag: Boolean indicating if force thresholds are met
- output: Array [Sn, Stx, Sty, Stor] containing force components
"""
if valid_mask is not None:
return self.get_force_cpu(vfield, verbose, valid_mask=valid_mask)
if not CUPY_AVAILABLE:
if verbose:
print("CuPy not available, falling back to CPU get_force")
return self.get_force(vfield, verbose)
return self.get_force_cpu(vfield, verbose)
# Decompose the vector field
self.decompose_cuda(vfield, verbose)
+34
View File
@@ -0,0 +1,34 @@
from dataclasses import dataclass
@dataclass
class MaskedPoissonConfig:
"""Configuration for masked Poisson solvers on irregular UV domains."""
method: str = "sparse" # "sparse" | "capfft"
max_boundary: int = 96
boundary_stride: int = 1
max_iters: int = 3
tol: float = 1e-4
reg: float = 1e-8
def __post_init__(self):
if self.method not in ("sparse", "capfft"):
raise ValueError("method must be one of: 'sparse', 'capfft'")
self.max_boundary = max(0, int(self.max_boundary))
self.boundary_stride = max(1, int(self.boundary_stride))
self.max_iters = max(1, int(self.max_iters))
self.tol = float(self.tol)
self.reg = float(self.reg)
@classmethod
def from_kwargs(cls, kwargs):
"""Build config from PoissonSolver / nHHD constructor kwargs."""
return cls(
method=kwargs.get("masked_poisson_solver", "sparse"),
max_boundary=kwargs.get("masked_poisson_max_boundary", 96),
boundary_stride=kwargs.get("masked_poisson_boundary_stride", 1),
max_iters=kwargs.get("masked_poisson_max_iters", 3),
tol=kwargs.get("masked_poisson_tol", 1e-4),
reg=kwargs.get("masked_poisson_reg", 1e-8),
)
+213
View File
@@ -0,0 +1,213 @@
import numpy
from scipy import sparse
from scipy.sparse import linalg as sparse_linalg
from .config import MaskedPoissonConfig
def normalize_mask(mask, shape):
"""Normalize optional mask to bool array matching grid shape."""
if mask is None:
return None
mask_arr = numpy.asarray(mask, dtype=bool)
if mask_arr.shape != shape:
raise ValueError("Mask dimensions should match the Poisson grid")
return mask_arr
def compute_boundary_mask(active):
"""Return active pixels with at least one inactive 4-neighbor."""
up = numpy.zeros_like(active)
down = numpy.zeros_like(active)
left = numpy.zeros_like(active)
right = numpy.zeros_like(active)
up[1:, :] = active[:-1, :]
down[:-1, :] = active[1:, :]
left[:, 1:] = active[:, :-1]
right[:, :-1] = active[:, 1:]
fully_surrounded = up & down & left & right
return active & ~fully_surrounded
def sample_boundary_coords(boundary_coords, config: MaskedPoissonConfig):
"""Subsample boundary pixel coordinates for capacitance correction."""
if boundary_coords.size == 0:
return boundary_coords
sampled = boundary_coords
if config.boundary_stride > 1:
sampled = sampled[:: config.boundary_stride]
if config.max_boundary > 0 and len(sampled) > config.max_boundary:
picks = numpy.linspace(0, len(sampled) - 1, num=config.max_boundary, dtype=numpy.int32)
sampled = sampled[picks]
return numpy.ascontiguousarray(sampled, dtype=numpy.int32)
class MaskedDomain:
"""Reduced Laplacian system on active mask pixels (2D only)."""
__slots__ = (
"shape",
"mask",
"index_map",
"factor",
"lap",
"num_active",
"pinned",
"active_coords",
)
def __init__(
self,
shape,
mask,
index_map,
factor,
lap,
num_active,
pinned,
active_coords,
):
self.shape = shape
self.mask = mask
self.index_map = index_map
self.factor = factor
self.lap = lap
self.num_active = num_active
self.pinned = pinned
self.active_coords = active_coords
@classmethod
def build(cls, mask, spacings):
"""Build masked Laplacian domain from bool mask and grid spacings (dy, dx).
mask: (H, W) bool; True = active domain pixels
spacings: (dy, dx) grid spacing tuple
returns: MaskedDomain with factorized sparse Laplacian
"""
active = numpy.asarray(mask, dtype=bool)
h, w = active.shape
active_coords = numpy.argwhere(active)
num_active = int(active.sum())
if num_active == 0:
return cls(
shape=active.shape,
mask=active.copy(),
index_map=numpy.full(active.shape, -1, dtype=numpy.int32),
factor=None,
lap=None,
num_active=0,
pinned=0,
active_coords=active_coords,
)
index_map = numpy.full(active.shape, -1, dtype=numpy.int32)
index_map[active] = numpy.arange(num_active, dtype=numpy.int32)
dy, dx = float(spacings[0]), float(spacings[1])
coeff_y = 1.0 / (dy * dy)
coeff_x = 1.0 / (dx * dx)
rows = []
cols = []
data = []
pinned = 0
for row, col in active_coords:
idx = int(index_map[row, col])
if idx == pinned:
rows.append(idx)
cols.append(idx)
data.append(1.0)
continue
diag = 0.0
for drow, dcol, coeff in ((-1, 0, coeff_y), (1, 0, coeff_y), (0, -1, coeff_x), (0, 1, coeff_x)):
nrow = row + drow
ncol = col + dcol
if nrow < 0 or nrow >= h or ncol < 0 or ncol >= w or not active[nrow, ncol]:
continue
rows.append(idx)
cols.append(int(index_map[nrow, ncol]))
data.append(-coeff)
diag += coeff
if diag == 0.0:
rows.append(idx)
cols.append(idx)
data.append(1.0)
else:
rows.append(idx)
cols.append(idx)
data.append(diag)
lap = sparse.csc_matrix((data, (rows, cols)), shape=(num_active, num_active), dtype=numpy.float64)
factor = sparse_linalg.factorized(lap)
return cls(
shape=active.shape,
mask=active.copy(),
index_map=index_map,
factor=factor,
lap=lap,
num_active=num_active,
pinned=pinned,
active_coords=active_coords,
)
def matches(self, mask):
"""True if this domain was built for the same mask array."""
mask_arr = numpy.asarray(mask, dtype=bool)
return self.shape == mask_arr.shape and numpy.array_equal(self.mask, mask_arr)
def prepare_rhs(self, f):
"""Extract gauge-fixed RHS on active pixels.
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
rhs -= rhs.mean()
rhs[self.pinned] = 0.0
return rhs
def scatter(self, sol_active, template):
"""Map reduced solution back to full grid.
sol_active: (N,) potential on active pixels
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 sol_active.size:
out[self.mask] = numpy.asarray(sol_active, dtype=numpy.float32)
return out
def residual(self, field, rhs_active):
"""Compute masked Laplacian residual L u - rhs on active pixels.
field: (H, W) full-grid field
rhs_active: (N,) RHS on active pixels (or None for homogeneous)
returns: (N,) residual vector
"""
vec = numpy.asarray(field, dtype=numpy.float64)[self.mask]
residual = self.lap.dot(vec)
if rhs_active is not None:
residual = residual - rhs_active
return residual
def to_cache_dict(self, capfft=None):
"""Legacy dict view for tests and facade compatibility."""
return {
"shape": self.shape,
"mask": self.mask,
"index_map": self.index_map,
"factor": self.factor,
"lap": self.lap,
"num_active": self.num_active,
"pinned": self.pinned,
"active_coords": self.active_coords,
"capfft": capfft,
}
+121 -7
View File
@@ -53,9 +53,87 @@ class StructuredGrid(object):
print(' Done!', end ='')
mtimer.end()
def _normalize_mask(self, mask):
if mask is None:
return None
mask_arr = numpy.asarray(mask, dtype=bool)
if tuple(mask_arr.shape) != tuple(self.dims):
raise ValueError("Mask dimensions should match the grid")
return mask_arr
def _masked_derivative_2d(self, field, mask, axis, spacing):
deriv = numpy.zeros_like(field, dtype=field.dtype)
active = mask.astype(bool)
if axis == 0:
prev_active = numpy.zeros_like(active)
next_active = numpy.zeros_like(active)
prev_active[1:, :] = active[:-1, :]
next_active[:-1, :] = active[1:, :]
both = active & prev_active & next_active
forward = active & ~prev_active & next_active
backward = active & prev_active & ~next_active
centered_view = deriv[1:-1, :]
centered_mask = both[1:-1, :]
centered_view[centered_mask] = (
field[2:, :][centered_mask] - field[:-2, :][centered_mask]
) / (2.0 * spacing)
forward_view = deriv[:-1, :]
forward_mask = forward[:-1, :]
forward_view[forward_mask] = (
field[1:, :][forward_mask] - field[:-1, :][forward_mask]
) / spacing
backward_view = deriv[1:, :]
backward_mask = backward[1:, :]
backward_view[backward_mask] = (
field[1:, :][backward_mask] - field[:-1, :][backward_mask]
) / spacing
elif axis == 1:
prev_active = numpy.zeros_like(active)
next_active = numpy.zeros_like(active)
prev_active[:, 1:] = active[:, :-1]
next_active[:, :-1] = active[:, 1:]
both = active & prev_active & next_active
forward = active & ~prev_active & next_active
backward = active & prev_active & ~next_active
centered_view = deriv[:, 1:-1]
centered_mask = both[:, 1:-1]
centered_view[centered_mask] = (
field[:, 2:][centered_mask] - field[:, :-2][centered_mask]
) / (2.0 * spacing)
forward_view = deriv[:, :-1]
forward_mask = forward[:, :-1]
forward_view[forward_mask] = (
field[:, 1:][forward_mask] - field[:, :-1][forward_mask]
) / spacing
backward_view = deriv[:, 1:]
backward_mask = backward[:, 1:]
backward_view[backward_mask] = (
field[:, 1:][backward_mask] - field[:, :-1][backward_mask]
) / spacing
else:
raise ValueError("axis must be 0 or 1 for 2D masked derivatives")
deriv[~active] = 0
return deriv
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
def divcurl(self, vfield, verbose=False):
def divcurl(self, vfield, verbose=False, mask=None):
"""Compute divergence and curl of a vector field.
vfield: (H, W, 2) or (Z, Y, X, 3) velocity/displacement field
mask: optional (H, W) bool; True = active domain pixels (2D only)
returns: 2D -> (div, curlw); 3D -> (div, curlu, curlv, curlw); zeros outside mask
"""
# compute div and curl of vfield (vector field)
#if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] - self.dims).any():
if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
@@ -66,7 +144,19 @@ class StructuredGrid(object):
sys.stdout.flush()
mtimer = Timer()
mask_arr = self._normalize_mask(mask)
if self.dim == 2:
if mask_arr is not None:
dudy = self._masked_derivative_2d(vfield[:, :, 0], mask_arr, axis=0, spacing=self.dx[0])
dudx = self._masked_derivative_2d(vfield[:, :, 0], mask_arr, axis=1, spacing=self.dx[1])
dvdy = self._masked_derivative_2d(vfield[:, :, 1], mask_arr, axis=0, spacing=self.dx[0])
dvdx = self._masked_derivative_2d(vfield[:, :, 1], mask_arr, axis=1, spacing=self.dx[1])
numpy.add(dudx, dvdy, dudx)
numpy.subtract(dvdx, dudy, dvdx)
dudx[~mask_arr] = 0
dvdx[~mask_arr] = 0
return (dudx, dvdx)
# self.dx = (dy,dx)
dudy, dudx = numpy.gradient(vfield[:,:,0], self.dx[0], self.dx[1])
@@ -242,8 +332,13 @@ class StructuredGrid(object):
return (curlu, curlv, curlw)
def rotated_gradient(self, sfield, verbose=False):
def rotated_gradient(self, sfield, verbose=False, mask=None):
"""Compute rotated gradient (-d/dy, d/dx) of a scalar potential.
sfield: (H, W) scalar potential field
mask: optional (H, W) bool; True = active domain pixels
returns: (H, W, 2) rotational component field; zero outside mask
"""
if (sfield.shape != self.dims):
#if (sfield.shape - self.dims).any():
raise ValueError("Dimensions of scalar field should match that of the grid")
@@ -256,10 +351,17 @@ class StructuredGrid(object):
sys.stdout.flush()
mtimer = Timer()
ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1])
mask_arr = self._normalize_mask(mask)
if mask_arr is not None:
ddy = self._masked_derivative_2d(sfield, mask_arr, axis=0, spacing=self.dx[0])
ddx = self._masked_derivative_2d(sfield, mask_arr, axis=1, spacing=self.dx[1])
else:
ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1])
ddy *= -1.0
grad = numpy.stack((ddy, ddx), axis=-1)
if mask_arr is not None:
grad[~mask_arr] = 0
if verbose:
print(' Done!', end ='')
@@ -305,8 +407,13 @@ class StructuredGrid(object):
return grad
def gradient(self, sfield, verbose=False):
def gradient(self, sfield, verbose=False, mask=None):
"""Compute gradient (d/dx, d/dy) of a scalar potential.
sfield: (H, W) or (Z, Y, X) scalar potential field
mask: optional (H, W) bool; True = active domain pixels (2D only)
returns: (H, W, 2) or (Z, Y, X, 3) gradient field; zero outside mask
"""
if (sfield.shape != self.dims):
#if (sfield.shape - self.dims).any():
raise ValueError("Dimensions of scalar field should match that of the grid")
@@ -316,10 +423,14 @@ class StructuredGrid(object):
sys.stdout.flush()
mtimer = Timer()
mask_arr = self._normalize_mask(mask)
if self.dim == 2:
# self.dx = (dy,dx)
ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1])
if mask_arr is not None:
ddy = self._masked_derivative_2d(sfield, mask_arr, axis=0, spacing=self.dx[0])
ddx = self._masked_derivative_2d(sfield, mask_arr, axis=1, spacing=self.dx[1])
else:
# self.dx = (dy,dx)
ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1])
grad = numpy.stack((ddx, ddy), axis = -1)
elif self.dim == 3:
@@ -328,6 +439,9 @@ class StructuredGrid(object):
ddz, ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1], self.dx[2])
grad = numpy.stack((ddx, ddy, ddz), axis = -1)
if mask_arr is not None:
grad[~mask_arr] = 0
if verbose:
print(' Done!', end ='')
mtimer.end()
+209 -822
View File
File diff suppressed because it is too large Load Diff
+5
View File
@@ -0,0 +1,5 @@
from .capacitance import CapacitanceFFTSolver
from .masked_sparse import MaskedSparseSolver
from .rectangular import RectangularFFTSolver
__all__ = ["RectangularFFTSolver", "MaskedSparseSolver", "CapacitanceFFTSolver"]
+3
View File
@@ -0,0 +1,3 @@
"""Shared solver protocol and helpers."""
__all__ = []
+113
View File
@@ -0,0 +1,113 @@
import numpy
from scipy import linalg as dense_linalg
from ..config import MaskedPoissonConfig
from ..domain import MaskedDomain, compute_boundary_mask, sample_boundary_coords
class CapacitanceFFTSolver:
"""Masked Poisson via rectangular FFT bulk solve + capacitance boundary correction."""
def __init__(self, config: MaskedPoissonConfig, bulk_solver):
self.config = config
self.bulk_solver = bulk_solver
self._capfft_cache = None
self._capfft_domain = None
def prepare(self, domain: MaskedDomain):
"""Precompute capacitance matrix for boundary correction.
domain: MaskedDomain with built Laplacian
returns: capacitance cache dict (also stored internally)
"""
self._capfft_domain = domain
self._capfft_cache = self._build_capfft_cache(domain)
return self._capfft_cache
def _build_capfft_cache(self, domain: MaskedDomain):
active = domain.mask
boundary_coords = numpy.argwhere(compute_boundary_mask(active))
boundary_coords = sample_boundary_coords(boundary_coords, self.config)
if boundary_coords.size == 0:
pinned_rc = numpy.array(domain.active_coords[domain.pinned], dtype=numpy.int32)
boundary_coords = pinned_rc.reshape(1, 2)
boundary_active_indices = domain.index_map[boundary_coords[:, 0], boundary_coords[:, 1]]
boundary_count = int(len(boundary_active_indices))
cap_matrix = numpy.zeros((boundary_count, boundary_count), dtype=numpy.float64)
zero_rhs = numpy.zeros(domain.num_active, dtype=numpy.float64)
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)
residual = domain.residual(response, zero_rhs)
cap_matrix[:, col] = residual[boundary_active_indices]
reg_scale = self.config.reg
if reg_scale <= 0:
reg_scale = 1e-8
reg = reg_scale * max(1.0, float(numpy.linalg.norm(cap_matrix, ord=numpy.inf)))
cap_matrix_reg = cap_matrix + reg * numpy.eye(boundary_count, dtype=numpy.float64)
lu, piv = dense_linalg.lu_factor(cap_matrix_reg)
return {
"boundary_coords": boundary_coords,
"boundary_active_indices": boundary_active_indices,
"boundary_count": boundary_count,
"cap_matrix": cap_matrix,
"cap_matrix_reg": cap_matrix_reg,
"cap_lu": lu,
"cap_piv": piv,
"reg": reg,
}
def get_cache(self, domain: MaskedDomain):
"""Return cached capacitance data, building if needed."""
if self._capfft_cache is None or self._capfft_domain is not domain:
self.prepare(domain)
return self._capfft_cache
def solve(self, f, domain: MaskedDomain):
"""Solve masked Poisson with capacitance-corrected rectangular FFT.
f: (H, W) scalar source field
domain: MaskedDomain with prebuilt Laplacian
returns: (H, W) scalar potential; zero outside mask
"""
rhs_active = domain.prepare_rhs(f)
if rhs_active.size == 0:
return numpy.zeros_like(f, dtype=numpy.float32)
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)
cap_cache = self.get_cache(domain)
if not cap_cache or cap_cache["boundary_count"] == 0:
field_active = field[domain.mask]
if field_active.size:
field_active = field_active.copy()
field_active -= field_active[domain.pinned]
return domain.scatter(field_active, f)
boundary_rows = cap_cache["boundary_active_indices"]
rhs_scale = max(1.0, float(numpy.linalg.norm(rhs_active)))
for _ in range(self.config.max_iters):
residual = domain.residual(field, rhs_active)
boundary_residual = residual[boundary_rows]
if float(numpy.linalg.norm(boundary_residual)) <= self.config.tol * rhs_scale:
break
alpha = dense_linalg.lu_solve((cap_cache["cap_lu"], cap_cache["cap_piv"]), boundary_residual)
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)
field -= correction
field_active = field[domain.mask]
if field_active.size:
field_active = field_active.copy()
field_active -= field_active[domain.pinned]
return domain.scatter(field_active, f)
+21
View File
@@ -0,0 +1,21 @@
import numpy
from ..domain import MaskedDomain
class MaskedSparseSolver:
"""Reference masked Poisson solver via factorized sparse Laplacian."""
def solve(self, f, domain: MaskedDomain):
"""Solve masked Poisson on active domain pixels.
f: (H, W) scalar source field
domain: MaskedDomain with prebuilt Laplacian
returns: (H, W) scalar potential; zero outside mask
"""
rhs = domain.prepare_rhs(f)
if rhs.size == 0:
return numpy.zeros_like(f, dtype=numpy.float32)
sol_active = domain.factor(rhs) # (N,) reduced-system potential
return domain.scatter(sol_active, f)
+530
View File
@@ -0,0 +1,530 @@
import numpy
import timeit
from scipy import signal
from ...backend_availability import (
CUPY_AVAILABLE,
OPENCL_AVAILABLE,
cp,
cp_signal,
cl,
cla,
cl_rfftn,
cl_irfftn,
)
from ..timing import Timer
class RectangularFFTSolver:
"""FFT Poisson solver on a full structured rectangular grid."""
def __init__(self, grid, spacings, poisson_backend="cpu", dtype=numpy.float32, stype="F"):
self.gdims = grid
self.gdx = spacings
self.dim = len(self.gdims)
self.dtype = dtype
self.stype = stype
self.poisson_backend = poisson_backend
self.ready = False
self.G = None
if self.poisson_backend == "cupy":
self._gpu_cache = None
if self.poisson_backend == "opencl":
self._opencl_cache = None
self._opencl_queue = None
self._opencl_context = None
self._opencl_crop_kernel = None
self._opencl_scatter_kernel = None
self._opencl_fused_mul_kernel = None
self._opencl_scatter_batch_kernel = None
self._opencl_crop_batch_kernel = None
self._opencl_batch_mul_kernel = None
self._opencl_last_timing = None
self._opencv_cache = None
def _create_radialGrid(self, verbose):
mtimer = Timer()
if verbose:
print(" - creating distance kernel:", end="")
rdims = tuple(2 * d - 1 for d in self.gdims)
R = numpy.indices(rdims, dtype=self.dtype)
for d in range(self.dim):
R[d] += 1 - self.gdims[d]
R[d] *= self.gdx[d]
if self.dim == 1:
numpy.absolute(R[0], R[0])
R = R[0]
elif self.dim == 2:
numpy.hypot(R[1], R[0], R[0])
R = R[0]
zval = 0.5 * (self.gdx[0] + self.gdx[1]) / 2.0
R[self.gdims[0] - 1, self.gdims[1] - 1] = zval
elif self.dim == 3:
numpy.hypot(R[2], R[1], R[1])
numpy.hypot(R[1], R[0], R[0])
R = R[0]
zval = 0.5 * (self.gdx[0] + self.gdx[1] + self.gdx[2]) / 3.0
R[self.gdims[0] - 1, self.gdims[1] - 1, self.gdims[2] - 1] = zval
if verbose:
print(R.shape, R.min(), R.max())
mtimer.end()
return R
def prepare(self, verbose=False):
"""Precompute Green's function for FFT convolution.
returns: None (sets self.G and self.ready)
"""
gtimer = Timer()
if verbose:
print("\nInitializing rectangular FFT Poisson solver, type =", self.stype)
self.G = None
if self.stype == "F":
self.G = self._create_radialGrid(verbose)
elif self.stype == "S":
raise SyntaxError("Spatial solver not suggested for structured grids")
ltimer = Timer()
if verbose:
print(" - computing the Green's function:")
if self.dim == 1:
numpy.multiply(self.G, 0.5, self.G)
elif self.dim == 2:
numpy.log(self.G, self.G)
numpy.multiply(self.G, (0.5 / numpy.pi), self.G)
elif self.dim == 3:
numpy.reciprocal(self.G, self.G)
numpy.multiply(self.G, (-0.25 / numpy.pi), self.G)
if verbose:
print(self.G.shape, self.G.min(), self.G.max(), end="")
ltimer.end()
self.ready = True
if verbose:
print("Rectangular FFT Poisson solver initialized", end="")
gtimer.end()
def solve(self, f, verbose=False):
"""Solve Poisson on the full rectangular grid via FFT convolution.
f: (H, W) or (Z, Y, X) scalar source field
returns: (H, W) scalar potential, same shape as f
"""
if not self.ready:
self.prepare(verbose)
fshape = f.shape
if self.gdims != fshape:
print(self.gdims, fshape)
raise ValueError("Shape of function should match shape of grid")
gtimer = Timer()
if verbose:
print("\nSolving rectangular Poisson Eq.")
if self.stype == "F":
p = signal.fftconvolve(f, self.G, mode="same")
numpy.multiply(p, numpy.prod(self.gdx), p)
else:
raise SyntaxError("Spatial solver not suggested for structured grids")
if verbose:
print("Rectangular Poisson solver finished")
gtimer.end()
return p
def solve_gpu(self, f, verbose=False, return_gpu=False):
"""GPU FFT Poisson solve using CuPy.
f: (H, W) scalar source field on CPU
return_gpu: if True, return CuPy array instead of numpy
returns: (H, W) scalar potential
"""
if not CUPY_AVAILABLE:
if verbose:
print("CuPy not available, falling back to CPU solver")
return self.solve(f, verbose)
if not self.ready:
self.prepare(verbose)
fshape = f.shape
gtimer = Timer()
if verbose:
print("\nSolving Poisson Eq. on GPU")
if self.gdims != fshape:
print(self.gdims, fshape)
raise ValueError("Shape of function should match shape of grid")
if self.stype == "F":
if self._gpu_cache is None:
if verbose:
print(" - Transferring Green's function to GPU...")
self._gpu_cache = {
"G_gpu": cp.asarray(self.G, dtype=cp.float32),
"gdx_prod": numpy.prod(self.gdx),
}
f_gpu = cp.asarray(f, dtype=cp.float32)
p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache["G_gpu"], mode="same")
p_gpu *= self._gpu_cache["gdx_prod"]
if verbose:
print("Poisson solver (GPU) finished")
gtimer.end()
if return_gpu:
return p_gpu
return cp.asnumpy(p_gpu)
raise SyntaxError("Spatial solver not suggested for structured grids")
def _solve_gpu_internal(self, f_gpu):
"""Internal GPU solve; f_gpu and return value are CuPy arrays."""
if self.stype == "F":
if self._gpu_cache is None:
self._gpu_cache = {
"G_gpu": cp.asarray(self.G, dtype=cp.float32),
"gdx_prod": numpy.prod(self.gdx),
}
p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache["G_gpu"], mode="same")
p_gpu *= self._gpu_cache["gdx_prod"]
return p_gpu
raise SyntaxError("Spatial solver not suggested for structured grids")
def clear_gpu_cache(self):
"""Release CuPy FFT cache buffers."""
if self._gpu_cache is not None:
self._gpu_cache = None
if CUPY_AVAILABLE:
cp.get_default_memory_pool().free_all_blocks()
def _get_opencl_queue(self):
if self._opencl_queue is not None:
return self._opencl_queue
if not OPENCL_AVAILABLE:
raise RuntimeError("OpenCL/pyvkfft is not available")
ctx = None
for platform in cl.get_platforms():
for device in platform.get_devices():
if device.type & cl.device_type.GPU == 0:
continue
ctx = cl.Context(devices=(device,))
break
if ctx is not None:
break
if ctx is None:
raise RuntimeError("No OpenCL GPU device found")
self._opencl_queue = cl.CommandQueue(ctx)
self._opencl_context = ctx
return self._opencl_queue
def _get_opencl_scatter_batch_kernel(self, queue):
if self._opencl_scatter_batch_kernel is not None:
return self._opencl_scatter_batch_kernel
kernel_code = """
__kernel void scatter_batch(
__global const float* src,
__global float* dst,
const int src_width,
const int src_height,
const int dst_width,
const int src_plane,
const int dst_plane
) {
int gid_x = get_global_id(0);
int gid_y = get_global_id(1);
int gid_z = get_global_id(2);
if (gid_x < src_width && gid_y < src_height) {
int src_idx = gid_z * src_plane + gid_y * src_width + gid_x;
int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
dst[dst_idx] = src[src_idx];
}
}
"""
program = cl.Program(self._opencl_context, kernel_code).build()
self._opencl_scatter_batch_kernel = program.scatter_batch
return self._opencl_scatter_batch_kernel
def _get_opencl_crop_batch_kernel(self, queue):
if self._opencl_crop_batch_kernel is not None:
return self._opencl_crop_batch_kernel
kernel_code = """
__kernel void crop_batch(
__global const float* src,
__global float* dst,
const int src_width,
const int dst_width,
const int dst_height,
const int start_x,
const int start_y,
const float scale,
const int src_plane,
const int dst_plane
) {
int gid_x = get_global_id(0);
int gid_y = get_global_id(1);
int gid_z = get_global_id(2);
if (gid_x < dst_width && gid_y < dst_height) {
int src_idx = gid_z * src_plane + (start_y + gid_y) * src_width + (start_x + gid_x);
int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
dst[dst_idx] = src[src_idx] * scale;
}
}
"""
program = cl.Program(self._opencl_context, kernel_code).build()
self._opencl_crop_batch_kernel = program.crop_batch
return self._opencl_crop_batch_kernel
def _get_opencl_batch_mul_kernel(self, queue):
if self._opencl_batch_mul_kernel is not None:
return self._opencl_batch_mul_kernel
kernel_code = """
__kernel void batch_complex_mul(
__global float2* fft_batch,
__global const float2* g,
const int n
) {
const int b = get_global_id(0);
const int base = b * n;
const int i0 = get_global_id(1) * 4;
const int i_end = min(i0 + 4, n);
for (int i = i0; i < i_end; i++) {
float2 gv = g[i];
float2 a = fft_batch[base + i];
fft_batch[base + i] = (float2)(fma(a.x, gv.x, -(a.y * gv.y)),
fma(a.x, gv.y, a.y * gv.x));
}
}
"""
program = cl.Program(self._opencl_context, kernel_code).build()
self._opencl_batch_mul_kernel = program.batch_complex_mul
return self._opencl_batch_mul_kernel
def _get_opencl_alloc_host_ptr(self):
ctx = self._opencl_context
mf = cl.mem_flags
def allocator(nbytes):
return cl.Buffer(ctx, mf.READ_WRITE | mf.ALLOC_HOST_PTR, size=nbytes)
return allocator
@staticmethod
def _center_crop_same(full_conv, target_shape):
starts = [(full_conv.shape[i] - target_shape[i]) // 2 for i in range(len(target_shape))]
slices = tuple(slice(starts[i], starts[i] + target_shape[i]) for i in range(len(target_shape)))
return full_conv[slices]
def solve_batch_opencl(self, f1, f2, verbose=False, profile=False):
"""Batched OpenCL FFT Poisson solve for two 2D fields.
f1, f2: (H, W) scalar source fields, same shape
returns: (p1, p2) each (H, W) scalar potential
"""
if not OPENCL_AVAILABLE:
if verbose:
print("OpenCL/pyvkfft not available, falling back to CPU solver")
return self.solve(f1, verbose), self.solve(f2, verbose)
if not self.ready:
self.prepare(verbose)
fshape = f1.shape
if fshape != f2.shape:
raise ValueError("solve_batch_opencl requires inputs with identical shapes")
if self.gdims != fshape:
print(self.gdims, fshape)
raise ValueError("Shape of function should match shape of grid")
timing = {"cache_prepare_ms": 0.0}
t_total = timeit.default_timer()
t0 = timeit.default_timer()
queue = self._get_opencl_queue()
timing["queue_ms"] = (timeit.default_timer() - t0) * 1000.0
t0 = timeit.default_timer()
full_shape = tuple(2 * s - 1 for s in fshape)
timing["host_prepare_ms"] = (timeit.default_timer() - t0) * 1000.0
cache_valid = (
self._opencl_cache is not None
and self._opencl_cache["fshape"] == fshape
and self._opencl_cache["full_shape"] == full_shape
and "f_batch_dev" in self._opencl_cache
)
if not cache_valid:
t0 = timeit.default_timer()
g_host = numpy.asarray(self.G, dtype=numpy.float32)
g_pad = numpy.zeros(full_shape, dtype=numpy.float32)
g_slices = tuple(slice(0, s) for s in g_host.shape)
g_pad[g_slices] = g_host
g_dev = cla.to_device(queue, g_pad)
g_fft_dev = cl_rfftn(g_dev)
if profile:
queue.finish()
same_starts = [fshape[i] - 1 for i in range(self.dim)]
fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1])
fft_n_elements = 1
for d in fft_shape:
fft_n_elements *= d
src_plane = fshape[0] * fshape[1]
dst_plane = full_shape[0] * full_shape[1]
alloc_pinned = self._get_opencl_alloc_host_ptr()
self._opencl_cache = {
"fshape": fshape,
"full_shape": full_shape,
"fft_shape": fft_shape,
"fft_n_elements": fft_n_elements,
"g_fft_dev": g_fft_dev,
"same_starts": same_starts,
"gdx_prod": numpy.prod(self.gdx),
"src_plane": src_plane,
"dst_plane": dst_plane,
"f_batch_dev": cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32),
"host_stage": numpy.empty((2,) + fshape, dtype=numpy.float32),
"f_stage_dev": cla.empty(queue, (2,) + fshape, dtype=numpy.float32),
"f_fft_batch_dev": cla.empty(queue, (2,) + fft_shape, dtype=numpy.complex64),
"conv_batch_dev": cla.empty(queue, (2,) + full_shape, dtype=numpy.float32),
"p_batch_dev": cla.empty(queue, (2,) + fshape, dtype=numpy.float32, allocator=alloc_pinned),
"p1_host": numpy.empty(fshape, dtype=numpy.float32),
"p2_host": numpy.empty(fshape, dtype=numpy.float32),
}
timing["cache_prepare_ms"] = (timeit.default_timer() - t0) * 1000.0
c = self._opencl_cache
fshape_c = c["fshape"]
full_shape_c = c["full_shape"]
t0 = timeit.default_timer()
H, W = fshape_c
f1_f32 = f1 if (f1.flags["C_CONTIGUOUS"] and f1.dtype == numpy.float32) else numpy.ascontiguousarray(f1, dtype=numpy.float32)
f2_f32 = f2 if (f2.flags["C_CONTIGUOUS"] and f2.dtype == numpy.float32) else numpy.ascontiguousarray(f2, dtype=numpy.float32)
host_stage = c["host_stage"]
numpy.copyto(host_stage[0], f1_f32)
numpy.copyto(host_stage[1], f2_f32)
f_stage_dev = c["f_stage_dev"]
cl.enqueue_copy(queue, f_stage_dev.data, host_stage, is_blocking=False)
c["f_batch_dev"].fill(0, queue=queue)
scatter_batch = self._get_opencl_scatter_batch_kernel(queue)
scatter_batch(
queue,
(W, H, 2),
None,
f_stage_dev.data,
c["f_batch_dev"].data,
numpy.int32(W),
numpy.int32(H),
numpy.int32(full_shape_c[1]),
numpy.int32(H * W),
numpy.int32(c["dst_plane"]),
)
if profile:
queue.finish()
timing["h2d_set_ms"] = (timeit.default_timer() - t0) * 1000.0
t0 = timeit.default_timer()
fft_batch = cl_rfftn(c["f_batch_dev"], c["f_fft_batch_dev"], ndim=2)
if profile:
queue.finish()
timing["fft_f_ms"] = (timeit.default_timer() - t0) * 1000.0
t0 = timeit.default_timer()
n_fft = c["fft_n_elements"]
global_n = max(1, (n_fft + 3) // 4)
local_n = min(256, global_n)
global_n = ((global_n + local_n - 1) // local_n) * local_n
batch_mul = self._get_opencl_batch_mul_kernel(queue)
batch_mul(
queue,
(2, global_n),
(1, local_n),
fft_batch.data,
c["g_fft_dev"].data,
numpy.int32(n_fft),
)
if profile:
queue.finish()
timing["mul_ms"] = (timeit.default_timer() - t0) * 1000.0
t0 = timeit.default_timer()
cl_irfftn(fft_batch, c["conv_batch_dev"], ndim=2)
if profile:
queue.finish()
timing["ifft_ms"] = (timeit.default_timer() - t0) * 1000.0
t0 = timeit.default_timer()
crop_batch = self._get_opencl_crop_batch_kernel(queue)
same_starts = c["same_starts"]
crop_batch(
queue,
(fshape_c[1], fshape_c[0], 2),
None,
c["conv_batch_dev"].data,
c["p_batch_dev"].data,
numpy.int32(full_shape_c[1]),
numpy.int32(fshape_c[1]),
numpy.int32(fshape_c[0]),
numpy.int32(same_starts[1]),
numpy.int32(same_starts[0]),
numpy.float32(c["gdx_prod"]),
numpy.int32(c["dst_plane"]),
numpy.int32(c["src_plane"]),
)
if profile:
queue.finish()
timing["crop_scale_ms"] = (timeit.default_timer() - t0) * 1000.0
t0 = timeit.default_timer()
p1 = c["p1_host"]
p2 = c["p2_host"]
mapped_p = c["p_batch_dev"].map_to_host(queue=queue, flags=cl.map_flags.READ, is_blocking=True)
numpy.copyto(p1, mapped_p[0])
numpy.copyto(p2, mapped_p[1])
mapped_p.base.release(queue=queue)
timing["d2h_get_ms"] = (timeit.default_timer() - t0) * 1000.0
timing["total_ms"] = (timeit.default_timer() - t_total) * 1000.0
self._opencl_last_timing = timing
if verbose:
print("Poisson solver (OpenCL FFT batch2) finished")
if profile:
print(
"OpenCL batch2 timing [ms]: "
f"total={timing['total_ms']:.3f}, "
f"queue={timing['queue_ms']:.3f}, "
f"cache={timing['cache_prepare_ms']:.3f}, "
f"host={timing['host_prepare_ms']:.3f}, "
f"h2d={timing['h2d_set_ms']:.3f}, "
f"fft={timing['fft_f_ms']:.3f}, "
f"mul={timing['mul_ms']:.3f}, "
f"ifft={timing['ifft_ms']:.3f}, "
f"d2h={timing['d2h_get_ms']:.3f}, "
f"crop={timing['crop_scale_ms']:.3f}"
)
return p1, p2
def solve_opencl_batch2(self, f1, f2, verbose=False, profile=False):
"""Alias for solve_batch_opencl (backward compatible name)."""
return self.solve_batch_opencl(f1, f2, verbose=verbose, profile=profile)
+177 -19
View File
@@ -132,6 +132,14 @@ class Sensor:
self.centroid_found = False
self.contour = np.empty(0)
self.depth_map = np.zeros_like(self.img)
self.valid_mask_raw = None
self.valid_mask = None
self.valid_mask_hhd = None
self.mask_stats = {
"coverage": 0.0,
"flow_mean_valid": 0.0,
"flow_mean_invalid": 0.0,
}
# Timing
self.last_log_time = time.perf_counter()
@@ -165,6 +173,12 @@ class Sensor:
self._hhd_grid_shape = None
cfg_backend = self.__config_tracker.get("backend", None)
self.poisson_backend = cfg_backend if cfg_backend is not None else backend
self.masked_poisson_solver = self.__config_tracker.get("masked_poisson_solver", "sparse")
self.masked_poisson_max_boundary = self.__config_tracker.get("masked_poisson_max_boundary", 96)
self.masked_poisson_boundary_stride = self.__config_tracker.get("masked_poisson_boundary_stride", 1)
self.masked_poisson_max_iters = self.__config_tracker.get("masked_poisson_max_iters", 3)
self.masked_poisson_tol = self.__config_tracker.get("masked_poisson_tol", 1e-4)
self.masked_poisson_reg = self.__config_tracker.get("masked_poisson_reg", 1e-8)
self._use_pyav = False
self._av_container = None
self._av_stream = None
@@ -204,7 +218,72 @@ class Sensor:
# Public API - Core Methods (Most Frequently Used)
# =========================================================================
def get_img(self, frame=None, already_processed=False):
def _resize_mask(self, mask, size):
if mask is None:
return None
mask_u8 = np.asarray(mask, dtype=np.uint8)
resized = cv2.resize(mask_u8, size, interpolation=cv2.INTER_NEAREST)
return resized.astype(bool)
def _prepare_valid_mask(self, valid_mask, raw_shape, stitched_shape):
if valid_mask is None:
self.valid_mask_raw = None
self.valid_mask = None
self.valid_mask_hhd = None
self.mask_stats["coverage"] = 0.0
return
mask_raw = np.asarray(valid_mask, dtype=bool)
if mask_raw.shape != raw_shape:
mask_raw = self._resize_mask(mask_raw, (raw_shape[1], raw_shape[0]))
if self.mirror:
mask_raw = np.fliplr(mask_raw)
self.valid_mask_raw = mask_raw
self.valid_mask = self._resize_mask(mask_raw, (stitched_shape[1], stitched_shape[0]))
self._update_hhd_mask_cache()
self.mask_stats["coverage"] = float(np.mean(self.valid_mask)) if self.valid_mask.size else 0.0
def _hhd_grid_size(self):
if self.valid_mask is None:
return None
h, w = self.valid_mask.shape[:2]
k = max(1, int(self._hhd_downsample))
if k <= 1:
return (w, h)
return (max(1, w // k), max(1, h // k))
def _update_hhd_mask_cache(self):
"""Precompute valid_mask at HHD resolution when the UV domain is fixed."""
if self.valid_mask is None:
self.valid_mask_hhd = None
return
ds_size = self._hhd_grid_size()
if ds_size is None:
self.valid_mask_hhd = None
return
if self.valid_mask.shape[1] == ds_size[0] and self.valid_mask.shape[0] == ds_size[1]:
self.valid_mask_hhd = self.valid_mask.astype(bool)
else:
self.valid_mask_hhd = self._resize_mask(self.valid_mask, ds_size)
def _compute_flow_means(self, flow):
if flow is None:
return
mag = np.linalg.norm(flow, axis=-1)
mask = self.valid_mask
if mask is None or mask.shape != mag.shape:
self.mask_stats["flow_mean_valid"] = float(mag.mean()) if mag.size else 0.0
self.mask_stats["flow_mean_invalid"] = 0.0
return
valid = mask.astype(bool)
invalid = ~valid
self.mask_stats["flow_mean_valid"] = float(mag[valid].mean()) if np.any(valid) else 0.0
self.mask_stats["flow_mean_invalid"] = float(mag[invalid].mean()) if np.any(invalid) else 0.0
def get_img(self, frame=None, already_processed=False, valid_mask=None, already_gray=False):
"""
Capture and preprocess a frame.
@@ -238,23 +317,39 @@ class Sensor:
# Start processing timing
t_processing_start = time.perf_counter()
self.frame = frame.copy() if isinstance(frame, np.ndarray) else frame
self.frame = frame if (already_processed and isinstance(frame, np.ndarray)) else (
frame.copy() if isinstance(frame, np.ndarray) else frame
)
if already_processed:
img_bgr = frame
if img_bgr.ndim == 2:
img_bgr = cv2.cvtColor(img_bgr, cv2.COLOR_GRAY2BGR)
elif img_bgr.ndim == 3 and img_bgr.shape[2] == 1:
img_bgr = cv2.cvtColor(img_bgr[:, :, 0], cv2.COLOR_GRAY2BGR)
gray = None
if already_gray:
if frame.ndim == 2:
gray = frame
elif frame.ndim == 3 and frame.shape[2] == 1:
gray = frame[:, :, 0]
else:
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
img_bgr = None
else:
if img_bgr.ndim == 2:
img_bgr = cv2.cvtColor(img_bgr, cv2.COLOR_GRAY2BGR)
elif img_bgr.ndim == 3 and img_bgr.shape[2] == 1:
img_bgr = cv2.cvtColor(img_bgr[:, :, 0], cv2.COLOR_GRAY2BGR)
else:
img_bgr = self.frame_processor.process(frame)
gray = None
t_after_extract = time.perf_counter()
t_after_stitch = time.perf_counter()
self.img_size_raw = img_bgr.shape
self._img_bgr_raw = img_bgr
gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
if gray is None:
self.img_size_raw = img_bgr.shape
self._img_bgr_raw = img_bgr
gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
else:
self.img_size_raw = gray.shape
self._img_bgr_raw = None
res = tuple(self._config_resolution_stitch)
if self._img_gray_buf is None or self._img_gray_buf.shape != (res[1], res[0]):
self._img_gray_buf = np.empty((res[1], res[0]), dtype=np.uint8)
@@ -265,6 +360,7 @@ class Sensor:
cv2.flip(self._img_gray_buf, 1, dst=self._img_gray_buf)
self.img = self._img_gray_buf
self._prepare_valid_mask(valid_mask, gray.shape, self.img.shape)
t_processing_end = time.perf_counter()
self._last_processing_time = t_processing_end - t_processing_start
@@ -281,7 +377,9 @@ class Sensor:
def _get_img_bgr(self):
"""Produce a resized+mirrored BGR image on demand (for display)."""
if self._img_bgr_raw is None:
return None
if self.img is None:
return None
return cv2.cvtColor(self.img, cv2.COLOR_GRAY2BGR)
res = tuple(self._config_resolution_stitch)
bgr = cv2.resize(self._img_bgr_raw, res, interpolation=cv2.INTER_LINEAR)
if self.mirror:
@@ -292,8 +390,12 @@ class Sensor:
# detect motion by optical flow magnitude
mag = np.linalg.norm(self.flow, axis=-1)
# print(f"mag: {mag.max()}, {mag.min()}, {mag.mean()}")
return mag.mean() > threshold
if self.valid_mask is not None and self.valid_mask.shape == mag.shape:
active = self.valid_mask.astype(bool)
motion_value = mag[active].mean() if np.any(active) else 0.0
else:
motion_value = mag.mean()
return motion_value > threshold
def compute_deformation(self, decompose=True, check_motion=True, threshold=6, tail_frames=10):
"""
@@ -352,15 +454,24 @@ class Sensor:
return False
def warmup_deformation(self):
"""Build optical-flow and masked-HHD caches using the current image."""
if self.img is None:
raise Exception("Warning: Get image first")
self.compute_flow()
self.compute_hhd()
return True
def compute_flow(self):
if self.img is None:
raise Exception("Warning: Get image first")
flow = self.__tracker.track(self.img)
flow = self.__tracker.track(self.img, valid_mask=self.valid_mask)
# flow = self.__tracker.track_raft(self.img)
if hasattr(flow, 'download'):
flow = flow.download()
self.flow = flow
self._compute_flow_means(self.flow)
def compute_hhd(self):
@@ -374,20 +485,50 @@ class Sensor:
except (TypeError, ValueError):
k = 1
mask_ds = None
if self.valid_mask is not None:
ds_size = self._hhd_grid_size()
if (
self.valid_mask_hhd is not None
and ds_size is not None
and self.valid_mask_hhd.shape[1] == ds_size[0]
and self.valid_mask_hhd.shape[0] == ds_size[1]
):
mask_ds = self.valid_mask_hhd
else:
mask_ds = self.valid_mask.astype(bool)
if k > 1:
ds_w = max(1, orig_w // k)
ds_h = max(1, orig_h // k)
flow_ds = cv2.resize(flow, (ds_w, ds_h), interpolation=cv2.INTER_AREA)
if mask_ds is not None and (
mask_ds.shape[1] != ds_w or mask_ds.shape[0] != ds_h
):
mask_ds = self._resize_mask(mask_ds, (ds_w, ds_h))
else:
flow_ds = flow
ds_shape = flow_ds.shape[:2]
self.valid_mask_hhd = mask_ds
if self.__hhd is None or self._hhd_grid_shape != ds_shape:
spacings = (18 * k, 18 * k)
self.__hhd = nHHD(cuda=self.cuda, grid=ds_shape, spacings=spacings, poisson_backend=self.poisson_backend, debug=self.debug)
self.__hhd = nHHD(
cuda=self.cuda,
grid=ds_shape,
spacings=spacings,
poisson_backend=self.poisson_backend,
debug=self.debug,
masked_poisson_solver=self.masked_poisson_solver,
masked_poisson_max_boundary=self.masked_poisson_max_boundary,
masked_poisson_boundary_stride=self.masked_poisson_boundary_stride,
masked_poisson_max_iters=self.masked_poisson_max_iters,
masked_poisson_tol=self.masked_poisson_tol,
masked_poisson_reg=self.masked_poisson_reg,
)
self._hhd_grid_shape = ds_shape
_, forces = self.__hhd.get_force(flow_ds)
_, forces = self.__hhd.get_force(flow_ds, valid_mask=mask_ds)
self.fnormal, self.fshearx, self.fsheary, self.ftor = forces
vnormal_ds = self.__hhd.get_normal()
@@ -400,6 +541,13 @@ class Sensor:
self.vnormal = vnormal_ds
self.vshear = vshear_ds
if self.valid_mask is not None:
active = self.valid_mask.astype(bool)
self.vnormal = self.vnormal.copy()
self.vshear = self.vshear.copy()
self.vnormal[~active] = 0
self.vshear[~active] = 0
self.compute_deformation_valid = True
return True
@@ -423,12 +571,12 @@ class Sensor:
raise Exception("Warning: vnormal is required for contact computation")
if multiple_foe:
result = self.__contact3d.compute_multiple_foe(vnormal_to_use)
result = self.__contact3d.compute_multiple_foe(vnormal_to_use, valid_mask=self.valid_mask)
else:
if self.cuda:
result = self.__contact3d.compute_foe_gpu(vnormal_to_use)
result = self.__contact3d.compute_foe_gpu(vnormal_to_use, valid_mask=self.valid_mask)
else:
result = self.__contact3d.compute_foe(vnormal_to_use)
result = self.__contact3d.compute_foe(vnormal_to_use, valid_mask=self.valid_mask)
self.centroid = result['centroid']
self.centroid_found = result['centroid_found']
@@ -436,6 +584,16 @@ class Sensor:
self.contour = result['contour']
self.depth_map = result['divergence_map']
if self.valid_mask is not None and self.centroid_found:
cx, cy = int(self.centroid[0]), int(self.centroid[1])
inside = (
0 <= cy < self.valid_mask.shape[0]
and 0 <= cx < self.valid_mask.shape[1]
and bool(self.valid_mask[cy, cx])
)
if not inside and self.verbose:
print("Warning: contact centroid lies outside the valid UV mask.")
self.compute_contact_valid = True
return result
+38 -10
View File
@@ -82,16 +82,39 @@ class Tracker:
def track(self, frame):
def _prepare_valid_mask(self, valid_mask, shape, erode_iters=1):
if valid_mask is None:
return None
mask = np.asarray(valid_mask, dtype=np.uint8)
if mask.shape != shape:
mask = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST)
if erode_iters > 0:
kernel = np.ones((3, 3), dtype=np.uint8)
mask = cv2.erode(mask, kernel, iterations=erode_iters)
return mask.astype(bool)
def apply_valid_mask(self, flow, valid_mask, erode_iters=1):
"""Zero flow outside the valid UV domain."""
if flow is None or valid_mask is None:
return flow
mask = self._prepare_valid_mask(valid_mask, flow.shape[:2], erode_iters=erode_iters)
flow_out = flow.copy()
flow_out[~mask] = 0
return flow_out
def track(self, frame, valid_mask=None):
"""Main tracking function that dispatches to CUDA or CPU implementation"""
if self.__algorithm == "raft":
return self.track_raft(frame.copy())
flow = self.track_raft(frame.copy())
elif self.__algorithm == "dis":
if frame.ndim == 3:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
return self.track_cpu(frame.copy())
flow = self.track_cpu(frame.copy(), valid_mask=valid_mask)
else:
flow = None
return self.apply_valid_mask(flow, valid_mask)
def track_cpu(self, frame):
def track_cpu(self, frame, valid_mask=None):
"""CPU-based optical flow tracking"""
frame_height, frame_width = frame.shape[:2]
@@ -114,7 +137,7 @@ class Tracker:
# Only check loop closure periodically (expensive operation)
if self.__frame_count % self.__loop_closure_interval == 0:
self.__checkLoopClosure(frame)
self.__checkLoopClosure(frame, valid_mask=valid_mask)
# CPU optical flow computation
self.__prev_fine_flow = self.__tracking_inst.calc(
@@ -126,7 +149,7 @@ class Tracker:
# Use fine flow for error estimation
error = self.__getForwardBackwardError(
self.__fine_base_frame_queue[-1], frame,
self.__prev_fine_flow)
self.__prev_fine_flow, valid_mask=valid_mask)
if error > self.__upper_threshold:
self.__fine_base_frame_queue.append(frame)
@@ -148,7 +171,7 @@ class Tracker:
# flag,output = decompose_obj.get_force(flow, verbose=0)
# return output
def __getForwardBackwardError(self, I0, I1, flow):
def __getForwardBackwardError(self, I0, I1, flow, valid_mask=None):
# if self.cuda:
# flow_map = cv2.cuda.add(self.map, flow)
# map_x, map_y = cv2.cuda.split(flow_map)
@@ -175,15 +198,20 @@ class Tracker:
diff = z.astype('float32') - I0.astype('float32')
else:
diff = z - I0
rev = np.mean(np.abs(diff))
diff_abs = np.abs(diff)
mask = self._prepare_valid_mask(valid_mask, diff_abs.shape, erode_iters=0)
if mask is not None and np.any(mask):
rev = float(diff_abs[mask].mean())
else:
rev = float(np.mean(diff_abs))
return rev
def __checkLoopClosure(self, frame): # update (prune) the frame queue with loop closure detection
def __checkLoopClosure(self, frame, valid_mask=None): # update (prune) the frame queue with loop closure detection
for i in range(len(self.__fine_base_frame_queue) - 2, -1, -1):
f = self.__tracking_inst.calc(
self.__fine_base_frame_queue[i], frame, None)
error = self.__getForwardBackwardError(
self.__fine_base_frame_queue[i], frame, f)
self.__fine_base_frame_queue[i], frame, f, valid_mask=valid_mask)
if error < self.__lower_threshold:
self.__fine_base_frame_queue = self.__fine_base_frame_queue[:i+1]
self.__fine_flow_queue = self.__fine_flow_queue[:i + 1]
+37
View File
@@ -0,0 +1,37 @@
"""UV-based finger mesh unroll / flatten pipeline.
Offline: Blender export -> build_or_load_domain -> unroll_map_uv.npz
Runtime: live camera frame -> UVUnrollFrameProcessor -> Sensor (with valid_mask)
3D display: UVDomain.surface.world_xyz
"""
from .domain import UVDomain
from .mesh_uv import align_frame_to_uv_map, barycentric_grid, face_uvs_from_loop_data, uv_triangle_to_image
from .precompute import (
build_or_load_domain,
build_uv_domain,
compute_uv_distortion_metrics,
load_calib_json,
load_exported_data,
recommend_edge_expand,
)
from .remap import UVRemap, flatten_frame, remap_valid_ratio
from .surface import UVSurfaceMap
__all__ = [
"UVDomain",
"UVRemap",
"UVSurfaceMap",
"align_frame_to_uv_map",
"barycentric_grid",
"build_or_load_domain",
"build_uv_domain",
"compute_uv_distortion_metrics",
"face_uvs_from_loop_data",
"flatten_frame",
"load_calib_json",
"load_exported_data",
"recommend_edge_expand",
"remap_valid_ratio",
"uv_triangle_to_image",
]
+54
View File
@@ -0,0 +1,54 @@
"""Cached UV domain: 2D remap plus optional 3D surface.
UVDomain is the canonical bundle saved to unroll_map_uv.npz and loaded by
build_or_load_domain, tests, and 3D viewer tools.
"""
from dataclasses import dataclass
from typing import Optional
import numpy as np
from .remap import UVRemap
from .surface import UVSurfaceMap
def _cache_has_surface(data: dict) -> bool:
return "world_xyz" in data or "flat_xyz" in data
@dataclass
class UVDomain:
"""Image-to-UV remap plus optional UV-to-world surface map."""
remap: UVRemap
surface: Optional[UVSurfaceMap] = None
@classmethod
def from_dict(cls, data: dict) -> "UVDomain":
"""Build domain from flat npz dict."""
remap = UVRemap.from_dict(data)
surface = UVSurfaceMap.from_dict(data) if _cache_has_surface(data) else None
return cls(remap=remap, surface=surface)
@classmethod
def load(cls, path: str) -> "UVDomain":
"""Load cached domain from .npz file."""
data = np.load(path)
return cls.from_dict({key: data[key] for key in data.files})
def to_dict(self) -> dict:
"""Serialize remap and surface into one npz-compatible dict."""
out = self.remap.to_dict()
if self.surface is not None:
out.update(self.surface.to_dict())
return out
def save(self, path: str) -> None:
"""Save domain cache to .npz file."""
np.savez_compressed(path, **self.to_dict())
@property
def valid_mask(self):
"""(H, W) bool active UV domain; same as remap.valid_mask."""
return self.remap.valid_mask
+101
View File
@@ -0,0 +1,101 @@
"""Low-level mesh/UV math for the unroll pipeline.
Rasterization helpers (barycentrics, UV-to-image warp, world projection) used
by offline precompute and by demo debug visualizations.
"""
import numpy as np
def face_uvs_from_loop_data(faces, uv_data, loop_to_vertex):
if len(uv_data) == 0:
raise ValueError("mesh.npz does not contain UV data; UV flattening is unavailable.")
if len(uv_data) != len(faces) * 3:
raise ValueError("Expected a triangulated mesh with exactly 3 UV loops per face.")
face_loop_vertices = loop_to_vertex.reshape(-1, 3)
if face_loop_vertices.shape != faces.shape or not np.array_equal(face_loop_vertices, faces):
raise ValueError("UV loop order does not match exported face vertex order.")
return uv_data.reshape(-1, 3, 2)
def project_world_to_image(points_world, K, world_to_cam):
"""Project world points into the image using Blender/OpenCV-compatible intrinsics."""
points_h = np.column_stack((points_world, np.ones(len(points_world), dtype=np.float64)))
points_cam = (world_to_cam @ points_h.T).T[:, :3]
z = -points_cam[:, 2]
x = K[0, 0] * (points_cam[:, 0] / z) + K[0, 2]
y = K[1, 1] * (points_cam[:, 1] / z) + K[1, 2]
return np.column_stack((x, y)), z
def barycentric_grid(px, py, tri):
a, b, c = tri
denom = (b[1] - c[1]) * (a[0] - c[0]) + (c[0] - b[0]) * (a[1] - c[1])
if abs(denom) < 1e-12:
return None
w0 = ((b[1] - c[1]) * (px - c[0]) + (c[0] - b[0]) * (py - c[1])) / denom
w1 = ((c[1] - a[1]) * (px - c[0]) + (a[0] - c[0]) * (py - c[1])) / denom
w2 = 1.0 - w0 - w1
return w0, w1, w2
def _triangle_edge_lengths(tri):
return np.array(
[
np.linalg.norm(tri[1] - tri[0]),
np.linalg.norm(tri[2] - tri[1]),
np.linalg.norm(tri[0] - tri[2]),
],
dtype=np.float64,
)
def triangle_area_3d(tri):
return 0.5 * np.linalg.norm(np.cross(tri[1] - tri[0], tri[2] - tri[0]))
def triangle_area_2d(tri):
return 0.5 * abs(
(tri[1, 0] - tri[0, 0]) * (tri[2, 1] - tri[0, 1])
- (tri[1, 1] - tri[0, 1]) * (tri[2, 0] - tri[0, 0])
)
def warp_uv_axis(values, edge_band, edge_expand):
"""Allocate more normalized output span to UV edge regions."""
u = np.asarray(values, dtype=np.float64)
e = float(np.clip(edge_band, 0.0, 0.49))
s = float(max(edge_expand, 1.0))
if e <= 0.0 or s <= 1.0:
return u.copy()
middle = 1.0 - 2.0 * e
total = middle + 2.0 * e * s
out = np.empty_like(u)
left = u < e
right = u > (1.0 - e)
center = ~(left | right)
out[left] = (u[left] * s) / total
out[center] = (e * s + (u[center] - e)) / total
out[right] = (e * s + middle + (u[right] - (1.0 - e)) * s) / total
return out
def uv_triangle_to_image(uv_tri, out_width, out_height, flip_v=True, edge_band=0.12, edge_expand=1.0):
"""Map UV triangle into output image coordinates with optional edge expansion."""
tri = np.empty_like(uv_tri)
u_warp = warp_uv_axis(uv_tri[:, 0], edge_band=edge_band, edge_expand=edge_expand)
v_warp = warp_uv_axis(uv_tri[:, 1], edge_band=edge_band, edge_expand=edge_expand)
tri[:, 0] = u_warp * (out_width - 1)
tri[:, 1] = ((1.0 - v_warp) if flip_v else v_warp) * (out_height - 1)
return tri
def align_frame_to_uv_map(frame_bgr, flip_vertical=True):
"""Match live-frame orientation to the Blender-exported UV remap convention."""
if not flip_vertical:
return frame_bgr
return np.flipud(frame_bgr).copy()
+349
View File
@@ -0,0 +1,349 @@
"""Offline UV domain precompute: build, cache, and load.
Loads Blender-exported mesh/intrinsics/pose, estimates UV distortion, rasterizes
each UV triangle into map_x/map_y (and optional world_xyz), then saves or loads
unroll_map_uv.npz. Overlapping UV islands use mesh-order depth buffering and
may resolve incorrectly where faces overlap.
"""
import json
import os
import time
import numpy as np
from .domain import UVDomain
from .mesh_uv import (
_triangle_edge_lengths,
barycentric_grid,
face_uvs_from_loop_data,
project_world_to_image,
triangle_area_2d,
triangle_area_3d,
uv_triangle_to_image,
)
from .remap import UVRemap, _cache_matches, remap_valid_ratio
from .surface import UVSurfaceMap
def load_calib_json(json_path):
"""Load OpenCV camera calibration from JSON.
returns: (K, dist_coeffs, (width, height))
"""
with open(json_path, "r", encoding="utf-8") as f:
calib = json.load(f)
K = np.array(calib["camera_matrix"], dtype=np.float64)
dist = np.array(calib["distortion"], dtype=np.float64).flatten()
img_size = tuple(int(v) for v in calib["image_size"])
return K, dist, img_size
def load_exported_data(mesh_path, intrinsics_path, pose_path):
"""Load Blender-exported mesh, intrinsics, and camera pose."""
mesh = np.load(mesh_path)
cam = np.load(intrinsics_path)
return {
"vertices": mesh["vertices"].astype(np.float64),
"faces": mesh["faces"].astype(np.int32),
"uv_data": mesh["uv_data"].astype(np.float64),
"loop_to_vertex": mesh["loop_to_vertex"].astype(np.int32),
"K": cam["K"].astype(np.float64),
"width": int(cam["width"]),
"height": int(cam["height"]),
"world_to_cam": np.load(pose_path).astype(np.float64),
}
def _summarize_uv_region(relative_scale, anisotropy, mask):
mask_scale = mask & np.isfinite(relative_scale) & (relative_scale > 0)
mask_aniso = mask & np.isfinite(anisotropy) & (anisotropy > 0)
if not np.any(mask_scale):
return {
"face_count": 0,
"relative_scale_p05": float("nan"),
"relative_scale_p50": float("nan"),
"relative_scale_p95": float("nan"),
"relative_scale_min": float("nan"),
"relative_scale_max": float("nan"),
"anisotropy_p50": float("nan"),
"anisotropy_p95": float("nan"),
"anisotropy_max": float("nan"),
}
return {
"face_count": int(np.count_nonzero(mask_scale)),
"relative_scale_p05": float(np.percentile(relative_scale[mask_scale], 5)),
"relative_scale_p50": float(np.percentile(relative_scale[mask_scale], 50)),
"relative_scale_p95": float(np.percentile(relative_scale[mask_scale], 95)),
"relative_scale_min": float(np.min(relative_scale[mask_scale])),
"relative_scale_max": float(np.max(relative_scale[mask_scale])),
"anisotropy_p50": float(np.percentile(anisotropy[mask_aniso], 50)) if np.any(mask_aniso) else float("nan"),
"anisotropy_p95": float(np.percentile(anisotropy[mask_aniso], 95)) if np.any(mask_aniso) else float("nan"),
"anisotropy_max": float(np.max(anisotropy[mask_aniso])) if np.any(mask_aniso) else float("nan"),
}
def compute_uv_distortion_metrics(data, edge_band=0.12):
"""Estimate local UV stretch/compression and anisotropy per face."""
faces = data["faces"]
vertices = data["vertices"]
face_uvs = face_uvs_from_loop_data(faces, data["uv_data"], data["loop_to_vertex"])
local_scale = np.full(len(faces), np.nan, dtype=np.float64)
relative_scale = np.full(len(faces), np.nan, dtype=np.float64)
anisotropy = np.full(len(faces), np.nan, dtype=np.float64)
uv_centers = np.full((len(faces), 2), np.nan, dtype=np.float64)
edge_mask = np.zeros(len(faces), dtype=bool)
for idx, (face, uv_tri) in enumerate(zip(faces, face_uvs)):
tri_3d = vertices[face]
tri_uv = uv_tri
uv_centers[idx] = np.mean(tri_uv, axis=0)
edges_3d = _triangle_edge_lengths(tri_3d)
edges_uv = _triangle_edge_lengths(tri_uv)
valid_edges = edges_uv > 1e-9
if np.count_nonzero(valid_edges) < 2:
continue
edge_scales = edges_3d[valid_edges] / edges_uv[valid_edges]
local_scale[idx] = float(np.exp(np.mean(np.log(np.clip(edge_scales, 1e-12, None)))))
anisotropy[idx] = float(np.max(edge_scales) / max(np.min(edge_scales), 1e-12))
area_uv = triangle_area_2d(tri_uv)
area_3d = triangle_area_3d(tri_3d)
if area_uv > 1e-12 and area_3d > 1e-12:
local_scale[idx] = float(np.sqrt(area_3d / area_uv))
valid_scale = np.isfinite(local_scale) & (local_scale > 0)
if not np.any(valid_scale):
raise RuntimeError("Could not compute UV distortion metrics; UV triangles may be degenerate.")
median_scale = float(np.median(local_scale[valid_scale]))
relative_scale[valid_scale] = local_scale[valid_scale] / median_scale
min_edge_distance = np.min(
np.column_stack(
[
uv_centers[:, 0],
1.0 - uv_centers[:, 0],
uv_centers[:, 1],
1.0 - uv_centers[:, 1],
]
),
axis=1,
)
edge_mask = min_edge_distance < edge_band
middle_mask = ~edge_mask
valid_aniso = np.isfinite(anisotropy) & (anisotropy > 0)
stats = {
"median_mm_per_uv": median_scale,
"relative_scale_p05": float(np.percentile(relative_scale[valid_scale], 5)),
"relative_scale_p50": float(np.percentile(relative_scale[valid_scale], 50)),
"relative_scale_p95": float(np.percentile(relative_scale[valid_scale], 95)),
"relative_scale_min": float(np.min(relative_scale[valid_scale])),
"relative_scale_max": float(np.max(relative_scale[valid_scale])),
"anisotropy_p50": float(np.percentile(anisotropy[valid_aniso], 50)) if np.any(valid_aniso) else float("nan"),
"anisotropy_p95": float(np.percentile(anisotropy[valid_aniso], 95)) if np.any(valid_aniso) else float("nan"),
"anisotropy_max": float(np.max(anisotropy[valid_aniso])) if np.any(valid_aniso) else float("nan"),
"edge_band": float(edge_band),
}
return {
"relative_scale": relative_scale,
"anisotropy": anisotropy,
"uv_centers": uv_centers,
"edge_mask": edge_mask,
"middle_mask": middle_mask,
"stats": stats,
"middle_stats": _summarize_uv_region(relative_scale, anisotropy, middle_mask),
"edge_stats": _summarize_uv_region(relative_scale, anisotropy, edge_mask),
}
def recommend_edge_expand(distortion, max_expand=3.0):
"""Suggest how much more UV space to allocate to the edge band."""
middle_p95 = distortion["middle_stats"]["relative_scale_p95"]
edge_p95 = distortion["edge_stats"]["relative_scale_p95"]
if not np.isfinite(middle_p95) or not np.isfinite(edge_p95) or middle_p95 <= 0:
return 1.0
return float(np.clip(np.sqrt(edge_p95 / middle_p95), 1.0, max_expand))
def build_uv_domain(
data,
source_K,
source_width,
source_height,
out_width,
out_height,
flip_v=True,
edge_band=0.12,
edge_expand=1.0,
include_surface=True,
verbose=True,
):
"""Build image-to-UV remap and optional UV-to-world surface map."""
faces = data["faces"]
vertices = data["vertices"]
face_uvs = face_uvs_from_loop_data(faces, data["uv_data"], data["loop_to_vertex"])
map_x = np.full((out_height, out_width), -1.0, dtype=np.float32)
map_y = np.full((out_height, out_width), -1.0, dtype=np.float32)
valid_mask = np.zeros((out_height, out_width), dtype=bool)
depth_buffer = np.full((out_height, out_width), np.inf, dtype=np.float64)
world_xyz = np.full((out_height, out_width, 3), np.nan, dtype=np.float32) if include_surface else None
if verbose:
print("Building UV domain (remap + surface) ..." if include_surface else "Building UV remap ...")
t_start = time.time()
for face_idx, (face, uv_tri) in enumerate(zip(faces, face_uvs)):
tri = uv_triangle_to_image(
uv_tri,
out_width,
out_height,
flip_v=flip_v,
edge_band=edge_band,
edge_expand=edge_expand,
)
x0 = max(0, int(np.floor(np.min(tri[:, 0]))))
x1 = min(out_width - 1, int(np.ceil(np.max(tri[:, 0]))))
y0 = max(0, int(np.floor(np.min(tri[:, 1]))))
y1 = min(out_height - 1, int(np.ceil(np.max(tri[:, 1]))))
if x1 < x0 or y1 < y0:
continue
yy, xx = np.mgrid[y0 : y1 + 1, x0 : x1 + 1]
bary = barycentric_grid(xx + 0.5, yy + 0.5, tri)
if bary is None:
continue
w0, w1, w2 = bary
inside = (w0 >= -1e-6) & (w1 >= -1e-6) & (w2 >= -1e-6)
if not np.any(inside):
continue
weights = np.stack((w0[inside], w1[inside], w2[inside]), axis=1)
points_world = weights @ vertices[face]
pixels, depth = project_world_to_image(points_world, source_K, data["world_to_cam"])
src_x = pixels[:, 0]
src_y = pixels[:, 1]
target_y = yy[inside]
target_x = xx[inside]
valid = (
(depth > 0)
& (src_x >= 0)
& (src_x <= source_width - 1)
& (src_y >= 0)
& (src_y <= source_height - 1)
& (depth < depth_buffer[target_y, target_x])
)
if not np.any(valid):
continue
tx = target_x[valid]
ty = target_y[valid]
map_x[ty, tx] = src_x[valid].astype(np.float32)
map_y[ty, tx] = src_y[valid].astype(np.float32)
valid_mask[ty, tx] = True
depth_buffer[ty, tx] = depth[valid]
if include_surface:
world_xyz[ty, tx] = points_world[valid].astype(np.float32)
if verbose and (face_idx % 200) == 0:
print(f" face {face_idx}/{len(faces)}")
if verbose:
print(f"Domain build took {time.time() - t_start:.1f} s")
remap = UVRemap(
map_x=map_x,
map_y=map_y,
valid_mask=valid_mask,
metadata={
"source_width": np.array(source_width, dtype=np.int32),
"source_height": np.array(source_height, dtype=np.int32),
"out_width": np.array(out_width, dtype=np.int32),
"out_height": np.array(out_height, dtype=np.int32),
"flip_v": np.array(int(flip_v), dtype=np.int32),
"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),
"method": np.array("uv"),
},
)
surface = UVSurfaceMap(world_xyz=world_xyz) if include_surface else None
return UVDomain(remap=remap, surface=surface)
def build_or_load_domain(
mesh_path,
intrinsics_path,
pose_path,
map_cache_path,
source_K,
source_width,
source_height,
out_width,
out_height,
flip_v=True,
edge_band=0.12,
edge_expand=0.0,
include_surface=True,
verbose=True,
):
"""Build or load full UV domain (2D remap + optional 3D surface).
edge_expand: use directly when > 1.0; otherwise auto-resolve from distortion.
returns: UVDomain
"""
data = load_exported_data(mesh_path, intrinsics_path, pose_path)
if edge_expand > 1.0:
resolved_edge_expand = float(edge_expand)
else:
distortion = compute_uv_distortion_metrics(data, edge_band=edge_band)
resolved_edge_expand = recommend_edge_expand(distortion)
if map_cache_path and os.path.exists(map_cache_path):
cached = UVDomain.load(map_cache_path)
cache_ok = _cache_matches(
cached.remap,
source_K=source_K,
source_width=source_width,
source_height=source_height,
out_width=out_width,
out_height=out_height,
flip_v=flip_v,
edge_band=edge_band,
edge_expand=resolved_edge_expand,
)
surface_ok = (not include_surface) or (cached.surface is not None)
if cache_ok and surface_ok:
if verbose:
print(f"Loaded cached UV domain. Valid ratio: {remap_valid_ratio(cached.remap):.3f}")
return cached
if verbose:
print("Cached UV domain is stale or missing surface map. Regenerating.")
domain = build_uv_domain(
data,
source_K=source_K,
source_width=source_width,
source_height=source_height,
out_width=out_width,
out_height=out_height,
flip_v=flip_v,
edge_band=edge_band,
edge_expand=resolved_edge_expand,
include_surface=include_surface,
verbose=verbose,
)
valid_ratio = remap_valid_ratio(domain.remap)
if verbose:
print(f"Generated UV domain valid ratio: {valid_ratio:.3f}")
if valid_ratio == 0.0:
raise RuntimeError("Generated UV domain contains no valid samples.")
if map_cache_path:
domain.save(map_cache_path)
return domain
+190
View File
@@ -0,0 +1,190 @@
"""2D UV flatten remap tables and runtime frame remapping.
Stores precomputed map_x/map_y/valid_mask and applies cv2.remap at runtime to
convert source camera frames into the UV flatten domain.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, Optional, Tuple
import cv2
import numpy as np
@dataclass
class UVRemap:
"""Precomputed UV flatten remap tables."""
map_x: np.ndarray
map_y: np.ndarray
valid_mask: np.ndarray
metadata: Dict[str, Any] = field(default_factory=dict)
_map1_fixed: Optional[np.ndarray] = field(default=None, repr=False, compare=False)
_map2_fixed: Optional[np.ndarray] = field(default=None, repr=False, compare=False)
_border_only_invalid: bool = field(default=False, repr=False, compare=False)
@property
def out_height(self) -> int:
return int(self.map_x.shape[0])
@property
def out_width(self) -> int:
return int(self.map_x.shape[1])
@classmethod
def from_dict(cls, remap_dict: dict) -> "UVRemap":
"""Build UVRemap from a flat dict (e.g. np.load result)."""
meta_keys = {
"source_width",
"source_height",
"out_width",
"out_height",
"flip_v",
"edge_band",
"edge_expand",
"projection_K",
"method",
}
metadata = {k: remap_dict[k] for k in meta_keys if k in remap_dict}
remap = cls(
map_x=np.asarray(remap_dict["map_x"], dtype=np.float32),
map_y=np.asarray(remap_dict["map_y"], dtype=np.float32),
valid_mask=np.asarray(remap_dict["valid_mask"], dtype=bool),
metadata=metadata,
)
remap.ensure_fast_maps()
return remap
@classmethod
def load(cls, path: str) -> "UVRemap":
"""Load cached remap from .npz file."""
data = np.load(path)
return cls.from_dict({key: data[key] for key in data.files})
def to_dict(self) -> dict:
"""Serialize to dict suitable for np.savez_compressed."""
out = {
"map_x": self.map_x.astype(np.float32),
"map_y": self.map_y.astype(np.float32),
"valid_mask": self.valid_mask.astype(bool),
}
out.update(self.metadata)
return out
def save(self, path: str) -> None:
"""Save remap cache to .npz file."""
np.savez_compressed(path, **self.to_dict())
def ensure_fast_maps(self) -> None:
"""Build OpenCV CV_16SC2 remap tables once for faster cv2.remap."""
if self._map1_fixed is not None:
return
self._map1_fixed, self._map2_fixed = cv2.convertMaps(
self.map_x,
self.map_y,
cv2.CV_16SC2,
)
invalid = ~self.valid_mask
if not np.any(invalid):
self._border_only_invalid = True
else:
self._border_only_invalid = bool(
np.all(self.map_x[invalid] <= 0.0) and np.all(self.map_y[invalid] <= 0.0)
)
def flatten(self, frame_bgr, border_value=0, dst=None):
"""Remap a source BGR frame into UV flatten coordinates.
frame_bgr: (H_src, W_src, 3) source camera image
border_value: fill value for invalid/out-of-domain pixels
returns: (H_out, W_out, 3) flattened BGR image
"""
return flatten_frame(frame_bgr, self, border_value=border_value, dst=dst)
def remap_valid_ratio(remap) -> float:
"""Fraction of output pixels covered by the valid UV domain."""
if isinstance(remap, UVRemap):
mask = remap.valid_mask
else:
mask = remap["valid_mask"]
return float(np.mean(mask))
def _cache_matches(
remap,
source_K,
source_width,
source_height,
out_width,
out_height,
flip_v,
edge_band,
edge_expand,
) -> bool:
"""True if cached remap matches the requested build parameters."""
if isinstance(remap, UVRemap):
rd = remap.to_dict()
else:
rd = remap
required = {"map_x", "map_y", "valid_mask", "source_width", "source_height", "out_width", "out_height"}
if not required.issubset(rd.keys()):
return False
if rd["map_x"].shape != (out_height, out_width):
return False
if int(rd["source_width"]) != source_width or int(rd["source_height"]) != source_height:
return False
if int(rd["out_width"]) != out_width or int(rd["out_height"]) != out_height:
return False
if "flip_v" in rd and int(rd["flip_v"]) != int(flip_v):
return False
if "edge_band" in rd and not np.isclose(float(rd["edge_band"]), float(edge_band), rtol=1e-6, atol=1e-6):
return False
if "edge_expand" in rd and not np.isclose(float(rd["edge_expand"]), float(edge_expand), rtol=1e-6, atol=1e-6):
return False
if "projection_K" in rd and not np.allclose(rd["projection_K"], source_K, rtol=1e-6, atol=1e-6):
return False
return True
def _remap_maps(remap) -> Tuple[np.ndarray, np.ndarray, np.ndarray, bool]:
if isinstance(remap, UVRemap):
remap.ensure_fast_maps()
valid_mask = remap.valid_mask
skip_invalid_fill = remap._border_only_invalid
if remap._map1_fixed is not None:
return remap._map1_fixed, remap._map2_fixed, valid_mask, skip_invalid_fill
return remap.map_x, remap.map_y, valid_mask, skip_invalid_fill
return (
remap["map_x"].astype(np.float32),
remap["map_y"].astype(np.float32),
remap["valid_mask"].astype(bool),
False,
)
def flatten_frame(frame, remap, border_value=0, dst=None):
"""Remap frame using map_x/map_y; zeros invalid UV pixels.
frame: (H, W, 3) or (H, W) source image
remap: UVRemap or dict with map_x, map_y, valid_mask
returns: flattened image, same channels as input
"""
map1, map2, valid_mask, skip_invalid_fill = _remap_maps(remap)
flat = cv2.remap(
frame,
map1,
map2,
interpolation=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=border_value,
dst=dst,
)
if skip_invalid_fill and border_value == 0:
return flat
if not flat.flags.writeable:
flat = flat.copy()
flat[~valid_mask] = border_value
return flat
+50
View File
@@ -0,0 +1,50 @@
"""UV pixel to world-space 3D coordinates.
Each flatten pixel maps to a world XYZ point for 3D visualization. Loaded
alongside the 2D remap as part of a cached UVDomain.
"""
from dataclasses import dataclass
import numpy as np
@dataclass
class UVSurfaceMap:
"""World-space 3D coordinates for each UV flatten pixel."""
world_xyz: np.ndarray
@property
def shape(self):
return self.world_xyz.shape[:2]
@property
def flat_xyz(self):
"""Legacy alias used by 3D viewer tools."""
return self.world_xyz
@classmethod
def from_dict(cls, data: dict):
"""Load from npz dict; accepts world_xyz or legacy flat_xyz key."""
if "world_xyz" in data:
xyz = np.asarray(data["world_xyz"], dtype=np.float32)
elif "flat_xyz" in data:
xyz = np.asarray(data["flat_xyz"], dtype=np.float32)
else:
raise KeyError("Surface map requires world_xyz or flat_xyz in cache")
return cls(world_xyz=xyz)
def to_dict(self) -> dict:
"""Serialize surface map for np.savez_compressed."""
return {"world_xyz": self.world_xyz.astype(np.float32)}
def valid_points(self, valid_mask):
"""Sample world XYZ at valid UV pixels.
valid_mask: (H, W) bool
returns: (N, 3) float32 world coordinates
"""
mask = np.asarray(valid_mask, dtype=bool)
xyz = self.world_xyz[mask]
return xyz[np.all(np.isfinite(xyz), axis=1)]
+110
View File
@@ -0,0 +1,110 @@
import numpy as np
from orisys import Sensor
from orisys.get_depth import Contact3D
from orisys.hhd import nHHD
def _circle_mask(shape, radius_ratio=0.35):
h, w = shape
yy, xx = np.mgrid[0:h, 0:w]
cy = (h - 1) / 2.0
cx = (w - 1) / 2.0
radius = min(h, w) * radius_ratio
return ((yy - cy) ** 2 + (xx - cx) ** 2) <= radius ** 2
def _radial_flow(shape, scale=0.05):
h, w = shape
yy, xx = np.mgrid[0:h, 0:w]
cy = (h - 1) / 2.0
cx = (w - 1) / 2.0
flow = np.zeros((h, w, 2), dtype=np.float32)
flow[:, :, 0] = (xx - cx) * scale
flow[:, :, 1] = (yy - cy) * scale
return flow
def test_masked_hhd_zeroes_invalid_region():
shape = (48, 56)
mask = _circle_mask(shape)
flow = _radial_flow(shape, scale=0.03)
flow[~mask] = 0
solver = nHHD(cuda=False, grid=shape, spacings=(1.0, 1.0), poisson_backend="cpu")
_, forces = solver.get_force(flow, valid_mask=mask)
vnormal = solver.get_normal()
vshear = solver.get_shear()
assert np.isfinite(forces).all()
assert np.allclose(vnormal[~mask], 0.0, atol=1e-5)
assert np.allclose(vshear[~mask], 0.0, atol=1e-5)
def test_masked_hhd_ignores_invalid_background_spike():
shape = (48, 56)
mask = _circle_mask(shape)
base = _radial_flow(shape, scale=0.02)
base[~mask] = 0
contaminated = base.copy()
contaminated[~mask, 0] = 25.0
contaminated[~mask, 1] = -15.0
solver = nHHD(cuda=False, grid=shape, spacings=(1.0, 1.0), poisson_backend="cpu")
_, forces_base = solver.get_force(base, valid_mask=mask)
_, forces_contaminated = solver.get_force(contaminated, valid_mask=mask)
assert np.allclose(forces_base, forces_contaminated, atol=1e-3, rtol=1e-3)
def test_sensor_detect_motion_ignores_invalid_region():
sensor = Sensor(
vid_src="external",
cuda=False,
isstitch=False,
verbose=False,
config_name="dd02-ov",
mirror=False,
backend="cpu",
)
try:
frame = np.zeros((728, 672, 3), dtype=np.uint8)
mask = np.zeros((728, 672), dtype=bool)
mask[180:540, 160:512] = True
sensor.get_img(frame=frame, already_processed=True, valid_mask=mask)
sensor.flow = np.zeros((728, 672, 2), dtype=np.float32)
sensor.flow[~mask] = 10.0
assert not sensor.detect_motion(threshold=1.0)
sensor.flow[:] = 0.0
sensor.flow[mask] = 10.0
assert sensor.detect_motion(threshold=1.0)
finally:
sensor.disconnect()
def test_contact_centroid_respects_domain_mask():
config_contact = {
"min_contour_area": 25,
"flow_threshold": 0.01,
"div_threshold": 0.0,
"expansion_threshold": 0.0,
"kernel_size": [3, 3],
"denoise_sigma": 0.0,
}
shape = (96, 96)
mask = _circle_mask(shape, radius_ratio=0.28)
flow = _radial_flow(shape, scale=0.05)
flow[~mask] = 0.0
flow[:, :8, 0] = 100.0
flow[:, :8, 1] = 100.0
contact = Contact3D(config_contact=config_contact, verbose=False)
result = contact.compute_foe(flow, valid_mask=mask)
assert result["centroid_found"]
cx, cy = result["centroid"]
assert mask[cy, cx]
+162
View File
@@ -0,0 +1,162 @@
from pathlib import Path
import numpy as np
from orisys.hhd import PoissonSolver, nHHD
REPO_ROOT = Path(__file__).resolve().parents[1]
def _circle_mask(shape, radius_ratio=0.35):
h, w = shape
yy, xx = np.mgrid[0:h, 0:w]
cy = (h - 1) / 2.0
cx = (w - 1) / 2.0
radius = min(h, w) * radius_ratio
return ((yy - cy) ** 2 + (xx - cx) ** 2) <= radius ** 2
def _smooth_rhs(shape):
h, w = shape
yy, xx = np.mgrid[0:h, 0:w]
field = (
np.sin(np.pi * xx / max(1, w - 1))
+ 0.5 * np.cos(2.0 * np.pi * yy / max(1, h - 1))
)
return field.astype(np.float32)
def _radial_flow(shape, scale=0.05):
h, w = shape
yy, xx = np.mgrid[0:h, 0:w]
cy = (h - 1) / 2.0
cx = (w - 1) / 2.0
flow = np.zeros((h, w, 2), dtype=np.float32)
flow[:, :, 0] = (xx - cx) * scale
flow[:, :, 1] = (yy - cy) * scale
return flow
def _masked_rhs_vector(rhs, mask, pinned):
vec = rhs[mask].astype(np.float64)
vec -= vec.mean()
vec[pinned] = 0.0
return vec
def _masked_residual_norm(solver, field, rhs, mask):
cache = solver._masked_cache if solver._masked_cache_matches(mask) else solver._build_masked_cache(mask)
rhs_vec = _masked_rhs_vector(rhs, mask, cache["pinned"])
residual = cache["lap"].dot(field[mask].astype(np.float64)) - rhs_vec
return float(np.linalg.norm(residual))
def test_capfft_solver_improves_over_raw_fft_on_synthetic_mask():
shape = (40, 48)
mask = _circle_mask(shape)
rhs = _smooth_rhs(shape)
rhs_masked = np.where(mask, rhs, 0).astype(np.float32)
sparse_solver = PoissonSolver(
grid=shape,
spacings=(1.0, 1.0),
poisson_backend="cpu",
masked_poisson_solver="sparse",
)
capfft_solver = PoissonSolver(
grid=shape,
spacings=(1.0, 1.0),
poisson_backend="cpu",
masked_poisson_solver="capfft",
masked_poisson_max_boundary=48,
masked_poisson_max_iters=3,
)
sparse = sparse_solver.solve(rhs_masked, mask=mask)
capfft = capfft_solver.solve(rhs_masked, mask=mask)
raw_fft = sparse_solver.solve(rhs_masked, mask=None)
assert np.allclose(capfft[~mask], 0.0, atol=1e-6)
err_cap = float(np.linalg.norm((capfft - sparse)[mask]))
err_raw = float(np.linalg.norm((raw_fft - sparse)[mask]))
assert err_cap < err_raw
res_cap = _masked_residual_norm(sparse_solver, capfft, rhs_masked, mask)
res_raw = _masked_residual_norm(sparse_solver, raw_fft, rhs_masked, mask)
assert res_cap < res_raw
def test_capfft_hhd_forces_track_sparse_reference():
shape = (36, 44)
mask = _circle_mask(shape)
flow = _radial_flow(shape, scale=0.03)
flow[~mask] = 0.0
sparse_hhd = nHHD(
cuda=False,
grid=shape,
spacings=(1.0, 1.0),
poisson_backend="cpu",
masked_poisson_solver="sparse",
)
capfft_hhd = nHHD(
cuda=False,
grid=shape,
spacings=(1.0, 1.0),
poisson_backend="cpu",
masked_poisson_solver="capfft",
masked_poisson_max_boundary=40,
masked_poisson_max_iters=3,
)
_, sparse_forces = sparse_hhd.get_force(flow, valid_mask=mask)
_, capfft_forces = capfft_hhd.get_force(flow, valid_mask=mask)
denom = max(1.0, float(np.linalg.norm(sparse_forces)))
rel = float(np.linalg.norm(capfft_forces - sparse_forces) / denom)
assert rel < 0.6
vnormal = capfft_hhd.get_normal()
vshear = capfft_hhd.get_shear()
assert np.allclose(vnormal[~mask], 0.0, atol=1e-5)
assert np.allclose(vshear[~mask], 0.0, atol=1e-5)
def test_capfft_uv_mask_fixture_improves_over_raw_fft():
uv_path = REPO_ROOT / "unroll_map_uv.npz"
if not uv_path.exists():
return
full_mask = np.load(uv_path)["valid_mask"].astype(bool)
mask = full_mask[::24, ::24]
if mask.sum() < 16:
mask = full_mask[::16, ::16]
shape = mask.shape
rhs = _smooth_rhs(shape)
rhs_masked = np.where(mask, rhs, 0).astype(np.float32)
sparse_solver = PoissonSolver(
grid=shape,
spacings=(1.0, 1.0),
poisson_backend="cpu",
masked_poisson_solver="sparse",
)
capfft_solver = PoissonSolver(
grid=shape,
spacings=(1.0, 1.0),
poisson_backend="cpu",
masked_poisson_solver="capfft",
masked_poisson_max_boundary=24,
masked_poisson_max_iters=3,
)
sparse = sparse_solver.solve(rhs_masked, mask=mask)
capfft = capfft_solver.solve(rhs_masked, mask=mask)
raw_fft = sparse_solver.solve(rhs_masked, mask=None)
err_cap = float(np.linalg.norm((capfft - sparse)[mask]))
err_raw = float(np.linalg.norm((raw_fft - sparse)[mask]))
assert err_cap < err_raw
+96
View File
@@ -0,0 +1,96 @@
from pathlib import Path
import numpy as np
import pytest
from orisys.unroll import UVDomain, UVRemap, UVSurfaceMap, flatten_frame, remap_valid_ratio
REPO_ROOT = Path(__file__).resolve().parents[1]
UV_MAP_PATH = REPO_ROOT / "unroll_map_uv.npz"
@pytest.fixture
def uv_domain():
if not UV_MAP_PATH.is_file():
pytest.skip(f"UV remap cache not found: {UV_MAP_PATH}")
return UVDomain.load(str(UV_MAP_PATH))
@pytest.fixture
def uv_remap(uv_domain):
return uv_domain.remap
def test_uv_remap_load_and_flatten(uv_remap):
h, w = uv_remap.out_height, uv_remap.out_width
src_h = int(uv_remap.metadata.get("source_height", h))
src_w = int(uv_remap.metadata.get("source_width", w))
source = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
flat = uv_remap.flatten(source, border_value=0)
assert flat.shape == (h, w, 3)
assert flat.dtype == np.uint8
invalid = ~uv_remap.valid_mask
if np.any(invalid):
assert np.all(flat[invalid] == 0)
def test_flatten_frame_dict_compat(uv_remap):
rd = uv_remap.to_dict()
source = np.zeros((int(rd["source_height"]), int(rd["source_width"]), 3), dtype=np.uint8)
flat = flatten_frame(source, rd, border_value=0)
assert flat.shape == uv_remap.map_x.shape + (3,)
def test_remap_valid_ratio(uv_remap):
ratio = remap_valid_ratio(uv_remap)
assert 0.0 < ratio <= 1.0
def test_uv_domain_load(uv_domain):
assert isinstance(uv_domain.remap, UVRemap)
assert uv_domain.remap.valid_mask.shape == (
uv_domain.remap.out_height,
uv_domain.remap.out_width,
)
def test_uv_remap_fast_maps_match_float(uv_remap):
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
source = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
uv_remap._map1_fixed = None
uv_remap._map2_fixed = None
float_flat = flatten_frame(source, {
"map_x": uv_remap.map_x,
"map_y": uv_remap.map_y,
"valid_mask": uv_remap.valid_mask,
}, border_value=0)
uv_remap.ensure_fast_maps()
fast_flat = flatten_frame(source, uv_remap, border_value=0)
assert np.array_equal(fast_flat, float_flat)
def test_uv_remap_reuses_provided_dst_buffer(uv_remap):
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
source = np.random.randint(0, 255, (src_h, src_w), dtype=np.uint8)
dst = np.empty(uv_remap.valid_mask.shape, dtype=np.uint8)
flat = flatten_frame(source, uv_remap, border_value=0, dst=dst)
assert flat is dst
assert flat.shape == uv_remap.valid_mask.shape
def test_uv_surface_map_when_present(uv_domain):
if uv_domain.surface is None:
pytest.skip("Cache has no world_xyz yet; regenerate domain to test surface")
surface = uv_domain.surface
assert isinstance(surface, UVSurfaceMap)
assert surface.world_xyz.shape[:2] == uv_domain.remap.valid_mask.shape
assert surface.flat_xyz is surface.world_xyz
points = surface.valid_points(uv_domain.valid_mask)
assert points.ndim == 2 and points.shape[1] == 3
assert np.all(np.isfinite(points))
+93
View File
@@ -0,0 +1,93 @@
from pathlib import Path
import cv2
import numpy as np
import pytest
from orisys.frame_processors import UVUnrollFrameProcessor
from orisys.unroll import UVDomain
REPO_ROOT = Path(__file__).resolve().parents[1]
UV_MAP_PATH = REPO_ROOT / "unroll_map_uv.npz"
@pytest.fixture
def uv_domain():
if not UV_MAP_PATH.is_file():
pytest.skip(f"UV remap cache not found: {UV_MAP_PATH}")
return UVDomain.load(str(UV_MAP_PATH))
def test_uv_unroll_processor_output_shape(uv_domain):
uv_remap = uv_domain.remap
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
processor = UVUnrollFrameProcessor(
uv_remap,
source_width=src_w,
source_height=src_h,
flip_source_vertical=True,
border_value=(0, 0, 0),
)
frame = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
unrolled = processor.process(frame)
assert unrolled.shape == (uv_remap.out_height, uv_remap.out_width, 3)
assert processor.valid_mask.shape == (uv_remap.out_height, uv_remap.out_width)
assert unrolled.dtype == np.uint8
def test_uv_unroll_processor_prepare_source(uv_domain):
uv_remap = uv_domain.remap
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
processor = UVUnrollFrameProcessor(
uv_remap,
source_width=src_w,
source_height=src_h,
flip_source_vertical=False,
)
frame = np.ones((src_h, src_w, 3), dtype=np.uint8) * 128
prepared = processor.prepare_source(frame)
assert prepared.shape == (src_h, src_w, 3)
def test_uv_unroll_processor_flatten_prepared_matches_process(uv_domain):
uv_remap = uv_domain.remap
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
processor = UVUnrollFrameProcessor(
uv_remap,
source_width=src_w,
source_height=src_h,
flip_source_vertical=True,
border_value=(0, 0, 0),
)
frame = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
prepared = processor.prepare_source(frame)
assert np.array_equal(processor.flatten_prepared(prepared), processor.process(frame))
def test_uv_unroll_processor_gray_fast_path_matches_bgr_result(uv_domain):
uv_remap = uv_domain.remap
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
processor = UVUnrollFrameProcessor(
uv_remap,
source_width=src_w,
source_height=src_h,
flip_source_vertical=True,
border_value=(0, 0, 0),
)
frame = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
prepared_gray = processor.prepare_source_gray(frame)
gray_unrolled = processor.flatten_prepared_gray(prepared_gray)
bgr_unrolled = processor.process(frame)
assert gray_unrolled.shape == (uv_remap.out_height, uv_remap.out_width)
assert np.allclose(gray_unrolled, cv2.cvtColor(bgr_unrolled, cv2.COLOR_BGR2GRAY), atol=1)
@@ -8,14 +8,70 @@ applied to each live frame with cv2.remap().
"""
import argparse
import json
import os
import sys
import time
import cv2
import numpy as np
import orisys
from orisys.frame_processors import UVUnrollFrameProcessor
from orisys.unroll import (
barycentric_grid,
build_or_load_domain,
compute_uv_distortion_metrics,
face_uvs_from_loop_data,
load_calib_json,
load_exported_data,
recommend_edge_expand,
uv_triangle_to_image,
)
class StepTimer:
"""Accumulate per-step wall times for rolling runtime averages."""
def __init__(self):
self._starts = {}
self.totals = {}
self.counts = {}
def start(self, name):
self._starts[name] = time.perf_counter()
def stop(self, name):
dt = time.perf_counter() - self._starts.pop(name)
self.totals[name] = self.totals.get(name, 0.0) + dt
self.counts[name] = self.counts.get(name, 0) + 1
return dt
def avg_ms(self, name):
if self.counts.get(name, 0) == 0:
return 0.0
return 1000.0 * self.totals[name] / self.counts[name]
def benchmark_step(label, func, *args, **kwargs):
"""Run func once and print its wall time."""
t0 = time.perf_counter()
result = func(*args, **kwargs)
ms = (time.perf_counter() - t0) * 1000.0
print(f"[benchmark] {label}: {ms:.1f} ms")
return result
def print_runtime_benchmark(timer):
"""Format rolling average timings for the live loop."""
parts = [
f"capture={timer.avg_ms('capture'):.1f}ms",
f"uv_prep={timer.avg_ms('uv_prepare'):.1f}ms",
f"uv_remap={timer.avg_ms('uv_remap'):.1f}ms",
f"get_img={timer.avg_ms('get_img'):.1f}ms",
f"flow={timer.avg_ms('flow'):.1f}ms",
f"hhd={timer.avg_ms('hhd'):.1f}ms",
f"contact={timer.avg_ms('contact'):.1f}ms",
f"viz={timer.avg_ms('viz'):.1f}ms",
]
return "bench[" + ", ".join(parts) + "]"
def resize_to_height(image, target_height):
@@ -58,6 +114,52 @@ def add_panel_title(image, title):
return out
def apply_display_mask(image, valid_mask, border_value=0):
"""Zero invalid UV pixels in visualization panels."""
if image is None or valid_mask is None:
return image
out = image.copy()
mask = valid_mask.astype(np.uint8)
if mask.shape != out.shape[:2]:
mask = cv2.resize(mask, (out.shape[1], out.shape[0]), interpolation=cv2.INTER_NEAREST)
invalid = ~mask.astype(bool)
out[invalid] = border_value
return out
def contact_inside_mask(centroid, valid_mask):
if centroid is None or valid_mask is None:
return None
if len(centroid) < 2:
return None
cx, cy = int(centroid[0]), int(centroid[1])
if cy < 0 or cy >= valid_mask.shape[0] or cx < 0 or cx >= valid_mask.shape[1]:
return False
return bool(valid_mask[cy, cx])
def overlay_runtime_status(image, sensor, centroid=None):
"""Annotate a panel with mask-aware runtime diagnostics."""
out = image.copy()
stats = getattr(sensor, "mask_stats", {}) or {}
coverage = 100.0 * float(stats.get("coverage", 0.0))
flow_valid = float(stats.get("flow_mean_valid", 0.0))
flow_invalid = float(stats.get("flow_mean_invalid", 0.0))
centroid_ok = contact_inside_mask(centroid, getattr(sensor, "valid_mask", None))
centroid_label = "n/a" if centroid_ok is None else ("inside" if centroid_ok else "outside")
lines = [
f"mask {coverage:.1f}%",
f"flow valid {flow_valid:.3f}",
f"flow invalid {flow_invalid:.3f}",
f"contact {centroid_label}",
]
y = 54
for line in lines:
cv2.putText(out, line, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 2, cv2.LINE_AA)
y += 24
return out
def make_tiled_dashboard(panels, tile_width=640, tile_height=360):
"""Arrange four panels into a 2x2 dashboard."""
tiles = [
@@ -69,382 +171,6 @@ def make_tiled_dashboard(panels, tile_width=640, tile_height=360):
return np.vstack([top_row, bottom_row])
def load_calib_json(json_path):
with open(json_path, "r", encoding="utf-8") as f:
calib = json.load(f)
K = np.array(calib["camera_matrix"], dtype=np.float64)
dist = np.array(calib["distortion"], dtype=np.float64).flatten()
img_size = tuple(int(v) for v in calib["image_size"])
return K, dist, img_size
def load_exported_data(mesh_path, intrinsics_path, pose_path):
mesh = np.load(mesh_path)
cam = np.load(intrinsics_path)
return {
"vertices": mesh["vertices"].astype(np.float64),
"faces": mesh["faces"].astype(np.int32),
"uv_data": mesh["uv_data"].astype(np.float64),
"loop_to_vertex": mesh["loop_to_vertex"].astype(np.int32),
"K": cam["K"].astype(np.float64),
"width": int(cam["width"]),
"height": int(cam["height"]),
"world_to_cam": np.load(pose_path).astype(np.float64),
}
def face_uvs_from_loop_data(faces, uv_data, loop_to_vertex):
if len(uv_data) == 0:
raise ValueError("mesh.npz does not contain UV data; UV flattening is unavailable.")
if len(uv_data) != len(faces) * 3:
raise ValueError("Expected a triangulated mesh with exactly 3 UV loops per face.")
face_loop_vertices = loop_to_vertex.reshape(-1, 3)
if face_loop_vertices.shape != faces.shape or not np.array_equal(face_loop_vertices, faces):
raise ValueError("UV loop order does not match exported face vertex order.")
return uv_data.reshape(-1, 3, 2)
def project_world_to_image(points_world, K, world_to_cam):
"""Project world points into the image using Blender/OpenCV-compatible intrinsics."""
points_h = np.column_stack((points_world, np.ones(len(points_world), dtype=np.float64)))
points_cam = (world_to_cam @ points_h.T).T[:, :3]
z = -points_cam[:, 2]
x = K[0, 0] * (points_cam[:, 0] / z) + K[0, 2]
y = K[1, 1] * (points_cam[:, 1] / z) + K[1, 2]
return np.column_stack((x, y)), z
def barycentric_grid(px, py, tri):
a, b, c = tri
denom = (b[1] - c[1]) * (a[0] - c[0]) + (c[0] - b[0]) * (a[1] - c[1])
if abs(denom) < 1e-12:
return None
w0 = ((b[1] - c[1]) * (px - c[0]) + (c[0] - b[0]) * (py - c[1])) / denom
w1 = ((c[1] - a[1]) * (px - c[0]) + (a[0] - c[0]) * (py - c[1])) / denom
w2 = 1.0 - w0 - w1
return w0, w1, w2
def triangle_edge_lengths_3d(tri):
return np.array(
[
np.linalg.norm(tri[1] - tri[0]),
np.linalg.norm(tri[2] - tri[1]),
np.linalg.norm(tri[0] - tri[2]),
],
dtype=np.float64,
)
def triangle_edge_lengths_2d(tri):
return np.array(
[
np.linalg.norm(tri[1] - tri[0]),
np.linalg.norm(tri[2] - tri[1]),
np.linalg.norm(tri[0] - tri[2]),
],
dtype=np.float64,
)
def triangle_area_3d(tri):
return 0.5 * np.linalg.norm(np.cross(tri[1] - tri[0], tri[2] - tri[0]))
def triangle_area_2d(tri):
return 0.5 * abs(
(tri[1, 0] - tri[0, 0]) * (tri[2, 1] - tri[0, 1])
- (tri[1, 1] - tri[0, 1]) * (tri[2, 0] - tri[0, 0])
)
def summarize_uv_region(relative_scale, anisotropy, mask):
mask_scale = mask & np.isfinite(relative_scale) & (relative_scale > 0)
mask_aniso = mask & np.isfinite(anisotropy) & (anisotropy > 0)
if not np.any(mask_scale):
return {
"face_count": 0,
"relative_scale_p05": float("nan"),
"relative_scale_p50": float("nan"),
"relative_scale_p95": float("nan"),
"relative_scale_min": float("nan"),
"relative_scale_max": float("nan"),
"anisotropy_p50": float("nan"),
"anisotropy_p95": float("nan"),
"anisotropy_max": float("nan"),
}
return {
"face_count": int(np.count_nonzero(mask_scale)),
"relative_scale_p05": float(np.percentile(relative_scale[mask_scale], 5)),
"relative_scale_p50": float(np.percentile(relative_scale[mask_scale], 50)),
"relative_scale_p95": float(np.percentile(relative_scale[mask_scale], 95)),
"relative_scale_min": float(np.min(relative_scale[mask_scale])),
"relative_scale_max": float(np.max(relative_scale[mask_scale])),
"anisotropy_p50": float(np.percentile(anisotropy[mask_aniso], 50)) if np.any(mask_aniso) else float("nan"),
"anisotropy_p95": float(np.percentile(anisotropy[mask_aniso], 95)) if np.any(mask_aniso) else float("nan"),
"anisotropy_max": float(np.max(anisotropy[mask_aniso])) if np.any(mask_aniso) else float("nan"),
}
def compute_uv_distortion_metrics(data, edge_band=0.12):
"""Estimate local UV stretch/compression and anisotropy per face."""
faces = data["faces"]
vertices = data["vertices"]
face_uvs = face_uvs_from_loop_data(faces, data["uv_data"], data["loop_to_vertex"])
local_scale = np.full(len(faces), np.nan, dtype=np.float64)
relative_scale = np.full(len(faces), np.nan, dtype=np.float64)
anisotropy = np.full(len(faces), np.nan, dtype=np.float64)
uv_centers = np.full((len(faces), 2), np.nan, dtype=np.float64)
edge_mask = np.zeros(len(faces), dtype=bool)
for idx, (face, uv_tri) in enumerate(zip(faces, face_uvs)):
tri_3d = vertices[face]
tri_uv = uv_tri
uv_centers[idx] = np.mean(tri_uv, axis=0)
edges_3d = triangle_edge_lengths_3d(tri_3d)
edges_uv = triangle_edge_lengths_2d(tri_uv)
valid_edges = edges_uv > 1e-9
if np.count_nonzero(valid_edges) < 2:
continue
edge_scales = edges_3d[valid_edges] / edges_uv[valid_edges]
local_scale[idx] = float(np.exp(np.mean(np.log(np.clip(edge_scales, 1e-12, None)))))
anisotropy[idx] = float(np.max(edge_scales) / max(np.min(edge_scales), 1e-12))
area_uv = triangle_area_2d(tri_uv)
area_3d = triangle_area_3d(tri_3d)
if area_uv > 1e-12 and area_3d > 1e-12:
local_scale[idx] = float(np.sqrt(area_3d / area_uv))
valid_scale = np.isfinite(local_scale) & (local_scale > 0)
if not np.any(valid_scale):
raise RuntimeError("Could not compute UV distortion metrics; UV triangles may be degenerate.")
median_scale = float(np.median(local_scale[valid_scale]))
relative_scale[valid_scale] = local_scale[valid_scale] / median_scale
min_edge_distance = np.min(
np.column_stack(
[
uv_centers[:, 0],
1.0 - uv_centers[:, 0],
uv_centers[:, 1],
1.0 - uv_centers[:, 1],
]
),
axis=1,
)
edge_mask = min_edge_distance < edge_band
middle_mask = ~edge_mask
valid_aniso = np.isfinite(anisotropy) & (anisotropy > 0)
stats = {
"median_mm_per_uv": median_scale,
"relative_scale_p05": float(np.percentile(relative_scale[valid_scale], 5)),
"relative_scale_p50": float(np.percentile(relative_scale[valid_scale], 50)),
"relative_scale_p95": float(np.percentile(relative_scale[valid_scale], 95)),
"relative_scale_min": float(np.min(relative_scale[valid_scale])),
"relative_scale_max": float(np.max(relative_scale[valid_scale])),
"anisotropy_p50": float(np.percentile(anisotropy[valid_aniso], 50)) if np.any(valid_aniso) else float("nan"),
"anisotropy_p95": float(np.percentile(anisotropy[valid_aniso], 95)) if np.any(valid_aniso) else float("nan"),
"anisotropy_max": float(np.max(anisotropy[valid_aniso])) if np.any(valid_aniso) else float("nan"),
"edge_band": float(edge_band),
}
return {
"relative_scale": relative_scale,
"anisotropy": anisotropy,
"uv_centers": uv_centers,
"edge_mask": edge_mask,
"middle_mask": middle_mask,
"stats": stats,
"middle_stats": summarize_uv_region(relative_scale, anisotropy, middle_mask),
"edge_stats": summarize_uv_region(relative_scale, anisotropy, edge_mask),
}
def recommend_edge_expand(distortion, max_expand=3.0):
"""Suggest how much more UV space to allocate to the edge band."""
middle_p95 = distortion["middle_stats"]["relative_scale_p95"]
edge_p95 = distortion["edge_stats"]["relative_scale_p95"]
if not np.isfinite(middle_p95) or not np.isfinite(edge_p95) or middle_p95 <= 0:
return 1.0
# Use a square-root response so large edge distortion does not over-expand.
return float(np.clip(np.sqrt(edge_p95 / middle_p95), 1.0, max_expand))
def warp_uv_axis(values, edge_band, edge_expand):
"""Allocate more normalized output span to UV edge regions."""
u = np.asarray(values, dtype=np.float64)
e = float(np.clip(edge_band, 0.0, 0.49))
s = float(max(edge_expand, 1.0))
if e <= 0.0 or s <= 1.0:
return u.copy()
middle = 1.0 - 2.0 * e
total = middle + 2.0 * e * s
out = np.empty_like(u)
left = u < e
right = u > (1.0 - e)
center = ~(left | right)
out[left] = (u[left] * s) / total
out[center] = (e * s + (u[center] - e)) / total
out[right] = (e * s + middle + (u[right] - (1.0 - e)) * s) / total
return out
def uv_triangle_to_image(uv_tri, out_width, out_height, flip_v=True, edge_band=0.12, edge_expand=1.0):
"""Map UV triangle into output image coordinates with optional edge expansion."""
tri = np.empty_like(uv_tri)
u_warp = warp_uv_axis(uv_tri[:, 0], edge_band=edge_band, edge_expand=edge_expand)
v_warp = warp_uv_axis(uv_tri[:, 1], edge_band=edge_band, edge_expand=edge_expand)
tri[:, 0] = u_warp * (out_width - 1)
tri[:, 1] = ((1.0 - v_warp) if flip_v else v_warp) * (out_height - 1)
return tri
def build_uv_remap(
data,
source_K,
source_width,
source_height,
out_width,
out_height,
flip_v=True,
edge_band=0.12,
edge_expand=1.0,
):
faces = data["faces"]
vertices = data["vertices"]
face_uvs = face_uvs_from_loop_data(faces, data["uv_data"], data["loop_to_vertex"])
map_x = np.full((out_height, out_width), -1.0, dtype=np.float32)
map_y = np.full((out_height, out_width), -1.0, dtype=np.float32)
valid_mask = np.zeros((out_height, out_width), dtype=bool)
depth_buffer = np.full((out_height, out_width), np.inf, dtype=np.float64)
print("Building UV remap ...")
t_start = time.time()
for face_idx, (face, uv_tri) in enumerate(zip(faces, face_uvs)):
tri = uv_triangle_to_image(
uv_tri,
out_width,
out_height,
flip_v=flip_v,
edge_band=edge_band,
edge_expand=edge_expand,
)
x0 = max(0, int(np.floor(np.min(tri[:, 0]))))
x1 = min(out_width - 1, int(np.ceil(np.max(tri[:, 0]))))
y0 = max(0, int(np.floor(np.min(tri[:, 1]))))
y1 = min(out_height - 1, int(np.ceil(np.max(tri[:, 1]))))
if x1 < x0 or y1 < y0:
continue
yy, xx = np.mgrid[y0:y1 + 1, x0:x1 + 1]
bary = barycentric_grid(xx + 0.5, yy + 0.5, tri)
if bary is None:
continue
w0, w1, w2 = bary
inside = (w0 >= -1e-6) & (w1 >= -1e-6) & (w2 >= -1e-6)
if not np.any(inside):
continue
weights = np.stack((w0[inside], w1[inside], w2[inside]), axis=1)
points_world = weights @ vertices[face]
pixels, depth = project_world_to_image(points_world, source_K, data["world_to_cam"])
src_x = pixels[:, 0]
src_y = pixels[:, 1]
target_y = yy[inside]
target_x = xx[inside]
valid = (
(depth > 0)
& (src_x >= 0)
& (src_x <= source_width - 1)
& (src_y >= 0)
& (src_y <= source_height - 1)
& (depth < depth_buffer[target_y, target_x])
)
if not np.any(valid):
continue
tx = target_x[valid]
ty = target_y[valid]
map_x[ty, tx] = src_x[valid].astype(np.float32)
map_y[ty, tx] = src_y[valid].astype(np.float32)
valid_mask[ty, tx] = True
depth_buffer[ty, tx] = depth[valid]
if (face_idx % 200) == 0:
print(f" face {face_idx}/{len(faces)}")
print(f"Remap build took {time.time() - t_start:.1f} s")
return {
"map_x": map_x,
"map_y": map_y,
"valid_mask": valid_mask,
"source_width": np.array(source_width, dtype=np.int32),
"source_height": np.array(source_height, dtype=np.int32),
"out_width": np.array(out_width, dtype=np.int32),
"out_height": np.array(out_height, dtype=np.int32),
"flip_v": np.array(int(flip_v), dtype=np.int32),
"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),
"method": np.array("uv"),
}
def remap_valid_ratio(remap):
return float(np.mean(remap["valid_mask"]))
def load_cached_remap(path):
data = np.load(path)
return {key: data[key] for key in data.files}
def cache_matches(remap, source_K, source_width, source_height, out_width, out_height, flip_v, edge_band, edge_expand):
required = {"map_x", "map_y", "valid_mask", "source_width", "source_height", "out_width", "out_height"}
if not required.issubset(remap.keys()):
return False
if remap["map_x"].shape != (out_height, out_width):
return False
if int(remap["source_width"]) != source_width or int(remap["source_height"]) != source_height:
return False
if int(remap["out_width"]) != out_width or int(remap["out_height"]) != out_height:
return False
if "flip_v" in remap and int(remap["flip_v"]) != int(flip_v):
return False
if "edge_band" in remap and not np.isclose(float(remap["edge_band"]), float(edge_band), rtol=1e-6, atol=1e-6):
return False
if "edge_expand" in remap and not np.isclose(float(remap["edge_expand"]), float(edge_expand), rtol=1e-6, atol=1e-6):
return False
if "projection_K" in remap and not np.allclose(remap["projection_K"], source_K, rtol=1e-6, atol=1e-6):
return False
return True
def flatten_frame(frame, remap, border_value=0):
flat = cv2.remap(
frame,
remap["map_x"].astype(np.float32),
remap["map_y"].astype(np.float32),
interpolation=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=border_value,
)
flat = flat.copy()
flat[~remap["valid_mask"].astype(bool)] = border_value
return flat
def make_uv_debug_image(data, out_width, out_height, flip_v=True, edge_band=0.12, edge_expand=1.0):
"""Create a UV-domain reference image to verify orientation of the remap."""
face_uvs = face_uvs_from_loop_data(data["faces"], data["uv_data"], data["loop_to_vertex"])
@@ -576,22 +302,59 @@ def open_source(video_path, camera_id):
return cv2.VideoCapture(video_path) if video_path else cv2.VideoCapture(camera_id)
def parse_video_source(value):
try:
return int(value)
except (TypeError, ValueError):
return value
PERF_TUNING_EPILOG = """
Performance tuning (config JSON + CLI):
tracker.dis_preset medium -> fast/ultrafast faster DIS optical flow
tracker.loop_closure_interval 1 -> 3-5 fewer loop-closure remaps
tracker.hhd_downsample 3 -> 4 smaller HHD Poisson grid
contact.denoise_sigma 12 -> 6 or 0 faster FoE contact step
--masked_poisson_solver sparse vs capfft benchmark HHD solve time
--no-viz skip dashboard/imshow ~20 ms/frame savings
"""
def align_frame_to_uv_map(frame_bgr, flip_vertical=True):
"""Match live-frame orientation to the Blender-exported UV remap convention."""
if not flip_vertical:
return frame_bgr
return np.flipud(frame_bgr).copy()
def run_deformation_timed(sensor, timer, check_motion=True, threshold=0, tail_frames=10):
"""Run flow + gated HHD with separate benchmark buckets."""
timer.start("flow")
sensor.compute_flow()
timer.stop("flow")
timer.start("hhd")
run_hhd = True
if check_motion:
is_motion = sensor.detect_motion(threshold=threshold)
try:
tail_frames = max(0, int(tail_frames))
except (TypeError, ValueError):
tail_frames = 0
if is_motion:
sensor.tail_frames_count = 0
elif sensor.tail_frames_count < tail_frames:
sensor.tail_frames_count += 1
else:
run_hhd = False
sensor.tail_frames_count = tail_frames
if run_hhd:
sensor.compute_hhd()
timer.stop("hhd")
def configure_masked_poisson(sensor, args):
"""Apply masked-HHD solver settings to the externally fed SDK sensor."""
sensor.masked_poisson_solver = args.masked_poisson_solver
sensor.masked_poisson_max_boundary = args.masked_poisson_max_boundary
sensor.masked_poisson_boundary_stride = args.masked_poisson_boundary_stride
sensor.masked_poisson_max_iters = args.masked_poisson_max_iters
sensor.masked_poisson_tol = args.masked_poisson_tol
sensor.masked_poisson_reg = args.masked_poisson_reg
def main():
parser = argparse.ArgumentParser(description="Real-time UV-based finger unrolling")
parser = argparse.ArgumentParser(
description="Real-time UV-based finger unrolling",
epilog=PERF_TUNING_EPILOG,
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("--mesh", required=True, help="Exported mesh.npz path")
parser.add_argument("--K", dest="intrinsics_path", required=True, help="Exported camera_intrinsics.npz path")
parser.add_argument("--pose", required=True, help="Exported world_to_cam.npy path")
@@ -599,8 +362,8 @@ def main():
# With a 12% UV edge band on each side, only about 76% of each axis is the
# lower-distortion middle region. Use a slightly larger full remap so the
# useful center keeps roughly the same sampling density as the prior 506x550.
parser.add_argument("--out_width", type=int, default=672)
parser.add_argument("--out_height", type=int, default=728)
parser.add_argument("--out_width", type=int, default=504)
parser.add_argument("--out_height", type=int, default=506)
parser.add_argument("--camera_id", type=int, default=0)
parser.add_argument("--video", help="Video file path (overrides camera)")
parser.add_argument("--map_cache", default="unroll_map_uv.npz")
@@ -610,13 +373,30 @@ def main():
parser.add_argument("--edge_expand", type=float, default=0.0, help="Extra UV allocation for edge band; 0 chooses automatically from distortion stats")
parser.add_argument("--sdk_config", default="./config/dd02-ov.json", help="Orisys SDK config path/name")
parser.add_argument("--sdk_source", default=None, help="Dummy source used to initialize Sensor; defaults to --video or --camera_id")
parser.add_argument("--masked_poisson_solver", choices=["sparse", "capfft"], default="sparse", help="Masked HHD Poisson solver used inside the SDK")
parser.add_argument("--masked_poisson_max_boundary", type=int, default=96, help="Maximum number of UV-boundary samples used by the capfft solver")
parser.add_argument("--masked_poisson_boundary_stride", type=int, default=1, help="Subsample stride applied before capfft boundary sampling")
parser.add_argument("--masked_poisson_max_iters", type=int, default=3, help="Maximum capfft correction iterations per Poisson solve")
parser.add_argument("--masked_poisson_tol", type=float, default=1e-4, help="Residual tolerance used by the capfft correction loop")
parser.add_argument("--masked_poisson_reg", type=float, default=1e-8, help="Diagonal regularization applied to the capacitance matrix")
parser.add_argument("--verbose", action="store_true", help="Enable verbose SDK logs")
parser.add_argument(
"--no-viz",
action="store_true",
help="Skip OpenCV dashboard; print forces and benchmark timings only",
)
args = parser.parse_args()
data = load_exported_data(args.mesh, args.intrinsics_path, args.pose)
print("[benchmark] Startup timing:")
data = benchmark_step("load_exported_data", load_exported_data, args.mesh, args.intrinsics_path, args.pose)
flip_v = not args.no_flip_v
flip_source_vertical = not args.no_flip_source_vertical
distortion = compute_uv_distortion_metrics(data, edge_band=args.edge_band)
distortion = benchmark_step(
"compute_uv_distortion_metrics",
compute_uv_distortion_metrics,
data,
edge_band=args.edge_band,
)
distortion_stats = distortion["stats"]
middle_stats = distortion["middle_stats"]
edge_stats = distortion["edge_stats"]
@@ -648,7 +428,7 @@ def main():
)
if args.calib:
source_K, dist_coeffs, img_size = load_calib_json(args.calib)
source_K, dist_coeffs, img_size = benchmark_step("load_calib_json", load_calib_json, args.calib)
source_width, source_height = img_size
print(f"Calibration loaded. Live image size: {(source_width, source_height)}")
else:
@@ -658,44 +438,30 @@ def main():
source_height = data["height"]
print("No calibration file provided; using exported intrinsics for projection.")
remap = None
if os.path.exists(args.map_cache):
cached = load_cached_remap(args.map_cache)
if cache_matches(
cached,
source_K=source_K,
source_width=source_width,
source_height=source_height,
out_width=args.out_width,
out_height=args.out_height,
flip_v=flip_v,
edge_band=args.edge_band,
edge_expand=edge_expand,
):
remap = cached
print(f"Loaded cached UV remap. Valid ratio: {remap_valid_ratio(remap):.3f}")
else:
print("Cached UV remap is stale. Regenerating.")
domain = benchmark_step(
"build_or_load_domain",
build_or_load_domain,
args.mesh,
args.intrinsics_path,
args.pose,
args.map_cache,
source_K,
source_width,
source_height,
args.out_width,
args.out_height,
flip_v=flip_v,
edge_band=args.edge_band,
edge_expand=edge_expand,
include_surface=True,
)
remap = domain.remap
remap.ensure_fast_maps()
edge_expand = float(remap.metadata.get("edge_expand", edge_expand))
if remap is None:
remap = build_uv_remap(
data,
source_K=source_K,
source_width=source_width,
source_height=source_height,
out_width=args.out_width,
out_height=args.out_height,
flip_v=flip_v,
edge_band=args.edge_band,
edge_expand=edge_expand,
)
valid_ratio = remap_valid_ratio(remap)
print(f"Generated UV remap valid ratio: {valid_ratio:.3f}")
if valid_ratio == 0.0:
raise RuntimeError("Generated UV remap contains no valid samples.")
np.savez_compressed(args.map_cache, **remap)
uv_debug = make_uv_debug_image(
uv_debug = benchmark_step(
"make_uv_debug_image",
make_uv_debug_image,
data,
out_width=args.out_width,
out_height=args.out_height,
@@ -703,7 +469,9 @@ def main():
edge_band=args.edge_band,
edge_expand=edge_expand,
)
uv_distortion = make_uv_distortion_image(
uv_distortion = benchmark_step(
"make_uv_distortion_image",
make_uv_distortion_image,
data,
distortion,
out_width=args.out_width,
@@ -712,11 +480,14 @@ def main():
edge_band=args.edge_band,
edge_expand=edge_expand,
)
uv_reference = np.hstack(
[
resize_to_height(uv_debug, args.out_height),
resize_to_height(uv_distortion, args.out_height),
]
uv_reference = benchmark_step(
"uv_reference_compose",
lambda: np.hstack(
[
resize_to_height(uv_debug, args.out_height),
resize_to_height(uv_distortion, args.out_height),
]
),
)
cap = open_source(args.video, args.camera_id)
@@ -727,80 +498,150 @@ def main():
cap.set(cv2.CAP_PROP_FRAME_WIDTH, source_width)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, source_height)
sensor = orisys.Sensor(
"external",
config_name=args.sdk_config,
verbose=args.verbose,
mirror=False,
sensor = benchmark_step(
"sensor_init",
lambda: orisys.Sensor(
"external",
config_name=args.sdk_config,
verbose=args.verbose,
mirror=False,
),
)
configure_masked_poisson(sensor, args)
print(
"Masked HHD solver: "
f"{sensor.masked_poisson_solver} "
f"(max_boundary={sensor.masked_poisson_max_boundary}, "
f"stride={sensor.masked_poisson_boundary_stride}, "
f"iters={sensor.masked_poisson_max_iters}, "
f"tol={sensor.masked_poisson_tol:.1e})"
)
print("Press 'q' to quit.")
frame_processor = UVUnrollFrameProcessor(
domain.remap,
camera_matrix=source_K if dist_coeffs is not None else None,
dist_coeffs=dist_coeffs,
source_width=source_width,
source_height=source_height,
flip_source_vertical=flip_source_vertical,
border_value=(0, 0, 0),
)
ok, warmup_frame = cap.read()
if ok:
warmup_prepared = frame_processor.prepare_source(warmup_frame)
warmup_unrolled = frame_processor.flatten_prepared(warmup_prepared)
sensor.get_img(
frame=warmup_unrolled,
already_processed=True,
valid_mask=frame_processor.valid_mask,
)
benchmark_step("hhd_warmup", sensor.warmup_deformation)
else:
print("[benchmark] hhd_warmup: skipped (no frame from source)")
if args.no_viz:
print("Running with --no-viz (dashboard disabled).")
else:
print("Press 'q' to quit.")
runtime_timer = StepTimer()
try:
while True:
runtime_timer.start("capture")
ok, frame = cap.read()
runtime_timer.stop("capture")
if not ok:
break
if frame.shape[1] != source_width or frame.shape[0] != source_height:
frame = cv2.resize(frame, (source_width, source_height), interpolation=cv2.INTER_LINEAR)
runtime_timer.start("uv_prepare")
frame_source = frame_processor.prepare_source(frame)
runtime_timer.stop("uv_prepare")
if dist_coeffs is not None:
frame_source = cv2.undistort(frame, source_K, dist_coeffs)
else:
frame_source = frame
frame_source = align_frame_to_uv_map(frame_source, flip_vertical=flip_source_vertical)
runtime_timer.start("uv_remap")
unrolled = frame_processor.flatten_prepared(frame_source)
runtime_timer.stop("uv_remap")
unrolled = flatten_frame(frame_source, remap, border_value=(0, 0, 0))
runtime_timer.start("get_img")
sensor.get_img(
frame=unrolled,
already_processed=True,
valid_mask=frame_processor.valid_mask,
)
runtime_timer.stop("get_img")
sensor.get_img(frame=unrolled, already_processed=True)
sensor.compute_deformation(check_motion=True, threshold=0)
run_deformation_timed(sensor, runtime_timer, check_motion=True, threshold=0)
runtime_timer.start("contact")
sensor.compute_contact()
runtime_timer.stop("contact")
fps, flow, vnormal, vshear, img, contour, centroid, depth_map, fnormal, fshearx, fsheary = sensor.read_info(
sensor.info.FPS,
sensor.info.VRAW,
sensor.info.VNORMAL,
sensor.info.VSHEAR,
sensor.info.IMG,
sensor.info.CONTOUR,
sensor.info.CENTROID,
sensor.info.DEPTH,
sensor.info.FNORMAL,
sensor.info.FSHEARX,
sensor.info.FSHEARY,
)
fps = sensor._fps
fnormal = sensor.fnormal
fshearx = sensor.fshearx
fsheary = sensor.fsheary
contact_ok = contact_inside_mask(sensor.centroid, sensor.valid_mask)
arrows = orisys.util.draw_arrows(
img,
flow,
threshold=2,
grid_spacing=20,
arrow_scale=1.0,
)
contact_img = orisys.util.draw_contact(img, contour, centroid)
if not args.no_viz:
runtime_timer.start("viz")
fps, flow, vnormal, vshear, img, contour, centroid, depth_map, fnormal, fshearx, fsheary = sensor.read_info(
sensor.info.FPS,
sensor.info.VRAW,
sensor.info.VNORMAL,
sensor.info.VSHEAR,
sensor.info.IMG,
sensor.info.CONTOUR,
sensor.info.CENTROID,
sensor.info.DEPTH,
sensor.info.FNORMAL,
sensor.info.FSHEARX,
sensor.info.FSHEARY,
)
depth_norm = depth_map
if depth_norm.size > 0 and depth_norm.max() > depth_norm.min():
depth_u8 = ((depth_norm - depth_norm.min()) / (depth_norm.max() - depth_norm.min()) * 255).astype(np.uint8)
else:
depth_u8 = np.zeros_like(depth_norm, dtype=np.uint8)
depth_colored = cv2.applyColorMap(depth_u8, cv2.COLORMAP_JET)
arrows = orisys.util.draw_arrows(
img,
flow,
threshold=1,
grid_spacing=15,
arrow_scale=1.0,
)
contact_img = orisys.util.draw_contact(img, contour, centroid)
dashboard = make_tiled_dashboard(
[
("Original", frame_source),
("UV Map | Distortion", uv_reference),
("UV Unrolled", unrolled),
("Flow", arrows),
],
tile_width=640,
tile_height=360,
)
cv2.imshow("UV Unroll Dashboard", dashboard)
depth_norm = depth_map
if depth_norm.size > 0 and depth_norm.max() > depth_norm.min():
depth_u8 = ((depth_norm - depth_norm.min()) / (depth_norm.max() - depth_norm.min()) * 255).astype(np.uint8)
else:
depth_u8 = np.zeros_like(depth_norm, dtype=np.uint8)
depth_colored = cv2.applyColorMap(depth_u8, cv2.COLORMAP_JET)
display_mask = sensor.valid_mask
arrows = apply_display_mask(arrows, display_mask, border_value=(0, 0, 0))
contact_img = apply_display_mask(contact_img, display_mask, border_value=(0, 0, 0))
depth_colored = apply_display_mask(depth_colored, display_mask, border_value=(0, 0, 0))
arrows = overlay_runtime_status(arrows, sensor, centroid=centroid)
contact_ok = contact_inside_mask(centroid, display_mask)
dashboard = make_tiled_dashboard(
[
("Original", frame_source),
("UV Map | Distortion", uv_reference),
("UV Unrolled", unrolled),
("Flow", arrows),
],
tile_width=640,
tile_height=360,
)
cv2.imshow("UV Unroll Dashboard", dashboard)
runtime_timer.stop("viz")
print(
f"FPS={fps:.2f}, 法向力={fnormal:.4f}, 切向力X={fshearx:.4f}, "
f"切向力Y={fsheary:.4f}, 展开图尺寸={unrolled.shape[:2]}",
f"切向力Y={fsheary:.4f}, 展开图尺寸={unrolled.shape[:2]}, "
f"solver={sensor.masked_poisson_solver}, "
f"mask={sensor.mask_stats['coverage']:.3f}, "
f"flow(valid/invalid)=({sensor.mask_stats['flow_mean_valid']:.3f}/"
f"{sensor.mask_stats['flow_mean_invalid']:.3f}), "
f"contact_in_mask={contact_ok}, "
f"{print_runtime_benchmark(runtime_timer)}",
end="\r",
)
+8 -6
View File
@@ -14,17 +14,19 @@ if SRC_DIR not in sys.path:
sys.path.insert(0, SRC_DIR)
import orisys
from orisys.unroll import UVDomain
def load_flatten_remap(path):
remap = np.load(path)
"""Load flatten cache; returns dict compatible with legacy viewer code."""
domain = UVDomain.load(path)
result = {
"map_x": remap["map_x"].astype(np.float32),
"map_y": remap["map_y"].astype(np.float32),
"valid_mask": remap["valid_mask"].astype(bool),
"map_x": domain.remap.map_x,
"map_y": domain.remap.map_y,
"valid_mask": domain.valid_mask,
}
if "flat_xyz" in remap:
result["flat_xyz"] = remap["flat_xyz"].astype(np.float32)
if domain.surface is not None:
result["flat_xyz"] = domain.surface.flat_xyz
return result
BIN
View File
Binary file not shown.