diff --git a/.gitignore b/.gitignore index e986e81..b8ce174 100644 --- a/.gitignore +++ b/.gitignore @@ -34,4 +34,5 @@ __pycache__/ /config/compound_eye.png # Legacy / scratch folders -/old_orisys/ \ No newline at end of file +/old_orisys/ +/tools/calibration/ \ No newline at end of file diff --git a/config/dd02-ov.json b/config/dd02-ov.json index 0252e68..23ea55e 100644 --- a/config/dd02-ov.json +++ b/config/dd02-ov.json @@ -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, diff --git a/demo_wuxi.py b/demo_wuxi.py deleted file mode 100644 index 04ad844..0000000 --- a/demo_wuxi.py +++ /dev/null @@ -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() diff --git a/docs/api.md b/docs/api.md index ed310cf..e07c23f 100644 --- a/docs/api.md +++ b/docs/api.md @@ -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 版本选择):** diff --git a/docs/get_started.md b/docs/get_started.md index fbc124e..cc11ec2 100644 --- a/docs/get_started.md +++ b/docs/get_started.md @@ -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),则说明安装成功。 # 使用方法 ## 运行示例程序 diff --git a/examples/finger_pipeline.py b/examples/finger_pipeline.py new file mode 100644 index 0000000..00f0cc4 --- /dev/null +++ b/examples/finger_pipeline.py @@ -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() diff --git a/setup.py b/setup.py index 91fca3c..f72959d 100644 --- a/setup.py +++ b/setup.py @@ -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"}, diff --git a/src/orisys/__init__.py b/src/orisys/__init__.py index d5e2e87..562c61d 100644 --- a/src/orisys/__init__.py +++ b/src/orisys/__init__.py @@ -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 diff --git a/src/orisys/frame_processors.py b/src/orisys/frame_processors.py index f3e646f..034169d 100644 --- a/src/orisys/frame_processors.py +++ b/src/orisys/frame_processors.py @@ -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}") diff --git a/src/orisys/get_depth.py b/src/orisys/get_depth.py index f8de03b..b596c26 100644 --- a/src/orisys/get_depth.py +++ b/src/orisys/get_depth.py @@ -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( diff --git a/src/orisys/hhd/__init__.py b/src/orisys/hhd/__init__.py index 5a2f11c..81abd69 100644 --- a/src/orisys/hhd/__init__.py +++ b/src/orisys/hhd/__init__.py @@ -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 diff --git a/src/orisys/hhd/api.py b/src/orisys/hhd/api.py index e03c2c4..42d1a0d 100644 --- a/src/orisys/hhd/api.py +++ b/src/orisys/hhd/api.py @@ -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) diff --git a/src/orisys/hhd/config.py b/src/orisys/hhd/config.py new file mode 100644 index 0000000..5b5f2ed --- /dev/null +++ b/src/orisys/hhd/config.py @@ -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), + ) diff --git a/src/orisys/hhd/domain.py b/src/orisys/hhd/domain.py new file mode 100644 index 0000000..ade666e --- /dev/null +++ b/src/orisys/hhd/domain.py @@ -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, + } diff --git a/src/orisys/hhd/grid.py b/src/orisys/hhd/grid.py index cdcf20b..dc768ba 100644 --- a/src/orisys/hhd/grid.py +++ b/src/orisys/hhd/grid.py @@ -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() diff --git a/src/orisys/hhd/poisson.py b/src/orisys/hhd/poisson.py index a1a6a17..984b91c 100644 --- a/src/orisys/hhd/poisson.py +++ b/src/orisys/hhd/poisson.py @@ -1,14 +1,12 @@ import numpy -import timeit -import sys -# import cv2 -from scipy import signal, spatial +from scipy import spatial -from ..backend_availability import ( - CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE, - cp, cp_signal, _torque_kernel, - cl, cla, cl_rfftn, cl_irfftn -) +from ..backend_availability import CUPY_AVAILABLE +from .config import MaskedPoissonConfig +from .domain import MaskedDomain, normalize_mask +from .solvers.capacitance import CapacitanceFFTSolver +from .solvers.masked_sparse import MaskedSparseSolver +from .solvers.rectangular import RectangularFFTSolver from .timing import Timer ''' @@ -51,904 +49,293 @@ from .timing import Timer # ------------------------------------------------------------------------------ ''' + class PoissonSolver(object): - # constructor def __init__(self, **kwargs): - args = kwargs.keys() - if ('grid' in args) == ('points' in args): - raise SyntaxError("Poisson Solver needs either shape of a regular grid or the points in an rectilinear/unstructured grid") + if ("grid" in args) == ("points" in args): + raise SyntaxError( + "Poisson Solver needs either shape of a regular grid or the points in an rectilinear/unstructured grid" + ) - if ('grid' in args) != ('spacings' in args): + if ("grid" in args) != ("spacings" in args): raise SyntaxError("Poisson Solver needs spacings for regular grid") - if ('points' in args) != ('pvolumes' in args): + if ("points" in args) != ("pvolumes" in args): raise SyntaxError("Poisson Solver needs point volumes for unstructured points") - # (default) data type self.dtype = numpy.float32 - if 'dtype' in args: - self.dtype = kwargs['dtype'] + if "dtype" in args: + self.dtype = kwargs["dtype"] - if 'grid' in args: - self.ptype = 'G' - self.stype = 'F' # default - self.gdims = kwargs['grid'] - self.gdx = kwargs['spacings'] + if "grid" in args: + self.ptype = "G" + self.stype = "F" + self.gdims = kwargs["grid"] + self.gdx = kwargs["spacings"] self.dim = len(self.gdims) if self.dim != len(self.gdx): raise ValueError("Dimensions of spacings should match that of the grid") - - elif 'points' in args: - self.ptype = 'P' - self.stype = 'S' # only option - self.points = kwargs['points'] - self.pvolumes = kwargs['pvolumes'] + elif "points" in args: + self.ptype = "P" + self.stype = "S" + self.points = kwargs["points"] + self.pvolumes = kwargs["pvolumes"] self.dim = len(self.points.shape) if self.pvolumes.shape[0] != self.points.shape[0]: raise ValueError("Number of pvolumes should match that of the points") - if 'solver' in args: - if (kwargs['solver'] != 'S' and kwargs['solver'] != 'F'): - raise SyntaxError('Solver can be only of type S(patial) or F(requency)') + if "solver" in args: + if kwargs["solver"] != "S" and kwargs["solver"] != "F": + raise SyntaxError("Solver can be only of type S(patial) or F(requency)") - if (kwargs['solver'] == 'F') and ('points' in args): - raise SyntaxError('Only Spatial solver supported for unstructured points') + if (kwargs["solver"] == "F") and ("points" in args): + raise SyntaxError("Only Spatial solver supported for unstructured points") - if (kwargs['solver'] == 'S') and ('grid' in args): - raise SyntaxError('Frequncy solver strongly suggested for structured grids') + if (kwargs["solver"] == "S") and ("grid" in args): + raise SyntaxError("Frequncy solver strongly suggested for structured grids") - self.stype = kwargs['solver'] - if 'poisson_backend' in args: - self.poisson_backend = kwargs['poisson_backend'] - if self.poisson_backend not in ('cpu', 'opencl', 'cupy'): + self.stype = kwargs["solver"] + + if "poisson_backend" in args: + self.poisson_backend = kwargs["poisson_backend"] + if self.poisson_backend not in ("cpu", "opencl", "cupy"): raise ValueError("poisson_backend must be one of: 'cpu', 'opencl', 'cupy'") + else: + self.poisson_backend = "cpu" - if (self.dim > 3): + self.masked_config = MaskedPoissonConfig.from_kwargs(kwargs) + self.masked_poisson_solver = self.masked_config.method + self.masked_poisson_max_boundary = self.masked_config.max_boundary + self.masked_poisson_boundary_stride = self.masked_config.boundary_stride + self.masked_poisson_max_iters = self.masked_config.max_iters + self.masked_poisson_tol = self.masked_config.tol + self.masked_poisson_reg = self.masked_config.reg + + if self.dim > 3: raise ValueError("Poisson Solver works only for 1, 2, or 3 dimensions") self.ready = False - if self.poisson_backend == 'cupy': - self._gpu_cache = None # Cache for GPU arrays - if self.poisson_backend == 'opencl': - self._opencl_cache = None # Cache for OpenCL FFT buffers - self._opencl_queue = None # OpenCL command queue - self._opencl_context = None # OpenCL context (set when queue is created) - self._opencl_crop_kernel = None # Cached OpenCL crop kernel - self._opencl_scatter_kernel = None # Cached scatter (reverse-crop) kernel - self._opencl_fused_mul_kernel = None # Cached fused complex multiply kernel - self._opencl_scatter_batch_kernel = None # Cached batch scatter kernel - self._opencl_crop_batch_kernel = None # Cached batch crop kernel - self._opencl_batch_mul_kernel = None # Cached batch complex multiply kernel - self._opencl_last_timing = None # Last per-stage OpenCL timing in ms - # Cache for OpenCV FFT (CPU backend) - self._opencv_cache = None # Cache for OpenCV FFT buffers - # if (self.ptype == 'G'): - # print( 'PoissonSolver:', self.dim, 'D grid =', self.gdims, 'with spacings', self.gdx) - # elif (self.ptype == 'P'): - # print( 'PoissonSolver', self.dim, 'D points =', self.points.shape) + self.G = None + self._masked_cache = None + self._masked_domain = None + self._opencv_cache = None - # -------------------------------------------------------------------------- - # create the Green's function in appropriate representation - # -------------------------------------------------------------------------- + self._rect_solver = None + self._sparse_solver = MaskedSparseSolver() + self._cap_solver = None - # generate a grid -(N-1) to (N-1) to define distance function - # where N is the size of the original grid - def _create_radialGrid(self, verbose): + if self.ptype == "G": + self._rect_solver = RectangularFFTSolver( + grid=self.gdims, + spacings=self.gdx, + poisson_backend=self.poisson_backend, + dtype=self.dtype, + stype=self.stype, + ) + self._cap_solver = CapacitanceFFTSolver(self.masked_config, self._rect_solver) - 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] - - # half of the average grid spacing - 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] - - # half of the average grid spacing - 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 - - # -------------------------------------------------------------------------- - # scale a function by point volumes for correct integration - def _scale(self, func): - - if self.ptype == 'P': - return numpy.multiply(func, self.pvolumes) - - elif self.ptype == 'G': - sfunc = func * numpy.prod(self.gdx) - - ''' - # volume of the boundary elements is only half the voxel - if self.dim == 1: - sfunc[ 0] *= 0.5 - sfunc[-1] *= 0.5 - elif self.dim == 2: - sfunc[ 0, :] *= 0.5 - sfunc[-1, :] *= 0.5 - sfunc[ :, 0] *= 0.5 - sfunc[ :,-1] *= 0.5 - elif self.dim == 3: - sfunc[ 0, :, :] *= 0.5 - sfunc[-1, :, :] *= 0.5 - sfunc[ :, 0, :] *= 0.5 - sfunc[ :,-1, :] *= 0.5 - sfunc[ :, :, 0] *= 0.5 - sfunc[ :, :,-1] *= 0.5 - ''' - return sfunc + 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_last_timing = None def prepare(self, verbose=False): + """Initialize Green's function / distance matrix for the domain. - # ---------------------------------------------------------------------- + verbose: print timing and setup info + returns: None (sets self.ready) + """ gtimer = Timer() if verbose: - print( '\nInitializing Poisson solver, type =', self.stype) + print("\nInitializing Poisson solver, type =", self.stype) - # ---------------------------------------------------------------------- - # generate a grid to store pairwise distances to store Green's function - self.G = None # definition depends upon solver type + if self.ptype == "G" and self.stype == "F": + self._rect_solver.prepare(verbose) + self.G = self._rect_solver.G + self.ready = self._rect_solver.ready + elif self.stype == "F" and self.ptype == "P": + raise SyntaxError("Frequncy solver not supported for unstructured points") + elif self.stype == "S" and self.ptype == "G": + raise SyntaxError("Spatial solver not suggested for structured grids") + elif self.stype == "S" and self.ptype == "P": + if self.dim == 1 and len(self.points.shape) == 1: + self.points = self.points[:, numpy.newaxis] - # ---------------------------------------------------------------------- - # for Fourier solution, - if self.stype == 'F': + ltimer = Timer() + self.G = spatial.distance.cdist(self.points, self.points) + zval = numpy.power(10.0, -10.0) + self.G += zval * numpy.identity(self.G.shape[0]) - if self.ptype == 'G': - self.G = self._create_radialGrid(verbose) + if verbose: + print(" - created pairwise distance matrix:", self.G.shape, self.G.min(), self.G.max(), end="") + ltimer.end() - elif self.ptype == 'P': - raise SyntaxError('Frequncy solver not supported for unstructured points') - #self.G = create_distArray() + 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) - # ---------------------------------------------------------------------- - # for spatial solution, - # use the pairwise distance matrix, N x N - # where N is the number of points - elif self.stype == 'S': - - if self.ptype == 'G': - raise SyntaxError('Spatial solver not suggested for structured grids') - #self.G = self.create_points(verbose) - - elif self.ptype == 'P': - # for 1D points, add an extra axis - if (self.dim == 1) and (len(self.points.shape) == 1): - self.points = self.points[:,numpy.newaxis] - - # create pairwise distance matrix now - ltimer = Timer() - self.G = spatial.distance.cdist(self.points, self.points) - - # TODO: use a good approximation here - zval = numpy.power(10.0, -10.0) - self.G += zval * numpy.identity(self.G.shape[0]) - - if verbose: - print(' - created pairwise distance matrix:', self.G.shape, self.G.min(), self.G.max(), end = "") - ltimer.end() - - # ---------------------------------------------------------------------- - # compute the Green's function - 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) + self.ready = True if verbose: - print( self.G.shape, self.G.min(), self.G.max(), end = "") - ltimer.end() - - # ---------------------------------------------------------------------- - self.ready = True - - if verbose: - print( 'Poisson solver initialized', end ='') + print("Poisson solver initialized", end="") gtimer.end() - # -------------------------------------------------------------------------- - # compute the integral solution - # -------------------------------------------------------------------------- - def solve(self, f, verbose=False): + def _normalize_mask(self, mask, shape): + return normalize_mask(mask, shape) + def _masked_cache_matches(self, mask): + return self._masked_domain is not None and self._masked_domain.matches(mask) + + def _get_or_build_domain(self, mask): + if self._masked_cache_matches(mask): + return self._masked_domain + + if self.dim != 2: + raise NotImplementedError("Masked Poisson solve is currently implemented for 2D grids only") + + domain = MaskedDomain.build(mask, self.gdx) + capfft = None + if self.masked_config.method == "capfft": + capfft = self._cap_solver.prepare(domain) + + self._masked_domain = domain + self._masked_cache = domain.to_cache_dict(capfft=capfft) + return domain + + def _build_masked_cache(self, mask): + """Legacy wrapper; returns dict cache for tests.""" + self._get_or_build_domain(mask) + return self._masked_cache + + def _compute_masked_residual(self, field, rhs_active, cache): + """Legacy wrapper for capacitance/residual tests.""" + domain = self._masked_domain + if domain is None or not domain.matches(cache["mask"]): + domain = MaskedDomain.build(cache["mask"], self.gdx) + return domain.residual(field, rhs_active) + + def _solve_masked(self, f, mask, verbose=False): + mask_arr = self._normalize_mask(mask, f.shape) + if mask_arr is None: + return None + if not numpy.any(mask_arr): + return numpy.zeros_like(f, dtype=numpy.float32) + + domain = self._get_or_build_domain(mask_arr) + if self.masked_config.method == "capfft": + return self._cap_solver.solve(f, domain) + return self._sparse_solver.solve(f, domain) + + def solve(self, f, verbose=False, mask=None): + """Solve Poisson equation on the prepared grid. + + f: (H, W) scalar source field (divergence or curl weight) + mask: optional (H, W) bool; True = active UV/domain pixels + returns: (H, W) scalar potential; zero outside mask when mask is set + """ if not self.ready: self.prepare(verbose) fshape = f.shape + mask_arr = self._normalize_mask(mask, fshape) gtimer = Timer() if verbose: - print( '\nSolving Poisson Eq.') + print("\nSolving Poisson Eq.") - # ---------------------------------------------------------------------- - # regular grid - if self.ptype == 'G': - - #if (self.gdims - fshape).any(): + if self.ptype == "G": if self.gdims != fshape: - print( self.gdims, fshape) + print(self.gdims, fshape) raise ValueError("Shape of function should match shape of grid") - # convolution in frequency domain using OpenCV FFT (2D) or numpy FFT (3D) - if self.stype == 'F': - if self.dim == 2: - # # Use OpenCV DFT for 2D (optimized with caching) - # # Convert to float32 for OpenCV (use ascontiguousarray for better performance) - # f_32 = numpy.ascontiguousarray(f, dtype=numpy.float32) - - # # For 'same' mode convolution, pad to (2*N-1) size - # full_shape = tuple(2 * s - 1 for s in f_32.shape) - - # # Initialize or validate cache - # cache_valid = ( - # self._opencv_cache is not None and - # self._opencv_cache['fshape'] == fshape and - # self._opencv_cache['full_shape'] == full_shape - # ) - - # if not cache_valid: - # # Cache G_fft (G doesn't change, so compute once) - # G_32 = numpy.ascontiguousarray(self.G, dtype=numpy.float32) - # G_fft = cv2.dft(G_32, flags=cv2.DFT_COMPLEX_OUTPUT) - - # # Pre-compute crop slices for 'same' mode extraction - # # For circular convolution at size (2*H-1, 2*W-1), the center corresponds to - # # position [H-1, W-1], which matches scipy's fftconvolve mode='same' extraction - # starts = [fshape[i] - 1 for i in range(len(fshape))] - # slices = tuple(slice(starts[i], starts[i] + fshape[i]) for i in range(len(fshape))) - - # # Pre-allocate reusable buffers - # self._opencv_cache = { - # 'fshape': fshape, - # 'full_shape': full_shape, - # 'G_fft': G_fft, # Cached: G doesn't change, so compute FFT once - # 'slices': slices, # Cached: crop slices pre-computed - # 'gdx_prod': numpy.prod(self.gdx), # Cached: grid spacing product - # # Reusable buffers (avoid reallocation on each call) - # 'f_pad': numpy.zeros(full_shape, dtype=numpy.float32), # Reused for padding - # 'conv_fft': numpy.empty((full_shape[0], full_shape[1], 2), dtype=numpy.float32), # Output for mulSpectrums - # } - - # cache = self._opencv_cache - - # # Zero-pad f into pre-allocated buffer (faster than creating new array) - # f_pad = cache['f_pad'] - # f_pad.fill(0.0) # Clear previous data efficiently - # f_slices = tuple(slice(0, s) for s in f_32.shape) - # f_pad[f_slices] = f_32 - - # # Compute FFT of f (cv2.dft returns new array, but we can reuse for mulSpectrums) - # f_fft = cv2.dft(f_pad, flags=cv2.DFT_COMPLEX_OUTPUT) - - # # Multiply in frequency domain using OpenCV's optimized mulSpectrums - # # This is faster than manual complex multiplication and supports output buffer - # # Signature: mulSpectrums(a, b, flags, c=None, conjB=False) - # cv2.mulSpectrums(f_fft, cache['G_fft'], cv2.DFT_COMPLEX_OUTPUT, - # cache['conv_fft'], False) - - # # Compute inverse FFT (cv2.idft returns new array) - # conv_full = cv2.idft(cache['conv_fft'], flags=cv2.DFT_SCALE | cv2.DFT_REAL_OUTPUT) - - # # Extract 'same' mode result (use cached slices, scale in-place) - # p = conv_full[cache['slices']].copy() - # numpy.multiply(p, cache['gdx_prod'], p) - # from scipy import signal - p = signal.fftconvolve(f, self.G, mode='same') - # Scale by grid spacing - numpy.multiply(p, numpy.prod(self.gdx), p) - else: - # For 3D, use numpy FFT (OpenCV DFT only supports 2D) - # from scipy import signal - p = signal.fftconvolve(f, self.G, mode='same') - # Scale by grid spacing - numpy.multiply(p, numpy.prod(self.gdx), p) + if mask_arr is not None: + p = self._solve_masked(f, mask_arr, verbose) + if verbose: + print("Poisson solver finished") + gtimer.end() + return p - # convolution in spatial domain - elif self.stype == 'S': - raise SyntaxError('Spatial solver not suggested for structured grids') - #p = (self.G).dot(f.reshape(-1)).reshape(fshape) - - # ---------------------------------------------------------------------- - # unstructured points - elif self.ptype == 'P': + if self.stype == "F": + p = self._rect_solver.solve(f, verbose=False) + else: + raise SyntaxError("Spatial solver not suggested for structured grids") + elif self.ptype == "P": if (self.points.shape[0] != fshape[0]) or (1 != len(fshape)): raise ValueError("Shape of function should be (N,1) with N as the number of points") - # convolution in spatial domain - # Dec 13, 2017 -- pre scaling works for unstructured grids! - if self.stype == 'S': + if self.stype == "S": p = numpy.multiply(f, self.pvolumes) numpy.dot(self.G, p, p) - #numpy.multiply(p, self.pvolumes, p) + else: + raise SyntaxError("Frequncy solver not supported for unstructured points") - elif self.stype == 'F': - raise SyntaxError('Frequncy solver not supported for unstructured points') - - # ---------------------------------------------------------------------- if verbose: - print( 'Poisson solver finished') + print("Poisson solver finished") gtimer.end() return p - # -------------------------------------------------------------------------- - # GPU-accelerated solve using CuPy - # -------------------------------------------------------------------------- def solve_gpu(self, f, verbose=False, return_gpu=False): - """ - GPU-accelerated Poisson solver using CuPy. - Falls back to CPU if CuPy is not available. - - Args: - f: Input field (numpy array) - verbose: Print timing information - - Returns: - p: Solution (numpy array on CPU) + """GPU FFT Poisson solve (full rectangular grid only). + + f: (H, W) scalar source field + return_gpu: return CuPy array if True + 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') - - # ---------------------------------------------------------------------- - # regular grid - if self.ptype == 'G': - if self.gdims != fshape: - print(self.gdims, fshape) - raise ValueError("Shape of function should match shape of grid") - - # convolution in frequency domain using GPU - if self.stype == 'F': - # Initialize GPU cache if needed - 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) - } - - # Transfer input to GPU - f_gpu = cp.asarray(f, dtype=cp.float32) - - # Perform FFT convolution on GPU - # CuPy's fftconvolve is GPU-accelerated - p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same') - - # Scale by grid spacing (on GPU) - p_gpu *= self._gpu_cache['gdx_prod'] - if verbose: - print('Poisson solver (GPU) finished') - gtimer.end() - - # Transfer result back to CPU - if return_gpu: - return p_gpu # Return CuPy array - return cp.asnumpy(p_gpu) - # p = cp.asnumpy(p_gpu) - - elif self.stype == 'S': - raise SyntaxError('Spatial solver not suggested for structured grids') - + if self.ptype != "G" or self.stype != "F": + raise SyntaxError("GPU solver applies to structured FFT grids only") + + result = self._rect_solver.solve_gpu(f, verbose=verbose, return_gpu=return_gpu) + self._gpu_cache = self._rect_solver._gpu_cache + return result + def _solve_gpu_internal(self, f_gpu): - """Internal GPU solve that accepts and returns GPU arrays.""" - """ - GPU-accelerated Poisson solver using CuPy. - Falls back to CPU if CuPy is not available. - - Args: - f_gpu: Input field (numpy array) - - Returns: - p_gpu: Solution (CuPy array) - """ - - # ---------------------------------------------------------------------- - # regular grid - if self.ptype == 'G': - - # convolution in frequency domain using GPU - if self.stype == 'F': - # Initialize GPU cache if needed - if self._gpu_cache is None: - self._gpu_cache = { - 'G_gpu': cp.asarray(self.G, dtype=cp.float32), - 'gdx_prod': numpy.prod(self.gdx) - } - - # Perform FFT convolution on GPU - # CuPy's fftconvolve is GPU-accelerated - p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same') - - # Scale by grid spacing (on GPU) - p_gpu *= self._gpu_cache['gdx_prod'] - - return p_gpu # Return CuPy array + """Internal GPU solve; input and output are CuPy arrays.""" + if not self.ready: + self.prepare(False) + return self._rect_solver._solve_gpu_internal(f_gpu) def clear_gpu_cache(self): - """Clear GPU cache to free GPU memory""" - if self._gpu_cache is not None: - self._gpu_cache = None - if CUPY_AVAILABLE: - cp.get_default_memory_pool().free_all_blocks() - - # ---------------------------------------------------------------------- - # unstructured points - not optimized for GPU - # elif self.ptype == 'P': - # if verbose: - # print(' - GPU solver not optimized for unstructured points, using CPU') - # return self.solve(f, verbose) - - # ---------------------------------------------------------------------- - - # return p - - def _get_opencl_queue(self): - """Create and cache an OpenCL queue on first available GPU.""" - 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_crop_kernel(self, queue): - """Create and cache OpenCL kernel for cropping convolution result.""" - if hasattr(self, '_opencl_crop_kernel') and self._opencl_crop_kernel is not None: - return self._opencl_crop_kernel - - # Simple 2D copy kernel - kernel_code = """ - __kernel void crop_and_scale( - __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 - ) { - int gid_x = get_global_id(0); - int gid_y = get_global_id(1); - - if (gid_x < dst_width && gid_y < dst_height) { - int src_idx = (start_y + gid_y) * src_width + (start_x + gid_x); - int dst_idx = 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_kernel = program.crop_and_scale - return self._opencl_crop_kernel - - def _get_opencl_scatter_kernel(self, queue): - """Kernel that copies unpadded input into a zero-padded device buffer (reverse of crop).""" - if self._opencl_scatter_kernel is not None: - return self._opencl_scatter_kernel - - kernel_code = """ - __kernel void scatter_to_padded( - __global const float* src, - __global float* dst, - const int src_width, - const int src_height, - const int dst_width - ) { - int gid_x = get_global_id(0); - int gid_y = get_global_id(1); - if (gid_x < src_width && gid_y < src_height) { - dst[gid_y * dst_width + gid_x] = src[gid_y * src_width + gid_x]; - } - } - """ - program = cl.Program(self._opencl_context, kernel_code).build() - self._opencl_scatter_kernel = program.scatter_to_padded - return self._opencl_scatter_kernel - - def _get_opencl_fused_mul_kernel(self, queue): - """Kernel that multiplies two complex spectra by one Green's spectrum in a single pass.""" - if self._opencl_fused_mul_kernel is not None: - return self._opencl_fused_mul_kernel - - kernel_code = """ - __kernel void fused_complex_mul( - __global float2* fft1, - __global float2* fft2, - __global const float2* g, - const int n - ) { - int i = get_global_id(0); - if (i < n) { - float2 gv = g[i]; - float2 a = fft1[i]; - float2 b = fft2[i]; - fft1[i] = (float2)(a.x * gv.x - a.y * gv.y, - a.x * gv.y + a.y * gv.x); - fft2[i] = (float2)(b.x * gv.x - b.y * gv.y, - b.x * gv.y + b.y * gv.x); - } - } - """ - program = cl.Program(self._opencl_context, kernel_code).build() - self._opencl_fused_mul_kernel = program.fused_complex_mul - return self._opencl_fused_mul_kernel - - def _get_opencl_scatter_batch_kernel(self, queue): - """Batched scatter: copies two unpadded inputs into a stacked zero-padded device buffer.""" - 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): - """Batched crop+scale: extracts 'same' region from two stacked full-conv results.""" - 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): - """Batched complex multiply: multiplies each batch plane by the same Green's spectrum. - Uses multiple elements per work item, __constant for broadcast-read Green's spectrum, - and fma for better throughput/precision.""" - 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): - """Return an allocator that creates buffers with ALLOC_HOST_PTR for pinned host access (reduces H2D/D2H transfer cost).""" - 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): - """Crop full convolution result to scipy 'same' output 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] - + """Release CuPy FFT cache buffers.""" + if self._rect_solver is not None: + self._rect_solver.clear_gpu_cache() + self._gpu_cache = None def solve_opencl_batch2(self, f1, f2, verbose=False, profile=False): - """ - Batched OpenCL Poisson solver for two 2D fields. - Optimized: scatter H2D (~4x smaller transfer), fused multiply - (reads g_fft once), async D2H (one sync instead of two). - When profile=True, records and optionally prints per-stage timings. - """ - 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 self.dim != 2: - # return self.solve_opencl(f1, verbose, profile), self.solve_opencl(f2, verbose, profile) + """Batched OpenCL FFT Poisson solve for two 2D fields. + f1, f2: (H, W) scalar source fields + returns: (p1, p2) each (H, W) scalar potential + """ if not self.ready: self.prepare(verbose) - fshape = f1.shape - if fshape != f2.shape: - raise ValueError("solve_opencl_batch2 requires inputs with identical shapes") - - if self.ptype != 'G' or self.stype != 'F': - return self.solve_opencl(f1, verbose, profile), self.solve_opencl(f2, verbose, profile) - - 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() - fshape = f1.shape - 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() - - # For circular conv at size (2s-1) with G of shape (2s-1) and f of shape s, - # alias-free positions are [s-1, 2s-2], matching scipy 'same' extraction. - 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] - - # Pinned allocator for staging (H2D) and output (D2H) buffers - 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'] - - # H2D: compact transfer of unpadded data, then scatter to padded buffer on device - 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 - - # Single batched 2D FFT for both fields (ndim=2 treats dim-0 as batch) - 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 - - # Batched multiply: each batch plane * same Green's spectrum (4 elements/work item, explicit local size) - 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 - - # Single batched 2D iFFT for both fields - t0 = timeit.default_timer() - conv_batch = cl_irfftn(fft_batch, c['conv_batch_dev'], ndim=2) - if profile: - queue.finish() - timing['ifft_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # Batch crop + scale on device - 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, - conv_batch.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 - - # D2H: map pinned buffer, copy into pre-allocated host arrays (zero allocation) - 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}" - ) + p1, p2 = self._rect_solver.solve_opencl_batch2(f1, f2, verbose=verbose, profile=profile) + self._opencl_cache = self._rect_solver._opencl_cache + self._opencl_queue = self._rect_solver._opencl_queue + self._opencl_context = self._rect_solver._opencl_context + self._opencl_last_timing = self._rect_solver._opencl_last_timing return p1, p2 - \ No newline at end of file + + # Note: api.py may call solve_opencl() for single-field OpenCL path; + # only solve_opencl_batch2 is implemented. Phase 2 should add or remove that call. diff --git a/src/orisys/hhd/solvers/__init__.py b/src/orisys/hhd/solvers/__init__.py new file mode 100644 index 0000000..565d16f --- /dev/null +++ b/src/orisys/hhd/solvers/__init__.py @@ -0,0 +1,5 @@ +from .capacitance import CapacitanceFFTSolver +from .masked_sparse import MaskedSparseSolver +from .rectangular import RectangularFFTSolver + +__all__ = ["RectangularFFTSolver", "MaskedSparseSolver", "CapacitanceFFTSolver"] diff --git a/src/orisys/hhd/solvers/base.py b/src/orisys/hhd/solvers/base.py new file mode 100644 index 0000000..0c3f444 --- /dev/null +++ b/src/orisys/hhd/solvers/base.py @@ -0,0 +1,3 @@ +"""Shared solver protocol and helpers.""" + +__all__ = [] diff --git a/src/orisys/hhd/solvers/capacitance.py b/src/orisys/hhd/solvers/capacitance.py new file mode 100644 index 0000000..7452ba5 --- /dev/null +++ b/src/orisys/hhd/solvers/capacitance.py @@ -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) diff --git a/src/orisys/hhd/solvers/masked_sparse.py b/src/orisys/hhd/solvers/masked_sparse.py new file mode 100644 index 0000000..634691a --- /dev/null +++ b/src/orisys/hhd/solvers/masked_sparse.py @@ -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) diff --git a/src/orisys/hhd/solvers/rectangular.py b/src/orisys/hhd/solvers/rectangular.py new file mode 100644 index 0000000..6782b70 --- /dev/null +++ b/src/orisys/hhd/solvers/rectangular.py @@ -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) diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index 9020696..89cdddc 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -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 diff --git a/src/orisys/tracker.py b/src/orisys/tracker.py index fb8cc9b..68805c6 100644 --- a/src/orisys/tracker.py +++ b/src/orisys/tracker.py @@ -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] diff --git a/src/orisys/unroll/__init__.py b/src/orisys/unroll/__init__.py new file mode 100644 index 0000000..dba757e --- /dev/null +++ b/src/orisys/unroll/__init__.py @@ -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", +] diff --git a/src/orisys/unroll/domain.py b/src/orisys/unroll/domain.py new file mode 100644 index 0000000..97c6a28 --- /dev/null +++ b/src/orisys/unroll/domain.py @@ -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 diff --git a/src/orisys/unroll/mesh_uv.py b/src/orisys/unroll/mesh_uv.py new file mode 100644 index 0000000..3b31c52 --- /dev/null +++ b/src/orisys/unroll/mesh_uv.py @@ -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() diff --git a/src/orisys/unroll/precompute.py b/src/orisys/unroll/precompute.py new file mode 100644 index 0000000..b0403d6 --- /dev/null +++ b/src/orisys/unroll/precompute.py @@ -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 diff --git a/src/orisys/unroll/remap.py b/src/orisys/unroll/remap.py new file mode 100644 index 0000000..27aad68 --- /dev/null +++ b/src/orisys/unroll/remap.py @@ -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 diff --git a/src/orisys/unroll/surface.py b/src/orisys/unroll/surface.py new file mode 100644 index 0000000..947033d --- /dev/null +++ b/src/orisys/unroll/surface.py @@ -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)] diff --git a/tests/test_masked_hhd.py b/tests/test_masked_hhd.py new file mode 100644 index 0000000..7340d67 --- /dev/null +++ b/tests/test_masked_hhd.py @@ -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] diff --git a/tests/test_masked_poisson_capfft.py b/tests/test_masked_poisson_capfft.py new file mode 100644 index 0000000..11a3976 --- /dev/null +++ b/tests/test_masked_poisson_capfft.py @@ -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 diff --git a/tests/test_uv_remap.py b/tests/test_uv_remap.py new file mode 100644 index 0000000..19f2f57 --- /dev/null +++ b/tests/test_uv_remap.py @@ -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)) diff --git a/tests/test_uv_unroll_processor.py b/tests/test_uv_unroll_processor.py new file mode 100644 index 0000000..ba498e7 --- /dev/null +++ b/tests/test_uv_unroll_processor.py @@ -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) diff --git a/tools/3d_patch_generation_blender/unroll_realtime.py b/tools/3d_patch_generation_blender/unroll_realtime.py index b580c25..5193a9f 100644 --- a/tools/3d_patch_generation_blender/unroll_realtime.py +++ b/tools/3d_patch_generation_blender/unroll_realtime.py @@ -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", ) diff --git a/tools/dd02_3dviewer/sdk_mapto3d_flatten.py b/tools/dd02_3dviewer/sdk_mapto3d_flatten.py index 822b739..919c712 100644 --- a/tools/dd02_3dviewer/sdk_mapto3d_flatten.py +++ b/tools/dd02_3dviewer/sdk_mapto3d_flatten.py @@ -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 diff --git a/unroll_map_uv.npz b/unroll_map_uv.npz index c16da64..a64028d 100644 Binary files a/unroll_map_uv.npz and b/unroll_map_uv.npz differ