orisys-tech/orisys-sdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'stable'

Browse Source
This commit is contained in:
ykcheungab
2026-05-14 21:02:23 +08:00
parent 436d8478cb a8cb4203fd
commit 7224fea411
82 changed files with 5812 additions and 5386 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+30 -10
View File
@@ -1,17 +1,37 @@
# Python bytecode / caches
__pycache__/
*.py[cod]
.pytest_cache/
# Packaging / build outputs
/build/ /build/
/dist/ /dist/
/dist_obfuscated/ /dist_obfuscated/
/release/
__pycache__/
orisys.egg-info/
/.cursor/
/dist_pyd/ /dist_pyd/
/dist_single_pyd/ /dist_single_pyd/
/old_orisys/
/dist_nuitka/ /dist_nuitka/
/build_nuitka /build_nuitka/
/data/ /release/
/SDK/
*.egg-info/
# Local tools / editor state
/.cursor/
/.vscode/
# Local data / models
/models/ /models/
/doc/.quarto/ /data/*
*.npy !/data/sample/
*.png !/data/sample/**
# Generated docs / tooling outputs
/docs/.quarto/
/docs/exported_data/
/tools/3d_patch_generation_blender/exported_data/
# Local or generated assets kept out of git
/config/compound_eye.png
# Legacy / scratch folders
/old_orisys/
API文档-0.3.0.pdf
BIN
View File
Binary file not shown.
RAFT_test.py
-156
View File
@@ -1,156 +0,0 @@
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import torch
import torchvision.transforms.functional as F
import torchvision.transforms as T
from torchvision.models.optical_flow import raft_large, raft_small
import numpy as np
import time
import orisys
def preprocess(batch, resize=None):
"""
Preprocess image batch for RAFT model.
Args:
batch: Tensor of shape (N, C, H, W)
resize: Optional tuple (height, width) to resize. If None, keeps original size.
Returns:
Preprocessed tensor
"""
transform_list = [
T.ConvertImageDtype(torch.float32),
T.Normalize(mean=0.5, std=0.5), # map [0, 1] into [-1, 1]
]
if resize is not None:
transform_list.append(T.Resize(size=resize))
transforms = T.Compose(transform_list)
batch = transforms(batch)
return batch
def setup_simple_raft():
"""
Setup a simpler version using available models
"""
model = raft_small(pretrained=True, progress=False).to('cuda')
model = model.eval()
# Alternative: Use LiteFlowNet (lighter alternative)
print("RAFT loaded successfully")
return model
def track_raft(prev_frame, frame, model, device):
gray = frame.copy()
start_time = time.time()
prev_tensor = torch.from_numpy(prev_frame.transpose(2, 0, 1).astype(np.float32) / 255.0)
curr_tensor = torch.from_numpy(gray.transpose(2, 0, 1).astype(np.float32) / 255.0)
prev_tensor = preprocess(prev_tensor.unsqueeze(0), resize=None).to(device)
curr_tensor = preprocess(curr_tensor.unsqueeze(0), resize=None).to(device)
t0 = time.perf_counter()
with torch.no_grad():
flow = model(prev_tensor, curr_tensor)
predicted_flows = flow[-1]
t1 = time.perf_counter()
print(f"RAFT time: {t1 - t0}")
flow_np = predicted_flows[0].cpu().numpy() # or .squeeze(0)
# Transpose to (H, W, 2) format expected by draw_arrows
flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2)
# torch.onnx.export(
# model, # model to export
# (prev_tensor, curr_tensor), # inputs of the model,
# "my_model.onnx", # filename of the ONNX model
# input_names=["prev_tensor", "curr_tensor"], # Rename inputs for the ONNX model
# dynamo=True # True or False to select the exporter to use
# )
# print("ONNX model exported successfully")
return flow_np
def main_simple():
# Initialize camera
# sensor = orisys.Sensor('s1.mp4', config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
model = setup_simple_raft()
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
cap = cv2.VideoCapture("s1.mp4")
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
prev_gray = None
fps_buffer = []
while True:
ret, frame = cap.read()
if not ret:
break
dx = 300
dy = 0
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame1 = frame[420:660, 840:1080, :]
# frame2 = frame[420+dy:660+dy, 840+dx:1080+dx, :]
# frame3 = frame[420-dy:660-dy, 840-dx:1080-dx, :]
# frame = np.concatenate([frame3, frame1, frame2], axis=1)
frame = frame1
# cv2.imshow("frame", frame3)
# cv2.waitKey(0)
if prev_gray is None:
prev_gray = frame.copy()
gray = frame
# Compute optical flow
start_time = time.time()
flow_np = track_raft(prev_gray, gray, model, device)
fps = 1.0 / (time.time() - start_time)
fps_buffer.append(fps)
if len(fps_buffer) > 30:
fps_buffer.pop(0)
avg_fps = np.mean(fps_buffer)
print(f"avg_fps: {avg_fps}")
# # Remove batch dimension: (1, 2, H, W) -> (2, H, W)
# flow_np = predicted_flows[0].cpu().numpy() # or .squeeze(0)
# # Transpose to (H, W, 2) format expected by draw_arrows
# flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2)
# Flow is already in pixel coordinates, no denormalization needed
# Keep as float32 for draw_arrows (it can handle float values)
print(f"Flow numpy shape: {flow_np.shape}, dtype: {flow_np.dtype}")
print(f"Flow range: x=[{flow_np[:,:,0].min():.2f}, {flow_np[:,:,0].max():.2f}], "
f"y=[{flow_np[:,:,1].min():.2f}, {flow_np[:,:,1].max():.2f}]")
arrows = orisys.util.draw_arrows(
gray,
flow_np,
threshold=2,
grid_spacing=40,
arrow_scale=1.0
)
arrows = cv2.resize(arrows, (400, 400))
cv2.imshow('Optical Flow', arrows)
# prev_gray = gray.copy()
#步骤 3.5: 键盘控制
key = cv2.waitKey(1) & 0xFF
if key == ord("q"):
print("\n退出程序...")
break
if key == ord("r"):
# 重置光流追踪器（当追踪漂移时使用）
prev_gray = None
print("传感器已重置")
if __name__ == "__main__":
main_simple()
README.md
+85
View File
@@ -0,0 +1,85 @@
# Orisys SDK
Orisys SDK is a Python toolkit for optical tactile sensing. It provides image capture,
stitching, dense optical flow, Helmholtz-Hodge decomposition, contact/depth estimation,
and several runnable examples for camera and video-based workflows.
## Repository Layout
- `src/orisys/`: installable SDK package
- `config/`: default and lab/runtime JSON configurations
- `examples/`: user-facing example scripts
- `scripts/`: maintenance, calibration, packaging, and demo scripts
- `tools/`: non-package tooling such as patch-generation utilities
- `tests/`: automated pytest suite
- `docs/`: usage and API documentation
## Quick Start
### 1. Create an environment
```bash
conda create -n orisys python=3.10
conda activate orisys
```
### 2. Install the SDK in editable mode
```bash
pip install -e .
```
Optional GPU extras:
```bash
pip install -e .[cuda12]
```
or
```bash
pip install -e .[cuda13]
```
### 3. Run tests
```bash
pytest
```
### 4. Run an example
```bash
python examples/basic_pipeline.py -v 0 -c ./config/ddjx01.json
```
## Configuration
The canonical configuration files now live in the repo-level `config/` folder.
The package loader `orisys.configs` resolves configs from:
1. an explicit JSON file path
2. a config file stem under `config/`
3. a backward-compatible JSON `"name"` match, such as `"prototype"`
Examples:
```python
import orisys
sensor = orisys.Sensor(0, config_name="ddjx01")
cfg = orisys.export_config("ddjx01", "./config/exported_config.json")
```
## Examples and Tools
- Example entry points: `examples/README.md`
- User documentation: `docs/get_started.md`
- Packaging helper: `scripts/build_nuitka.py`
- Specialized utilities: `tools/`
## Notes for Collaborators
- Generated build outputs such as `build/`, `dist/`, and `SDK/` are intentionally ignored.
- Large or generated tooling outputs should stay under ignored output folders rather than be committed.
- If you add new images or arrays that should be versioned, review `.gitignore` first.
Readme.docx
BIN
View File
Binary file not shown.
build_nuitka.py
-155
View File
@@ -1,155 +0,0 @@
import os
import sys
import shutil
import subprocess
from pathlib import Path
from importlib.machinery import EXTENSION_SUFFIXES
def install_nuitka():
print("Checking Nuitka installation...")
try:
import nuitka # noqa: F401
print("Nuitka is installed.")
except ImportError:
print("Installing Nuitka...")
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "nuitka", "ordered-set", "zstandard"]
)
def get_version() -> str:
init_file = Path("src/orisys/__init__.py")
if not init_file.exists():
return "0.3.2"
for line in init_file.read_text(encoding="utf-8").splitlines():
if line.startswith("__version__"):
return line.split("=", 1)[1].strip().strip('"').strip("'")
return "0.3.2"
def find_compiled_extension(output_dir: Path) -> Path:
candidates = []
for suffix in EXTENSION_SUFFIXES:
candidates.extend(output_dir.glob(f"*{suffix}"))
if not candidates:
all_files = ", ".join(p.name for p in output_dir.iterdir()) if output_dir.exists() else "(none)"
raise RuntimeError(f"Nuitka did not produce a compiled extension. Found: {all_files}")
# Prefer the shortest filename (usually the package extension itself)
return sorted(candidates, key=lambda p: len(p.name))[0]
def build():
install_nuitka()
version = get_version()
src_dir = Path("src/orisys")
output_dir = Path("build_nuitka")
dist_dir = Path("dist_nuitka")
if output_dir.exists():
shutil.rmtree(output_dir)
if dist_dir.exists():
shutil.rmtree(dist_dir)
output_dir.mkdir()
dist_dir.mkdir()
print("\n=== Compiling with Nuitka ===")
env = os.environ.copy()
env["PYTHONPATH"] = "src" + os.pathsep + env.get("PYTHONPATH", "")
cmd = [
sys.executable, "-m", "nuitka",
"--module",
"--include-package=orisys",
"--output-dir=" + str(output_dir),
"--remove-output",
"--no-pyi-file",
"src/orisys",
]
print(f"Running: {' '.join(cmd)}")
subprocess.check_call(cmd, env=env)
ext_file = find_compiled_extension(output_dir)
ext_suffix = "".join(ext_file.suffixes)
print(f"Generated compiled extension: {ext_file}")
print("\n=== Preparing Wheel Structure ===")
wheel_root = dist_dir / "wheel_root"
pkg_dir = wheel_root / "orisys"
pkg_dir.mkdir(parents=True)
dest_ext = pkg_dir / f"__init__{ext_suffix}"
shutil.copy(ext_file, dest_ext)
print(f"Copied compiled package to {dest_ext}")
configs_src = src_dir / "configs"
configs_dst = pkg_dir / "configs"
configs_dst.mkdir()
json_count = 0
for json_file in configs_src.glob("*.json"):
shutil.copy(json_file, configs_dst)
json_count += 1
print(f"Copied {json_count} JSON config files")
examples_dst = pkg_dir / "examples"
examples_dst.mkdir(exist_ok=True)
(examples_dst / "__init__.py").write_text("# Examples for orisys package\n", encoding="utf-8")
example_files = [Path("example1_basic_pipeline.py"), Path("example2_two_sensor_pipeline.py")]
example_count = 0
for example_file in example_files:
if example_file.exists():
shutil.copy(example_file, examples_dst)
example_count += 1
print(f"Copied example: {example_file.name}")
print(f"Copied {example_count} example files to orisys/examples/")
setup_py_path = dist_dir / "setup.py"
setup_content = f"""
from setuptools import setup, Distribution
class BinaryDistribution(Distribution):
def has_ext_modules(self):
return True
setup(
name="orisys",
version="{version}",
description="Orisys - Optical Tactile Sensor SDK (Nuitka Compiled)",
author="Orisys",
packages=['orisys', 'orisys.examples'],
package_dir={{'': 'wheel_root'}},
package_data={{
'orisys': ['*{ext_suffix}', 'configs/*.json'],
'orisys.examples': ['*.py'],
}},
include_package_data=True,
install_requires=[
"numpy>=1.24.4",
"opencv-python>=4.10.0.84",
"scipy>=1.10.1",
],
extras_require={{
"cuda12": ["cupy-cuda12x>=12.3.0"],
"cuda13": ["cupy-cuda13x>=13.0.0"],
"cuda": ["cupy-cuda12x>=12.3.0"],
}},
distclass=BinaryDistribution,
zip_safe=False,
)
"""
setup_py_path.write_text(setup_content, encoding="utf-8")
print("\n=== Building Wheel ===")
subprocess.check_call([sys.executable, "setup.py", "bdist_wheel"], cwd=dist_dir)
wheels = list((dist_dir / "dist").glob("*.whl"))
if wheels:
print(f"\n[SUCCESS] Wheel created: {wheels[0]}")
else:
print("\n[ERROR] Wheel not found.")
if __name__ == "__main__":
build()
calibration_data_collection.py
-195
View File
@@ -1,195 +0,0 @@
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import orisys
import numpy as np
import argparse
def main():
parser = argparse.ArgumentParser(description='Calibration Data Collection')
parser.add_argument('-i','--vid_source', type=str, default=0, help='Video Source')
parser.add_argument('-id','--device_id', type=int, default=0, help='Device ID')
args = parser.parse_args()
device_id = args.device_id
vid_source = args.vid_source
sample_name = f"dd01a_{device_id}" #dd01 (平面款), dd01a (平面款 手搓), dd02 (手指款)
test_type = "normal" # shearX, shearY, normal
print("\n")
print(f"Collecting data for device ID: {device_id} on test type: {test_type}, video source: {vid_source}")
print("\n")
sensor = orisys.Sensor(vid_source, config_name="./config/prototype_of3.json", cal_path= f"./config/{sample_name}.npy", cuda=False, verbose=True)
print(f"\n按 'q' 退出程序，按 'r' 重置光流追踪器，按 's' 开始保存，按 'e' 停止保存，按 'j' 减小力，按 'l' 增大力")
force = 0 # kgf
force_step = 0.1 # kgf
force_min = 0 # kgf
force_max = 3 # kgf
is_save = False
save_mode = ["raw", "stitched", "arrows", "contact", "depth"]
save_mode_index = 0
frame_count = 0
writer = None
save_path = None
while True:
sensor.get_img()
if is_save and writer is not None and sensor.frame is not None:
frame_count += 1
if save_mode_index == 0:
save_img = sensor.frame
elif save_mode_index == 1:
save_img = sensor.img
elif save_mode_index == 2:
save_img = arrows
elif save_mode_index == 3:
save_img = image_with_foe
elif save_mode_index == 4:
save_img = div_normalized
# Convert grayscale (1 channel) to BGR (3 channels) if needed
if len(save_img.shape) == 2 or (len(save_img.shape) == 3 and save_img.shape[2] == 1):
if len(save_img.shape) == 3:
save_img = save_img.squeeze(2)
save_img = cv2.cvtColor(save_img, cv2.COLOR_GRAY2BGR)
writer.write(save_img)
sensor.compute_deformation()
sensor.compute_contact()
fps, flow, vnormal, img, contour, centroid, depth_map, fnormal, fshearx, fsheary = sensor.read_info(
sensor.info.FPS,
sensor.info.VRAW,
sensor.info.VNORMAL,
sensor.info.IMG,
sensor.info.CONTOUR,
sensor.info.CENTROID,
sensor.info.DEPTH,
sensor.info.FNORMAL,
sensor.info.FSHEARX,
sensor.info.FSHEARY
)
arrows = orisys.util.draw_arrows(
img,
flow,
threshold=4,
grid_spacing=10,
arrow_scale=1.0
)
cv2.putText(arrows, f"FPS: {fps:.0f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
cv2.imshow("Deformation", arrows)
image_with_foe = orisys.util.draw_contact(img, contour, centroid)
cv2.imshow("Contact Centroid", image_with_foe)
div_abs = depth_map
if div_abs.max() > div_abs.min():
div_normalized = ((div_abs - div_abs.min()) / (div_abs.max() - div_abs.min()) * 255).astype(np.uint8)
else:
div_normalized = np.zeros_like(div_abs, dtype=np.uint8)
# 应用颜色映射（使用JET：蓝色=低散度，红色=高散度）
cv2.imshow("Depth Map", cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET))
key = cv2.waitKey(1) & 0xFF
if key == ord("s"):
save_name = f"{sample_name}_{test_type}_{force:.1f}kgf"
save_path = f"./data/{sample_name}/{test_type}/{save_name}.mp4"
# Ensure directory exists
os.makedirs(os.path.dirname(save_path), exist_ok=True)
# Get frame dimensions from the actual frame
save_img = None
if save_mode_index == 0:
save_img = sensor.frame
elif save_mode_index == 1:
save_img = sensor.img
elif save_mode_index == 2:
save_img = arrows
elif save_mode_index == 3:
save_img = image_with_foe
elif save_mode_index == 4:
save_img = div_normalized
if save_img is not None:
# Convert grayscale (1 channel) to BGR (3 channels) if needed for dimension check
temp_img = save_img.copy()
if len(temp_img.shape) == 2 or (len(temp_img.shape) == 3 and temp_img.shape[2] == 1):
if len(temp_img.shape) == 3:
temp_img = temp_img.squeeze(2)
temp_img = cv2.cvtColor(temp_img, cv2.COLOR_GRAY2BGR)
frame_height, frame_width = temp_img.shape[:2]
writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (frame_width, frame_height))
if writer.isOpened():
is_save = True
print(f"Start writing to {save_path}")
else:
print(f"Failed to open video writer for {save_path}")
writer = None
else:
print("No frame available to determine video dimensions")
if key == ord("e"):
is_save = False
if writer is not None:
writer.release()
writer = None
if save_path is not None:
print(f"Stop writing to {save_path}, total frames: {frame_count}")
else:
print(f"Stop writing, total frames: {frame_count}")
frame_count = 0
if key == ord("p"):
print(f"Normal Force: {fnormal:.1f} kgf, Shear ForceX: {fshearx:.1f} kgf, Shear ForceY: {fsheary:.1f} kgf")
if key == ord("i"):
save_mode_index += 1
if save_mode_index >= len(save_mode):
save_mode_index = 0
print(f"Save mode: {save_mode[save_mode_index]}")
if key == ord("k"):
save_mode_index -= 1
if save_mode_index < 0:
save_mode_index = len(save_mode) - 1
print(f"Save mode: {save_mode[save_mode_index]}")
if key == ord("j"):
force -= force_step
if force < force_min:
force = force_min
print(f"Force: {force:.1f} kgf")
if key == ord("l"):
force += force_step
if force > force_max:
force = force_max
print(f"Force: {force:.1f} kgf")
if key == ord("q"):
print("\n退出程序...")
if writer is not None:
writer.release()
writer = None
break
if key == ord("r"):
# 重置光流追踪器（当追踪漂移时使用）
sensor.reset()
print("传感器已重置")
# =========================================================================
# 4. 清理资源
# =========================================================================
if writer is not None:
writer.release()
sensor.disconnect() # 释放视频源和相关资源
if __name__ == '__main__':
main()
config/compound_eye.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 1014 KiB

config/ddjx01.json
+1 -1
View File
@@ -11,7 +11,7 @@
}, },
"stitcher": { "stitcher": {
"min_area_frac": 0.01, "min_area_frac": 0.01,
"max_area_frac": 0.04, "max_area_frac": 0.05,
"max_aspect_ratio": 50, "max_aspect_ratio": 50,
"index": [1,2,3,6,7,8,11,12,13], "index": [1,2,3,6,7,8,11,12,13],
"margin": 10, "margin": 10,
data/RGB/c0.5mm.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.7 MiB

data/RGB/init_2.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.7 MiB

data/jx01/jx01_hex1.3_arrows.mp4
BIN
View File
Binary file not shown.
data/jx01/jx01_hex1.3_depth.mp4
BIN
View File
Binary file not shown.
data/jx01/jx01_hex1.3_pyqt.mp4
BIN
View File
Binary file not shown.
data/jx01/jx01_hex1.5_2_arrows.mp4
BIN
View File
Binary file not shown.
data/jx01/jx01_hex1.5_2_depth.mp4
BIN
View File
Binary file not shown.
data/jx01/jx01_hex1.5_2_pyqt.mp4
BIN
View File
Binary file not shown.
patch_generation/s1.jpg → data/sample/s1.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 429 KiB

After

Width:  |  Height:  |  Size: 429 KiB

s1_480p.mp4 → data/sample/s1_480p.mp4
View File
s1_720p.mp4 → data/sample/s1_720p.mp4
View File
doc/get_started.docx
BIN
View File
Binary file not shown.
doc/软件使用说明-0.3.0.pdf
BIN
View File
Binary file not shown.
Readme.md → docs/api.md
+7 -5
View File
@@ -72,8 +72,8 @@ conda install conda-forge::pyvkfft
### 输入参数 ### 输入参数
* **vid_src** (`int | str`, 必填): 传感器 ID、相机编号或视频路径。默认为 `0`。 * **vid_src** (`int | str`, 必填): 传感器 ID、相机编号或视频路径。默认为 `0`。
* **config_name** (`str`，可选): 传感器配置`json`文件路径，用于加载对应的参数配置，默认为 `./config/ddjx01`。 * **config_name** (`str`，可选): 传感器配置文件名称或 `json` 文件路径，用于加载对应的参数配置；例如 `ddjx01` 或 `./config/ddjx01.json`。
* **cal_path** (`str`，可选): 复眼拼接校准文件路径，默认为 `./config/cal.npy`。有多个传感器同时运行时建议传感器类创建时建议填写，避免校准文件覆写。在首次运行时程序会在找不到文件时自动创建。 * **cal_path** (`str`，可选，已弃用): 该参数仅为兼容旧代码保留，当前版本中不再使用，传入后也会被忽略。
* **backend** (`str`，可选): 光流场解耦计算后端，默认为 `auto`，可选`cupy`,`opencl`,`cpu`。 * **backend** (`str`，可选): 光流场解耦计算后端，默认为 `auto`，可选`cupy`,`opencl`,`cpu`。
* **verbose** (`bool`，可选): 是否输出初始化及运行过程中的详细日志信息，默认为 `False`。 * **verbose** (`bool`，可选): 是否输出初始化及运行过程中的详细日志信息，默认为 `False`。
@@ -99,12 +99,12 @@ sensor = Sensor("/data/data.mp4")
采集并预处理一帧图像并存到`Sensor` 类别中。 采集并预处理一帧图像并存到`Sensor` 类别中。
### 返回 ### 返回
* `np.ndarray`：处理后的BGR格式图像(`uint8`，尺寸为json文件中`resolution_stitch`对应分辨率)。 * `bool`：是否成功读取图像。
* `None`：当输入源读取失败（例如视频结束、摄像头断开）时返回。
### 状态更新 ### 状态更新
* 成功时更新`sensor`类里面的图像。 * 成功时更新`sensor`类里面的`sensor.img`。处理后的BGR格式图像(`uint8`，尺寸为json文件中`resolution_stitch`对应分辨率)。
### 示例 ### 示例
```python ```python
@@ -123,6 +123,8 @@ sensor.get_img()
### 输入参数 ### 输入参数
* **check_motion**(`bool`，可选): 为 `True` 时先根据光流判断是否有运动，有运动才做 HHD 分解；为 `False` 时每帧都做 HHD。默认 `True`。 * **check_motion**(`bool`，可选): 为 `True` 时先根据光流判断是否有运动，有运动才做 HHD 分解；为 `False` 时每帧都做 HHD。默认 `True`。
* **threshold**(`float`，可选): **check_motion** 功能的阈值，默认为`6`。 * **threshold**(`float`，可选): **check_motion** 功能的阈值，默认为`6`。
* **tail_frames** (`int`，可选): 光流变化低于阈值后继续解算多少帧，确保传感器回复，，默认为 `10`.
### 返回 ### 返回
* `bool`：是否接触到物体，如果**check_motion**为`False`时返回`None`。 * `bool`：是否接触到物体，如果**check_motion**为`False`时返回`None`。
doc/get_started.md → docs/get_started.md
+3 -3
View File
@@ -1,6 +1,6 @@
叠动科技软件使用手册 叠动科技软件使用手册
版本： 0.3.2 版本： 0.3.2
更新日期： 2026年3月6日 更新日期： 2026年4月16日
# 简介 # 简介
本文件介绍如何上手和运行叠动视触觉传感器配套的开发套件。本套件是配合叠动视触觉传感器进行图像处理，计算形变、力以及接触信息。它可以从摄像头或视频文件中读取图像，处理后打印和展示处理后信息。 本文件介绍如何上手和运行叠动视触觉传感器配套的开发套件。本套件是配合叠动视触觉传感器进行图像处理，计算形变、力以及接触信息。它可以从摄像头或视频文件中读取图像，处理后打印和展示处理后信息。
@@ -120,7 +120,7 @@ print(orisys.__version__)
3. 如果第一步没有输出，传感器的ID就是“0”，如果第一步有一个设备，传感器的ID就是“1”，如此类推。 3. 如果第一步没有输出，传感器的ID就是“0”，如果第一步有一个设备，传感器的ID就是“1”，如此类推。
4. 在确定了传感器的ID（下方例子假设为“0”）后，打开Anaconda Prompt, 切换到工作路径，然后在命令行中输入： 4. 在确定了传感器的ID（下方例子假设为“0”）后，打开Anaconda Prompt, 切换到工作路径，然后在命令行中输入：
``` bash ``` bash
python example1_basic_pipeline.py -v 0 python examples/basic_pipeline.py -v 0
``` ```
5. 程序会运行，并自动打开可视化窗口和打印资料，在窗口上按键盘”q”键可以停止运行，按”r”键可以重设数据 5. 程序会运行，并自动打开可视化窗口和打印资料，在窗口上按键盘”q”键可以停止运行，按”r”键可以重设数据
@@ -130,7 +130,7 @@ print(orisys.__version__)
3. 根据新增的输出，例如”/dev/video0”，确定传感器的ID为“0”, 如果是”/dev/video3”的话ID为“3” ，如此类推。 3. 根据新增的输出，例如”/dev/video0”，确定传感器的ID为“0”, 如果是”/dev/video3”的话ID为“3” ，如此类推。
4. 在确定了传感器的ID（下方例子假设为“0”）后，打开终端, 激活相应的python环境，切换到工作路径，然后在命令行中输入： 4. 在确定了传感器的ID（下方例子假设为“0”）后，打开终端, 激活相应的python环境，切换到工作路径，然后在命令行中输入：
``` bash ``` bash
python example1_basic_pipeline.py -v 0 python examples/basic_pipeline.py -v 0
``` ```
5. 程序会运行，并自动打开可视化窗口和打印资料，在窗口上按键盘”q”键可以停止运行，按”r”键可以重设数据 5. 程序会运行，并自动打开可视化窗口和打印资料，在窗口上按键盘”q”键可以停止运行，按”r”键可以重设数据
res/orisys_logo.png → docs/images/orisys_logo.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 22 KiB

After

Width:  |  Height:  |  Size: 22 KiB

res/orisys_logo_notext.png → docs/images/orisys_logo_notext.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 48 KiB

After

Width:  |  Height:  |  Size: 48 KiB

doc/render_command.txt → docs/render_command.txt
View File
example1_basic_pipeline.py
-113
View File
@@ -1,113 +0,0 @@
"""
示例 1: 基础用法 - 单传感器同步处理
本示例演示 orisys 的基本使用流程：
1. 初始化传感器（视频文件或摄像头）
2. 循环获取图像并计算形变和触碰点
3. 读取各类传感器数据
4. 可视化显示
"""
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import orisys
import numpy as np
import argparse
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--video","-v", type=str, default="0", help="Camera ID / Video file path")
parser.add_argument("--config","-c", type=str, default="./config/ddjx01.json", help="Configuration file path")
parser.add_argument("--cal","-cal", type=str, default="./config/ddjx01.npy", help="Calibration file path")
parser.add_argument("--verbose","-verbose", type=bool, default=True, help="Print detailed information")
args = parser.parse_args()
# 步骤1：创建传感器对象 (必要步骤)，
# config_name为配置文件名（内置配置文件名）或路径（自定义配置文件路径）
# cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件
# 方式1: 从视频文件读取
try:
input = int(args.video)
is_camera = True
sensor = orisys.Sensor(input, config_name=args.config, cal_path= args.cal, verbose=args.verbose)
except:
is_camera = False
sensor = orisys.Sensor(args.video, config_name=args.config, cal_path= args.cal, verbose=args.verbose)
print("\n按 'q' 退出程序，按 'r' 重置传感器追踪器")
while True:
# 步骤2：信息提取
# 步骤 2.1: 获取并拼接图像 (必要步骤)
sensor.get_img()
# 步骤 2.2: 计算形变场 (必要步骤)
is_touched = sensor.compute_deformation(check_motion=True, threshold=6)
# 步骤 2.3: 计算触碰点
if is_touched:
sensor.compute_contact()
#步骤 2.4: 读取信息
fps, flow, vnormal, img, contour, centroid, depth_map, fnormal, fshearx, fsheary = sensor.read_info(
sensor.info.FPS,
sensor.info.VRAW,
sensor.info.VNORMAL,
sensor.info.IMG,
sensor.info.CONTOUR,
sensor.info.CENTROID,
sensor.info.DEPTH,
sensor.info.FNORMAL,
sensor.info.FSHEARX,
sensor.info.FSHEARY
)
#步骤 3: 可视化显示
# 步骤 3.1: 绘制形变场
arrows = orisys.util.draw_arrows(
img,
flow,
threshold=2,
grid_spacing=20,
arrow_scale=1.0
)
cv2.imshow("Deformation", arrows)
# 步骤 3.2: 在图像上触碰点和区域
image_with_foe = orisys.util.draw_contact(img, contour, centroid)
cv2.imshow("Contact Centroid", image_with_foe)
# 步骤 3.3: 绘制深度图
# 归一化散度值到0-255范围（使用绝对值，然后归一化）
div_abs = depth_map
if div_abs.max() > div_abs.min():
div_normalized = ((div_abs - div_abs.min()) / (div_abs.max() - div_abs.min()) * 255).astype(np.uint8)
else:
div_normalized = np.zeros_like(div_abs, dtype=np.uint8)
# 应用颜色映射（使用JET：蓝色=低散度，红色=高散度）
cv2.imshow("Depth Map", cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET))
#步骤 3.4: 打印资料
print(f"fps: {fps}, fnormal: {fnormal}, fshearx: {fshearx}, fsheary: {fsheary}, depth_shape: {depth_map.shape}, sensor.resolution_raw: {sensor.img_size_raw}")
#步骤 3.5: 键盘控制
key = cv2.waitKey(1) & 0xFF
if key == ord("q"):
print("\n退出程序...")
break
if key == ord("r"):
# 重置光流追踪器（当追踪漂移时使用）
sensor.reset()
print("传感器已重置")
# =========================================================================
# 4. 清理资源
# =========================================================================
sensor.disconnect() # 释放视频源和相关资源
if __name__ == '__main__':
main()
examples/README.md
+50
View File
@@ -0,0 +1,50 @@
# Examples
This folder contains user-facing example scripts for the Orisys SDK.
## Files
- `basic_pipeline.py`: single-sensor basic processing and visualization flow
- `two_sensor_pipeline.py`: dual-sensor multiprocessing example
- `plot_data.py`: real-time plotting example with PyQtGraph
## Typical usage
Run examples from the repository root so relative paths such as `./config/ddjx01.json`
resolve correctly.
### Basic pipeline
```bash
python examples/basic_pipeline.py -v 0 -c ./config/ddjx01.json
```
### Two-sensor pipeline
```bash
python examples/two_sensor_pipeline.py
```
### Real-time plotting
```bash
python examples/plot_data.py -v 0 -c ./config/ddjx01.json
```
## Requirements
- Python environment with the SDK dependencies installed
- a supported camera or test video
- valid JSON config in `./config/`
Some examples also require:
- `pyqtgraph`
- `PyQt5`
- multiple cameras for multi-sensor demos
## Notes
- Most examples assume a tactile sensor camera is connected.
- Press `q` in the display window to quit where applicable.
- If your camera index is not `0`, change `-v` accordingly.
examples/basic_pipeline.py
+115
View File
@@ -0,0 +1,115 @@
"""
示例 1：单传感器基础流程
本示例演示 Orisys SDK 的标准处理流程：
1. 初始化传感器（摄像头或视频文件）
2. 获取图像并计算形变
3. 按需计算接触区域
4. 读取并可视化结果
SDK version: 0.3.1
"""
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import orisys
import numpy as np
import argparse
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--video","-v", type=str, default="0", help="视频源：摄像头编号或视频文件路径")
parser.add_argument("--config","-c", type=str, default="./config/ddjx01.json", help="配置文件名称或路径")
parser.add_argument("--verbose","-verbose", type=bool, default=True, help="是否输出详细日志")
args = parser.parse_args()
# 步骤 1：创建传感器对象（必需）
# config_name 可以是内置配置名称，也可以是自定义配置文件路径
# cal_path 为标定文件路径；每台设备建议使用独立的标定文件
# 输入既可以是摄像头编号，也可以是视频文件路径
try:
input = int(args.video)
is_camera = True
sensor = orisys.Sensor(input, config_name=args.config, verbose=args.verbose)
except:
is_camera = False
sensor = orisys.Sensor(args.video, config_name=args.config, verbose=args.verbose)
print("\n按键说明：'q' 退出程序，'r' 重置追踪器。")
while True:
# 步骤 2：数据处理
# 步骤 2.1：获取并拼接图像（必需）
sensor.get_img()
# 步骤 2.2：计算形变场（必需）；返回值表示当前帧是否检测到接触/运动
is_contact_now = sensor.compute_deformation(check_motion=True, threshold=0)
# 步骤 2.3：仅在当前帧检测到接触/运动时计算接触区域
# if is_contact_now:
sensor.compute_contact()
# 步骤 2.4：读取结果
fps, flow, vnormal, img, contour, centroid, depth_map, fnormal, fshearx, fsheary = sensor.read_info(
sensor.info.FPS,
sensor.info.VRAW,
sensor.info.VNORMAL,
sensor.info.IMG,
sensor.info.CONTOUR,
sensor.info.CENTROID,
sensor.info.DEPTH,
sensor.info.FNORMAL,
sensor.info.FSHEARX,
sensor.info.FSHEARY
)
# 步骤 3：可视化显示
# 步骤 3.1：显示形变矢量场
arrows = orisys.util.draw_arrows(
img,
flow,
threshold=2,
grid_spacing=20,
arrow_scale=1.0
)
cv2.imshow("形变矢量场", arrows)
# 步骤 3.2：显示接触区域和质心
image_with_foe = orisys.util.draw_contact(img, contour, centroid)
cv2.imshow("接触区域与质心", image_with_foe)
# 步骤 3.3：显示深度图
# 将深度值归一化到 0-255 范围
div_abs = depth_map
if div_abs.max() > div_abs.min():
div_normalized = ((div_abs - div_abs.min()) / (div_abs.max() - div_abs.min()) * 255).astype(np.uint8)
else:
div_normalized = np.zeros_like(div_abs, dtype=np.uint8)
# 应用伪彩色映射
cv2.imshow("深度图", cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET))
# 步骤 3.4：打印关键结果
print(
f"FPS={fps:.2f}, 法向力={fnormal:.4f}, 切向力X={fshearx:.4f}, "
f"切向力Y={fsheary:.4f}, 深度图尺寸={depth_map.shape}, 原始分辨率={sensor.img_size_raw}"
)
# 步骤 3.5：键盘控制
key = cv2.waitKey(1) & 0xFF
if key == ord("q"):
print("\n正在退出示例程序...")
break
if key == ord("r"):
# 当光流追踪出现漂移时，可手动重置追踪器
sensor.reset()
print("追踪器已重置。")
# =========================================================================
# 4. 释放资源
# =========================================================================
sensor.disconnect() # 释放视频源及相关资源
if __name__ == '__main__':
main()
example3_plot_data.py → examples/plot_data.py
+107 -56
View File
@@ -1,3 +1,15 @@
"""
示例 3：实时曲线与幅值图显示
本示例演示如何在读取传感器结果的同时：
1. 显示形变矢量场
2. 显示流场幅值热力图
3. 使用 PyQtGraph 实时绘制法向力与切向力曲线
4. 按键控制录像开始与结束
SDK version: 0.3.1
"""
import os import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2 import cv2
@@ -5,10 +17,45 @@ import orisys
import numpy as np import numpy as np
import time import time
import argparse import argparse
import pyqtgraph as pg try:
import pyqtgraph as pg
except:
print("PyQtGraph is not installed. Please install it using 'pip install pyqt5 pyqtgraph'.")
exit()
def lp_filter(series, alpha=0.5):
series = np.asarray(series)
if series.ndim == 0:
return float(series)
if series.shape[0] < 2:
return float(series[-1])
return float(series[-1] * alpha + series[-2] * (1 - alpha))
def median_filter(series, window_size=5):
series = np.asarray(series)
if series.ndim == 0:
return float(series)
if series.shape[0] == 0:
return 0.0
try:
window_n = max(1, int(window_size))
except (TypeError, ValueError):
window_n = 1
return float(np.median(series[-window_n:]))
def moving_average(series, window_size=5):
series = np.asarray(series)
if series.ndim == 0:
return float(series)
if series.shape[0] == 0:
return 0.0
try:
window_n = max(1, int(window_size))
except (TypeError, ValueError):
window_n = 1
return float(np.mean(series[-window_n:]))
def create_plot(): def create_plot():
p1 = win.addPlot(title="法相力") p1 = win.addPlot(title="法向力")
p1.setLabel('bottom', '时间', units='s') p1.setLabel('bottom', '时间', units='s')
win.nextRow() win.nextRow()
p2 = win.addPlot(title="切向力 X") p2 = win.addPlot(title="切向力 X")
@@ -33,75 +80,75 @@ def update(plot_af, plot_tfx, plot_tfy, data, p1, p2, p3):
# 深度图 # 深度图
def draw_magnitude_map(vfield, threshold=3.0, colormap=cv2.COLORMAP_JET) -> np.ndarray: def draw_magnitude_map(vfield, threshold=3.0, colormap=cv2.COLORMAP_JET) -> np.ndarray:
""" """
Draw color map of pointwise magnitude of velocity field. 绘制流场幅值伪彩色图。
Args: Args:
vfield: Velocity field to visualize (2D vector field) vfield: 需要可视化的二维矢量场
threshold: Minimum magnitude to display (below threshold shows as zero/blue) threshold: 最小显示阈值；低于阈值的区域显示为低值颜色
colormap: OpenCV colormap to use (default: COLORMAP_JET) colormap: OpenCV 颜色映射类型（默认 COLORMAP_JET）
Returns: Returns:
np.ndarray: 400x400 BGR image with colored magnitude map np.ndarray: 400x400 的 BGR 伪彩图像
""" """
flow_field = vfield flow_field = vfield
# Compute magnitude at each point # 计算每个像素点的幅值
magnitude = np.sqrt(flow_field[:, :, 0]**2 + flow_field[:, :, 1]**2) magnitude = np.sqrt(flow_field[:, :, 0]**2 + flow_field[:, :, 1]**2)
# Smooth by downsampling and interpolating back # 先降采样再插值回原尺寸，使显示更平滑
downsample_factor = 10 downsample_factor = 10
h, w = magnitude.shape h, w = magnitude.shape
magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor] magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor]
magnitude_smoothed = cv2.resize(magnitude_downsampled, (w, h), interpolation=cv2.INTER_CUBIC) magnitude_smoothed = cv2.resize(magnitude_downsampled, (w, h), interpolation=cv2.INTER_CUBIC)
# Apply threshold - set below-threshold values to zero # 应用阈值，将低于阈值的区域置零
magnitude_thresholded = magnitude_smoothed.copy() magnitude_thresholded = magnitude_smoothed.copy()
magnitude_thresholded[magnitude_smoothed < threshold] = 0 magnitude_thresholded[magnitude_smoothed < threshold] = 0
# Normalize to 0-255 range # 归一化到 0-255
mag_min, mag_max = magnitude_thresholded.min(), magnitude_thresholded.max() mag_min, mag_max = magnitude_thresholded.min(), magnitude_thresholded.max()
if mag_max > mag_min: if mag_max > mag_min:
magnitude_norm = ((magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255).astype(np.uint8) magnitude_norm = ((magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255).astype(np.uint8)
else: else:
magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8) magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8)
# Apply colormap # 应用伪彩色映射
magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap) magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap)
# Resize to output size # 缩放到固定输出尺寸
out_img = cv2.resize(magnitude_colored, (400, 400), interpolation=cv2.INTER_CUBIC) out_img = cv2.resize(magnitude_colored, (400, 400), interpolation=cv2.INTER_CUBIC)
return out_img return out_img
# 实时图表 # 实时图表
npoints = 100 # 图表数据长度 npoints = 100 # 曲线缓冲区长度
pg.setConfigOption('background', 'w') pg.setConfigOption('background', 'w')
pg.setConfigOption('foreground', 'k') pg.setConfigOption('foreground', 'k')
app = pg.mkQApp("Plotting Example") app = pg.mkQApp("Orisys Plot Example")
win = pg.GraphicsLayoutWidget(show=True, title="Basic plotting examples") win = pg.GraphicsLayoutWidget(show=True, title="Orisys 实时力曲线")
win.setWindowTitle('pyqtgraph example: Plotting') win.setWindowTitle('Orisys 实时力曲线')
p1, p2, p3, plot_af, plot_tfx, plot_tfy = create_plot() p1, p2, p3, plot_af, plot_tfx, plot_tfy = create_plot()
data = np.zeros([npoints,2,3]) data = np.zeros([npoints,2,3])
start_time = None # Will be set when sensor starts start_time = None # Will be set when sensor starts
for i in range(3): for i in range(3):
data[:,0,i] = np.linspace(-npoints,0,npoints) data[:,0,i] = np.linspace(-npoints,0,npoints)/30
def main(): def main():
global start_time global start_time
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--video","-v", type=str, default="0", help="Camera ID / Video file path") parser.add_argument("--video","-v", type=str, default="0", help="视频源：摄像头编号或视频文件路径")
parser.add_argument("--config","-c", type=str, default="./config/ddjx01.json", help="Configuration file path") parser.add_argument("--config","-c", type=str, default="./config/ddjx01.json", help="配置文件名称或路径")
parser.add_argument("--cal","-cal", type=str, default="./config/ddjx01.npy", help="Calibration file path") parser.add_argument("--cal","-cal", type=str, default="./config/ddjx01.npy", help="标定文件路径")
parser.add_argument("--verbose","-verbose", type=bool, default=True, help="Print detailed information") parser.add_argument("--verbose","-verbose", type=bool, default=True, help="是否输出详细日志")
args = parser.parse_args() args = parser.parse_args()
# 步骤1：创建传感器对象 (必要步骤)， # 步骤 1：创建传感器对象（必需）
# config_name为配置文件名（内置配置文件名）或路径（自定义配置文件路径） # config_name 可以是内置配置名称，也可以是自定义配置文件路径
# cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件 # cal_path 为标定文件路径；每台设备建议使用独立的标定文件
# 方式1: 从视频文件读取 # 输入既可以是摄像头编号，也可以是视频文件路径
try: try:
input = int(args.video) input = int(args.video)
is_camera = True is_camera = True
@@ -109,23 +156,23 @@ def main():
except: except:
is_camera = False is_camera = False
sensor = orisys.Sensor(args.video, config_name=args.config, cal_path= args.cal, verbose=args.verbose) sensor = orisys.Sensor(args.video, config_name=args.config, cal_path= args.cal, verbose=args.verbose)
# 打开视频 # 下面保留了历史示例写法，便于参考
# sensor1 = Sensor("./s1.mp4",isstitch=True,cuda = True,verbose = True)
# 方式1: 从视频文件读取
# sensor1 = orisys.Sensor("./s1.mp4", config_name="./config/ddjx01.json", cal_path= "./config/ddjx01.npy", cuda=False, verbose=True) # sensor1 = orisys.Sensor("./s1.mp4", config_name="./config/ddjx01.json", cal_path= "./config/ddjx01.npy", cuda=False, verbose=True)
start_time = time.time() start_time = time.time()
is_save = False is_save = False
frame_count = 0 frame_count = 0
writer = None writer = None
save_path = None
print("\n按键说明：'q' 退出程序，'s' 开始录像，'e' 结束录像。")
while True: while True:
sensor.get_img() # 提取并拼接复眼图像 sensor.get_img() # 获取并拼接复眼图像
if is_save and writer is not None and sensor.frame is not None: if is_save and writer is not None and sensor.frame is not None:
frame_count += 1 frame_count += 1
save_img = sensor.frame save_img = sensor.frame
# Convert grayscale (1 channel) to BGR (3 channels) if needed # 若为灰度图，则先转换为 BGR 再写入视频
if len(save_img.shape) == 2 or (len(save_img.shape) == 3 and save_img.shape[2] == 1): if len(save_img.shape) == 2 or (len(save_img.shape) == 3 and save_img.shape[2] == 1):
if len(save_img.shape) == 3: if len(save_img.shape) == 3:
save_img = save_img.squeeze(2) save_img = save_img.squeeze(2)
@@ -133,44 +180,48 @@ def main():
writer.write(save_img) writer.write(save_img)
sensor.compute_deformation(check_motion=False) # 计算形变和分解 sensor.compute_deformation(check_motion=True, threshold=0) # 计算形变与分解结果
# 提取信息： # 读取信息：
# 向量场: VRAW, VNORMAL, VSHEAR: 原始，法向，切向 # 向量场：VRAW、VNORMAL、VSHEAR -> 原始 / 法向 / 切向
# 帧率，时间戳: FPS, TIMESTAMP # 标量：FNORMAL、FSHEARX、FSHEARY -> 法向力 / X 向切向力 / Y 向切向力
# 变形量总和: FNORMAL , FSHEARX, FSHEARY : 法向，切向（X)，切向（Y)
fps, fn, fx, fy, flow, vnormal, vshear = sensor.read_info(sensor.info.FPS, sensor.info.FNORMAL, sensor.info.FSHEARX, sensor.info.FSHEARY, sensor.info.VRAW, sensor.info.VNORMAL,sensor.info.VSHEAR) fps, fn, fx, fy, flow, vnormal, vshear = sensor.read_info(sensor.info.FPS, sensor.info.FNORMAL, sensor.info.FSHEARX, sensor.info.FSHEARY, sensor.info.VRAW, sensor.info.VNORMAL,sensor.info.VSHEAR)
print(fps,fn,fx,fy) print(f"FPS={fps:.2f}, 法向力={fn:.4f}, 切向力X={fx:.4f}, 切向力Y={fy:.4f}")
# 显示向量场 # 显示形变矢量场
arrows = orisys.util.draw_arrows(sensor.img, flow, threshold=5, grid_spacing=10, arrow_scale=1.0) arrows = orisys.util.draw_arrows(sensor.img, flow, threshold=5, grid_spacing=10, arrow_scale=1.0)
cv2.imshow("Deformation", arrows) cv2.imshow("形变矢量场", arrows)
# 显示深度图 # 显示流场幅值图
mag = draw_magnitude_map(flow) mag = draw_magnitude_map(flow)
cv2.imshow("Magnitude", mag) cv2.imshow("流场幅值图", mag)
# 显示图表 - 更新时间轴 # 更新实时曲线
current_time = time.time() - start_time current_time = time.time() - start_time
data_filtered = data.copy()
for i in range(3): for i in range(3):
data[:-1, 0, i] = data[1:, 0, i] # shift time data left data[:-1, 0, i] = data[1:, 0, i] # 时间轴左移
data[:-1, 1, i] = data[1:, 1, i] # shift force data left data[:-1, 1, i] = data[1:, 1, i] # 力数据左移
data[-1, 0, :] = current_time # set current time for all plots data[-1, 0, :] = current_time # 写入当前时间
data[-1, 1, 0] = fn data[-1, 1, 0] = fn
data[-1, 1, 1] = fx data[-1, 1, 1] = fx
data[-1, 1, 2] = fy data[-1, 1, 2] = fy
update(plot_af, plot_tfx, plot_tfy, data, p1, p2, p3) # Update plot # 对最新采样点做一阶低通，滤波对象是历史曲线而不是单个标量。
data_filtered[-1, 1, 0] = moving_average(data[:, 1, 0])
data_filtered[-1, 1, 1] = moving_average(data[:, 1, 1])
data_filtered[-1, 1, 2] = moving_average(data[:, 1, 2])
update(plot_af, plot_tfx, plot_tfy, data_filtered, p1, p2, p3) # 更新曲线
key = cv2.waitKey(1) key = cv2.waitKey(1)
if key & 0xFF == ord("q"): if key & 0xFF == ord("q"):
print("\n 停止程序") print("\n正在退出示例程序...")
break break
if key == ord("s"): if key == ord("s"):
save_path = f"./data/pulse_test.mp4" save_path = f"./data/pulse_test.mp4"
# Ensure directory exists # 确保输出目录存在
os.makedirs(os.path.dirname(save_path), exist_ok=True) os.makedirs(os.path.dirname(save_path), exist_ok=True)
# Get frame dimensions from the actual frame # 从当前帧获取视频尺寸
save_img = sensor.frame save_img = sensor.frame
if save_img is not None: if save_img is not None:
# Convert grayscale (1 channel) to BGR (3 channels) if needed for dimension check # 如为灰度图，先转换为 BGR 以获取视频尺寸
temp_img = save_img.copy() temp_img = save_img.copy()
if len(temp_img.shape) == 2 or (len(temp_img.shape) == 3 and temp_img.shape[2] == 1): if len(temp_img.shape) == 2 or (len(temp_img.shape) == 3 and temp_img.shape[2] == 1):
if len(temp_img.shape) == 3: if len(temp_img.shape) == 3:
@@ -181,24 +232,24 @@ def main():
writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (frame_width, frame_height)) writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (frame_width, frame_height))
if writer.isOpened(): if writer.isOpened():
is_save = True is_save = True
print(f"Start writing to {save_path}") print(f"开始录像：{save_path}")
else: else:
print(f"Failed to open video writer for {save_path}") print(f"无法创建录像文件：{save_path}")
writer = None writer = None
else: else:
print("No frame available to determine video dimensions") print("当前没有可用帧，无法确定录像尺寸。")
if key == ord("e"): if key == ord("e"):
is_save = False is_save = False
if writer is not None: if writer is not None:
writer.release() writer.release()
writer = None writer = None
if save_path is not None: if save_path is not None:
print(f"Stop writing to {save_path}, total frames: {frame_count}") print(f"录像结束：{save_path}，总帧数：{frame_count}")
else: else:
print(f"Stop writing, total frames: {frame_count}") print(f"录像结束，总帧数：{frame_count}")
frame_count = 0 frame_count = 0
sensor.disconnect() # 断开传感器连接 sensor.disconnect() # 断开传感器连接并释放资源
example2_two_sensor_pipeline.py → examples/two_sensor_pipeline.py
+65 -60
View File
@@ -1,22 +1,24 @@
""" """
示例 4: 多传感器并行处理（多进程模式 - 不使用 MultiSensor） 示例 2：双传感器并行处理（多进程）
本示例演示如何使用 multiprocessing 同时运行两个传感器： 本示例演示如何使用 multiprocessing 同时运行两个传感器：
- 每个传感器运行在独立的进程中 - 每个传感器运行在独立进程中
- 每个进程拥有独立的 CUDA 上下文 - 每个进程拥有独立的 CUDA 上下文
- 使用 multiprocessing.Process 和 Queue 进行进程间通信 - 通过 Process 和 Queue 在主进程中汇总结果并显示
架构说明： 架构说明：
进程 1 (传感器 0): 采集 → 光流 → HHD → 接触检测 → 队列 进程 1（传感器 0）：采集 → 光流 → HHD → 接触检测 → 队列
进程 2 (传感器 1): 采集 → 光流 → HHD → 接触检测 → 队列 进程 2（传感器 1）：采集 → 光流 → HHD → 接触检测 → 队列
↓ ↓
主进程: 显示 ← ───────────────┘ 主进程： 显示 ← ───────────────┘
适用场景：多指触觉传感器、多视角采集、需要 GPU 加速的多路处理 适用场景：多指触觉传感器、多视角采集，以及需要 GPU 加速的多路处理。
注意事项： 注意事项：
- Windows 下必须在 if __name__ == '__main__' 中调用 mp.freeze_support() - Windows 下必须在 if __name__ == '__main__' 中调用 mp.freeze_support()
- 每个传感器进程会占用独立的 GPU 显存 - 每个传感器进程都会占用独立的 GPU 显存
SDK version: 0.3.1
""" """
import os import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
@@ -26,20 +28,20 @@ import numpy as np
import multiprocessing as mp import multiprocessing as mp
from multiprocessing import Queue, Process from multiprocessing import Queue, Process
import time import time
import argparse
def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_event, target_fps=None): def sensor_worker(sensor_id, vid_src, config_name, result_queue, stop_event, target_fps=None):
""" """
传感器工作进程函数 传感器工作进程。
Args: Args:
sensor_id: 传感器编号 sensor_id: 传感器编号
vid_src: 视频源（文件路径或摄像头索引） vid_src: 视频源（文件路径或摄像头索引）
config_name: 配置名称 config_name: 配置名称
cal_path: 标定路径
result_queue: 结果队列，用于向主进程发送数据 result_queue: 结果队列，用于向主进程发送数据
stop_event: 停止事件，用于通知工作进程退出 stop_event: 停止事件，用于通知工作进程退出
target_fps: 目标帧率；若设置则用定时器限速，使多传感器 FPS 接近一致 target_fps: 目标帧率；设置后将对处理循环限速，使多路传感器 FPS 更接近
""" """
try: try:
# 在每个进程中创建独立的传感器实例 # 在每个进程中创建独立的传感器实例
@@ -48,27 +50,26 @@ def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_
isstitch=True, isstitch=True,
cuda=True, cuda=True,
verbose=False, # 减少子进程输出 verbose=False, # 减少子进程输出
config_name=config_name, config_name=config_name
cal_path=cal_path
) )
print(f"传感器 {sensor_id} 已启动" + (f" (target_fps={target_fps})" if target_fps else "")) print(f"传感器 {sensor_id} 已启动" + (f"（目标帧率：{target_fps} FPS）" if target_fps else ""))
frame_interval = 1.0 / target_fps if target_fps and target_fps > 0 else 0.0 frame_interval = 1.0 / target_fps if target_fps and target_fps > 0 else 0.0
while not stop_event.is_set(): while not stop_event.is_set():
t_start = time.perf_counter() t_start = time.perf_counter()
try: try:
# 步骤1: 获取图像 # 步骤 1：获取图像
sensor.get_img() sensor.get_img()
# 步骤2: 计算形变场 # 步骤 2：计算形变场
sensor.compute_deformation() sensor.compute_deformation()
# 步骤3: 计算触碰点 # 步骤 3：计算接触区域（如需要可取消注释）
# sensor.compute_contact() # sensor.compute_contact()
# 步骤4: 读取信息 # 步骤 4：读取结果
fps, flow, vnormal, img, contour, centroid, depth_map = sensor.read_info( fps, flow, vnormal, img, contour, centroid, depth_map = sensor.read_info(
sensor.info.FPS, sensor.info.FPS,
sensor.info.VRAW, sensor.info.VRAW,
@@ -79,7 +80,7 @@ def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_
sensor.info.DEPTH sensor.info.DEPTH
) )
# 步骤5: 准备数据字典（确保所有数据都是可序列化的） # 步骤 5：整理结果字典（确保数据可序列化）
data = { data = {
'sensor_id': sensor_id, 'sensor_id': sensor_id,
'fps': fps, 'fps': fps,
@@ -92,100 +93,104 @@ def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_
'timestamp': time.time() 'timestamp': time.time()
} }
# 步骤6: 发送数据到主进程（非阻塞） # 步骤 6：将结果发送到主进程（非阻塞）
try: try:
result_queue.put_nowait(data) result_queue.put_nowait(data)
except Exception: except Exception:
# 如果队列已满，跳过这一帧 # 如果队列已满，则丢弃当前帧，避免阻塞工作进程
pass pass
except Exception as e: except Exception as e:
print(f"传感器 {sensor_id} 处理错误: {e}") print(f"传感器 {sensor_id} 处理失败：{e}")
time.sleep(0.1) time.sleep(0.1)
# 步骤7: 按 target_fps 限速，使多传感器 FPS 接近 # 步骤 7：按 target_fps 限速，使多路传感器 FPS 更接近
if frame_interval > 0: if frame_interval > 0:
elapsed = time.perf_counter() - t_start elapsed = time.perf_counter() - t_start
sleep_time = frame_interval - elapsed sleep_time = frame_interval - elapsed
if sleep_time > 0.001: # 至少 1ms 才 sleep，避免无意义调度 if sleep_time > 0.001: # 至少 1ms 才 sleep，避免无意义调度
time.sleep(sleep_time) time.sleep(sleep_time)
# 清理资源 # 释放资源
sensor.disconnect() sensor.disconnect()
print(f"传感器 {sensor_id} 已停止") print(f"传感器 {sensor_id} 已停止。")
except Exception as e: except Exception as e:
print(f"传感器 {sensor_id} 初始化失败: {e}") print(f"传感器 {sensor_id} 初始化失败：{e}")
result_queue.put({'sensor_id': sensor_id, 'error': str(e)}) result_queue.put({'sensor_id': sensor_id, 'error': str(e)})
def main(): def main():
# ========================================================================= # =========================================================================
# 1. 配置多个传感器 # 1. 配置多路传感器
# ========================================================================= # =========================================================================
# 目标帧率：设为正数时，两路传感器会按此帧率限速，使 FPS 接近一致；设为 None 则不限速 # 目标帧率：设为正数时，两路传感器会按该帧率限速；设为 None 则不限速
parser = argparse.ArgumentParser()
parser.add_argument("--video1","-v1", type=str, default="0", help="视频文件路径")
parser.add_argument("--video2","-v2", type=str, default="0", help="视频文件路径")
args = parser.parse_args()
target_fps = 30 target_fps = 30
num_sensors = 2 num_sensors = 2
sensor_configs = [ sensor_configs = [
{ {
'sensor_id': 0, 'sensor_id': 0,
'vid_src': 2, 'vid_src': args.video1,
'config_name': './config/prototype_of3_720p.json', 'config_name': './config/ddjx01.json'
'cal_path': './config/ddjx01.npy'
}, },
{ {
'sensor_id': 1, 'sensor_id': 1,
'vid_src': 4, # 可以改为不同的视频源或摄像头索引 'vid_src': args.video2, # 可替换为其他视频源或摄像头索引
'config_name': './config/prototype_of3_720p.json', 'config_name': './config/ddjx01.json'
'cal_path': './config/ddjx01.npy'
}, },
] ]
# ========================================================================= # =========================================================================
# 2. 创建进程间通信对象 # 2. 创建进程间通信对象
# ========================================================================= # =========================================================================
# 增大队列大小以缓冲更多帧，避免因主进程处理延迟导致的丢帧 # 适当增大队列容量，以缓冲更多帧并减轻主进程偶发延迟带来的丢帧
result_queues = [Queue(maxsize=30) for _ in sensor_configs] result_queues = [Queue(maxsize=30) for _ in sensor_configs]
stop_events = [mp.Event() for _ in sensor_configs] stop_events = [mp.Event() for _ in sensor_configs]
# ========================================================================= # =========================================================================
# 3. 启动传感器工作进程 # 3. 启动工作进程
# ========================================================================= # =========================================================================
processes = [] processes = []
for i, config in enumerate(sensor_configs): for i, config in enumerate(sensor_configs):
p = Process( p = Process(
target=sensor_worker, target=sensor_worker,
args=(i, config['vid_src'], config['config_name'], config['cal_path'], args=(i, config['vid_src'], config['config_name'],
result_queues[i], stop_events[i], target_fps) result_queues[i], stop_events[i], target_fps)
) )
p.start() p.start()
processes.append(p) processes.append(p)
print(f"\n已启动 {len(processes)} 个传感器进程") print(f"\n已启动 {len(processes)} 个传感器进程。")
print("按 'q' 退出程序") print("按键说明：'q' 退出程序。")
# 用于存储每个传感器的最新数据 # 保存每个传感器的最新一帧数据
latest_data = [None] * len(sensor_configs) latest_data = [None] * len(sensor_configs)
# 用于控制打印频率，减少I/O开销 # 控制打印频率，减少 I/O 开销
print_counter = 0 print_counter = 0
print_interval = 30 # 每30帧打印一次，进一步减少I/O print_interval = 30 # 每 30 帧打印一次
# 用于控制按键检查频率 # 控制按键检查频率
key_check_counter = 0 key_check_counter = 0
key_check_interval = 5 # 每5帧检查一次按键 key_check_interval = 5 # 每 5 帧检查一次按键
try: try:
# ===================================================================== # =====================================================================
# 4. 主循环：从队列获取数据并显示 # 4. 主循环：从队列获取结果并显示
# ===================================================================== # =====================================================================
while True: while True:
# 从所有队列中快速获取最新数据（高效清空队列，只保留最新帧） # 从所有队列中快速获取最新数据，只保留最后一帧
for i, queue in enumerate(result_queues): for i, queue in enumerate(result_queues):
# 快速清空队列，只保留最后一帧 # 清空队列，仅保留最后一帧
data = None data = None
# 一次性获取所有可用数据，只保留最后一个 # 读取所有可用数据，只保留最新结果
while True: while True:
try: try:
data = queue.get_nowait() data = queue.get_nowait()
@@ -195,14 +200,14 @@ def main():
if data is not None: if data is not None:
latest_data[i] = data latest_data[i] = data
# 显示所有传感器的数据 # 显示所有传感器的最新结果
for data in latest_data: for data in latest_data:
if data is None or 'error' in data: if data is None or 'error' in data:
continue continue
sid = data['sensor_id'] sid = data['sensor_id']
# 显示形变场 # 显示形变矢量场
if data['img'] is not None and data['flow'] is not None: if data['img'] is not None and data['flow'] is not None:
arrows = orisys.util.draw_arrows( arrows = orisys.util.draw_arrows(
data['img'], data['img'],
@@ -211,7 +216,7 @@ def main():
grid_spacing=10, grid_spacing=10,
arrow_scale=1.0 arrow_scale=1.0
) )
cv2.imshow(f"Deformation {sid}", arrows) cv2.imshow(f"传感器 {sid} - 形变矢量场", arrows)
# # 显示触碰点和区域 # # 显示触碰点和区域
# if data['img'] is not None: # if data['img'] is not None:
@@ -233,33 +238,33 @@ def main():
# cv2.imshow(f"Depth Map {sid}", # cv2.imshow(f"Depth Map {sid}",
# cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET)) # cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET))
# 降低打印频率以减少I/O开销（在主循环级别控制） # 降低打印频率，减少主循环中的 I/O 开销
print_counter += 1 print_counter += 1
if print_counter >= print_interval: if print_counter >= print_interval:
for data in latest_data: for data in latest_data:
if data is not None and 'error' not in data: if data is not None and 'error' not in data:
sid = data['sensor_id'] sid = data['sensor_id']
print(f"传感器 {sid}: FPS={data['fps']:5.1f}") print(f"传感器 {sid}：FPS={data['fps']:5.1f}")
print_counter = 0 print_counter = 0
# 降低按键检查频率以减少开销 # 降低按键检查频率，减少额外开销
key_check_counter += 1 key_check_counter += 1
if key_check_counter >= key_check_interval: if key_check_counter >= key_check_interval:
key = cv2.waitKey(1) & 0xFF key = cv2.waitKey(1) & 0xFF
if key == ord('q'): if key == ord('q'):
print("\n退出程序...") print("\n正在退出示例程序...")
break break
key_check_counter = 0 key_check_counter = 0
finally: finally:
# ===================================================================== # =====================================================================
# 5. 清理资源 # 5. 释放资源
# ===================================================================== # =====================================================================
# 通知所有工作进程停止 # 通知所有工作进程停止
for event in stop_events: for event in stop_events:
event.set() event.set()
# 等待所有进程结束 # 等待所有工作进程结束
for p in processes: for p in processes:
p.join(timeout=2) p.join(timeout=2)
if p.is_alive(): if p.is_alive():
@@ -267,10 +272,10 @@ def main():
p.join() p.join()
cv2.destroyAllWindows() cv2.destroyAllWindows()
print("所有传感器进程已停止") print("所有传感器进程已停止。")
if __name__ == '__main__': if __name__ == '__main__':
# Windows 下多进程必需的初始化 # Windows 下运行多进程示例所需的初始化
mp.freeze_support() mp.freeze_support()
main() main()
five_sensors.py
-296
View File
@@ -1,296 +0,0 @@
"""
示例 4: 多传感器并行处理（多进程模式 - 不使用 MultiSensor）
本示例演示如何使用 multiprocessing 同时运行两个传感器：
- 每个传感器运行在独立的进程中
- 每个进程拥有独立的 CUDA 上下文
- 使用 multiprocessing.Process 和 Queue 进行进程间通信
架构说明：
进程 1 (传感器 0): 采集 → 光流 → HHD → 接触检测 → 队列
进程 2 (传感器 1): 采集 → 光流 → HHD → 接触检测 → 队列
↓
主进程: 显示 ← ───────────────┘
适用场景：多指触觉传感器、多视角采集、需要 GPU 加速的多路处理
注意事项：
- Windows 下必须在 if __name__ == '__main__' 中调用 mp.freeze_support()
- 每个传感器进程会占用独立的 GPU 显存
"""
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import orisys
import numpy as np
import multiprocessing as mp
from multiprocessing import Queue, Process
import time
def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_event, target_fps=None):
"""
传感器工作进程函数
Args:
sensor_id: 传感器编号
vid_src: 视频源（文件路径或摄像头索引）
config_name: 配置名称
cal_path: 标定路径
result_queue: 结果队列，用于向主进程发送数据
stop_event: 停止事件，用于通知工作进程退出
target_fps: 目标帧率；若设置则用定时器限速，使多传感器 FPS 接近一致
"""
try:
# 在每个进程中创建独立的传感器实例
sensor = orisys.Sensor(
vid_src,
isstitch=True,
cuda=True,
verbose=False, # 减少子进程输出
config_name=config_name,
cal_path=cal_path
)
print(f"传感器 {sensor_id} 已启动" + (f" (target_fps={target_fps})" if target_fps else ""))
frame_interval = 1.0 / target_fps if target_fps and target_fps > 0 else 0.0
while not stop_event.is_set():
t_start = time.perf_counter()
try:
# 步骤1: 获取图像
sensor.get_img()
# 步骤2: 计算形变场
sensor.compute_deformation()
# 步骤3: 计算触碰点
sensor.compute_contact()
# 步骤4: 读取信息
fps, flow, vnormal, img, contour, centroid, depth_map = sensor.read_info(
sensor.info.FPS,
sensor.info.VRAW,
sensor.info.VNORMAL,
sensor.info.IMG,
sensor.info.CONTOUR,
sensor.info.CENTROID,
sensor.info.DEPTH
)
# 步骤5: 准备数据字典（确保所有数据都是可序列化的）
data = {
'sensor_id': sensor_id,
'fps': fps,
'flow': flow.copy() if flow is not None else None,
'vnormal': vnormal.copy() if vnormal is not None else None,
'img': img.copy() if img is not None else None,
'contour': contour,
'centroid': centroid,
'depth_map': depth_map.copy() if depth_map is not None else None,
'timestamp': time.time()
}
# 步骤6: 发送数据到主进程（非阻塞）
try:
result_queue.put_nowait(data)
except Exception:
# 如果队列已满，跳过这一帧
pass
except Exception as e:
print(f"传感器 {sensor_id} 处理错误: {e}")
time.sleep(0.1)
# 步骤7: 精确限速（time.sleep 在 Windows 上精度仅 ~15.6ms，需留足余量后 spin-wait）
if frame_interval > 0:
deadline = t_start + frame_interval
remaining = deadline - time.perf_counter()
if remaining > 0.018:
time.sleep(remaining - 0.018) # coarse sleep; leave 18ms margin for Windows granularity
while time.perf_counter() < deadline:
pass # spin-wait for the last few ms
# 清理资源
sensor.disconnect()
print(f"传感器 {sensor_id} 已停止")
except Exception as e:
print(f"传感器 {sensor_id} 初始化失败: {e}")
result_queue.put({'sensor_id': sensor_id, 'error': str(e)})
def main():
# =========================================================================
# 1. 配置多个传感器
# =========================================================================
# 目标帧率：设为 None 不限速（可跑满 100+ FPS）；设为正数则 sleep 限速，使 FPS ≤ target_fps
target_fps = 30 # 不限速；设为 30 等可 cap 帧率
num_sensors = 5
sensor_configs = [
{
'sensor_id': 0,
'vid_src': './s1_480p.mp4',
'config_name': './config/prototype_of3_480p.json',
'cal_path': './config/dd01a_0_480p.npy'
},
{
'sensor_id': 1,
'vid_src': './s1_480p.mp4', # 可以改为不同的视频源或摄像头索引
'config_name': './config/prototype_of3_480p.json',
'cal_path': './config/dd01a_0_480p.npy'
},
{
'sensor_id': 2,
'vid_src': './s1_480p.mp4', # 可以改为不同的视频源或摄像头索引
'config_name': './config/prototype_of3_480p.json',
'cal_path': './config/dd01a_0_480p.npy'
},
{
'sensor_id': 3,
'vid_src': './s1_480p.mp4', # 可以改为不同的视频源或摄像头索引
'config_name': './config/prototype_of3_480p.json',
'cal_path': './config/dd01a_0_480p.npy'
},
{
'sensor_id': 4,
'vid_src': './s1_480p.mp4', # 可以改为不同的视频源或摄像头索引
'config_name': './config/prototype_of3_480p.json',
'cal_path': './config/dd01a_0_480p.npy'
},
]
# =========================================================================
# 2. 创建进程间通信对象
# =========================================================================
# 增大队列大小以缓冲更多帧，避免因主进程处理延迟导致的丢帧
result_queues = [Queue(maxsize=30) for _ in sensor_configs]
stop_events = [mp.Event() for _ in sensor_configs]
# =========================================================================
# 3. 启动传感器工作进程
# =========================================================================
processes = []
for i, config in enumerate(sensor_configs):
p = Process(
target=sensor_worker,
args=(i, config['vid_src'], config['config_name'], config['cal_path'],
result_queues[i], stop_events[i], target_fps)
)
p.start()
processes.append(p)
print(f"\n已启动 {len(processes)} 个传感器进程")
print("按 'q' 退出程序")
# 用于存储每个传感器的最新数据
latest_data = [None] * len(sensor_configs)
# 用于控制打印频率，减少I/O开销
print_counter = 0
print_interval = 30 # 每30帧打印一次，进一步减少I/O
# 用于控制按键检查频率
key_check_counter = 0
key_check_interval = 5 # 每5帧检查一次按键
try:
# =====================================================================
# 4. 主循环：从队列获取数据并显示
# =====================================================================
while True:
# 从所有队列中快速获取最新数据（高效清空队列，只保留最新帧）
for i, queue in enumerate(result_queues):
# 快速清空队列，只保留最后一帧
data = None
# 一次性获取所有可用数据，只保留最后一个
while True:
try:
data = queue.get_nowait()
except:
break
if data is not None:
latest_data[i] = data
# 显示所有传感器的数据
for data in latest_data:
if data is None or 'error' in data:
continue
sid = data['sensor_id']
# 显示形变场
if data['img'] is not None and data['flow'] is not None:
arrows = orisys.util.draw_arrows(
data['img'],
data['flow'],
threshold=5,
grid_spacing=10,
arrow_scale=1.0
)
cv2.imshow(f"Deformation {sid}", arrows)
# # 显示触碰点和区域
# if data['img'] is not None:
# image_with_foe = orisys.util.draw_contact(
# data['img'],
# data['contour'],
# data['centroid']
# )
# cv2.imshow(f"Contact Centroid {sid}", image_with_foe)
# # 显示深度图
# if data['depth_map'] is not None and data['depth_map'].size > 0:
# div_abs = data['depth_map']
# if div_abs.max() > div_abs.min():
# div_normalized = ((div_abs - div_abs.min()) /
# (div_abs.max() - div_abs.min()) * 255).astype(np.uint8)
# else:
# div_normalized = np.zeros_like(div_abs, dtype=np.uint8)
# cv2.imshow(f"Depth Map {sid}",
# cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET))
# 降低打印频率以减少I/O开销（在主循环级别控制）
print_counter += 1
if print_counter >= print_interval:
for data in latest_data:
if data is not None and 'error' not in data:
sid = data['sensor_id']
print(f"传感器 {sid}: FPS={data['fps']:5.1f}")
print_counter = 0
# 降低按键检查频率以减少开销
key_check_counter += 1
if key_check_counter >= key_check_interval:
key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
print("\n退出程序...")
break
key_check_counter = 0
finally:
# =====================================================================
# 5. 清理资源
# =====================================================================
# 通知所有工作进程停止
for event in stop_events:
event.set()
# 等待所有进程结束
for p in processes:
p.join(timeout=2)
if p.is_alive():
p.terminate()
p.join()
cv2.destroyAllWindows()
print("所有传感器进程已停止")
if __name__ == '__main__':
# Windows 下多进程必需的初始化
mp.freeze_support()
main()
hand_test.py
-354
View File
@@ -1,354 +0,0 @@
"""
示例 4: 多传感器并行处理（多进程模式 - 不使用 MultiSensor）
本示例演示如何使用 multiprocessing 同时运行两个传感器：
- 每个传感器运行在独立的进程中
- 每个进程拥有独立的 CUDA 上下文
- 使用 multiprocessing.Process 和 Queue 进行进程间通信
架构说明：
进程 1 (传感器 0): 采集 → 光流 → HHD → 接触检测 → 队列
进程 2 (传感器 1): 采集 → 光流 → HHD → 接触检测 → 队列
↓
主进程: 显示 ← ───────────────┘
适用场景：多指触觉传感器、多视角采集、需要 GPU 加速的多路处理
注意事项：
- Windows 下必须在 if __name__ == '__main__' 中调用 mp.freeze_support()
- 每个传感器进程会占用独立的 GPU 显存
"""
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import orisys
import numpy as np
import multiprocessing as mp
from multiprocessing import Queue, Process
import time
import ctypes
import platform
import ctypes
# Multi-sensor FPS stabilization options
STABILITY_MODE = True # True: prioritize similar FPS across sensors
TARGET_FPS = 30.0 # Used only when STABILITY_MODE=True
SEND_VISUALIZATION_DATA = False # True sends img/flow/depth to main process (higher IPC cost)
def _enable_high_res_sleep():
"""
Enable high-resolution sleep for accurate FPS pacing.
Note: time.perf_counter() is already high-resolution on all platforms (Windows, macOS, Linux).
This function only improves time.sleep() accuracy on Windows.
"""
if os.name == "nt": # Windows (x64 and ARM64)
try:
# timeBeginPeriod improves time.sleep() accuracy (default ~15ms -> 1ms)
# time.perf_counter() is already high-res on Windows (uses QueryPerformanceCounter)
winmm = ctypes.WinDLL("winmm")
winmm.timeBeginPeriod(1) # Request 1ms timer resolution for sleep()
return winmm
except Exception:
return None
# macOS and Linux: time.perf_counter() and time.sleep() are already high-resolution
# No special setup needed
return None
def _disable_high_res_sleep(handle):
"""Restore default scheduler timer resolution (Windows only)."""
if handle is None:
return
if os.name == "nt" and isinstance(handle, ctypes.WinDLL):
try:
handle.timeEndPeriod(1)
except Exception:
pass
# macOS and Linux don't need cleanup
def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_event):
"""
传感器工作进程函数
Args:
sensor_id: 传感器编号
vid_src: 视频源（文件路径或摄像头索引）
config_name: 配置名称
result_queue: 结果队列，用于向主进程发送数据
stop_event: 停止事件，用于通知工作进程退出
"""
timer_handle = None
try:
# Limit per-process CPU thread fan-out to reduce contention jitter.
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
cv2.setNumThreads(1)
if STABILITY_MODE and TARGET_FPS > 0:
timer_handle = _enable_high_res_sleep()
# 在每个进程中创建独立的传感器实例
sensor = orisys.Sensor(
vid_src,
isstitch=True,
cuda=True,
verbose=False, # 减少子进程输出
config_name=config_name,
cal_path=cal_path
)
print(f"传感器 {sensor_id} 已启动")
target_period = (1.0 / TARGET_FPS) if (STABILITY_MODE and TARGET_FPS > 0) else 0.0
next_tick = time.perf_counter()
while not stop_event.is_set():
try:
# 步骤1: 获取图像
sensor.get_img()
# 步骤2: 计算形变场
sensor.compute_deformation()
# 步骤3: 计算触碰点
sensor.compute_contact()
# 步骤4: 读取信息
if SEND_VISUALIZATION_DATA:
fps, flow, img, contour, centroid, depth_map = sensor.read_info(
sensor.info.FPS,
sensor.info.VRAW,
sensor.info.IMG,
sensor.info.CONTOUR,
sensor.info.CENTROID,
sensor.info.DEPTH
)
# Full payload (higher IPC cost)
data = {
'sensor_id': sensor_id,
'fps': fps,
'flow': flow.copy() if flow is not None else None,
'img': img.copy() if img is not None else None,
'contour': contour,
'centroid': centroid,
'depth_map': depth_map.copy() if depth_map is not None else None,
'timestamp': time.time()
}
else:
# Lightweight payload for stable multi-process throughput
fps, = sensor.read_info(sensor.info.FPS)
data = {
'sensor_id': sensor_id,
'fps': fps,
'timestamp': time.time()
}
# 步骤6: 发送数据到主进程（非阻塞）
try:
result_queue.put_nowait(data)
except:
# Queue full: keep latest data (drop oldest)
try:
result_queue.get_nowait()
result_queue.put_nowait(data)
except:
pass
# Pace workers to a common target so sensor FPS stays similar.
# Use drift-free scheduler: next_tick increments by fixed period.
if target_period > 0:
next_tick += target_period
now = time.perf_counter()
sleep_time = next_tick - now
if sleep_time > 0:
time.sleep(sleep_time)
else:
# If we are behind, reset schedule to current time.
next_tick = now
except Exception as e:
print(f"传感器 {sensor_id} 处理错误: {e}")
time.sleep(0.1)
# 清理资源
sensor.disconnect()
if timer_handle is not None:
_disable_high_res_sleep(timer_handle)
print(f"传感器 {sensor_id} 已停止")
except Exception as e:
if timer_handle is not None:
_disable_high_res_sleep(timer_handle)
print(f"传感器 {sensor_id} 初始化失败: {e}")
result_queue.put({'sensor_id': sensor_id, 'error': str(e)})
def main():
# =========================================================================
# 1. 配置多个传感器
# =========================================================================
num_sensors = 5
sensor_configs = [
{
'sensor_id': 0,
'vid_src': './s1_720p.mp4',
'config_name': './config/prototype_of3_720p.json',
'cal_path': './config/dd01a_0_720p.npy'
},
{
'sensor_id': 1,
'vid_src': './s1_720p.mp4',
'config_name': './config/prototype_of3_720p.json',
'cal_path': './config/dd01a_0_720p.npy'
},
{
'sensor_id': 2,
'vid_src': './s1_720p.mp4',
'config_name': './config/prototype_of3_720p.json',
'cal_path': './config/dd01a_0_720p.npy'
},
{
'sensor_id': 3,
'vid_src': './s1_720p.mp4',
'config_name': './config/prototype_of3_720p.json',
'cal_path': './config/dd01a_0_720p.npy'
},
{
'sensor_id': 4,
'vid_src': './s1_720p.mp4',
'config_name': './config/prototype_of3_720p.json',
'cal_path': './config/dd01a_0_720p.npy'
},
]
# =========================================================================
# 2. 创建进程间通信对象
# =========================================================================
# Keep queue short to prevent backlog-induced jitter (latest-data policy)
result_queues = [Queue(maxsize=2) for _ in sensor_configs]
stop_events = [mp.Event() for _ in sensor_configs] # 停止事件
# =========================================================================
# 3. 启动传感器工作进程
# =========================================================================
processes = []
for i, config in enumerate(sensor_configs):
p = Process(
target=sensor_worker,
args=(i, config['vid_src'], config['config_name'], config['cal_path'], result_queues[i], stop_events[i])
)
p.start()
processes.append(p)
print(f"\n已启动 {len(processes)} 个传感器进程")
print("按 'q' 退出程序")
# 用于存储每个传感器的最新数据
latest_data = [None] * len(sensor_configs)
# 用于控制打印频率，减少I/O开销
print_counter = 0
print_interval = 30 # 每30帧打印一次，进一步减少I/O
# 用于控制按键检查频率
key_check_counter = 0
key_check_interval = 5 # 每5帧检查一次按键
try:
# =====================================================================
# 4. 主循环：从队列获取数据并显示
# =====================================================================
while True:
# 从所有队列中快速获取最新数据（高效清空队列，只保留最新帧）
for i, queue in enumerate(result_queues):
# 快速清空队列，只保留最后一帧
data = None
# 一次性获取所有可用数据，只保留最后一个
while True:
try:
data = queue.get_nowait()
except:
break
if data is not None:
latest_data[i] = data
# 显示所有传感器的数据
for data in latest_data:
if data is None or 'error' in data:
continue
sid = data['sensor_id']
# 显示形变场
if data.get('img') is not None and data.get('flow') is not None:
arrows = orisys.util.draw_arrows(
data['img'],
data['flow'],
threshold=5,
grid_spacing=10,
arrow_scale=1.0
)
cv2.imshow(f"Deformation {sid}", arrows)
# # 显示触碰点和区域
# if data['img'] is not None:
# image_with_foe = orisys.util.draw_contact(
# data['img'],
# data['contour'],
# data['centroid']
# )
# cv2.imshow(f"Contact Centroid {sid}", image_with_foe)
# # 显示深度图
# if data['depth_map'] is not None and data['depth_map'].size > 0:
# div_abs = data['depth_map']
# if div_abs.max() > div_abs.min():
# div_normalized = ((div_abs - div_abs.min()) /
# (div_abs.max() - div_abs.min()) * 255).astype(np.uint8)
# else:
# div_normalized = np.zeros_like(div_abs, dtype=np.uint8)
# cv2.imshow(f"Depth Map {sid}",
# cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET))
# 降低打印频率以减少I/O开销（在主循环级别控制）
print_counter += 1
if print_counter >= print_interval:
for data in latest_data:
if data is not None and 'error' not in data:
sid = data['sensor_id']
print(f"传感器 {sid}: FPS={data['fps']:5.1f}")
print_counter = 0
# 降低按键检查频率以减少开销
key_check_counter += 1
if key_check_counter >= key_check_interval:
key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
print("\n退出程序...")
break
key_check_counter = 0
finally:
# =====================================================================
# 5. 清理资源
# =====================================================================
# 通知所有工作进程停止
for event in stop_events:
event.set()
# 等待所有进程结束
for p in processes:
p.join(timeout=2)
if p.is_alive():
p.terminate()
p.join()
cv2.destroyAllWindows()
print("所有传感器进程已停止")
if __name__ == '__main__':
# Windows 下多进程必需的初始化
mp.freeze_support()
main()
patch_generation/color_pattern.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 1.8 MiB

patch_generation/raw_1.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 344 KiB

patch_generation/red_pattern.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 912 KiB

patch_generation/test_hue_gray.py
-224
View File
@@ -1,224 +0,0 @@
import cv2
import numpy as np
import matplotlib.pyplot as plt
def calculate_local_contrast(image_input, window_size=5):
"""
Calculate Global Average Local Contrast (for fine texture detail evaluation in optical flow).
Args:
image_input: str (image file path) or numpy array (grayscale/RGB image).
window_size: int (odd number, default=5) — sliding window size for fine texture analysis.
Returns:
float: Global average local contrast (higher = better fine texture detail; target ≥0.2 for optical flow).
Raises:
ValueError: If image is invalid, window size is even/too small, or processing fails.
"""
# ---------------------- Step 1: Input Validation & Preprocessing ----------------------
# Validate window size (must be odd, ≥3 to capture meaningful texture)
if window_size % 2 == 0 or window_size < 3:
raise ValueError(f"Window size must be an odd integer ≥3 (got {window_size}). Use 3/5/7.")
# Load/convert image to grayscale (consistent with optical flow workflow)
if isinstance(image_input, str):
img = cv2.imread(image_input)
if img is None:
raise ValueError(f"Failed to load image from path: {image_input}")
elif isinstance(image_input, np.ndarray):
img = image_input.copy()
else:
raise ValueError("Input must be an image file path (str) or numpy array (image).")
# Convert to grayscale (handle RGB/BGR)
if len(img.shape) == 3:
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
else:
gray = img # Already grayscale
# Normalize to 0-255 (avoid non-standard pixel ranges, e.g., float32)
gray = cv2.normalize(gray, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
# ---------------------- Step 2: Sliding Window Local Contrast Calculation ----------------------
# Define padding (replicate edge pixels to avoid boundary distortion)
pad_size = window_size // 2
gray_padded = cv2.copyMakeBorder(
gray,
top=pad_size, bottom=pad_size, left=pad_size, right=pad_size,
borderType=cv2.BORDER_REPLICATE
)
# Initialize array to store local contrast for each window
local_contrast_vals = []
# Iterate over every pixel (center of the sliding window)
for y in range(pad_size, gray_padded.shape[0] - pad_size):
for x in range(pad_size, gray_padded.shape[1] - pad_size):
# Extract 5x5 (or custom) window around (x,y)
window = gray_padded[y-pad_size : y+pad_size+1, x-pad_size : x+pad_size+1]
# Compute Michelson contrast for the window (avoid division by zero)
window_max = np.max(window)
window_min = np.min(window)
if window_max + window_min == 0: # All pixels are 0 (black window)
lc = 0.0
else:
lc = (window_max - window_min) / (window_max + window_min)
local_contrast_vals.append(lc)
# ---------------------- Step 3: Global Average Local Contrast ----------------------
global_local_contrast = np.mean(local_contrast_vals) if local_contrast_vals else 0.0
# Round for readability (4 decimal places)
return global_local_contrast
def calculate_edge_strength_density(image_input):
"""
Calculate Edge Strength and Edge Density for optical flow evaluation.
Args:
image_input: Can be (1) str (image file path) or (2) numpy array (grayscale/RGB image).
Returns:
dict: Contains two metrics:
- edge_strength: Average gradient magnitude of valid edges (higher = sharper edges).
- edge_density: Percentage of valid edge pixels (higher = more texture/edges).
Raises:
ValueError: If image cannot be loaded/processed.
"""
# ---------------------- Step 1: Load/Preprocess Image ----------------------
img = image_input
# Convert to grayscale (handle RGB/BGR)
if len(img.shape) == 3:
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
else:
gray = img # Already grayscale
# Normalize to 0-255 (avoid float/non-standard pixel ranges)
gray = cv2.normalize(gray, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
# ---------------------- Step 2: Compute Gradient Magnitude (Sobel) ----------------------
# Sobel kernel size = 3 (balance between noise reduction and edge sharpness for optical flow)
sobel_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3) # Horizontal gradient
sobel_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3) # Vertical gradient
# Compute gradient magnitude (edge strength raw values)
gradient_magnitude = np.sqrt(sobel_x**2 + sobel_y**2)
# ---------------------- Step 3: Adaptive Threshold for Valid Edges ----------------------
# Use median gradient to filter noise (robust to outliers)
median_gradient = np.median(gradient_magnitude)
# Threshold = 1.5× median (filters weak noise, keeps meaningful edges for optical flow)
edge_threshold = 1.5 * median_gradient
# Create binary mask for valid edges (1 = valid edge, 0 = non-edge)
valid_edge_mask = (gradient_magnitude > edge_threshold).astype(np.uint8)
# ---------------------- Step 4: Calculate Metrics ----------------------
# Edge Strength: Average gradient magnitude of valid edges
valid_gradient_values = gradient_magnitude[valid_edge_mask == 1]
edge_strength = np.mean(valid_gradient_values) if len(valid_gradient_values) > 0 else 0.0
# Edge Density: Percentage of valid edge pixels (total pixels = height × width)
total_pixels = gray.shape[0] * gray.shape[1]
valid_edge_pixels = np.sum(valid_edge_mask)
edge_density = (valid_edge_pixels / total_pixels) * 100 if total_pixels > 0 else 0.0
# print(f"edge_strength: {edge_strength}, edge_density: {edge_density}")
# Round for readability (3 decimal places)
return edge_strength, edge_density
def custom_gray_conversion(img, r_weight=0.299, g_weight=0.587, b_weight=0.114,
red_boost=1.2, green_reduce=0.8):
"""
Custom grayscale calculation with tunable RGB weights + red/green contrast enhancement
Args:
image_path: Path to input image
r_weight: Weight for Red channel (default: 0.299)
g_weight: Weight for Green channel (default: 0.587)
b_weight: Weight for Blue channel (default: 0.114)
red_boost: Multiplier to boost red region grayscale (>1 = brighter)
green_reduce: Multiplier to reduce green region grayscale (<1 = darker)
Returns:
default_gray: Grayscale with standard weights
custom_gray: Grayscale with tunable weights + red/green enhancement
"""
# 1. Read image and validate
# img_bgr = cv2.imread(image_path)
img_bgr = img
if img_bgr is None:
raise ValueError("Failed to read image - check path!")
# 2. Convert BGR (OpenCV default) to RGB and split channels
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
# Convert to float32 to avoid integer overflow/precision loss
r = img_rgb[:, :, 0].astype(np.float32)
g = img_rgb[:, :, 1].astype(np.float32)
b = img_rgb[:, :, 2].astype(np.float32)
# 3. Calculate DEFAULT grayscale (standard weights)
default_gray = 0.299 * r + 0.587 * g + 0.114 * b
# 4. Calculate CUSTOM grayscale (tunable weights)
custom_gray_base = r_weight * r + g_weight * g + b_weight * b
# 5. Enhance red/green contrast (optional but effective)
# Define red/green regions using RGB value ratios (no HSV needed)
# red_mask = (r > g * 1.2) & (r > b * 1.2) & (r >= 40) # Red-dominant pixels
# green_mask = (g > r * 1.2) & (g > b * 1.2) & (g >= 40) # Green-dominant pixels
# Apply targeted adjustments to red/green regions
custom_gray = custom_gray_base.copy()
# custom_gray[red_mask] *= red_boost # Boost red region grayscale
# custom_gray[green_mask] *= green_reduce # Reduce green region grayscale
# 6. Normalize to 0-255 range and convert back to uint8 (OpenCV standard)
# Clip values first to avoid overflow, then scale if needed
default_gray = np.clip(default_gray, 0, 255).astype(np.uint8)
custom_gray = np.clip(custom_gray, 0, 255).astype(np.uint8)
return default_gray, custom_gray
# -------------------------- Test the Code --------------------------
if __name__ == "__main__":
# Replace with your image path
IMAGE_PATH = "./raw_1.jpg"
img = cv2.imread(IMAGE_PATH)
dx = 50
dy = 0
tile = img[410+dy:650+dy,940+dx:1180+dx,:]
# Example 1: Use standard weights (verify correctness)
default_gray, red_gray = custom_gray_conversion(
tile,
r_weight=1, g_weight=0, b_weight=0
)
default_gray, green_gray = custom_gray_conversion(
tile,
r_weight=0, g_weight=1, b_weight=0
)
default_gray, blue_gray = custom_gray_conversion(
tile,
r_weight=0, g_weight=0, b_weight=1
)
edge_strength, edge_density = calculate_edge_strength_density(default_gray)
local_contrast = calculate_local_contrast(default_gray)
print(f"标准灰度: 边缘强度: {edge_strength}, 边缘密度: {edge_density}, 局部对比度: {local_contrast}")
edge_strength, edge_density = calculate_edge_strength_density(red_gray)
local_contrast = calculate_local_contrast(red_gray)
print(f"红色灰度: 边缘强度: {edge_strength}, 边缘密度: {edge_density}, 局部对比度: {local_contrast}")
edge_strength, edge_density = calculate_edge_strength_density(green_gray)
local_contrast = calculate_local_contrast(green_gray)
print(f"绿色灰度: 边缘强度: {edge_strength}, 边缘密度: {edge_density}, 局部对比度: {local_contrast}")
edge_strength, edge_density = calculate_edge_strength_density(blue_gray)
local_contrast = calculate_local_contrast(blue_gray)
print(f"蓝色灰度: 边缘强度: {edge_strength}, 边缘密度: {edge_density}, 局部对比度: {local_contrast}")
plt.subplot(2, 2, 1)
plt.imshow(default_gray, cmap='gray')
plt.subplot(2, 2, 2)
plt.imshow(red_gray, cmap='gray')
plt.subplot(2, 2, 3)
plt.imshow(green_gray, cmap='gray')
plt.subplot(2, 2, 4)
plt.imshow(blue_gray, cmap='gray')
plt.show()
patch_generation/yellow_pattern.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 736 KiB

benchmark.py → scripts/benchmark.py
View File
scripts/build_nuitka.py
+228
View File
@@ -0,0 +1,228 @@
import os
import sys
import shutil
import subprocess
from pathlib import Path
from importlib.machinery import EXTENSION_SUFFIXES
INCLUDE_LAYOUT = {
"configs": {
"source_dir": Path("config"),
"target_subdir": Path("configs"),
"patterns": ["*.json"],
},
"examples": {
"source_dir": Path("examples"),
"target_subdir": Path("examples"),
"files": [
Path("examples/basic_pipeline.py"),
Path("examples/two_sensor_pipeline.py"),
Path("examples/plot_data.py"),
],
},
"docs": {
"source_dir": Path("docs"),
"target_subdir": Path("docs"),
"patterns": ["*.md", "*.pdf"],
},
}
def install_nuitka():
print("Checking Nuitka installation...")
try:
import nuitka # noqa: F401
print("Nuitka is installed.")
except ImportError:
print("Installing Nuitka...")
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "nuitka", "ordered-set", "zstandard"]
)
def get_version() -> str:
init_file = Path("src/orisys/__init__.py")
if not init_file.exists():
return "0.3.1"
for line in init_file.read_text(encoding="utf-8").splitlines():
if line.startswith("__version__"):
return line.split("=", 1)[1].strip().strip('"').strip("'")
return "0.3.1"
def find_compiled_extension(output_dir: Path, module_name: str) -> Path:
candidates = []
for suffix in EXTENSION_SUFFIXES:
candidates.extend(output_dir.glob(f"{module_name}*{suffix}"))
if not candidates:
all_files = ", ".join(p.name for p in output_dir.iterdir()) if output_dir.exists() else "(none)"
raise RuntimeError(
f"Nuitka did not produce a compiled extension for {module_name}. Found: {all_files}"
)
return sorted(candidates, key=lambda p: len(p.name))[0]
def move_wheel_to_sdk(wheel_path: Path, version: str) -> Path:
sdk_version_dir = Path("SDK") / version
sdk_version_dir.mkdir(parents=True, exist_ok=True)
dest_wheel = sdk_version_dir / wheel_path.name
if dest_wheel.exists():
dest_wheel.unlink()
shutil.move(str(wheel_path), str(dest_wheel))
return dest_wheel
def compile_orisys(build_root: Path, env: dict[str, str]) -> Path:
module_output_dir = build_root / "orisys"
if module_output_dir.exists():
shutil.rmtree(module_output_dir)
module_output_dir.mkdir(parents=True)
cmd = [
sys.executable,
"-m",
"nuitka",
"--module",
"--include-package=orisys",
"--output-dir=" + str(module_output_dir),
"--remove-output",
"--no-pyi-file",
"src/orisys",
]
print(f"Running: {' '.join(cmd)}")
subprocess.check_call(cmd, env=env)
ext_file = find_compiled_extension(module_output_dir, "orisys")
print(f"Generated compiled extension for orisys: {ext_file}")
return ext_file
def copy_compiled_orisys(pkg_dir: Path, ext_file: Path) -> str:
ext_suffix = "".join(ext_file.suffixes)
dest_ext = pkg_dir / f"__init__{ext_suffix}"
shutil.copy(ext_file, dest_ext)
print(f"Copied compiled module orisys to {dest_ext}")
return ext_suffix
def copy_included_files(pkg_dir: Path) -> dict[str, int]:
copied_counts = {}
for group_name, group_cfg in INCLUDE_LAYOUT.items():
src_dir = group_cfg["source_dir"]
dst_dir = pkg_dir / group_cfg["target_subdir"]
dst_dir.mkdir(exist_ok=True)
copied = 0
if "files" in group_cfg:
for file_path in group_cfg["files"]:
if file_path.exists():
shutil.copy(file_path, dst_dir)
copied += 1
print(f"Copied {group_name}: {file_path.name}")
else:
for pattern in group_cfg.get("patterns", []):
for src_file in src_dir.glob(pattern):
shutil.copy(src_file, dst_dir)
copied += 1
print(f"Copied {group_name}: {src_file.name}")
if group_name == "examples":
(dst_dir / "__init__.py").write_text("# Examples for orisys package\n", encoding="utf-8")
copied_counts[group_name] = copied
return copied_counts
def get_package_data_entries(ext_suffix: str) -> dict[str, list[str]]:
package_data = {
"orisys": [f"*{ext_suffix}"],
"orisys.examples": ["*.py"],
}
for group_name, group_cfg in INCLUDE_LAYOUT.items():
if group_name == "examples":
continue
target_subdir = str(group_cfg["target_subdir"]).replace("\\", "/")
for pattern in group_cfg.get("patterns", []):
package_data["orisys"].append(f"{target_subdir}/{pattern}")
return package_data
def build():
install_nuitka()
version = get_version()
output_dir = Path("build_nuitka")
dist_dir = Path("dist_nuitka")
if output_dir.exists():
shutil.rmtree(output_dir)
if dist_dir.exists():
shutil.rmtree(dist_dir)
output_dir.mkdir()
dist_dir.mkdir()
print("\n=== Compiling with Nuitka ===")
env = os.environ.copy()
env["PYTHONPATH"] = "src" + os.pathsep + env.get("PYTHONPATH", "")
ext_file = compile_orisys(output_dir, env)
print("\n=== Preparing Wheel Structure ===")
wheel_root = dist_dir / "wheel_root"
pkg_dir = wheel_root / "orisys"
pkg_dir.mkdir(parents=True)
ext_suffix = copy_compiled_orisys(pkg_dir, ext_file)
copied_counts = copy_included_files(pkg_dir)
print(f"Copied {copied_counts['configs']} JSON config files")
print(f"Copied {copied_counts['examples']} example files to orisys/examples/")
print(f"Copied {copied_counts['docs']} doc files to orisys/docs/")
package_data = get_package_data_entries(ext_suffix)
setup_py_path = dist_dir / "setup.py"
setup_content = f"""
from setuptools import setup, Distribution
class BinaryDistribution(Distribution):
def has_ext_modules(self):
return True
setup(
name="orisys",
version="{version}",
description="Orisys - Optical Tactile Sensor SDK (Nuitka Compiled)",
author="Orisys",
packages=['orisys', 'orisys.examples'],
package_dir={{'': 'wheel_root'}},
package_data={package_data},
include_package_data=True,
install_requires=[
"numpy>=1.24.4",
"opencv-python>=4.10.0.84",
"scipy>=1.10.1",
],
extras_require={{
"cuda12": ["cupy-cuda12x>=12.3.0"],
"cuda13": ["cupy-cuda13x>=13.0.0"],
"cuda": ["cupy-cuda12x>=12.3.0"],
}},
distclass=BinaryDistribution,
zip_safe=False,
)
"""
setup_py_path.write_text(setup_content, encoding="utf-8")
print("\n=== Building Wheel ===")
subprocess.check_call([sys.executable, "setup.py", "bdist_wheel"], cwd=dist_dir)
wheels = sorted((dist_dir / "dist").glob("*.whl"))
if wheels:
final_wheel = move_wheel_to_sdk(wheels[0], version)
print(f"\n[SUCCESS] Wheel created and moved to: {final_wheel}")
else:
print("\n[ERROR] Wheel not found.")
if __name__ == "__main__":
build()
calibration_script.py → scripts/calibration_script.py
View File
scripts/demo_wuxi.py
+961
View File
@@ -0,0 +1,961 @@
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
# Force pyqtgraph to bind to PyQt5 when multiple Qt bindings are installed.
os.environ.setdefault("PYQTGRAPH_QT_LIB", "PyQt5")
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, QHBoxLayout, QVBoxLayout, QWidget, QSizePolicy, QShortcut, QApplication
from pyqtgraph.Qt import QtGui
# Final desired appearance on the mirrored 1872x1560 (6:5) display.
# The desktop is first shown on a 1920x1080 (16:9) screen, then stretched by
# the OS onto the 6:5 mirror. We therefore pre-distort the source layout so
# the mirrored output looks correct.
TARGET_SINGLE_PANEL_W = 936
TARGET_ASPECT_W = TARGET_SINGLE_PANEL_W * 2
TARGET_VIDEO_FRAME_H = 1053
TARGET_BOTTOM_PANEL_H = 507
TARGET_ASPECT_H_TOTAL = TARGET_VIDEO_FRAME_H + TARGET_BOTTOM_PANEL_H
TARGET_MIRROR_W = 1872
TARGET_MIRROR_H = 1560
SINGLE_PANEL_W = TARGET_SINGLE_PANEL_W
ASPECT_W = TARGET_ASPECT_W
BOTTOM_PANEL_H = TARGET_BOTTOM_PANEL_H
VIDEO_FRAME_W = TARGET_ASPECT_W
VIDEO_FRAME_H = TARGET_VIDEO_FRAME_H
ASPECT_H_TOTAL = TARGET_ASPECT_H_TOTAL
UI_SCALE = 1.0
ARROW_STRETCH = 507
PLOT_STRETCH = 429
VIDEO_LOOP_PATH = "./叠动机械爪_成片.m4v"
npoints = 120 # 图表数据长度
dispaly_angle = -90 # 箭头图显示角度
touch_threshold = 5 # 触碰反应阈值
frame_to_recover = 10 # 触碰后等待回复帧数
# Opneing Hour for keep reading image
open_hour = 9
close_hour = 17
manual_open_enable = True # 控制是否可以手动切换营业状态 （按键G)
is_open = True
# 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 = False
def scale_ui_px(px, floor=1):
return max(floor, int(round(px * UI_SCALE)))
def compute_source_layout(screen_w, screen_h):
"""Pre-distort the app on the source screen so the mirrored monitor is correct."""
target_video_ratio = TARGET_VIDEO_FRAME_H / float(TARGET_ASPECT_H_TOTAL)
target_bottom_ratio = TARGET_BOTTOM_PANEL_H / float(TARGET_ASPECT_H_TOTAL)
return {
"window_w": max(320, int(screen_w)),
"window_h": max(240, int(screen_h)),
"video_w": max(320, int(screen_w)),
"video_h": max(1, int(round(screen_h * target_video_ratio))),
"bottom_h": max(1, int(round(screen_h * target_bottom_ratio))),
"single_panel_w": max(1, int(round(screen_w / 2.0))),
}
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(panel_widget):
"""
Pixel size of the square arrow canvas from the window client area and the
fixed horizontal split ARROW_STRETCH : PLOT_STRETCH (matches QHBoxLayout stretch).
"""
w = panel_widget.width()
h = panel_widget.height()
if w < 16 or h < 16:
return max(32, ASPECT_H_TOTAL)
w_arrow = int(round(w * ARROW_STRETCH / float(ARROW_STRETCH + PLOT_STRETCH)))
# Keep the rendered arrow image square inside the left panel.
return max(32, min(h, w_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 using source-screen pre-distortion."""
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):
pm = pm.scaled(lw, lh, Qt.IgnoreAspectRatio, Qt.SmoothTransformation)
label.setPixmap(pm)
def set_video_label_bgr(label, bgr_image):
"""Show loop video in the top frame using the fitted display size."""
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):
pm = pm.scaled(lw, lh, Qt.IgnoreAspectRatio, Qt.FastTransformation)
label.setPixmap(pm)
def read_loop_video_frame_realtime(video_cap, video_clock):
"""Read loop-video frame paced to wall clock; skip frames when main loop is slow."""
if video_cap is None or not video_cap.isOpened():
return False, None
now = time.perf_counter()
fps = max(1.0, float(video_clock.get("fps", 30.0)))
frame_dt = 1.0 / fps
next_ts = video_clock.get("next_ts")
if next_ts is None:
next_ts = now
frames_due = 1
if now > next_ts + frame_dt:
frames_due += int((now - next_ts) / frame_dt)
# Cap skip burst to avoid long catch-up loops after pauses.
frames_due = min(frames_due, int(max(2.0, fps)))
for _ in range(frames_due - 1):
if not video_cap.grab():
video_cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
if not video_cap.grab():
return False, None
vret, vframe = video_cap.read()
if not vret:
video_cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
vret, vframe = video_cap.read()
if not vret:
return False, None
next_ts = next_ts + frames_due * frame_dt
if next_ts < now - 2.0 * frame_dt:
next_ts = now
video_clock["next_ts"] = next_ts
return True, vframe
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", scale_ui_px(11, floor=8))
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]))
def build_sensor_panel(parent_layout):
panel = QWidget()
panel_layout = QHBoxLayout(panel)
panel_layout.setContentsMargins(0, 0, 0, 0)
panel_layout.setSpacing(0)
# Ignored: do not let pixmap/content drive QLabel size — layout stretch sets arrow vs plot widths.
_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)
panel_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)
panel_layout.addWidget(plot_widget, stretch=PLOT_STRETCH)
disconnected_label = QLabel("已断开")
disconnected_label.setAlignment(Qt.AlignCenter)
disconnected_label.setStyleSheet(
f"background-color: #0B1016; color: #FF4D4D; font-size: {scale_ui_px(54, floor=24)}px; "
"font-weight: bold; letter-spacing: 2px; margin: 0; padding: 0; border: none;"
)
disconnected_label.hide()
panel_layout.addWidget(disconnected_label, stretch=SINGLE_PANEL_W)
parent_layout.addWidget(panel, stretch=1)
p1, p2, p3, plot_af, plot_tfx, plot_tfy = create_plot(plot_widget)
data = np.zeros([npoints, 2, 3])
for i in range(3):
data[:, 0, i] = np.linspace(0.0, npoints / 30.0, npoints)
return {
"panel": panel,
"arrow_label": arrow_label,
"plot_widget": plot_widget,
"disconnected_label": disconnected_label,
"p1": p1,
"p2": p2,
"p3": p3,
"plot_af": plot_af,
"plot_tfx": plot_tfx,
"plot_tfy": plot_tfy,
"data": data,
}
# 实时图表：单窗口，两套并排（左 sensor1，右 sensor2）
pg.setConfigOption("background", PLOT_THEME_BG)
pg.setConfigOption("foreground", PLOT_THEME_FG)
app = pg.mkQApp("Wuxi Demo")
screen = app.primaryScreen()
screen_geometry = screen.availableGeometry() if screen is not None else None
screen_w = screen_geometry.width() if screen_geometry is not None else 1920
screen_h = screen_geometry.height() if screen_geometry is not None else 1080
source_layout = compute_source_layout(screen_w, screen_h)
ASPECT_W = source_layout["window_w"]
ASPECT_H_TOTAL = source_layout["window_h"]
VIDEO_FRAME_W = source_layout["video_w"]
VIDEO_FRAME_H = source_layout["video_h"]
BOTTOM_PANEL_H = source_layout["bottom_h"]
SINGLE_PANEL_W = source_layout["single_panel_w"]
UI_SCALE = max(0.1, screen_h / float(TARGET_ASPECT_H_TOTAL))
combined_win = CombinedTactileWindow()
combined_win.setWindowTitle("Wuxi tactile — dual flow + forces")
main_layout = QVBoxLayout(combined_win)
main_layout.setContentsMargins(0, 0, 0, 0)
main_layout.setSpacing(0)
video_label = QLabel()
video_label.setAlignment(Qt.AlignCenter)
video_label.setStyleSheet("background-color: #000000; margin: 0; padding: 0; border: none;")
video_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
video_label.setMinimumSize(0, 0)
main_layout.addWidget(video_label, stretch=VIDEO_FRAME_H)
bottom_container = QWidget()
bottom_layout = QHBoxLayout(bottom_container)
bottom_layout.setContentsMargins(0, 0, 0, 0)
bottom_layout.setSpacing(0)
main_layout.addWidget(bottom_container, stretch=BOTTOM_PANEL_H)
panel1_ui = build_sensor_panel(bottom_layout)
panel2_ui = build_sensor_panel(bottom_layout)
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.showFullScreen()
closed_hour_overlay = QLabel(combined_win)
closed_hour_overlay.setAlignment(Qt.AlignCenter)
closed_hour_overlay.setStyleSheet(
f"background-color: {PLOT_THEME_BG}; color: #A0A0A0; "
f"font-size: {scale_ui_px(72, floor=32)}px; font-weight: bold; "
"margin: 0; padding: 0; border: none;"
)
closed_hour_overlay.setText("Closed hour\n非营业时间")
closed_hour_overlay.hide()
start_time = None # Will be set when sensor starts
fps = 30.0
def outside_opening_hours(now):
return now.hour < open_hour or now.hour >= close_hour
def main():
global start_time
parser = argparse.ArgumentParser()
parser.add_argument("--video1","-v1", type=str, default="0", help="视频文件路径")
parser.add_argument("--video2","-v2", type=str, 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
sensor2 = None
video_cap = None
try:
input1 = int(args.video1)
input2 = int(args.video2)
is_camera = True
except:
is_camera = False
try:
if is_camera:
sensor1 = orisys.Sensor(
input1,
config_name="./config/ddjx01_wuxi.json",
cal_path="./config/dd01a_0306.npy",
isstitch=True,
cuda=False,
verbose=False,
)
sensor2 = orisys.Sensor(
input2,
config_name="./config/ddjx01_wuxi.json",
cal_path="./config/dd01a_0306.npy",
isstitch=True,
cuda=False,
verbose=False,
)
else:
sensor1 = orisys.Sensor(
args.video1,
config_name="./config/ddjx01_wuxi.json",
cal_path="./config/dd01a_0306.npy",
isstitch=True,
cuda=False,
verbose=False,
)
sensor2 = orisys.Sensor(
args.video2,
config_name="./config/ddjx01_wuxi.json",
cal_path="./config/dd01a_0306.npy",
isstitch=True,
cuda=False,
verbose=False,
)
def _shortcut_reset_tracker():
sensor1.reset()
sensor2.reset()
print("传感器已重置 (快捷键 R)", flush=True)
def _shortcut_quit():
combined_win.close()
def _shortcut_toggle_open():
global is_open
is_open = not is_open
print(f"营业状态: {'营业' if is_open else '非营业'}", flush=True)
_sc_r = QShortcut(QKeySequence(Qt.Key_R), combined_win)
_sc_r.activated.connect(_shortcut_reset_tracker)
if manual_open_enable:
_sc_g = QShortcut(QKeySequence(Qt.Key_G), combined_win)
_sc_g.activated.connect(_shortcut_toggle_open)
_sc_q = QShortcut(QKeySequence(Qt.Key_Q), combined_win)
_sc_q.activated.connect(_shortcut_quit)
video_cap = cv2.VideoCapture(VIDEO_LOOP_PATH)
video_clock = {"fps": 30.0, "next_ts": None}
if not video_cap.isOpened():
print(f"[video] Failed to open loop video: {VIDEO_LOOP_PATH}", flush=True)
else:
src_fps = float(video_cap.get(cv2.CAP_PROP_FPS) or 0.0)
if np.isfinite(src_fps) and src_fps >= 1.0:
video_clock["fps"] = src_fps
start_time = time.time()
is_save = False
frame_count = 0
writer = None
sensor1_active = True
sensor2_active = True
# Local to main(): assignments inside the loop make this local; must exist
# before the first iteration that skips the closed branch.
start_open = False
while True:
now = datetime.now()
if outside_opening_hours(now) or not is_open:
start_open = False
closed_hour_overlay.setGeometry(0, 0, combined_win.width(), combined_win.height())
closed_hour_overlay.raise_()
closed_hour_overlay.show()
app.processEvents()
if combined_win.user_closed or not combined_win.isVisible():
print("\n窗口已关闭，退出程序")
if writer is not None:
writer.release()
writer = None
break
time.sleep(1)
continue
if not start_open:
start_open = True
sensor1.reset()
sensor2.reset()
print(f"传感器已重置: {datetime.now():%Y-%m-%d %H:%M:%S}", flush=True)
closed_hour_overlay.hide()
if video_cap is not None and video_cap.isOpened():
vret, vframe = read_loop_video_frame_realtime(video_cap, video_clock)
if vret:
if vframe.shape[1] != VIDEO_FRAME_W or vframe.shape[0] != VIDEO_FRAME_H:
vframe = cv2.resize(vframe, (VIDEO_FRAME_W, VIDEO_FRAME_H), interpolation=cv2.INTER_LINEAR)
set_video_label_bgr(video_label, vframe)
if sensor1_active:
img1_valid = sensor1.get_img()
if img1_valid is None:
panel1_ui["arrow_label"].hide()
panel1_ui["plot_widget"].hide()
panel1_ui["disconnected_label"].show()
sensor1_active = False
print(f"Sensor1 disconnected. Showing DISCONNECTED screen. {datetime.now():%Y-%m-%d %H:%M:%S}", flush=True)
else:
sensor1.compute_deformation(check_motion=True, tail_frames=frame_to_recover, threshold=touch_threshold)
_, fn, fx, fy, flow = sensor1.read_info(
sensor1.info.FPS,
sensor1.info.FNORMAL,
sensor1.info.FSHEARX,
sensor1.info.FSHEARY,
sensor1.info.VRAW,
)
_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]
lw, lh = panel1_ui["arrow_label"].width(), panel1_ui["arrow_label"].height()
if lw > 16 and lh > 16:
side = max(32, min(lw, lh))
else:
side = compute_arrow_canvas_side(panel1_ui["panel"])
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(panel1_ui["arrow_label"], arrows)
data1 = panel1_ui["data"]
for i in range(3):
data1[:-1, 0, i] = data1[1:, 0, i]
data1[:-1, 1, i] = data1[1:, 1, i]
data1[-1, 0, :] = frame_count / 15.0
data1[-1, 1, 0] = fn
data1[-1, 1, 1] = abs(fx)
data1[-1, 1, 2] = abs(fy)
update(
panel1_ui["plot_af"],
panel1_ui["plot_tfx"],
panel1_ui["plot_tfy"],
data1,
panel1_ui["p1"],
panel1_ui["p2"],
panel1_ui["p3"],
)
if sensor2_active:
img2_valid = sensor2.get_img()
if img2_valid is None:
panel2_ui["arrow_label"].hide()
panel2_ui["plot_widget"].hide()
panel2_ui["disconnected_label"].show()
sensor2_active = False
print(f"Sensor2 disconnected. Showing DISCONNECTED screen. {datetime.now():%Y-%m-%d %H:%M:%S}", flush=True)
else:
sensor2.compute_deformation(check_motion=True, tail_frames=frame_to_recover, threshold=touch_threshold)
_, fn2, fx2, fy2, flow2 = sensor2.read_info(
sensor2.info.FPS,
sensor2.info.FNORMAL,
sensor2.info.FSHEARX,
sensor2.info.FSHEARY,
sensor2.info.VRAW,
)
_bg_bgr2 = (22, 16, 11)
_img2 = sensor2.img
_h2, _w2 = _img2.shape[:2]
black_img2 = np.empty((_h2, _w2, 3), dtype=_img2.dtype)
black_img2[:, :, 0] = _bg_bgr2[0]
black_img2[:, :, 1] = _bg_bgr2[1]
black_img2[:, :, 2] = _bg_bgr2[2]
lw2, lh2 = panel2_ui["arrow_label"].width(), panel2_ui["arrow_label"].height()
if lw2 > 16 and lh2 > 16:
side2 = max(32, min(lw2, lh2))
else:
side2 = compute_arrow_canvas_side(panel2_ui["panel"])
arrows2 = draw_arrows_wuxi(
black_img2,
flow2,
threshold=2,
grid_spacing=20,
arrow_scale=0.5,
below_threshold_color=(160, 160, 160),
target_size=int(side2),
)
set_arrow_label_bgr(panel2_ui["arrow_label"], arrows2)
data2 = panel2_ui["data"]
for i in range(3):
data2[:-1, 0, i] = data2[1:, 0, i]
data2[:-1, 1, i] = data2[1:, 1, i]
data2[-1, 0, :] = frame_count / 15.0
data2[-1, 1, 0] = fn2
data2[-1, 1, 1] = abs(fx2)
data2[-1, 1, 2] = abs(fy2)
update(
panel2_ui["plot_af"],
panel2_ui["plot_tfx"],
panel2_ui["plot_tfy"],
data2,
panel2_ui["p1"],
panel2_ui["p2"],
panel2_ui["p3"],
)
app.processEvents()
now = datetime.now()
if now.hour == 9 and now.minute == 0 and now.second == 0:
if sensor1_active:
sensor1.reset()
if sensor2_active:
sensor2.reset()
print(
f"[session] 传感器已重置: {now:%Y-%m-%d %H:%M:%S} ",
flush=True,
)
if not sensor1_active and not sensor2_active:
while not combined_win.user_closed and combined_win.isVisible():
app.processEvents()
time.sleep(0.05)
break
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))
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} {datetime.now():%Y-%m-%d %H:%M:%S}", flush=True)
if 'sensor2' in locals() and sensor2 is not None:
try:
sensor2.disconnect()
except Exception as disconnect_err:
print(f"[session] disconnect error (sensor2): {disconnect_err} {datetime.now():%Y-%m-%d %H:%M:%S}", flush=True)
if video_cap is not None:
video_cap.release()
if __name__ == '__main__':
main()
promote_video.py → scripts/promote_video.py
+10 -3
View File
@@ -153,7 +153,9 @@ def main():
frame_count = 0 frame_count = 0
writer = None writer = None
while True: while True:
sensor1.get_img() # 提取并拼接复眼图像 isimg = sensor1.get_img() # 提取并拼接复眼图像
if isimg is None:
break
sensor1.compute_deformation() # 计算形变和分解 sensor1.compute_deformation() # 计算形变和分解
@@ -165,8 +167,13 @@ def main():
fps, fn, fx, fy, flow, depth_map = sensor1.read_info(sensor1.info.FPS, sensor1.info.FNORMAL, sensor1.info.FSHEARX, sensor1.info.FSHEARY, sensor1.info.VRAW, sensor1.info.DEPTH) fps, fn, fx, fy, flow, depth_map = sensor1.read_info(sensor1.info.FPS, sensor1.info.FNORMAL, sensor1.info.FSHEARX, sensor1.info.FSHEARY, sensor1.info.VRAW, sensor1.info.DEPTH)
print(fps,fn,fx,fy) print(fps,fn,fx,fy)
# 显示向量场 # 显示向量场
black_img = np.zeros_like(sensor1.img) _bg_bgr = (22, 16, 11)
arrows = orisys.util.draw_arrows(sensor1.img, flow, threshold=2, grid_spacing=20, arrow_scale=0.5) black_img = np.empty_like(sensor1.img)
black_img[:, :, 0] = _bg_bgr[0]
black_img[:, :, 1] = _bg_bgr[1]
black_img[:, :, 2] = _bg_bgr[2]
arrows = orisys.util.draw_arrows(black_img, flow, threshold=2, grid_spacing=20, arrow_scale=0.5,below_threshold_color=(160, 160, 160),edge_margin=15,colormap=cv2.COLORMAP_PARULA)
cv2.imshow("Deformation", arrows) cv2.imshow("Deformation", arrows)
# 显示深度图 # 显示深度图
scripts/test_ffmpeg.py
+159
View File
@@ -0,0 +1,159 @@
import argparse
import av
import cv2
import platform
import subprocess
import sys
import time
TARGET_FPS = 30
WIDTH = 1920
HEIGHT = 1080
def list_windows_dshow_devices():
"""
Print DirectShow devices using ffmpeg, if available.
"""
cmd = ["ffmpeg", "-hide_banner", "-list_devices", "true", "-f", "dshow", "-i", "dummy"]
print("Listing Windows DirectShow devices via ffmpeg...")
try:
result = subprocess.run(cmd, capture_output=True, text=True, check=False)
except FileNotFoundError as exc:
raise RuntimeError("ffmpeg was not found in PATH; cannot list DirectShow devices.") from exc
output = (result.stderr or "") + (result.stdout or "")
print(output.strip())
def open_pyav_capture(system_name, device, pixel_format):
"""
Open a camera device via PyAV on Linux or Windows.
"""
if system_name == "linux":
options = {
"video_size": f"{WIDTH}x{HEIGHT}",
"framerate": str(TARGET_FPS),
"input_format": pixel_format or "mjpeg",
}
container = av.open(device or "/dev/video4", format="v4l2", options=options)
elif system_name == "windows":
if not device:
raise RuntimeError(
"On Windows, PyAV/dshow requires a camera name, for example "
'--device "USB Camera". Use --list-devices first.'
)
options = {
"video_size": f"{WIDTH}x{HEIGHT}",
"framerate": str(TARGET_FPS),
}
if pixel_format:
options["pixel_format"] = pixel_format
container = av.open(f"video={device}", format="dshow", options=options)
else:
raise RuntimeError(
f"Unsupported OS: {platform.system()}. Only Linux and Windows are supported."
)
stream = container.streams.video[0]
return container, stream
def main():
parser = argparse.ArgumentParser(
description="Capture from camera via PyAV on Linux or Windows, optionally save to MP4."
)
parser.add_argument(
"-d",
"--device",
type=str,
default="",
help=(
"Linux: V4L2 device path like /dev/video4. "
"Windows: DirectShow camera name like 'USB Camera'. "
"If empty, Linux uses /dev/video4; Windows requires --device."
),
)
parser.add_argument(
"--pixel-format",
type=str,
default="",
help=(
"Optional input pixel format / format hint. "
"Examples: Linux 'mjpeg'; Windows may accept values such as 'mjpeg' if supported."
),
)
parser.add_argument(
"--list-devices",
action="store_true",
help="On Windows, list DirectShow device names via ffmpeg and exit.",
)
parser.add_argument(
"-o",
"--output",
type=str,
default="",
help="Output MP4 path (e.g. capture.mp4). If empty, only preview, no file.",
)
args = parser.parse_args()
system_name = platform.system().lower()
if args.list_devices:
if system_name != "windows":
raise RuntimeError("--list-devices is currently only implemented for Windows.")
list_windows_dshow_devices()
return
container, stream = open_pyav_capture(system_name, args.device, args.pixel_format)
print(
f"Using PyAV backend on {platform.system()} "
f"with device={args.device or '/dev/video4'}"
)
frame_count = 0
t_start = time.perf_counter()
writer = None
if args.output:
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
writer = cv2.VideoWriter(
args.output, fourcc, float(TARGET_FPS), (WIDTH, HEIGHT)
)
if not writer.isOpened():
container.close()
raise RuntimeError(f"Could not open VideoWriter for {args.output}")
print(f"Recording to {args.output} @ {TARGET_FPS} fps, {WIDTH}x{HEIGHT}")
try:
for frame in container.decode(stream):
t0 = time.perf_counter()
img = frame.to_ndarray(format="bgr24") # OpenCV-compatible
t1 = time.perf_counter()
if writer is not None:
writer.write(img)
cv2.imshow("pyav_cam", img)
key = cv2.waitKey(1) & 0xFF
t2 = time.perf_counter()
frame_count += 1
if frame_count % 30 == 0:
elapsed = t2 - t_start
print(
f"Frame {frame_count}: convert={1000*(t1-t0):.1f}ms, "
f"imshow+waitKey={1000*(t2-t1):.1f}ms, avg_fps={frame_count/elapsed:.1f}"
)
if key == ord("q"):
break
finally:
container.close()
if writer is not None:
writer.release()
print(f"Saved {frame_count} frames to {args.output}")
cv2.destroyAllWindows()
if __name__ == "__main__":
main()
setup.py
+7 -8
View File
@@ -1,11 +1,12 @@
from setuptools import setup, find_packages from setuptools import setup, find_packages
from pathlib import Path from pathlib import Path
# Find all JSON files in the configs directory config_dir = Path("config")
configs_dir = Path("src/orisys/configs") data_files = []
json_files = [] if config_dir.exists():
if configs_dir.exists(): json_files = [str(p) for p in sorted(config_dir.glob("*.json"))]
json_files = [str(p.relative_to("src")) for p in configs_dir.glob("*.json")] if json_files:
data_files.append(("orisys/configs", json_files))
setup( setup(
name="orisys", name="orisys",
@@ -14,10 +15,8 @@ setup(
author="Your Name", author="Your Name",
package_dir={"": "src"}, package_dir={"": "src"},
packages=find_packages(where="src"), packages=find_packages(where="src"),
package_data={
"orisys": ["configs/*.json"],
},
include_package_data=True, include_package_data=True,
data_files=data_files,
install_requires=[ install_requires=[
"numpy>=1.24.4", "numpy>=1.24.4",
"opencv-python>=4.10.0.84", "opencv-python>=4.10.0.84",
src/orisys/HHD.py
-2187
View File
File diff suppressed because it is too large Load Diff
src/orisys/configs.py
+126
View File
@@ -0,0 +1,126 @@
import json
from pathlib import Path
def _candidate_config_dirs() -> list[Path]:
"""
Candidate config directories for editable installs and packaged wheels.
"""
here = Path(__file__).resolve().parent
candidates = [
here / "configs", # packaged wheel layout: orisys/configs/*.json
here.parent.parent / "config", # editable install from repo: repo_root/config/*.json
here.parent / "config", # fallback
Path.cwd() / "config", # script launched from repo root
]
unique: list[Path] = []
for path in candidates:
if path not in unique:
unique.append(path)
return unique
def _get_config_dir() -> Path:
for candidate in _candidate_config_dirs():
if candidate.exists() and any(candidate.glob("*.json")):
return candidate
# Default to the first candidate so error messages are stable.
return _candidate_config_dirs()[0]
CONFIG_DIR = _get_config_dir()
def _config_files() -> list[Path]:
if not CONFIG_DIR.exists():
return []
return sorted(CONFIG_DIR.glob("*.json"))
def _load_json(path: Path) -> dict:
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def load_config(config_name_or_path):
"""
Load sensor configuration from a JSON file path or config name.
Config lookup order:
1. Exact file path (absolute or relative)
2. File stem under the config directory (e.g. ``ddjx01`` -> ``config/ddjx01.json``)
3. JSON field ``"name"`` matching the requested name (keeps ``prototype`` compatible)
"""
config_path = Path(config_name_or_path)
if config_path.exists() and config_path.is_file():
if config_path.suffix.lower() != ".json":
raise ValueError(f"Config file must be a JSON file: {config_path}")
return _load_json(config_path)
if config_path.is_absolute():
raise FileNotFoundError(f"Config file not found at absolute path: {config_path}")
direct = CONFIG_DIR / f"{config_name_or_path}.json"
if direct.exists():
return _load_json(direct)
# Backward-compatible alias resolution: support old names via JSON "name" field.
for candidate in _config_files():
try:
data = _load_json(candidate)
except Exception:
continue
if data.get("name") == config_name_or_path:
return data
available = [f.stem for f in _config_files()]
available_str = ", ".join(available) if available else "none"
raise FileNotFoundError(
f"Input config file not found with name or path: {config_name_or_path}\n"
f" Tried as file path: {config_path} (not found)\n"
f" Tried as config file: {direct} (not found)\n"
f" Available config files: {available_str}"
)
def list_config():
"""
List all available JSON configuration files by file stem.
"""
return [f.stem for f in _config_files()]
def export_config(config_name=None, path="./config/exported_config.json"):
"""
Export a configuration to a JSON file with readable formatting.
"""
config = load_config(config_name)
Path(path).parent.mkdir(parents=True, exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
formatted = _format_json_value(config)
f.write(formatted)
f.write("\n")
print(f"Config name: {config_name} exported to {path}")
return config
def _format_json_value(value, indent_level=0, indent_str=" "):
if isinstance(value, dict):
if not value:
return "{}"
items = []
for k, v in value.items():
key_str = json.dumps(k, ensure_ascii=False)
value_str = _format_json_value(v, indent_level + 1, indent_str)
items.append(f"{indent_str * (indent_level + 1)}{key_str}: {value_str}")
return "{\n" + ",\n".join(items) + "\n" + indent_str * indent_level + "}"
if isinstance(value, list):
if not value:
return "[]"
items = [json.dumps(item, ensure_ascii=False) for item in value]
return "[" + ", ".join(items) + "]"
return json.dumps(value, ensure_ascii=False)
src/orisys/configs/__init__.py
-172
View File
@@ -1,172 +0,0 @@
import json
import sys
from pathlib import Path
def _get_config_dir():
"""Resolve config directory, robustly handling Nuitka compiled paths."""
# 1. Start with the directory of this file
current_path = Path(__file__).parent
# 2. Check strict relative location (standard case)
# If we are in src/orisys/configs/__init__.py, configs is current_path
if (current_path / "prototype.json").exists():
return current_path
# 3. Walk up the directory tree to find a 'configs' folder containing 'prototype.json'
# This handles nested virtual paths like .../orisys/orisys/configs
# We look for a physical directory that matches our expectation
search_path = current_path
# Limit depth to avoid infinite loops, though unlikely with .parent
for _ in range(5):
# Check if 'configs' exists in this directory
candidate = search_path / "configs"
if candidate.exists() and candidate.is_dir():
if (candidate / "prototype.json").exists():
return candidate
# Also check if the current directory ITSELF is the configs dir we missed
# (e.g. if we are somehow inside configs but missed check 1)
if search_path.name == 'configs' and (search_path / "prototype.json").exists():
return search_path
if search_path.parent == search_path: # Reached root
break
search_path = search_path.parent
# 4. Try finding 'orisys' package root via sys.modules (if loaded)
try:
if 'orisys' in sys.modules:
mod = sys.modules['orisys']
if hasattr(mod, '__file__') and mod.__file__:
pkg_root = Path(mod.__file__).parent
# pkg_root is .../site-packages/orisys
candidate = pkg_root / "configs"
if candidate.exists() and (candidate / "prototype.json").exists():
return candidate
except Exception:
pass
# Fallback to current path if nothing else works
return current_path
CONFIG_DIR = _get_config_dir()
def load_config(config_name_or_path):
"""
Load sensor configuration from JSON file.
Accepts either:
1. A file path (absolute or relative) to a JSON config file
2. A config name (without .json extension) to load from package configs directory
Args:
config_name_or_path: Either:
- A file path string (e.g., "/path/to/config.json" or "./config.json")
- A config name string (e.g., "prototype") to load from package configs
Returns:
Dictionary containing the configuration
Raises:
FileNotFoundError: If the config file doesn't exist
ValueError: If the file path is not a JSON file
"""
# Convert to Path object for easier handling
config_path = Path(config_name_or_path)
# Check if it's an existing file (absolute or relative path)
if config_path.exists() and config_path.is_file():
# User provided a file path
if config_path.suffix.lower() != '.json':
raise ValueError(f"Config file must be a JSON file: {config_path}")
with open(config_path, 'r', encoding='utf-8') as f:
return json.load(f)
# If it's an absolute path but doesn't exist, raise error immediately
if config_path.is_absolute():
raise FileNotFoundError(f"Config file not found at absolute path: {config_path}")
# Otherwise, treat it as a config name and look in package directory
package_config_path = CONFIG_DIR / f"{config_name_or_path}.json"
if not package_config_path.exists():
# List available configs for better error message
available = [f.stem for f in CONFIG_DIR.glob("*.json")]
available_str = ", ".join(available) if available else "none"
raise FileNotFoundError(
f"Input config file not found with name or path: {config_name_or_path}\n"
f" Tried as file path: {config_path} (not found)\n"
f" Tried as package config: {package_config_path} (not found)\n"
f" Available package configs: {available_str}"
)
with open(package_config_path, 'r', encoding='utf-8') as f:
return json.load(f)
def list_config():
"""
List all available JSON configuration files in the configs directory.
Returns:
List of configuration names (without .json extension), sorted alphabetically.
Returns an empty list if no JSON files are found.
"""
config_files = sorted(CONFIG_DIR.glob("*.json"))
return [f.stem for f in config_files]
def export_config(config_name = None, path="./config/exported_config.json"):
"""
Export configuration to JSON file with readable formatting.
Arrays/lists are kept on a single line, while nested objects are indented.
Args:
config: Configuration dictionary to export
path: Output file path (default: "./config/exported_config.json")
Returns:
The exported config dictionary
"""
config = load_config(config_name)
Path(path).parent.mkdir(parents=True, exist_ok=True)
with open(path, 'w', encoding='utf-8') as f:
formatted = _format_json_value(config)
f.write(formatted)
f.write('\n') # Add trailing newline
print(f"Config name: {config_name} exported to {path}")
return config
def _format_json_value(value, indent_level=0, indent_str=" "):
"""
Recursively format JSON value with arrays/lists on single line.
Args:
value: The value to format (dict, list, or primitive)
indent_level: Current indentation level
indent_str: String used for indentation (default: 4 spaces)
Returns:
Formatted string representation
"""
if isinstance(value, dict):
if not value:
return "{}"
items = []
for k, v in value.items():
key_str = json.dumps(k, ensure_ascii=False)
value_str = _format_json_value(v, indent_level + 1, indent_str)
items.append(f"{indent_str * (indent_level + 1)}{key_str}: {value_str}")
return "{\n" + ",\n".join(items) + "\n" + indent_str * indent_level + "}"
elif isinstance(value, list):
if not value:
return "[]"
# Format arrays/lists on a single line
items = [json.dumps(item, ensure_ascii=False) for item in value]
return "[" + ", ".join(items) + "]"
else:
# Primitive types (string, number, bool, None)
return json.dumps(value, ensure_ascii=False)
src/orisys/configs/ddjx01.json
-43
View File
@@ -1,43 +0,0 @@
{
"name": "prototype",
"version": "1.0",
"camera": {
"resolution_raw": [1920, 1080],
"fps": 30,
"units": 15,
"units_used": 9,
"resolution_stitch": [365, 365],
"stitch_res_list": [365,263,123,63]
},
"stitcher": {
"min_area_frac": 0.01,
"max_area_frac": 0.04,
"max_aspect_ratio": 100,
"index": [1,2,3,6,7,8,11,12,13],
"margin": 10,
"nominal_ov_pos": 100,
"n": [3,3],
"row_tol": 0.6,
"box_size": [220, 220]
},
"tracker": {
"adaptive": true,
"lower_threshold": 12.0,
"upper_threshold": 18.0,
"loop_closure_interval": 1,
"algorithm": "dis",
"dis_preset": "medium",
"dis_mode": 3,
"algorithm_list": ["raft", "dis"]
},
"contact": {
"min_contour_area": 3000,
"flow_threshold":1,
"div_threshold":0.02,
"expansion_threshold": 0.01,
"kernel_size": [5,5],
"denoise_sigma": 15
}
}
src/orisys/configs/prototype.json
-39
View File
@@ -1,39 +0,0 @@
{
"name": "prototype",
"version": "1.0",
"camera": {
"resolution_raw": [1920, 1080],
"fps": 30,
"units": 15,
"units_used": 9,
"resolution_stitch": [400, 400]
},
"stitcher": {
"min_area_frac": 0.01,
"max_area_frac": 0.04,
"max_aspect_ratio": 100,
"index": [1,2,3,6,7,8,11,12,13],
"margin": 10,
"nominal_ov_pos": 100,
"n": [3,3],
"row_tol": 0.6
},
"tracker": {
"adaptive": true,
"lower_threshold": 12.0,
"upper_threshold": 18.0,
"loop_closure_interval": 1,
"dis_preset": "medium",
"finest_scale": 3
},
"contact": {
"min_contour_area": 3000,
"flow_threshold":1,
"div_threshold":0.02,
"expansion_threshold": 0.01,
"kernel_size": [5,5],
"denoise_sigma": 15
}
}
src/orisys/hhd/__init__.py
+20
View File
@@ -0,0 +1,20 @@
"""
HHD subsystem package.
This package is intended to hold the refactored implementation currently living in
`orisys.HHD`. The old monolithic file can remain as a compatibility layer while code
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
"""
from .api import nHHD
from .grid import StructuredGrid
from .poisson import PoissonSolver
from .timing import Timer
__all__ = ["nHHD", "StructuredGrid", "PoissonSolver", "Timer"]
src/orisys/hhd/api.py
+456
View File
@@ -0,0 +1,456 @@
import numpy
from ..backend_availability import (
CUPY_AVAILABLE,
OPENCL_AVAILABLE,
KERNEL_AVAILABLE,
cp,
_torque_kernel,
)
from .grid import StructuredGrid
from .poisson import PoissonSolver
class nHHD(object):
# constructor
def __init__(self, cuda=True, **kwargs):
# Select Poisson backend: 'auto' (default), 'cupy', 'opencl', or 'cpu'
backend = kwargs.get('poisson_backend', 'auto').lower()
if backend not in ('auto', 'cupy', 'opencl', 'cpu'):
raise ValueError("poisson_backend must be one of: 'auto', 'cupy', 'opencl', 'cpu'")
# Resolve 'auto' to a specific backend
if backend == 'auto':
backend = 'cupy' if CUPY_AVAILABLE else ('opencl' if OPENCL_AVAILABLE else 'cpu')
# Validate backend availability and fallback to 'cpu' if needed
if backend == 'cupy' and not CUPY_AVAILABLE:
backend = 'cpu'
elif backend == 'opencl' and not OPENCL_AVAILABLE:
backend = 'cpu'
self.debug = kwargs.get('debug', False)
self.poisson_backend = backend
# Initialize CUDA streams for cupy backend
if self.poisson_backend == 'cupy':
self.stream1 = cp.cuda.Stream()
self.stream2 = cp.cuda.Stream()
self.cuda = cuda
args = kwargs.keys()
if ('grid' in args) == ('points' in args):
raise SyntaxError("nHHD Solver needs either shape of a regular grid or the points in an rectilinear/unstructured grid")
if ('grid' in args) != ('spacings' in args):
raise SyntaxError("nHHD needs spacings for regular grid")
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.dim = self.mesh.dim
if (self.dim != 2) and (self.dim != 3):
raise ValueError("nHHD Solver works only for 2 and 3 dimensions")
self._force_cpu_cache = None
self.psolver.prepare(False)
def get_force(self, vfield, verbose=False):
if self.poisson_backend == 'cupy':
return self.get_force_cuda(vfield, verbose)
else:
return self.get_force_cpu(vfield, verbose)
def get_shear(self):
if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
return self.h
elif self.poisson_backend == 'cupy':
# Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps)
if hasattr(self, 'h') and self.h is not None:
return self.h
return cp.asnumpy(self._h_gpu)
def get_normal(self):
if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
return self.d
elif self.poisson_backend == 'cupy':
# Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps)
if hasattr(self, 'd') and self.d is not None:
return self.d
return cp.asnumpy(self._d_gpu)
def _solve_poisson_backend(self, field, verbose=False):
"""
Dispatch Poisson solve using the selected backend.
"""
if self.poisson_backend == 'opencl':
return self.psolver.solve_opencl(field, verbose, profile=self.debug)
if self.poisson_backend == 'cupy':
return self.psolver.solve_gpu(field, verbose)
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
if vfield.shape[-1] != self.dim:
raise ValueError("nHHD.decompose requires a valid-dimensional vector field")
if verbose > 0:
print('\n ==> nhhd.decompose', vfield.shape)
if self.dim == 2:
# compute div curl
(self.div, self.curlw) = self.mesh.divcurl(vfield)
if self.poisson_backend == 'opencl':
self.div = numpy.ascontiguousarray(self.div, dtype=numpy.float32)
self.curlw = numpy.ascontiguousarray(self.curlw, dtype=numpy.float32)
if verbose > 0:
print('div =', self.div.shape, self.div.min(), self.div.max())
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):
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)
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)
if verbose > 0:
print('d =', self.d.shape, self.d.min(), self.d.max())
self.r = self.mesh.rotated_gradient(self.nRu)
if verbose > 0:
print('r =', self.r.shape, self.r.min(), self.r.max())
else:
# compute div curl
(self.div, self.curlu, self.curlv, self.curlw) = self.mesh.divcurl(vfield)
if verbose > 0:
print('div =', self.div.shape, self.div.min(), self.div.max())
print('curlu =', self.curlu.shape, self.curlu.min(), self.curlu.max())
print('curlv =', self.curlv.shape, self.curlv.min(), self.curlv.max())
print('curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# compute potentials
self.nD = self._solve_poisson_backend(self.div, verbose > 0)
if verbose > 0:
print('nD =', self.nD.shape, self.nD.min(), self.nD.max())
self.nRu = self._solve_poisson_backend(self.curlu, verbose > 0)
self.nRu *= -1.0
if verbose > 0:
print('nRu =', self.nRu.shape, self.nRu.min(), self.nRu.max())
self.nRv = self._solve_poisson_backend(self.curlv, verbose > 0)
self.nRv *= -1.0
if verbose > 0:
print('nRv =', self.nRv.shape, self.nRv.min(), self.nRv.max())
self.nRw = self._solve_poisson_backend(self.curlw, verbose > 0)
self.nRw *= -1.0
if verbose > 0:
print('nRw =', self.nRw.shape, self.nRw.min(), self.nRw.max())
# compute divergent field as gradient of D potential
self.d = self.mesh.gradient(self.nD)
if verbose > 0:
print('d =', self.d.shape, self.d.min(), self.d.max())
# temporary use of r
self.r = numpy.stack((self.nRu, self.nRv, self.nRw), axis=-1)
# compute rotational field as curl of R potential
(cu, cv, cw) = self.mesh.curl3D(self.r)
self.r = numpy.stack((cu, cv, cw), axis=-1)
if verbose > 0:
print('r =', self.r.shape, self.r.min(), self.r.max())
self.h = numpy.add(self.d, self.r)
numpy.subtract(vfield, self.h, self.h)
if verbose > 0:
print('h =', self.h.shape, self.h.min(), self.h.max())
def decompose_cuda(self, vfield, verbose=0):
"""
GPU-optimized decomposition that keeps intermediate results on GPU.
Returns GPU arrays to avoid CPU-GPU transfers.
Args:
vfield: Vector field (numpy array or CuPy array)
verbose: Verbosity level (default: 0)
Returns:
Tuple of GPU arrays: (d_gpu, h_gpu, r_gpu, nRu_gpu) for 2D fields
"""
if not CUPY_AVAILABLE:
if verbose > 0:
print("CuPy not available, falling back to CPU decompose")
self.decompose(vfield, verbose)
return None
if vfield.shape[-1] != self.dim:
raise ValueError("nHHD.decompose requires a valid-dimensional vector field")
if verbose > 0:
print('\n ==> nhhd.decompose_cuda', vfield.shape)
# Convert input to GPU if needed
if isinstance(vfield, numpy.ndarray):
vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
else:
vfield_gpu = vfield
if self.dim == 2:
# Compute div curl on GPU
dx0, dx1 = self.mesh.dx[0], self.mesh.dx[1]
dudy, dudx = cp.gradient(vfield_gpu[:, :, 0], dx0, dx1)
dvdy, dvdx = cp.gradient(vfield_gpu[:, :, 1], dx0, dx1)
div_gpu = cp.add(dudx, dvdy)
curlw_gpu = cp.subtract(dvdx, dudy)
if verbose > 0:
print('div =', div_gpu.shape, float(div_gpu.min()), float(div_gpu.max()))
print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max()))
# Use a single stream for both solves. Running concurrent FFT-based
# solves on separate streams can cause occasional long-run jitter.
nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
cp.cuda.get_current_stream().synchronize()
if verbose > 0:
print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max()))
print('nR =', nRu_gpu.shape, float(nRu_gpu.min()), float(nRu_gpu.max()))
# Compute gradients on GPU
ddy, ddx = cp.gradient(nD_gpu, dx0, dx1)
d_gpu = cp.stack((ddx, ddy), axis=-1)
if verbose > 0:
print('d =', d_gpu.shape, float(d_gpu.min()), float(d_gpu.max()))
# Compute rotated gradient on GPU
rdy, rdx = cp.gradient(nRu_gpu, dx0, dx1)
rdy = cp.multiply(rdy, -1.0)
r_gpu = cp.stack((rdy, rdx), axis=-1)
if verbose > 0:
print('r =', r_gpu.shape, float(r_gpu.min()), float(r_gpu.max()))
# Compute h on GPU
h_gpu = cp.add(d_gpu, r_gpu)
h_gpu = cp.subtract(vfield_gpu, h_gpu)
if verbose > 0:
print('h =', h_gpu.shape, float(h_gpu.min()), float(h_gpu.max()))
# Ensure all default-stream work (gradient, etc.) is complete before use
cp.cuda.get_current_stream().synchronize()
# Store GPU arrays
self._d_gpu = d_gpu
self._h_gpu = h_gpu
self._r_gpu = r_gpu
self._nRu_gpu = nRu_gpu
return (d_gpu, h_gpu, r_gpu, nRu_gpu)
else:
# For 3D, fall back to CPU for now
if verbose > 0:
print("3D GPU decompose not yet implemented, using CPU")
self.decompose(vfield, verbose)
return None
def get_force_cpu(self, vfield, verbose=False):
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']
self.decompose(vfield, verbose)
# Optimized: use abs instead of absolute (slightly faster)
self.Sn = numpy.sum(numpy.abs(self.d))
# 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])
# Optimized: compute cross products directly without creating intermediate arrays
# For 2D: cross(a, b) = a[0]*b[1] - a[1]*b[0]
# Extract components once for reuse
r_x = self.r[:, :, 0]
r_y = self.r[:, :, 1]
grid_x = grid[:, :, 0]
grid_y = grid[:, :, 1]
# Compute cross products directly (faster than numpy.cross for 2D)
# cross(r_p, r) where r_p = grid - nR_max_loc
max_y, max_x = nR_max_loc[0], nR_max_loc[1]
cross_p = (grid_x - max_x) * r_y - (grid_y - max_y) * r_x
# cross(r_n, r) where r_n = grid - nR_min_loc
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
self.Stor = numpy.sum(cross_p) + numpy.sum(cross_n)
# Optimized: pre-allocate output array with explicit dtype
output = numpy.array([self.Sn, self.Stx, self.Sty, self.Stor], dtype=numpy.float32)
# Optimized: combine threshold checks (short-circuit evaluation)
flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000)
return flag, output
def _init_force_cpu_cache(self, shape):
"""
Initialize cached CPU grid arrays for force computation.
Avoids per-frame mgrid/stack allocations and reduces jitter.
"""
mgrid = numpy.mgrid[0:shape[0], 0:shape[1]]
self._force_cpu_cache = {
'shape': shape,
'grid': numpy.stack((mgrid[0], mgrid[1]), axis=2),
}
def _init_force_gpu_cache(self, shape):
"""
Initialize cached GPU arrays for force computation.
Call once when shape is known to avoid repeated allocations.
Args:
shape: Tuple (height, width) of the vector field
"""
if not CUPY_AVAILABLE:
return
# Cache grid arrays (only need to create once per shape)
self._force_gpu_cache = {
'shape': shape,
'mgrid_0': cp.arange(shape[0], dtype=cp.float32)[:, None],
'mgrid_1': cp.arange(shape[1], dtype=cp.float32)[None, :],
}
# Pre-broadcast for efficiency
self._force_gpu_cache['mgrid_0_bc'] = cp.broadcast_to(
self._force_gpu_cache['mgrid_0'], shape
).copy() # .copy() to make contiguous
self._force_gpu_cache['mgrid_1_bc'] = cp.broadcast_to(
self._force_gpu_cache['mgrid_1'], shape
).copy()
def get_force_cuda(self, vfield, verbose=False):
"""
Further optimized GPU force computation using CuPy ElementwiseKernel.
Fuses all operations into a single kernel for maximum performance.
Args:
vfield: Vector field (numpy array)
verbose: Print debug information (default: False)
Returns:
tuple: (flag, output) where:
- flag: Boolean indicating if force thresholds are met
- output: Array [Sn, Stx, Sty, Stor] containing force components
"""
if not CUPY_AVAILABLE:
if verbose:
print("CuPy not available, falling back to CPU get_force")
return self.get_force(vfield, verbose)
# Decompose the vector field
self.decompose_cuda(vfield, verbose)
d_gpu = self._d_gpu
h_gpu = self._h_gpu
r_gpu = self._r_gpu
nRu_gpu = self._nRu_gpu
# Use reduction kernels for better performance
# Compute Sn as L1 sum directly on GPU (faster than linalg.norm for this case)
Sn = float(cp.sum(cp.abs(d_gpu)))
# Compute Stx, Sty
Stx = float(cp.sum(h_gpu[:, :, 0]))
Sty = float(cp.sum(h_gpu[:, :, 1]))
# Find extrema indices
nR_max_idx = int(cp.argmax(nRu_gpu))
nR_min_idx = int(cp.argmin(nRu_gpu))
shape = nRu_gpu.shape
max_0, max_1 = nR_max_idx // shape[1], nR_max_idx % shape[1]
min_0, min_1 = nR_min_idx // shape[1], nR_min_idx % shape[1]
# Initialize cache if needed
if not hasattr(self, '_force_gpu_cache') or self._force_gpu_cache is None or self._force_gpu_cache['shape'] != shape:
self._init_force_gpu_cache(shape)
mgrid_0 = self._force_gpu_cache['mgrid_0_bc']
mgrid_1 = self._force_gpu_cache['mgrid_1_bc']
r_x = r_gpu[:, :, 0]
r_y = r_gpu[:, :, 1]
if KERNEL_AVAILABLE and _torque_kernel is not None:
torque = _torque_kernel(
r_x, r_y, mgrid_0, mgrid_1,
float(max_0), float(max_1), float(min_0), float(min_1)
)
Stor = float(cp.sum(torque))
else:
Stor = float(cp.sum(
(mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
(mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
))
self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor
output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32)
# Keep stable CPU copies for get_normal()/get_shear() so callers get this frame's
# result instead of reading GPU memory that can change (avoids CuPy-only jumps)
self.d = cp.asnumpy(self._d_gpu)
self.h = cp.asnumpy(self._h_gpu)
flag = (Sn >= 300 and Stx >= 300 and Sty >= 300 and abs(Stor) >= 1000)
return flag, output
def clear_force_gpu_cache(self):
"""Clear the GPU cache for force computation to free memory."""
if hasattr(self, '_force_gpu_cache'):
self._force_gpu_cache = None
if CUPY_AVAILABLE:
cp.get_default_memory_pool().free_all_blocks()
src/orisys/hhd/grid.py
+385
View File
@@ -0,0 +1,385 @@
import numpy
import timeit
import sys
# import cv2
from scipy import signal, spatial
from ..backend_availability import (
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
cp, cp_signal, _torque_kernel,
cl, cla, cl_rfftn, cl_irfftn
)
from .timing import Timer
class StructuredGrid(object):
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
def __init__(self, **kwargs):
'''
kwargs:
grid: ndarray of grid dimensions (Y,X) or (Z,Y,X)
spacings: ndarray of grid spacings (dy, dx) or (dz, dy, dx)
verbose: verbosity level
'''
args = kwargs.keys()
if ('grid' not in args) or ('spacings' not in args):
raise SyntaxError("Dimensions and spacings of the grid are required")
verbose = 0
if 'verbose' in args:
verbose = kwargs['verbose']
self.dims = kwargs['grid']
self.dx = kwargs['spacings']
self.dim = len(self.dims)
if self.dim != 2 and self.dim != 3:
raise ValueError("StructuredGrid works for 2D and 3D only")
if self.dim != len(self.dx):
raise ValueError("Dimensions of spacings should match that of the grid")
# Initialize GPU cache
self._gpu_divcurl_cache = None
if verbose > 0:
print(' Initializing', self.dim, 'D structured grid...', end ='')
sys.stdout.flush()
mtimer = Timer()
print(' Done!', end ='')
mtimer.end()
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
def divcurl(self, vfield, verbose=False):
# 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):
raise ValueError("Dimensions of vector field should match that of the grid")
if verbose:
print(' Computing div and curl...', end ='')
sys.stdout.flush()
mtimer = Timer()
if self.dim == 2:
# self.dx = (dy,dx)
dudy, dudx = numpy.gradient(vfield[:,:,0], self.dx[0], self.dx[1])
dvdy, dvdx = numpy.gradient(vfield[:,:,1], self.dx[0], self.dx[1])
numpy.add(dudx, dvdy, dudx)
numpy.subtract(dvdx, dudy, dvdx)
if verbose:
print(' Done!', end ='')
mtimer.end()
return (dudx, dvdx)
elif self.dim == 3:
# self.dx = (dz,dy,dx)
dudz, dudy, dudx = numpy.gradient(vfield[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
dvdz, dvdy, dvdx = numpy.gradient(vfield[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
dwdz, dwdy, dwdx = numpy.gradient(vfield[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
numpy.add(dudx, dvdy, dudx)
numpy.add(dudx, dvdz, dudx)
numpy.subtract(dwdy, dvdz, dwdy)
numpy.subtract(dudz, dwdx, dudz)
numpy.subtract(dvdx, dudy, dvdx)
if verbose:
print(' Done!', end ='')
mtimer.end()
return (dudx, dwdy, dudz, dvdx)
# def divcurl_gpu(self, vfield, verbose=False):
# """
# GPU-accelerated computation of divergence and curl of vector field.
# Uses CuPy for faster gradient computation on GPU.
# """
# if not CUPY_AVAILABLE:
# if verbose:
# print('CuPy not available, falling back to CPU')
# return self.divcurl(vfield, verbose)
# if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
# raise ValueError("Dimensions of vector field should match that of the grid")
# if verbose:
# print(' Computing div and curl on GPU...', end ='')
# sys.stdout.flush()
# mtimer = Timer()
# if self.dim == 2:
# # Transfer to GPU
# vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
# # Compute gradients on GPU
# dudy, dudx = cp.gradient(vfield_gpu[:,:,0], self.dx[0], self.dx[1])
# dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], self.dx[0], self.dx[1])
# # Compute div and curl on GPU
# div_gpu = dudx + dvdy
# curl_gpu = dvdx - dudy
# # Transfer back to CPU
# div = cp.asnumpy(div_gpu)
# curl = cp.asnumpy(curl_gpu)
# if verbose:
# print(' Done!', end ='')
# mtimer.end()
# return (div, curl)
# elif self.dim == 3:
# # Transfer to GPU
# vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
# # Compute gradients on GPU
# dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
# dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
# dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
# # Compute div and curl components on GPU
# div_gpu = dudx + dvdy + dwdz
# curlu_gpu = dwdy - dvdz
# curlv_gpu = dudz - dwdx
# curlw_gpu = dvdx - dudy
# # Transfer back to CPU
# div = cp.asnumpy(div_gpu)
# curlu = cp.asnumpy(curlu_gpu)
# curlv = cp.asnumpy(curlv_gpu)
# curlw = cp.asnumpy(curlw_gpu)
# if verbose:
# print(' Done!', end ='')
# mtimer.end()
# return (div, curlu, curlv, curlw)
def curl3D(self, vfield, verbose=False):
#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):
raise ValueError("Dimensions of vector field should match that of the grid")
if self.dim != 3:
raise ValueError("curl3D works only for 2D")
if verbose:
print(' Computing curl...', end ='')
sys.stdout.flush()
mtimer = Timer()
# self.dx = (dz,dy,dx)
dudz, dudy, dudx = numpy.gradient(vfield[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
dvdz, dvdy, dvdx = numpy.gradient(vfield[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
dwdz, dwdy, dwdx = numpy.gradient(vfield[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
numpy.subtract(dwdy, dvdz, dwdy)
numpy.subtract(dudz, dwdx, dudz)
numpy.subtract(dvdx, dudy, dvdx)
if verbose:
print(' Done!', end ='')
mtimer.end()
return (dwdy, dudz, dvdx)
def curl3D_gpu(self, vfield, verbose=False):
"""
GPU-accelerated computation of curl for 3D vector field.
Uses CuPy for faster gradient computation on GPU.
"""
if not CUPY_AVAILABLE:
if verbose:
print('CuPy not available, falling back to CPU')
return self.curl3D(vfield, verbose)
if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
raise ValueError("Dimensions of vector field should match that of the grid")
if self.dim != 3:
raise ValueError("curl3D works only for 3D")
if verbose:
print(' Computing curl on GPU...', end ='')
sys.stdout.flush()
mtimer = Timer()
# Transfer to GPU
vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
# Compute gradients on GPU
dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
# Compute curl components on GPU
curl_u = dwdy - dvdz
curl_v = dudz - dwdx
curl_w = dvdx - dudy
# Transfer back to CPU
curlu = cp.asnumpy(curl_u)
curlv = cp.asnumpy(curl_v)
curlw = cp.asnumpy(curl_w)
if verbose:
print(' Done!', end ='')
mtimer.end()
return (curlu, curlv, curlw)
def rotated_gradient(self, sfield, verbose=False):
if (sfield.shape != self.dims):
#if (sfield.shape - self.dims).any():
raise ValueError("Dimensions of scalar field should match that of the grid")
if self.dim != 2:
raise ValueError("rotated_gradient works only for 2D")
if verbose:
print(' Computing rotated gradient...', end ='')
sys.stdout.flush()
mtimer = Timer()
ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1])
ddy *= -1.0
grad = numpy.stack((ddy, ddx), axis=-1)
if verbose:
print(' Done!', end ='')
mtimer.end()
return grad
def rotated_gradient_gpu(self, sfield, verbose=False):
"""
GPU-accelerated computation of rotated gradient for 2D scalar field.
Uses CuPy for faster gradient computation on GPU.
"""
if not CUPY_AVAILABLE:
if verbose:
print('CuPy not available, falling back to CPU')
return self.rotated_gradient(sfield, verbose)
if (sfield.shape != self.dims):
raise ValueError("Dimensions of scalar field should match that of the grid")
if self.dim != 2:
raise ValueError("rotated_gradient works only for 2D")
if verbose:
print(' Computing rotated gradient on GPU...', end ='')
sys.stdout.flush()
mtimer = Timer()
# Transfer to GPU
sfield_gpu = cp.asarray(sfield, dtype=cp.float32)
# Compute gradients on GPU
ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1])
ddy *= -1.0
# Stack and transfer back to CPU
grad_gpu = cp.stack((ddy, ddx), axis=-1)
grad = cp.asnumpy(grad_gpu)
if verbose:
print(' Done!', end ='')
mtimer.end()
return grad
def gradient(self, sfield, verbose=False):
if (sfield.shape != self.dims):
#if (sfield.shape - self.dims).any():
raise ValueError("Dimensions of scalar field should match that of the grid")
if verbose:
print(' Computing gradient...', end ='')
sys.stdout.flush()
mtimer = Timer()
if self.dim == 2:
# 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:
# self.dx = (dz,dy,dx)
ddz, ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1], self.dx[2])
grad = numpy.stack((ddx, ddy, ddz), axis = -1)
if verbose:
print(' Done!', end ='')
mtimer.end()
return grad
def gradient_gpu(self, sfield, verbose=False):
"""
GPU-accelerated computation of gradient for scalar field.
Uses CuPy for faster gradient computation on GPU.
"""
if not CUPY_AVAILABLE:
if verbose:
print('CuPy not available, falling back to CPU')
return self.gradient(sfield, verbose)
if (sfield.shape != self.dims):
raise ValueError("Dimensions of scalar field should match that of the grid")
if verbose:
print(' Computing gradient on GPU...', end ='')
sys.stdout.flush()
mtimer = Timer()
# Transfer to GPU
sfield_gpu = cp.asarray(sfield, dtype=cp.float32)
if self.dim == 2:
# Compute gradients on GPU
ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1])
grad_gpu = cp.stack((ddx, ddy), axis=-1)
elif self.dim == 3:
# Compute gradients on GPU
ddz, ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1], self.dx[2])
grad_gpu = cp.stack((ddx, ddy, ddz), axis=-1)
# Transfer back to CPU
grad = cp.asnumpy(grad_gpu)
if verbose:
print(' Done!', end ='')
mtimer.end()
return grad
def clear_gpu_cache(self):
"""Clear GPU cache to free GPU memory"""
if self._gpu_divcurl_cache is not None:
self._gpu_divcurl_cache = None
if CUPY_AVAILABLE:
cp.get_default_memory_pool().free_all_blocks()
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
src/orisys/hhd/kernels.py
+16
View File
@@ -0,0 +1,16 @@
"""
Low-level kernels and shared numeric helpers for HHD.
Use this file for pieces that should not live inside the public `nHHD` API or the
`PoissonSolver` class, for example:
- finite-difference helper functions
- small backend-dispatch helpers
- reusable CuPy/OpenCL kernel preparation code
- shared CPU math that does not depend on object state
"""
# No standalone helpers from the old monolithic `HHD.py` are required yet.
# Keep this module in place so future extracted CUDA/OpenCL/CPU helpers have a
# stable home without changing the public package layout again.
__all__ = []
src/orisys/hhd/poisson.py
+954
View File
@@ -0,0 +1,954 @@
import numpy
import timeit
import sys
# import cv2
from scipy import signal, spatial
from ..backend_availability import (
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
cp, cp_signal, _torque_kernel,
cl, cla, cl_rfftn, cl_irfftn
)
from .timing import Timer
'''
# ------------------------------------------------------------------------------
Poisson Solver for 1D, 2D, and 3D scalar fields.
Data may be defined on a rectlinear grid or unstructured points.
1) create an object
PoissonSolver(args)
'solver' : type of solver
'S' for spatial convolution
only option for unstructured points
'F' for convolution in frequency space
default option for regular grids
--> for regular grid
'grid' : shape of the rectlinear grid
tuple of dimensions (Z,Y,X)
'spacings' : spacings of the rectlinear grid
tuple of spacings (dz,dy,dx)
--> for unstrucutred points
'points' : unstructured points
numpy array of shape (N, dim)
'pvolumes' : point volumes for each point
numpy array of shape (N, 1)
2) call prepare()
this initializes internal datastructures depending upon the mesh.
Once initialized, solve() can be called any number of times to
solve the Poisson equation on the same domain. If the domain
changes (e.g., a Lagrangian mesh), create a new object.
3) call solve(f), where, f is a numpy array defining a scalar field
for rectlinear grid, f.shape = grid.shape
for unstructured points, f.shape = (points.shape[0], 1)
# ------------------------------------------------------------------------------
'''
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) != ('spacings' in args):
raise SyntaxError("Poisson Solver needs spacings for regular grid")
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 'grid' in args:
self.ptype = 'G'
self.stype = 'F' # default
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']
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 (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')
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'")
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)
# --------------------------------------------------------------------------
# create the Green's function in appropriate representation
# --------------------------------------------------------------------------
# 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):
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
def prepare(self, verbose=False):
# ----------------------------------------------------------------------
gtimer = Timer()
if verbose:
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
# ----------------------------------------------------------------------
# for Fourier solution,
if self.stype == 'F':
if self.ptype == 'G':
self.G = self._create_radialGrid(verbose)
elif self.ptype == 'P':
raise SyntaxError('Frequncy solver not supported for unstructured points')
#self.G = create_distArray()
# ----------------------------------------------------------------------
# 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)
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 ='')
gtimer.end()
# --------------------------------------------------------------------------
# compute the integral solution
# --------------------------------------------------------------------------
def solve(self, f, verbose=False):
if not self.ready:
self.prepare(verbose)
fshape = f.shape
gtimer = Timer()
if verbose:
print( '\nSolving Poisson Eq.')
# ----------------------------------------------------------------------
# regular grid
if self.ptype == 'G':
#if (self.gdims - fshape).any():
if 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)
# 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.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':
p = numpy.multiply(f, self.pvolumes)
numpy.dot(self.G, p, p)
#numpy.multiply(p, self.pvolumes, p)
elif self.stype == 'F':
raise SyntaxError('Frequncy solver not supported for unstructured points')
# ----------------------------------------------------------------------
if verbose:
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)
"""
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')
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
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]
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)
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}"
)
return p1, p2
src/orisys/hhd/timing.py
+36
View File
@@ -0,0 +1,36 @@
import numpy
import timeit
class Timer(object):
def display(self):
tseconds = self.etime - self.stime
print(' [[ elapsed time:', end = "")
if tseconds < numpy.power(10.,-6.):
print("%.3f micro-sec." % (tseconds*numpy.power(10,6)),end = "")
elif tseconds < numpy.power(10.,-3.):
print("%.3f milli-sec." % (tseconds*numpy.power(10.,3.)),end = "")
elif tseconds < 60.0:
print("%.3f sec." % (tseconds),end = "")
else:
m = int(tseconds/ 60.0)
print("%d min. %.3f sec." % (m, (tseconds - 60*m)),end = "")
print(']]')
def start(self):
self.stime = timeit.default_timer()
def end(self, print_time=True):
self.etime = timeit.default_timer()
if print_time:
self.display()
def __init__(self):
self.start()
src/orisys/sensor.py
+51 -24
View File
@@ -19,7 +19,7 @@ from typing import Tuple
from .Stitcher import Stitcher from .Stitcher import Stitcher
from .tracker import Tracker from .tracker import Tracker
from .HHD import nHHD from .hhd import nHHD
from . import configs from . import configs
from .get_depth import Contact3D from .get_depth import Contact3D
from .backend_availability import (check_cupy_available, check_opencl_available) from .backend_availability import (check_cupy_available, check_opencl_available)
@@ -202,7 +202,7 @@ class Sensor:
Capture and preprocess a frame. Capture and preprocess a frame.
Returns: Returns:
np.ndarray: Processed grayscale image np.ndarray | None: Processed image, or None on capture failure.
""" """
# Simple direct read - buffer size of 1 should minimize frame accumulation # Simple direct read - buffer size of 1 should minimize frame accumulation
# Timing measurement helps identify if capture or processing is the bottleneck # Timing measurement helps identify if capture or processing is the bottleneck
@@ -265,39 +265,66 @@ class Sensor:
def detect_motion(self, threshold=6): def detect_motion(self, threshold=6):
# detect motion by optical flow magnitude # detect motion by optical flow magnitude
mag = np.linalg.norm(self.flow, axis=1) mag = np.linalg.norm(self.flow, axis=-1)
# print(f"mag: {mag.max()}, {mag.min()}, {mag.mean()}") # print(f"mag: {mag.max()}, {mag.min()}, {mag.mean()}")
return mag.mean() > threshold return mag.mean() > threshold
def compute_deformation(self, decompose=True, check_motion=False, threshold=6, tail_frames = 5): def compute_deformation(self, decompose=True, check_motion=True, threshold=6, tail_frames=10):
""" """
Compute optical flow and decompose into force components. Compute optical flow and optionally decompose it into force components.
Returns: Parameters
np.ndarray: Optical flow field ----------
decompose : bool
If True, run HHD decomposition (`compute_hhd`) after flow computation.
check_motion : bool
If True, gate decomposition with `detect_motion(threshold)`. When motion
falls below threshold, decomposition may still run for up to
`tail_frames` additional frames to avoid abrupt force drop-off.
threshold : float
Motion-detection threshold on mean flow magnitude.
tail_frames : int
Number of extra non-motion frames allowed to keep running HHD when
`check_motion=True`.
Returns
-------
bool | None
- `True` : motion is currently detected (`check_motion=True`)
- `False` : motion is currently not detected (`check_motion=True`)
- `None` : motion gating disabled (`check_motion=False`)
""" """
if self.img is None: if self.img is None:
raise Exception("Warning: Get image first") raise Exception("Warning: Get image first")
self.compute_flow() self.compute_flow()
# self.compute_hhd() if not check_motion:
if check_motion: self.tail_frames_count = 0
if self.detect_motion(threshold=threshold):
self.tail_frames_count = 0
self.compute_hhd()
return True
else:
if self.tail_frames_count <= tail_frames:
self.tail_frames_count +=1
self.compute_hhd()
return True
return False
else:
if decompose: if decompose:
self.compute_hhd() self.compute_hhd()
return None return None
# return True
is_motion = self.detect_motion(threshold=threshold)
try:
tail_frames = max(0, int(tail_frames))
except (TypeError, ValueError):
tail_frames = 0
if is_motion:
self.tail_frames_count = 0
if decompose:
self.compute_hhd()
return True
# No motion this frame. Keep HHD alive briefly if requested, but do not
# report this as "motion detected".
if decompose and self.tail_frames_count < tail_frames:
self.tail_frames_count += 1
self.compute_hhd()
else:
self.tail_frames_count = tail_frames
return False
def compute_flow(self): def compute_flow(self):
if self.img is None: if self.img is None:
test_3Dflow.py
-449
View File
@@ -1,449 +0,0 @@
"""
Prototype tactile 2.5D motion analysis on top of the existing dense flow output.
This script compares several motion representations derived from optical flow:
1. Raw flow arrows / magnitude
2. Radial and tangential flow around an estimated contact center
3. Differential flow features: divergence, curl, strain magnitude
4. Pseudo-height reconstructed from divergence (qualitative only)
"""
import os
import sys
import argparse
from pathlib import Path
import time
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import numpy as np
WINDOW_TITLE = "Tactile 2.5D Motion Prototype"
_local_src = os.path.join(os.path.dirname(__file__), "src")
if _local_src not in sys.path:
sys.path.insert(0, _local_src)
import orisys
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--video", "-v", type=str, default="0", help="Camera ID / video file path")
parser.add_argument("--config", "-c", type=str, default="./config/ddjx01.json", help="Configuration file path")
parser.add_argument("--cal", "-cal", type=str, default="./config/ddjx01.npy", help="Calibration file path")
parser.add_argument("--verbose", action="store_true", help="Print detailed sensor logs")
parser.add_argument("--no-display", action="store_true", help="Disable OpenCV windows")
parser.add_argument("--save", type=str, default="", help="Optional output video path for the comparison panel")
parser.add_argument("--max-frames", type=int, default=0, help="Stop after N frames (0 = unlimited)")
parser.add_argument("--grid-spacing", type=int, default=20, help="Arrow visualization grid spacing")
parser.add_argument("--flow-threshold", type=float, default=2.0, help="Arrow visualization magnitude threshold")
parser.add_argument("--print-every", type=int, default=1, help="Print scalar metrics every N frames")
parser.add_argument("--panel-size", type=int, default=320, help="Width/height of each panel tile")
parser.add_argument("--plane-width-mm", type=float, default=20.0, help="Physical tactile imaging plane width in mm")
parser.add_argument("--plane-height-mm", type=float, default=20.0, help="Physical tactile imaging plane height in mm")
parser.add_argument("--camera-distance-mm", type=float, default=10.0, help="Distance from camera plane to imaging plane in mm")
parser.add_argument("--display-max-width", type=int, default=1800, help="Maximum preview window width in pixels")
parser.add_argument("--display-max-height", type=int, default=1000, help="Maximum preview window height in pixels")
return parser.parse_args()
def create_sensor(args):
try:
video_input = int(args.video)
except ValueError:
video_input = args.video
return orisys.Sensor(
video_input,
config_name=args.config,
cal_path=args.cal,
verbose=args.verbose,
)
def ensure_bgr(image):
if image.ndim == 2:
return cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
return image
def resize_panel(image, panel_size):
return cv2.resize(ensure_bgr(image), (panel_size, panel_size), interpolation=cv2.INTER_LINEAR)
def fit_display_image(image, max_width, max_height):
h, w = image.shape[:2]
scale = min(max_width / max(float(w), 1.0), max_height / max(float(h), 1.0), 1.0)
if scale >= 0.999:
return image
new_size = (max(1, int(w * scale)), max(1, int(h * scale)))
return cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)
def signed_heatmap(field, colormap=cv2.COLORMAP_TURBO):
max_abs = float(np.max(np.abs(field)))
if max_abs <= 1e-9:
norm = np.full(field.shape, 127, dtype=np.uint8)
else:
norm = np.clip((field / max_abs + 1.0) * 127.5, 0, 255).astype(np.uint8)
return cv2.applyColorMap(norm, colormap)
def magnitude_heatmap(field, colormap=cv2.COLORMAP_TURBO):
max_val = float(np.max(field))
if max_val <= 1e-9:
norm = np.zeros(field.shape, dtype=np.uint8)
else:
norm = np.clip(field / max_val * 255.0, 0, 255).astype(np.uint8)
return cv2.applyColorMap(norm, colormap)
def add_title(image, title):
panel = image.copy()
cv2.rectangle(panel, (0, 0), (panel.shape[1], 28), (0, 0, 0), -1)
cv2.putText(panel, title, (8, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
return panel
def make_panel(image, title, panel_size):
return add_title(resize_panel(image, panel_size), title)
def stack_panels(panels, columns=4):
filler = np.zeros_like(panels[0])
rows = []
for idx in range(0, len(panels), columns):
row = panels[idx:idx + columns]
if len(row) < columns:
row = row + [filler] * (columns - len(row))
rows.append(cv2.hconcat(row))
return cv2.vconcat(rows)
def estimate_contact_center(divergence, flow_magnitude):
positive = np.maximum(divergence, 0.0)
weights = positive * (flow_magnitude + 1e-6)
total = float(weights.sum())
h, w = divergence.shape
if total <= 1e-9:
return (w // 2, h // 2)
ys, xs = np.indices(divergence.shape, dtype=np.float32)
cx = int(np.clip((weights * xs).sum() / total, 0, w - 1))
cy = int(np.clip((weights * ys).sum() / total, 0, h - 1))
return (cx, cy)
def prepare_support_mask(support_map):
if support_map is None:
return None
support = np.abs(support_map).astype(np.float32)
max_support = float(np.max(support))
if max_support > 0.0:
support /= max_support
support = (support > 0.05).astype(np.float32)
if float(np.mean(support)) < 0.002:
return None
return support
def decompose_radial_tangential(flow, center):
h, w = flow.shape[:2]
ys, xs = np.indices((h, w), dtype=np.float32)
cx, cy = center
rx = xs - float(cx)
ry = ys - float(cy)
r_norm = np.sqrt(rx * rx + ry * ry)
r_norm = np.where(r_norm < 1e-6, 1.0, r_norm)
rhat_x = rx / r_norm
rhat_y = ry / r_norm
that_x = -rhat_y
that_y = rhat_x
u = flow[:, :, 0]
v = flow[:, :, 1]
radial = u * rhat_x + v * rhat_y
tangential = u * that_x + v * that_y
return radial, tangential
def compute_flow_jacobian_features(flow):
u = flow[:, :, 0]
v = flow[:, :, 1]
du_dy, du_dx = np.gradient(u)
dv_dy, dv_dx = np.gradient(v)
divergence = du_dx + dv_dy
curl = dv_dx - du_dy
strain_xx = du_dx
strain_yy = dv_dy
strain_xy = 0.5 * (du_dy + dv_dx)
strain_mag = np.sqrt(strain_xx * strain_xx + strain_yy * strain_yy + 2.0 * strain_xy * strain_xy)
return {
"divergence": divergence,
"curl": curl,
"strain_mag": strain_mag,
}
def reconstruct_pseudo_height(divergence, flow_magnitude=None, support_map=None):
rhs = divergence.astype(np.float32).copy()
# No-contact guard: tiny flow noise becomes large low-frequency drift after Poisson solve.
if flow_magnitude is not None:
mean_mag = float(np.mean(flow_magnitude))
max_mag = float(np.max(flow_magnitude))
if mean_mag < 0.03 and max_mag < 0.15:
return np.zeros_like(rhs, dtype=np.float32)
support = prepare_support_mask(support_map)
if support_map is not None:
if support is None:
return np.zeros_like(rhs, dtype=np.float32)
rhs *= support
noise_floor = max(float(np.std(rhs)) * 1.5, 1e-4)
rhs[np.abs(rhs) < noise_floor] = 0.0
if not np.any(rhs):
return np.zeros_like(rhs, dtype=np.float32)
# Zero-pad before FFT solve to reduce periodic wrap-around ghosts, then crop back.
h, w = rhs.shape
pad_y = max(8, h // 2)
pad_x = max(8, w // 2)
rhs_padded = np.pad(rhs, ((pad_y, pad_y), (pad_x, pad_x)), mode="constant")
rhs_padded = rhs_padded - float(rhs_padded.mean())
hp, wp = rhs_padded.shape
fy = np.fft.fftfreq(hp).reshape(-1, 1)
fx = np.fft.fftfreq(wp).reshape(1, -1)
denom = (2.0 * np.pi * fx) ** 2 + (2.0 * np.pi * fy) ** 2
denom[0, 0] = 1.0
rhs_hat = np.fft.fft2(rhs_padded)
height_hat = -rhs_hat / denom
height_hat[0, 0] = 0.0
height_full = np.real(np.fft.ifft2(height_hat)).astype(np.float32)
height = height_full[pad_y:pad_y + h, pad_x:pad_x + w]
height = cv2.GaussianBlur(height, (0, 0), 1.0)
return height
def build_visualization(img, flow, center, features, radial, tangential, pseudo_height, depth_map, args):
arrows = orisys.util.draw_arrows(
img,
flow,
threshold=args.flow_threshold,
grid_spacing=args.grid_spacing,
arrow_scale=1.0,
)
contact_view = ensure_bgr(img)
cv2.circle(contact_view, center, 6, (0, 255, 0), -1)
cv2.circle(contact_view, center, 12, (0, 255, 0), 2)
panel_images = [
(contact_view, "Image / center"),
(arrows, "Raw flow"),
(signed_heatmap(radial), "Radial flow"),
(signed_heatmap(tangential), "Tangential flow"),
(signed_heatmap(features["divergence"]), "Divergence"),
(signed_heatmap(features["curl"]), "Curl"),
(magnitude_heatmap(features["strain_mag"]), "Strain mag"),
(signed_heatmap(depth_map), "Depth map"),
(signed_heatmap(pseudo_height), "Pseudo height"),
]
panels = [make_panel(image, title, args.panel_size) for image, title in panel_images]
return stack_panels(panels)
def make_writer(path, frame_size, fps):
if not path:
return None
parent = Path(path).parent
if parent and not parent.exists():
parent.mkdir(parents=True, exist_ok=True)
return cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), fps, frame_size)
def approximate_raw_flow_to_mm(flow_px, plane_width_mm, plane_height_mm, camera_distance_mm, image_shape):
h, w = image_shape[:2]
fx = camera_distance_mm * w / max(plane_width_mm, 1e-6)
fy = camera_distance_mm * h / max(plane_height_mm, 1e-6)
flow_mm = flow_px.copy().astype(np.float32)
flow_mm[:, :, 0] *= camera_distance_mm / max(fx, 1e-6)
flow_mm[:, :, 1] *= camera_distance_mm / max(fy, 1e-6)
return flow_mm
def compute_metrics(flow_magnitude, raw_flow_mm, radial, tangential, features, depth_map, pseudo_height):
return {
"flow_mean": float(flow_magnitude.mean()),
"flow_mean_mm": float(np.linalg.norm(raw_flow_mm, axis=2).mean()),
"radial_abs_mean": float(np.mean(np.abs(radial))),
"tangential_abs_mean": float(np.mean(np.abs(tangential))),
"divergence_mean": float(features["divergence"].mean()),
"curl_abs_mean": float(np.mean(np.abs(features["curl"]))),
"strain_mean": float(features["strain_mag"].mean()),
"depth_peak": float(np.max(np.abs(depth_map))),
"pseudo_height_peak": float(np.max(np.abs(pseudo_height))),
}
def format_metrics(frame_count, fps, metrics):
line = (
f"frame={frame_count:04d} fps={fps:6.2f} "
f"flow_mean={metrics['flow_mean']:.4f} "
f"flow_mean_mm={metrics['flow_mean_mm']:.4f} "
f"radial_abs={metrics['radial_abs_mean']:.4f} "
f"tangential_abs={metrics['tangential_abs_mean']:.4f} "
f"div_mean={metrics['divergence_mean']:.4f} "
f"curl_abs={metrics['curl_abs_mean']:.4f} "
f"strain={metrics['strain_mean']:.4f} "
f"depth_peak={metrics['depth_peak']:.4f} "
f"height_peak={metrics['pseudo_height_peak']:.4f}"
)
t = metrics.get("_timings")
if t:
line += (
f" | analyze {t['total']:.2f}ms"
f" (mag:{t['magnitude']:.2f} jac:{t['jacobian']:.2f}"
f" rad_tan:{t['rad_tan']:.2f} pseudo_h:{t['pseudo_h']:.2f}"
f" metrics:{t['metrics']:.2f} vis:{t['visualize']:.2f})"
)
return line
def read_sensor_frame(sensor):
sensor.compute_flow()
sensor.compute_hhd()
sensor.compute_contact()
return sensor.read_info(
sensor.info.FPS,
sensor.info.VRAW,
sensor.info.IMG,
sensor.info.DEPTH,
)
def analyze_frame(img, flow, depth_map, args):
t = time.perf_counter()
flow_magnitude = np.linalg.norm(flow, axis=2)
t_mag = time.perf_counter()
features = compute_flow_jacobian_features(flow)
t_jac = time.perf_counter()
center = estimate_contact_center(features["divergence"], flow_magnitude)
radial, tangential = decompose_radial_tangential(flow, center)
t_decomp = time.perf_counter()
pseudo_height = reconstruct_pseudo_height(
features["divergence"],
flow_magnitude=flow_magnitude,
support_map=depth_map,
)
t_height = time.perf_counter()
raw_flow_mm = approximate_raw_flow_to_mm(
flow,
args.plane_width_mm,
args.plane_height_mm,
args.camera_distance_mm,
img.shape,
)
metrics = compute_metrics(flow_magnitude, raw_flow_mm, radial, tangential, features, depth_map, pseudo_height)
t_metrics = time.perf_counter()
canvas = build_visualization(
img,
flow,
center,
features,
radial,
tangential,
pseudo_height,
depth_map,
args,
)
t_vis = time.perf_counter()
metrics["_timings"] = {
"magnitude": (t_mag - t) * 1000,
"jacobian": (t_jac - t_mag) * 1000,
"rad_tan": (t_decomp - t_jac) * 1000,
"pseudo_h": (t_height - t_decomp) * 1000,
"metrics": (t_metrics - t_height) * 1000,
"visualize": (t_vis - t_metrics) * 1000,
"total": (t_vis - t) * 1000,
}
return canvas, metrics
def main():
args = parse_args()
sensor = create_sensor(args)
writer = None
frame_count = 0
print_every = max(args.print_every, 1)
try:
print("\nPress 'q' to quit, 'r' to reset tracking.")
while True:
img = sensor.get_img()
if img is None:
print("No more frames returned from sensor.")
break
fps, flow, img, depth_map = read_sensor_frame(sensor)
canvas, metrics = analyze_frame(img, flow, depth_map, args)
if writer is None and args.save:
writer = make_writer(args.save, (canvas.shape[1], canvas.shape[0]), max(float(fps), 1.0))
if writer is not None:
writer.write(canvas)
frame_count += 1
if frame_count % print_every == 0:
print(format_metrics(frame_count, fps, metrics))
if not args.no_display:
display_canvas = fit_display_image(
canvas,
max_width=args.display_max_width,
max_height=args.display_max_height,
)
cv2.imshow(WINDOW_TITLE, display_canvas)
key = cv2.waitKey(1) & 0xFF
if key == ord("q"):
print("\nExiting...")
break
if key == ord("r"):
sensor.reset()
print("Sensor reset.")
if args.max_frames > 0 and frame_count >= args.max_frames:
print(f"Reached max_frames={args.max_frames}.")
break
finally:
if writer is not None:
writer.release()
sensor.disconnect()
if not args.no_display:
cv2.destroyAllWindows()
if __name__ == "__main__":
main()
test_cam.py
-75
View File
@@ -1,75 +0,0 @@
import cv2
import time
# Toggle these to test whether exposure is limiting delivered FPS.
USE_MANUAL_EXPOSURE = False
MANUAL_EXPOSURE = -6 # Typical V4L2/OpenCV manual range is negative values.
cap = cv2.VideoCapture(4, cv2.CAP_V4L2)
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M', 'J', 'P', 'G'))
cap.set(cv2.CAP_PROP_FPS, 30)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
# For V4L2 in OpenCV: auto-exposure is usually 3, manual is usually 1.
if USE_MANUAL_EXPOSURE:
ok_auto = cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 1)
ok_exp = cap.set(cv2.CAP_PROP_EXPOSURE, MANUAL_EXPOSURE)
else:
ok_auto = cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 3)
ok_exp = True
actual_fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))
print(f"Backend: {cap.getBackendName()}")
print(
f"FOURCC: {chr(actual_fourcc & 0xFF)}"
f"{chr((actual_fourcc >> 8) & 0xFF)}"
f"{chr((actual_fourcc >> 16) & 0xFF)}"
f"{chr((actual_fourcc >> 24) & 0xFF)}"
)
print(f"Resolution: {int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))}x{int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))}")
print(f"FPS: {cap.get(cv2.CAP_PROP_FPS)}")
print(
f"Exposure mode: {'manual' if USE_MANUAL_EXPOSURE else 'auto'} "
f"(set_auto_ok={ok_auto}, set_exp_ok={ok_exp})"
)
print(f"AutoExposure(readback): {cap.get(cv2.CAP_PROP_AUTO_EXPOSURE)}")
print(f"Exposure(readback): {cap.get(cv2.CAP_PROP_EXPOSURE)}")
frame_count = 0
start = time.perf_counter()
while cap.isOpened():
t0 = time.perf_counter()
ok = cap.grab()
t1 = time.perf_counter()
if not ok:
break
ret, frame = cap.retrieve()
t2 = time.perf_counter()
if not ret:
continue
cv2.imshow("webcam", frame)
t3 = time.perf_counter()
key = cv2.waitKey(1) & 0xFF
t4 = time.perf_counter()
frame_count += 1
if frame_count % 30 == 0:
elapsed = t4 - start
print(
f"Frame {frame_count}: "
f"grab={1000*(t1-t0):.1f}ms, "
f"retrieve={1000*(t2-t1):.1f}ms, "
f"imshow={1000*(t3-t2):.1f}ms, "
f"waitKey={1000*(t4-t3):.1f}ms, "
f"avg_fps={frame_count/elapsed:.1f}"
)
if key == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
test_downsample.py
-252
View File
@@ -1,252 +0,0 @@
"""
HHD Downsample Factor Comparison
Reads a video file and runs the full pipeline at multiple hhd_downsample factors.
Computes per-frame HHD timing and vnormal quality metrics against the baseline (factor=1),
then plots the results.
Usage:
python test_downsample.py -v jx01_hex1.5_2.mp4 -c ./config/prototype_of3_720p.json -f 1,2,3,4 -n 50 -w 30
"""
import os
import sys
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
_local_src = os.path.join(os.path.dirname(__file__), "src")
if _local_src not in sys.path:
sys.path.insert(0, _local_src)
import cv2
import orisys
import numpy as np
import time
import argparse
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
def run_sensor(vid_src, config, cal, downsample, n_warmup, n_measure):
"""Run sensor for n_warmup + n_measure frames. Capture timing and vnormal for every measurement frame."""
try:
vid = int(vid_src)
except ValueError:
vid = vid_src
sensor = orisys.Sensor(
vid, config_name=config, cal_path=cal,
isstitch=True, cuda=False, verbose=False,
hhd_downsample=downsample,
)
hhd_times = []
vnormals = []
frame_idx = 0
n_total = n_warmup + n_measure
while frame_idx < n_total:
try:
sensor.get_img()
except Exception:
break
sensor.compute_flow()
t0 = time.perf_counter()
sensor.compute_hhd()
t_hhd = time.perf_counter() - t0
if frame_idx >= n_warmup:
hhd_times.append(t_hhd)
vn = sensor.vnormal.copy() if sensor.vnormal is not None else None
vnormals.append(vn)
frame_idx += 1
sensor.disconnect()
return {"hhd_times": hhd_times, "vnormals": vnormals}
def per_frame_rmse(a, b):
if a is None or b is None:
return float("nan")
diff = a.astype(np.float64) - b.astype(np.float64)
return float(np.sqrt(np.mean(diff ** 2)))
def per_frame_max_abs_err(a, b):
if a is None or b is None:
return float("nan")
return float(np.max(np.abs(a.astype(np.float64) - b.astype(np.float64))))
def per_frame_cosine_sim(a, b):
if a is None or b is None:
return float("nan")
a64 = a.astype(np.float64)
b64 = b.astype(np.float64)
dot = np.sum(a64 * b64, axis=-1)
mag_a = np.linalg.norm(a64, axis=-1)
mag_b = np.linalg.norm(b64, axis=-1)
denom = mag_a * mag_b
valid = denom > 1e-12
if not np.any(valid):
return float("nan")
return float(np.mean(dot[valid] / denom[valid]))
def main():
parser = argparse.ArgumentParser(description="Compare HHD downsample factors")
parser.add_argument("--video", "-v", type=str, default="2", help="Camera ID or video file path")
parser.add_argument("--config", "-c", type=str, default="./config/prototype_of3_720p.json")
parser.add_argument("--cal", "-cal", type=str, default="./config/ddjx01.npy")
parser.add_argument("--factors", "-f", type=str, default="1,2,3", help="Comma-separated downsample factors")
parser.add_argument("--warmup", "-w", type=int, default=5, help="Warmup frames (skipped)")
parser.add_argument("--frames", "-n", type=int, default=30, help="Measurement frames")
parser.add_argument("--output", "-o", type=str, default="downsample_comparison.png", help="Output plot filename")
args = parser.parse_args()
factors = [int(f) for f in args.factors.split(",")]
print(f"Video: {args.video}")
print(f"Config: {args.config}")
print(f"Factors: {factors}")
print(f"Warmup: {args.warmup} frames, Measure: {args.frames} frames")
print()
# --- collect data ---
results = {}
for factor in factors:
print(f"--- Running factor {factor} ---")
results[factor] = run_sensor(
args.video, args.config, args.cal,
factor, args.warmup, args.frames,
)
n = len(results[factor]["hhd_times"])
if n > 0:
arr = np.array(results[factor]["hhd_times"]) * 1000
print(f" {n} frames: hhd mean={arr.mean():.2f}ms, median={np.median(arr):.2f}ms, "
f"min={arr.min():.2f}ms, max={arr.max():.2f}ms")
else:
print(" No frames captured")
print()
# --- timing summary table ---
print("=" * 72)
print(f"{'Factor':<10} {'Mean HHD (ms)':<16} {'Median (ms)':<14} {'Speedup':<10}")
print("-" * 72)
baseline_median = None
for factor in factors:
arr = np.array(results[factor]["hhd_times"]) * 1000
if len(arr) == 0:
continue
med = float(np.median(arr))
if baseline_median is None:
baseline_median = med
speedup = baseline_median / med if med > 0 else 0
print(f"{factor:<10} {arr.mean():<16.2f} {med:<14.2f} {speedup:<10.2f}x")
print("=" * 72)
print()
# --- compute per-frame quality metrics vs baseline (factor=1) ---
has_baseline = 1 in factors and len(factors) >= 2
metrics = {} # factor -> {"rmse": [...], "max_err": [...], "cosine": [...]}
if has_baseline:
base_vn = results[1]["vnormals"]
n_base = len(base_vn)
for factor in factors:
if factor == 1:
continue
test_vn = results[factor]["vnormals"]
n_common = min(n_base, len(test_vn))
rmse_arr, maxerr_arr, cos_arr = [], [], []
for i in range(n_common):
rmse_arr.append(per_frame_rmse(base_vn[i], test_vn[i]))
maxerr_arr.append(per_frame_max_abs_err(base_vn[i], test_vn[i]))
cos_arr.append(per_frame_cosine_sim(base_vn[i], test_vn[i]))
metrics[factor] = {
"rmse": np.array(rmse_arr),
"max_err": np.array(maxerr_arr),
"cosine": np.array(cos_arr),
}
# print per-factor summary
print(f"{'Factor':<10} {'Mean RMSE':<14} {'Mean MaxErr':<14} {'Mean CosSim':<14}")
print("-" * 52)
for factor in sorted(metrics):
m = metrics[factor]
print(f"{factor:<10} {np.nanmean(m['rmse']):<14.6f} {np.nanmean(m['max_err']):<14.4f} "
f"{np.nanmean(m['cosine']):<14.6f}")
print("=" * 52)
print()
# --- plot ---
n_quality_factors = len(metrics)
n_plots = 1 + (3 if has_baseline else 0) # timing + (rmse, max_err, cosine)
fig, axes = plt.subplots(n_plots, 1, figsize=(12, 3.5 * n_plots), sharex=True)
if n_plots == 1:
axes = [axes]
colors = plt.cm.tab10.colors
frame_nums = None
# 1) HHD timing per frame
ax = axes[0]
for i, factor in enumerate(factors):
arr = np.array(results[factor]["hhd_times"]) * 1000
frames = np.arange(len(arr))
if frame_nums is None:
frame_nums = frames
ax.plot(frames, arr, label=f"ds={factor}", color=colors[i % len(colors)], linewidth=1)
ax.set_ylabel("HHD time (ms)")
ax.set_title("Per-frame HHD computation time")
ax.legend(loc="upper right")
ax.grid(True, alpha=0.3)
if has_baseline:
# 2) RMSE
ax = axes[1]
for i, factor in enumerate(sorted(metrics)):
m = metrics[factor]
frames = np.arange(len(m["rmse"]))
ci = factors.index(factor) if factor in factors else i
ax.plot(frames, m["rmse"], label=f"ds={factor}", color=colors[ci % len(colors)], linewidth=1)
ax.set_ylabel("RMSE (vnormal)")
ax.set_title("Per-frame RMSE vs baseline (ds=1)")
ax.legend(loc="upper right")
ax.grid(True, alpha=0.3)
# 3) Max absolute error
ax = axes[2]
for i, factor in enumerate(sorted(metrics)):
m = metrics[factor]
frames = np.arange(len(m["max_err"]))
ci = factors.index(factor) if factor in factors else i
ax.plot(frames, m["max_err"], label=f"ds={factor}", color=colors[ci % len(colors)], linewidth=1)
ax.set_ylabel("Max Abs Error (vnormal)")
ax.set_title("Per-frame max absolute error vs baseline (ds=1)")
ax.legend(loc="upper right")
ax.grid(True, alpha=0.3)
# 4) Cosine similarity
ax = axes[3]
for i, factor in enumerate(sorted(metrics)):
m = metrics[factor]
frames = np.arange(len(m["cosine"]))
ci = factors.index(factor) if factor in factors else i
ax.plot(frames, m["cosine"], label=f"ds={factor}", color=colors[ci % len(colors)], linewidth=1)
ax.set_ylabel("Cosine Similarity")
ax.set_title("Per-frame cosine similarity vs baseline (ds=1)")
ax.legend(loc="upper right")
ax.grid(True, alpha=0.3)
ax.set_ylim(bottom=min(0.999, ax.get_ylim()[0]))
axes[-1].set_xlabel("Frame")
fig.tight_layout()
fig.savefig(args.output, dpi=150)
print(f"Plot saved to: {args.output}")
if __name__ == "__main__":
main()
test_ffmpeg.py
-81
View File
@@ -1,81 +0,0 @@
import argparse
import av
import cv2
import time
DEVICE = "/dev/video4"
TARGET_FPS = 30
WIDTH = 1920
HEIGHT = 1080
def main():
parser = argparse.ArgumentParser(description="Capture from V4L2 via PyAV, optionally save to MP4.")
parser.add_argument(
"-o",
"--output",
type=str,
default="",
help="Output MP4 path (e.g. capture.mp4). If empty, only preview, no file.",
)
args = parser.parse_args()
# Open v4l2 device with same style options as ffmpeg CLI
container = av.open(
DEVICE,
format="v4l2",
options={
"input_format": "mjpeg",
"video_size": f"{WIDTH}x{HEIGHT}",
"framerate": str(TARGET_FPS),
},
)
stream = container.streams.video[0]
frame_count = 0
t_start = time.perf_counter()
writer = None
if args.output:
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
writer = cv2.VideoWriter(
args.output, fourcc, float(TARGET_FPS), (WIDTH, HEIGHT)
)
if not writer.isOpened():
container.close()
raise RuntimeError(f"Could not open VideoWriter for {args.output}")
print(f"Recording to {args.output} @ {TARGET_FPS} fps, {WIDTH}x{HEIGHT}")
try:
for frame in container.decode(stream):
t0 = time.perf_counter()
img = frame.to_ndarray(format="bgr24") # OpenCV-compatible
t1 = time.perf_counter()
if writer is not None:
writer.write(img)
cv2.imshow("pyav_cam", img)
key = cv2.waitKey(1) & 0xFF
t2 = time.perf_counter()
frame_count += 1
if frame_count % 30 == 0:
elapsed = t2 - t_start
print(
f"Frame {frame_count}: convert={1000*(t1-t0):.1f}ms, "
f"imshow+waitKey={1000*(t2-t1):.1f}ms, avg_fps={frame_count/elapsed:.1f}"
)
if key == ord("q"):
break
finally:
container.close()
if writer is not None:
writer.release()
print(f"Saved {frame_count} frames to {args.output}")
cv2.destroyAllWindows()
if __name__ == "__main__":
main()
test_opencl.py
-33
View File
@@ -1,33 +0,0 @@
import pyopencl as cl
# Get all platforms
platforms = cl.get_platforms()
if not platforms:
raise RuntimeError("No OpenCL platforms found")
# Use the first platform (rusticl in your case)
platform = platforms[0]
print("Using platform:", platform.name)
# Get devices for that platform
devices = platform.get_devices()
if not devices:
raise RuntimeError("No devices found on platform")
for platform in cl.get_platforms():
for device in platform.get_devices():
print("Device:", device.name)
# if device.type & cl.device_type.GPU == 0:
# print("Skipping non-GPU device:", device.name)
# Use the first device (your Adreno GPU)
device = devices[0]
print("Using device:", device.name)
# Create a context with that device
ctx = cl.Context([device])
# Create a command queue
queue = cl.CommandQueue(ctx)
print("Context and command queue created successfully!")
test_stitch.py
-77
View File
@@ -1,77 +0,0 @@
"""
示例 1: 基础用法 - 单传感器同步处理
本示例演示 orisys 的基本使用流程：
1. 初始化传感器（视频文件或摄像头）
2. 循环获取图像并计算形变和触碰点
3. 读取各类传感器数据
4. 可视化显示
"""
import os
import argparse
import cv2
import orisys
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--video", "-v", type=str, default="0", help="Camera ID / Video file path")
parser.add_argument("--config", "-c", type=str, default="./config/ddjx01.json", help="Configuration file path")
parser.add_argument("--cal", "-cal", type=str, default="./config/ddjx01.npy", help="Calibration file path")
parser.add_argument("--verbose", "-verbose", type=bool, default=True, help="Print detailed information")
parser.add_argument(
"--match-brightness",
dest="match_brightness",
action="store_true",
help="Enable per-tile brightness matching",
)
parser.set_defaults(match_brightness=False)
args = parser.parse_args()
try:
vid_src = int(args.video)
except ValueError:
vid_src = args.video
# 步骤1：创建传感器对象 (必要步骤)，
# config_name为配置文件名（内置配置文件名）或路径（自定义配置文件路径）
# cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件
sensor = orisys.Sensor(vid_src, config_name=args.config, cal_path=args.cal, verbose=args.verbose)
sensor.stitcher.configure_feature_stitch(match_brightness=args.match_brightness)
print("\n按 'q' 退出程序，按 'r' 重置传感器追踪器")
while True:
# 步骤2：信息提取
# 步骤 2.1: 获取并拼接图像 (必要步骤)
img = sensor.get_img()
if img is None:
print("Failed to get image (end of stream or capture error). Exiting.")
break
if img is not None:
cv2.imshow("OpenCV Stitch (feature-based)", img)
# 步骤 2.2: 计算形变场 (必要步骤) — sensor.compute_deformation(...)
# 步骤 2.3: 计算触碰点
# 步骤 3: 键盘控制
key = cv2.waitKey(1) & 0xFF
if key == ord("q"):
print("\n退出程序...")
break
if key == ord("r"):
# 重置光流追踪器（当追踪漂移时使用）
sensor.reset()
print("传感器已重置")
# =========================================================================
# 4. 清理资源
# =========================================================================
sensor.disconnect() # 释放视频源和相关资源
if __name__ == "__main__":
main()
test_usb_fps.py
-213
View File
@@ -1,213 +0,0 @@
#!/usr/bin/env python3
"""
USB camera format and FPS tester for Jetson.
- Lists available formats/resolutions via v4l2-ctl.
- For selected modes, opens the camera with OpenCV and measures per-frame
capture time (ms) to show real delivered FPS/latency.
"""
import argparse
import subprocess
import sys
import time
import statistics
from typing import List, Tuple
import cv2
def run_v4l2_list_formats(device: str) -> str:
"""Return the raw output of v4l2-ctl --list-formats-ext."""
try:
out = subprocess.check_output(
["v4l2-ctl", f"--device={device}", "--list-formats-ext"],
stderr=subprocess.STDOUT,
text=True,
)
return out
except (subprocess.CalledProcessError, FileNotFoundError) as e:
print(f"[WARN] Failed to run v4l2-ctl: {e}")
return ""
def parse_formats(output: str) -> List[Tuple[str, int, int, List[float]]]:
"""
Parse v4l2-ctl --list-formats-ext output.
Returns list of (fourcc, width, height, fps_list).
"""
lines = output.splitlines()
formats: List[Tuple[str, int, int, List[float]]] = []
current_fourcc = None
for line in lines:
line = line.strip()
if line.startswith("[") and "):" in line and "'" in line:
# Example: [0]: 'MJPG' (Motion-JPEG, compressed)
try:
fourcc = line.split("'")[1]
current_fourcc = fourcc
except Exception:
current_fourcc = None
elif line.startswith("Size: Discrete") and current_fourcc:
# Example: Size: Discrete 1280x720
try:
parts = line.split()
wh = parts[-1]
w_str, h_str = wh.split("x")
width = int(w_str)
height = int(h_str)
formats.append((current_fourcc, width, height, []))
except Exception:
pass
elif "Interval:" in line and "fps" in line and formats:
# Example: Interval: Discrete 0.033s (30.000 fps)
try:
fps_part = line.split("(")[-1].split("fps")[0]
fps = float(fps_part)
fourcc, w, h, fps_list = formats[-1]
fps_list.append(fps)
formats[-1] = (fourcc, w, h, fps_list)
except Exception:
pass
return formats
def fourcc_to_int(code: str) -> int:
return cv2.VideoWriter_fourcc(*code)
def benchmark_mode(
device_index: int,
fourcc: str,
width: int,
height: int,
target_fps: float,
warmup: int = 10,
measure_frames: int = 60,
) -> None:
"""Open a specific mode and measure per-frame capture time."""
print(f"\n=== Testing {fourcc} {width}x{height} @ {target_fps:.1f} fps ===")
cap = cv2.VideoCapture(device_index, cv2.CAP_V4L2)
if not cap.isOpened():
print(" [ERROR] Failed to open device with V4L2 backend.")
return
cap.set(cv2.CAP_PROP_FOURCC, fourcc_to_int(fourcc))
cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
cap.set(cv2.CAP_PROP_FPS, target_fps)
# Query negotiated parameters
negotiated_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
negotiated_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
negotiated_fps = cap.get(cv2.CAP_PROP_FPS)
negotiated_fourcc_val = int(cap.get(cv2.CAP_PROP_FOURCC))
negotiated_fourcc = "".join(
chr((negotiated_fourcc_val >> 8 * i) & 0xFF) for i in range(4)
)
print(
f" Negotiated: {negotiated_w}x{negotiated_h} {negotiated_fourcc} "
f"@ {negotiated_fps:.1f} fps (backend={cap.getBackendName()})"
)
# Warm-up frames (to let auto-exposure, buffers, etc. settle)
for i in range(warmup):
ret, _ = cap.read()
if not ret:
print(f" [WARN] Warmup frame {i} failed.")
time.sleep(0.05)
times_ms: List[float] = []
for i in range(measure_frames):
t0 = time.perf_counter()
ret, _ = cap.read()
t1 = time.perf_counter()
if not ret:
print(f" [WARN] Capture failed at frame {i}.")
break
times_ms.append((t1 - t0) * 1000.0)
cap.release()
if not times_ms:
print(" [ERROR] No successful frames captured.")
return
mean_ms = statistics.mean(times_ms)
med_ms = statistics.median(times_ms)
min_ms = min(times_ms)
max_ms = max(times_ms)
eff_fps = 1000.0 / mean_ms
print(
f" Frames measured: {len(times_ms)}\n"
f" Capture time (ms): mean={mean_ms:.2f}, median={med_ms:.2f}, "
f"min={min_ms:.2f}, max={max_ms:.2f}\n"
f" Effective FPS from capture timing: {eff_fps:.2f}"
)
def main() -> None:
parser = argparse.ArgumentParser(
description="List USB camera formats and benchmark capture times."
)
parser.add_argument(
"--device",
default="/dev/video4",
help="V4L2 device path (default: /dev/video0)",
)
parser.add_argument(
"--index",
type=int,
default=0,
help="OpenCV device index (default: 0)",
)
parser.add_argument(
"--test-modes",
nargs="*",
default=["MJPG:640x480:30", "MJPG:1280x720:30"],
help=(
"Modes to benchmark as FOURCC:WxH:FPS. "
"Example: MJPG:640x480:30 YUYV:640x480:30"
),
)
args = parser.parse_args()
print(f"== v4l2-ctl formats for {args.device} ==")
raw = run_v4l2_list_formats(args.device)
if not raw:
print("No v4l2-ctl output; skipping format list.")
else:
print(raw)
parsed = parse_formats(raw)
if parsed:
print("\n== Parsed formats (FOURCC WxH [fps...]) ==")
for fourcc, w, h, fps_list in parsed:
fps_str = ", ".join(f"{f:.1f}" for f in fps_list) if fps_list else "?"
print(f" {fourcc} {w}x{h} [{fps_str}]")
print("\n== Benchmarking capture time with OpenCV ==")
for mode_str in args.test_modes:
try:
fourcc_s, wh_s, fps_s = mode_str.split(":")
w_s, h_s = wh_s.lower().split("x")
w = int(w_s)
h = int(h_s)
fps = float(fps_s)
except ValueError:
print(f"[WARN] Cannot parse test mode '{mode_str}', expected FOURCC:WxH:FPS")
continue
benchmark_mode(args.index, fourcc_s, w, h, fps)
if __name__ == "__main__":
try:
main()
except KeyboardInterrupt:
print("\nInterrupted by user.", file=sys.stderr)
tools/3d_patch_generation_blender/blender_export.py
+147
View File
@@ -0,0 +1,147 @@
import bpy
import numpy as np
import json
import os
from mathutils import Matrix, Vector
def export_camera_intrinsics(camera_obj, scene):
"""Return camera intrinsic matrix K (3x3) and distortion (none)."""
cam = camera_obj.data
render = scene.render
res_x = render.resolution_x
res_y = render.resolution_y
# Sensor size and focal length
if cam.sensor_fit == 'HORIZONTAL':
sensor_width = cam.sensor_width
sensor_height = sensor_width / (res_x / res_y)
else:
sensor_height = cam.sensor_height
sensor_width = sensor_height * (res_x / res_y)
focal = cam.lens # in mm
# Pixel sizes in mm
pixel_width = sensor_width / res_x
pixel_height = sensor_height / res_y
# Intrinsic matrix (OpenCV style: fx, fy in pixels, cx, cy principal point)
fx = focal / pixel_width
fy = focal / pixel_height
cx = res_x / 2.0
cy = res_y / 2.0
K = np.array([[fx, 0, cx],
[0, fy, cy],
[0, 0, 1]], dtype=np.float64)
return K, res_x, res_y
def export_mesh_data(obj, filepath):
"""Export mesh vertices, faces, and UVs to a compressed .npz file."""
me = obj.data
# Transform vertices to world space
vertices = np.array([obj.matrix_world @ v.co for v in me.vertices], dtype=np.float32)
faces = np.array([poly.vertices[:] for poly in me.polygons], dtype=np.int32)
# UVs (per loop, need mapping to vertices for simplicity we'll store per-face)
uv_layer = me.uv_layers.active
if uv_layer:
uv_data = np.array([uv_layer.data[i].uv for i in range(len(uv_layer.data))], dtype=np.float32)
loop_to_vertex = np.array([me.loops[i].vertex_index for i in range(len(me.loops))], dtype=np.int32)
else:
uv_data = np.zeros((0,2), dtype=np.float32)
loop_to_vertex = np.zeros(0, dtype=np.int32)
np.savez_compressed(filepath,
vertices=vertices,
faces=faces,
uv_data=uv_data,
loop_to_vertex=loop_to_vertex)
print(f"Mesh exported to {filepath}")
def export_camera_pose(camera_obj, filepath):
"""Export camera extrinsics (world-to-camera transform matrix)."""
# Camera world matrix: transforms from camera local to world
cam_to_world = np.array(camera_obj.matrix_world, dtype=np.float32)
# We need world-to-camera transform (inverse)
world_to_cam = np.linalg.inv(cam_to_world)
np.save(filepath, world_to_cam)
print(f"Camera pose exported to {filepath}")
def export_pixel_to_3d_lookup(obj, camera_obj, scene, filepath):
"""
Precompute for every pixel the 3D hit point (world coordinates)
and save as a 2D array (height, width, 3). This avoids raycasting later.
"""
depsgraph = bpy.context.evaluated_depsgraph_get()
render = scene.render
res_x = render.resolution_x
res_y = render.resolution_y
# Get camera data
cam = camera_obj.data
if cam.sensor_fit == 'HORIZONTAL':
sensor_width = cam.sensor_width
sensor_height = sensor_width / (res_x / res_y)
else:
sensor_height = cam.sensor_height
sensor_width = sensor_height * (res_x / res_y)
focal = cam.lens
frame_width = sensor_width / focal
frame_height = sensor_height / focal
# Preallocate array (height, width, 3) for XYZ world coordinates
xyz_map = np.full((res_y, res_x, 3), np.nan, dtype=np.float32)
origin = camera_obj.matrix_world.translation
cam_rot = camera_obj.matrix_world.to_quaternion()
for y in range(res_y):
# Normalized sensor coordinate v (0=bottom, 1=top)
v_ndc = (y + 0.5) / res_y
sy = (v_ndc - 0.5) * frame_height
for x in range(res_x):
u_ndc = (x + 0.5) / res_x
sx = (u_ndc - 0.5) * frame_width
direction_cam = Vector((sx, sy, -1.0))
direction_world = cam_rot @ direction_cam
direction_world.normalize()
hit, location, normal, face_index, hit_obj, matrix = scene.ray_cast(
depsgraph, origin, direction_world)
if hit and hit_obj == obj:
xyz_map[y, x, :] = location
np.savez_compressed(filepath, xyz_map=xyz_map)
print(f"Pixel-to-3D lookup saved to {filepath} (shape {res_y}x{res_x})")
def export_all():
# Configuration
output_dir = bpy.path.abspath("//exported_data/")
os.makedirs(output_dir, exist_ok=True)
obj = bpy.context.active_object # or specify name
cam = bpy.data.objects.get("Camera")
if not cam:
cam = bpy.context.scene.camera
if not obj or not cam:
raise RuntimeError("Select mesh and ensure camera exists.")
# 1. Mesh data
export_mesh_data(obj, os.path.join(output_dir, "mesh.npz"))
# 2. Camera intrinsics (K matrix) and resolution
K, w, h = export_camera_intrinsics(cam, bpy.context.scene)
np.savez(os.path.join(output_dir, "camera_intrinsics.npz"), K=K, width=w, height=h)
# 3. Camera pose (world-to-camera)
export_camera_pose(cam, os.path.join(output_dir, "world_to_cam.npy"))
# 4. Optional: precomputed pixel-to-3D map (fastest)
export_pixel_to_3d_lookup(obj, cam, bpy.context.scene,
os.path.join(output_dir, "pixel_to_xyz.npz"))
print("Export complete.")
if __name__ == "__main__":
export_all()
tools/3d_patch_generation_blender/checker_patch_dd02c.py
+585
View File
@@ -0,0 +1,585 @@
"""
Texture Generator for Arbitrary 3D Surfaces (Blender Script)
- Mode: "patches" (black/white checker patches) or "dots" (circular markers)
- Mesh UVs are filled white so that non‑visible parts become white.
- Texture background (no‑mesh areas) stays black.
"""
import bpy
import numpy as np
import os
import time
from mathutils import Vector
# ----------------------------------------------------------------------
# Common Utilities
# ----------------------------------------------------------------------
def draw_cross(pixels, width, height, cx_uv, cy_uv, length_px, thickness_px, color):
"""
Draw a cross (plus sign) in pixel space.
cx_uv, cy_uv: centre of cross in UV coordinates (0‑1), origin at bottom‑left.
length_px: total length of each arm (pixels).
thickness_px: line thickness (odd number gives symmetric result).
"""
# Convert UV centre to pixel coordinates (y from bottom)
cx = int(round(cx_uv * (width - 1)))
cy = int(round(cy_uv * (height - 1)))
half_len = length_px // 2
half_thick = thickness_px // 2
# Horizontal arm
x0 = max(0, cx - half_len)
x1 = min(width - 1, cx + half_len)
y0 = max(0, cy - half_thick)
y1 = min(height - 1, cy + half_thick)
for y in range(y0, y1 + 1):
for x in range(x0, x1 + 1):
idx = (y * width + x) * 4
pixels[idx:idx+4] = color
# Vertical arm
x0 = max(0, cx - half_thick)
x1 = min(width - 1, cx + half_thick)
y0 = max(0, cy - half_len)
y1 = min(height - 1, cy + half_len)
for y in range(y0, y1 + 1):
for x in range(x0, x1 + 1):
idx = (y * width + x) * 4
pixels[idx:idx+4] = color
def fill_background(pixels, width, height, color):
"""Fill the entire image with a solid RGBA color (0-1 floats)."""
r, g, b, a = color
for y in range(height):
for x in range(width):
idx = (y * width + x) * 4
pixels[idx] = r
pixels[idx+1] = g
pixels[idx+2] = b
pixels[idx+3] = a
def fill_polygon(image, pixels, width, height, polygon, color):
"""Fill a polygon in the pixel array (flattened, RGBA)."""
poly = [(int(round(x)), int(round(y))) for x, y in polygon]
if len(poly) < 3:
return
min_x = max(0, min(p[0] for p in poly))
max_x = min(width-1, max(p[0] for p in poly))
min_y = max(0, min(p[1] for p in poly))
max_y = min(height-1, max(p[1] for p in poly))
for y in range(min_y, max_y + 1):
intersections = []
for i in range(len(poly)):
x1, y1 = poly[i]
x2, y2 = poly[(i+1) % len(poly)]
if y1 == y2:
continue
if (y1 <= y < y2) or (y2 <= y < y1):
t = (y - y1) / (y2 - y1)
x = x1 + t * (x2 - x1)
intersections.append(x)
intersections.sort()
for k in range(0, len(intersections), 2):
if k+1 >= len(intersections):
break
x_start = int(round(intersections[k]))
x_end = int(round(intersections[k+1]))
x_start = max(0, min(width-1, x_start))
x_end = max(0, min(width-1, x_end))
for x in range(x_start, x_end + 1):
idx = (y * width + x) * 4
pixels[idx:idx+4] = color
def fill_polygon_value(grid, width, height, polygon, value):
"""Fill a polygon into a 2D scalar grid."""
poly = [(int(round(x)), int(round(y))) for x, y in polygon]
if len(poly) < 3:
return
min_x = max(0, min(p[0] for p in poly))
max_x = min(width - 1, max(p[0] for p in poly))
min_y = max(0, min(p[1] for p in poly))
max_y = min(height - 1, max(p[1] for p in poly))
for y in range(min_y, max_y + 1):
intersections = []
for i in range(len(poly)):
x1, y1 = poly[i]
x2, y2 = poly[(i + 1) % len(poly)]
if y1 == y2:
continue
if (y1 <= y < y2) or (y2 <= y < y1):
t = (y - y1) / (y2 - y1)
x = x1 + t * (x2 - x1)
intersections.append(x)
intersections.sort()
for k in range(0, len(intersections), 2):
if k + 1 >= len(intersections):
break
x_start = int(round(intersections[k]))
x_end = int(round(intersections[k + 1]))
x_start = max(0, min(width - 1, x_start))
x_end = max(0, min(width - 1, x_end))
grid[y, x_start:x_end + 1] = value
def draw_ellipse(pixels, width, height, cx, cy, a, b, angle, color):
"""Filled ellipse with semi‑axes a,b (pixels), rotation angle (rad)."""
cos_a = np.cos(angle)
sin_a = np.sin(angle)
rx = abs(a * cos_a) + abs(b * sin_a)
ry = abs(a * sin_a) + abs(b * cos_a)
x_min = int(max(0, cx - rx))
x_max = int(min(width-1, cx + rx))
y_min = int(max(0, cy - ry))
y_max = int(min(height-1, cy + ry))
for y in range(y_min, y_max + 1):
for x in range(x_min, x_max + 1):
dx = x - cx
dy = y - cy
u = dx * cos_a + dy * sin_a
v = -dx * sin_a + dy * cos_a
if (u/a)**2 + (v/b)**2 <= 1.0:
idx = (y * width + x) * 4
pixels[idx:idx+4] = color
def get_uv_from_face(obj, face_index, location):
"""World‑space location → UV coordinate on the mesh face."""
me = obj.data
local_loc = obj.matrix_world.inverted() @ location
poly = me.polygons[face_index]
vert_indices = poly.vertices
loop_indices = poly.loop_indices
if len(vert_indices) < 3:
return None
v0 = me.vertices[vert_indices[0]].co
v1 = me.vertices[vert_indices[1]].co
v2 = me.vertices[vert_indices[2]].co
uv_layer = me.uv_layers.active.data
uv0 = uv_layer[loop_indices[0]].uv
uv1 = uv_layer[loop_indices[1]].uv
uv2 = uv_layer[loop_indices[2]].uv
def tri_area(a, b, c):
return 0.5 * ((b - a).cross(c - a)).length
total_area = tri_area(v0, v1, v2)
if total_area == 0:
return None
w0 = tri_area(local_loc, v1, v2) / total_area
w1 = tri_area(local_loc, v2, v0) / total_area
w2 = 1.0 - w0 - w1
uv = w0 * uv0 + w1 * uv1 + w2 * uv2
return Vector((uv.x, uv.y))
def raycast_from_camera(scene, depsgraph, cam, obj, u_ndc, v_ndc, frame_width, frame_height):
"""Camera raycast through normalised sensor coordinates (0‑1)."""
sx = (u_ndc - 0.5) * frame_width
sy = (v_ndc - 0.5) * frame_height
direction_cam = Vector((sx, sy, -1.0))
direction_world = cam.matrix_world.to_quaternion() @ direction_cam
direction_world.normalize()
origin = cam.matrix_world.translation
hit, location, normal, face_index, hit_obj, matrix = scene.ray_cast(
depsgraph, origin, direction_world)
if hit and hit_obj == obj:
return location, face_index
return None, None
def rasterize_mesh_uvs(obj, pixels, tex_width, tex_height, color=(1,1,1,1)):
"""
Fill all triangles of the active UV map with 'color' (default white).
This marks every part of the mesh, visible or not, as white.
"""
me = obj.data
uv_layer = me.uv_layers.active
if not uv_layer:
print("Warning: No active UV layer – mesh rasterization skipped.")
return
uv_data = uv_layer.data
for poly in me.polygons:
# Get loop indices for this polygon
loops = poly.loop_indices
# Build UV polygon
uv_poly = []
for loop_idx in loops:
uv = uv_data[loop_idx].uv
uv_poly.append((uv.x, uv.y))
# Convert UV (0‑1) to pixel coordinates
pixel_poly = [(u * (tex_width - 1), v * (tex_height - 1)) for u, v in uv_poly]
fill_polygon(None, pixels, tex_width, tex_height, pixel_poly, color)
def rasterize_mesh_mask(obj, tex_width, tex_height):
"""Rasterize the active UV map to a boolean mesh mask."""
mask = np.zeros((tex_height, tex_width), dtype=bool)
me = obj.data
uv_layer = me.uv_layers.active
if not uv_layer:
return mask
uv_data = uv_layer.data
for poly in me.polygons:
uv_poly = []
for loop_idx in poly.loop_indices:
uv = uv_data[loop_idx].uv
uv_poly.append((uv.x, uv.y))
pixel_poly = [(u * (tex_width - 1), v * (tex_height - 1)) for u, v in uv_poly]
fill_polygon_value(mask, tex_width, tex_height, pixel_poly, True)
return mask
def paint_mask_color(pixels, width, height, mask, color):
"""Paint all True pixels of a boolean mask with an RGBA color."""
r, g, b, a = color
ys, xs = np.nonzero(mask)
for y, x in zip(ys, xs):
idx = (y * width + x) * 4
pixels[idx:idx+4] = (r, g, b, a)
def complete_checker_gaps(pixels, width, height, mesh_mask, patch_state, max_gap_px=6):
"""
Fill small white slivers inside the mesh by expanding nearby checker colors.
Only grows locally (max_gap_px) so non-visible mesh areas remain white.
patch_state: -1 = unassigned, 0 = black patch, 1 = white patch
"""
for _ in range(max_gap_px):
candidate = mesh_mask & (patch_state < 0)
if not np.any(candidate):
break
black_counts = np.zeros_like(patch_state, dtype=np.uint8)
white_counts = np.zeros_like(patch_state, dtype=np.uint8)
for dy in (-1, 0, 1):
for dx in (-1, 0, 1):
if dx == 0 and dy == 0:
continue
src_y0 = max(0, -dy)
src_y1 = height - max(0, dy)
src_x0 = max(0, -dx)
src_x1 = width - max(0, dx)
dst_y0 = max(0, dy)
dst_y1 = height - max(0, -dy)
dst_x0 = max(0, dx)
dst_x1 = width - max(0, -dx)
neighbor_vals = patch_state[src_y0:src_y1, src_x0:src_x1]
black_counts[dst_y0:dst_y1, dst_x0:dst_x1] += (neighbor_vals == 0)
white_counts[dst_y0:dst_y1, dst_x0:dst_x1] += (neighbor_vals == 1)
fillable = candidate & ((black_counts + white_counts) > 0)
if not np.any(fillable):
break
patch_state[fillable] = np.where(
black_counts[fillable] > white_counts[fillable], 0, 1
).astype(np.int8)
paint_mask_color(pixels, width, height, patch_state == 0, (0.0, 0.0, 0.0, 1.0))
paint_mask_color(pixels, width, height, patch_state == 1, (1.0, 1.0, 1.0, 1.0))
def checker_patch_color(i, j):
"""Return a black/white checker color for patch index (i, j)."""
if (i + j) % 2 == 0:
return np.array([1.0, 1.0, 1.0], dtype=np.float32)
return np.array([0.0, 0.0, 0.0], dtype=np.float32)
# ----------------------------------------------------------------------
# Mode 1: Colored Patches
# ----------------------------------------------------------------------
def generate_patch_texture(
obj, cam, scene, depsgraph,
nx=40, ny=30, r=0.05,
texture_width=2048, texture_height=2048,
output_path="//patch_texture.png",uv_centre=None, cross_length=40, cross_thickness=3, cross_color=(1,0,0,1),
gap_fill_px=6
):
cam_data = cam.data
render = scene.render
res_x = render.resolution_x
res_y = render.resolution_y
if cam_data.sensor_fit == 'HORIZONTAL':
sensor_width = cam_data.sensor_width
sensor_height = sensor_width / (res_x / res_y)
else:
sensor_height = cam_data.sensor_height
sensor_width = sensor_height * (res_x / res_y)
focal = cam_data.lens
frame_width = sensor_width / focal
frame_height = sensor_height / focal
patch_uv_polygons = []
patch_colors = np.zeros((ny, nx, 3), dtype=np.float32)
print(f"Raycasting {nx*ny} patches...")
for j in range(ny):
v_start = j / ny
v_end = (j+1) / ny
for i in range(nx):
u_start = i / nx
u_end = (i+1) / nx
corners_ndc = [(u_start, v_start), (u_end, v_start),
(u_end, v_end), (u_start, v_end)]
uv_poly = []
for (u_ndc, v_ndc) in corners_ndc:
loc, face_idx = raycast_from_camera(scene, depsgraph, cam, obj,
u_ndc, v_ndc, frame_width, frame_height)
if loc is not None:
uv = get_uv_from_face(obj, face_idx, loc)
if uv is not None:
uv_poly.append(uv)
if len(uv_poly) < 3:
patch_uv_polygons.append(None)
continue
patch_uv_polygons.append(uv_poly)
patch_colors[j, i] = checker_patch_color(i, j)
# Create image – start with black background
img_name = "PatchTexture"
if img_name in bpy.data.images:
bpy.data.images.remove(bpy.data.images[img_name])
image = bpy.data.images.new(img_name, width=texture_width, height=texture_height, alpha=True)
pixels = list(image.pixels)
mesh_mask = rasterize_mesh_mask(obj, texture_width, texture_height)
patch_state = np.full((texture_height, texture_width), -1, dtype=np.int8)
# Fill entire image with black (no mesh areas)
fill_background(pixels, texture_width, texture_height, (0,0,0,1))
# Rasterize all mesh UVs in white (visible + non‑visible parts)
paint_mask_color(pixels, texture_width, texture_height, mesh_mask, (1,1,1,1))
# Now draw the visible patches on top
idx = 0
for j in range(ny):
for i in range(nx):
uv_poly = patch_uv_polygons[idx]; idx += 1
if uv_poly is None: continue
pixel_poly = [(u*(texture_width-1), v*(texture_height-1)) for u,v in uv_poly]
color = patch_colors[j,i]
rgba = (color[0], color[1], color[2], 1.0)
fill_polygon(image, pixels, texture_width, texture_height, pixel_poly, rgba)
fill_polygon_value(
patch_state,
texture_width,
texture_height,
pixel_poly,
1 if np.mean(color) >= 0.5 else 0,
)
# Fill small steep-edge gaps from neighboring checker cells.
complete_checker_gaps(
pixels, texture_width, texture_height, mesh_mask, patch_state, max_gap_px=gap_fill_px
)
# (after drawing all patches)
if uv_centre is not None:
draw_cross(pixels, texture_width, texture_height,
uv_centre[0], uv_centre[1],
cross_length, cross_thickness, cross_color)
image.pixels[:] = pixels
image.update()
abs_path = bpy.path.abspath(output_path)
image.filepath_raw = abs_path
image.file_format = 'PNG'
image.save()
print(f"Patch texture saved to {abs_path}")
return image
# ----------------------------------------------------------------------
# Mode 2: Circular Dot Markers
# ----------------------------------------------------------------------
def generate_dot_texture(
obj, cam, scene, depsgraph,
dot_diameter_pixels=12,
spacing_pixels=40,
dot_color=(0,0,0,1), # RGBA 0-1
texture_width=2048, texture_height=2048,
output_path="//dot_texture.png",uv_centre=None, cross_length=40, cross_thickness=3, cross_color=(1,0,0,1)
):
render = scene.render
res_x = render.resolution_x
res_y = render.resolution_y
cam_data = cam.data
if cam_data.sensor_fit == 'HORIZONTAL':
sensor_width = cam_data.sensor_width
sensor_height = sensor_width / (res_x / res_y)
else:
sensor_height = cam_data.sensor_height
sensor_width = sensor_height * (res_x / res_y)
focal = cam_data.lens
frame_width = sensor_width / focal
frame_height = sensor_height / focal
margin = spacing_pixels // 2
x_centers = np.arange(margin, res_x, spacing_pixels)
y_centers = np.arange(margin, res_y, spacing_pixels)
offset_pix = max(2, dot_diameter_pixels // 4)
du_ndc = offset_pix / res_x
dv_ndc = offset_pix / res_y
dots = [] # (cx_uv, cy_uv, a_uv, b_uv, angle_uv)
print(f"Raycasting {len(x_centers)*len(y_centers)} dot centers...")
for y_pix in y_centers:
v_ndc = y_pix / res_y
for x_pix in x_centers:
u_ndc = x_pix / res_x
center_loc, face_idx = raycast_from_camera(scene, depsgraph, cam, obj,
u_ndc, v_ndc, frame_width, frame_height)
if center_loc is None: continue
center_uv = get_uv_from_face(obj, face_idx, center_loc)
if center_uv is None: continue
offsets = [
( u_ndc + du_ndc, v_ndc),
( u_ndc - du_ndc, v_ndc),
( u_ndc, v_ndc + dv_ndc),
( u_ndc, v_ndc - dv_ndc)
]
uv_offsets = []
valid = True
for (u_off, v_off) in offsets:
loc_off, f_idx_off = raycast_from_camera(scene, depsgraph, cam, obj,
u_off, v_off, frame_width, frame_height)
if loc_off is None:
valid = False; break
uv_off = get_uv_from_face(obj, f_idx_off, loc_off)
if uv_off is None:
valid = False; break
uv_offsets.append(uv_off)
if not valid: continue
scale_u = 1.0 / res_x
scale_v = 1.0 / res_y
du_dx = (uv_offsets[0] - uv_offsets[1]) / (2 * offset_pix * scale_u)
dv_dx = (uv_offsets[2] - uv_offsets[3]) / (2 * offset_pix * scale_v)
r_pix = dot_diameter_pixels / 2.0
a_uv = r_pix * du_dx.length * scale_u
b_uv = r_pix * dv_dx.length * scale_v
angle = np.arctan2(du_dx.y, du_dx.x) if du_dx.length > 0 else 0.0
dots.append((center_uv.x, center_uv.y, a_uv, b_uv, angle))
print(f"Found {len(dots)} valid dots.")
if uv_centre is not None:
draw_cross(pixels, texture_width, texture_height,
uv_centre[0], uv_centre[1],
cross_length, cross_thickness, cross_color)
# Create image
img_name = "DotTexture"
if img_name in bpy.data.images:
bpy.data.images.remove(bpy.data.images[img_name])
image = bpy.data.images.new(img_name, width=texture_width, height=texture_height, alpha=True)
pixels = list(image.pixels)
# Black background (no mesh)
fill_background(pixels, texture_width, texture_height, (0,0,0,1))
# White for all mesh areas
rasterize_mesh_uvs(obj, pixels, texture_width, texture_height, color=(1,1,1,1))
# Draw dots
for (cx_uv, cy_uv, a_uv, b_uv, angle_uv) in dots:
cx_px = cx_uv * (texture_width - 1)
cy_px = cy_uv * (texture_height - 1)
a_px = a_uv * texture_width
b_px = b_uv * texture_height
draw_ellipse(pixels, texture_width, texture_height,
cx_px, cy_px, a_px, b_px, angle_uv, dot_color)
image.pixels[:] = pixels
image.update()
abs_path = bpy.path.abspath(output_path)
image.filepath_raw = abs_path
image.file_format = 'PNG'
image.save()
print(f"Dot texture saved to {abs_path}")
return image
# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------
def main():
# ==================== CONFIGURATION ====================
MODE = "patches"
OBJ_NAME = "FingerMesh"
CAM_NAME = "Camera"
TEXTURE_WIDTH = 4096
TEXTURE_HEIGHT = 4096
OUTPUT_PATH = "//" + ("patch_texture.png" if MODE=="patches" else "dot_texture.png")
# Patch mode
PATCH_NX = 40; PATCH_NY = 40; COLOR_R = 0.05
# Dot mode
DOT_DIAMETER_PIX = 12; DOT_SPACING_PIX = 40; DOT_COLOR = (0,0,0,1)
# Cross
CROSS_LENGTH_PX = 40
CROSS_THICKNESS_PX = 3
CROSS_COLOR = (1, 0, 0, 1) # red cross
# ======================================================
start_time = time.time()
scene = bpy.context.scene
depsgraph = bpy.context.evaluated_depsgraph_get()
if OBJ_NAME:
obj = bpy.data.objects.get(OBJ_NAME)
else:
obj = bpy.context.active_object
if obj is None or obj.type != 'MESH':
raise ValueError("No valid mesh object found.")
cam = bpy.data.objects.get(CAM_NAME) if CAM_NAME else scene.camera
if cam is None or cam.type != 'CAMERA':
raise ValueError("No camera found.")
if not obj.data.uv_layers.active:
obj.data.uv_layers.new(name="UVMap")
# Compute UV centre
uv_layer = obj.data.uv_layers.active
uv_coords = np.array([uv_layer.data[i].uv.to_tuple() for i in range(len(uv_layer.data))])
uv_min = uv_coords.min(axis=0)
uv_max = uv_coords.max(axis=0)
uv_centre = (uv_min + uv_max) / 2.0
print(f"UV centre: ({uv_centre[0]:.4f}, {uv_centre[1]:.4f})")
# Compute UV bounding box centre
me = obj.data
uv_layer = me.uv_layers.active
uv_coords = np.array([uv_layer.data[i].uv.to_tuple() for i in range(len(uv_layer.data))])
uv_min = uv_coords.min(axis=0)
uv_max = uv_coords.max(axis=0)
uv_centre = (uv_min + uv_max) / 2.0
print(f"UV centre: {uv_centre[0]:.4f}, {uv_centre[1]:.4f}")
if MODE == "patches":
generate_patch_texture(
obj, cam, scene, depsgraph,
nx=PATCH_NX, ny=PATCH_NY, r=COLOR_R,
texture_width=TEXTURE_WIDTH, texture_height=TEXTURE_HEIGHT,
output_path=OUTPUT_PATH,
uv_centre=uv_centre,
cross_length=CROSS_LENGTH_PX,
cross_thickness=CROSS_THICKNESS_PX,
cross_color=CROSS_COLOR
)
elif MODE == "dots":
generate_dot_texture(
obj, cam, scene, depsgraph,
dot_diameter_pixels=DOT_DIAMETER_PIX,
spacing_pixels=DOT_SPACING_PIX,
dot_color=DOT_COLOR,
texture_width=TEXTURE_WIDTH, texture_height=TEXTURE_HEIGHT,
output_path=OUTPUT_PATH,
uv_centre=uv_centre,
cross_length=CROSS_LENGTH_PX,
cross_thickness=CROSS_THICKNESS_PX,
cross_color=CROSS_COLOR
)
print(f"Completed in {time.time()-start_time:.2f} seconds.")
if __name__ == "__main__":
main()
tools/3d_patch_generation_blender/dd02c.blend
BIN
View File
Binary file not shown.
tools/3d_patch_generation_blender/mapto3d.py
+177
View File
@@ -0,0 +1,177 @@
import numpy as np
import cv2
import os
def load_exported_data(data_dir):
"""Load all exported files."""
# Mesh
mesh = np.load(os.path.join(data_dir, "mesh.npz"))
vertices = mesh['vertices']
faces = mesh['faces']
uv_data = mesh['uv_data']
loop_to_vertex = mesh['loop_to_vertex']
# Camera intrinsics
cam_data = np.load(os.path.join(data_dir, "camera_intrinsics.npz"))
K = cam_data['K']
width = int(cam_data['width'])
height = int(cam_data['height'])
# Camera pose (world-to-camera)
world_to_cam = np.load(os.path.join(data_dir, "world_to_cam.npy"))
# Precomputed pixel-to-3D map (optional)
xyz_map_path = os.path.join(data_dir, "pixel_to_xyz.npz")
if os.path.exists(xyz_map_path):
xyz_map = np.load(xyz_map_path)['xyz_map']
else:
xyz_map = None
return {
'vertices': vertices,
'faces': faces,
'uv_data': uv_data,
'loop_to_vertex': loop_to_vertex,
'K': K,
'width': width,
'height': height,
'world_to_cam': world_to_cam,
'xyz_map': xyz_map
}
def ray_intersect_triangle(ray_origin, ray_dir, v0, v1, v2):
"""Möller–Trumbore intersection algorithm. Returns t (distance) or None."""
epsilon = 1e-8
edge1 = v1 - v0
edge2 = v2 - v0
h = np.cross(ray_dir, edge2)
a = np.dot(edge1, h)
if abs(a) < epsilon:
return None
f = 1.0 / a
s = ray_origin - v0
u = f * np.dot(s, h)
if u < 0.0 or u > 1.0:
return None
q = np.cross(s, edge1)
v = f * np.dot(ray_dir, q)
if v < 0.0 or u + v > 1.0:
return None
t = f * np.dot(edge2, q)
if t > epsilon:
return t
return None
def pixel_to_3d_ray(u, v, K, world_to_cam):
"""
Given pixel coordinates (u,v) and camera intrinsics K,
return a ray in world coordinates: origin and direction.
"""
# Invert K to get normalized image coordinates
fx = K[0,0]; fy = K[1,1]; cx = K[0,2]; cy = K[1,2]
x_norm = (u - cx) / fx
y_norm = (v - cy) / fy
# Ray in camera space: (x_norm, y_norm, 1.0)
ray_cam = np.array([x_norm, y_norm, 1.0], dtype=np.float64)
# Transform to world
cam_to_world = np.linalg.inv(world_to_cam)
R = cam_to_world[:3, :3]
t = cam_to_world[:3, 3]
ray_world_dir = R @ ray_cam
ray_world_dir /= np.linalg.norm(ray_world_dir)
ray_world_origin = t
return ray_world_origin, ray_world_dir
def map_image_to_3d(image_path, data, use_precomputed=True):
"""
Map each pixel of the image to a 3D world point.
If use_precomputed and xyz_map exists, return that directly.
Otherwise, perform raycasting on the mesh.
Returns:
points: (H, W, 3) array of 3D points (NaN for misses).
colors: (H, W, 3) array of RGB colors from image.
"""
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # to RGB
h, w = img.shape[:2]
if use_precomputed and data['xyz_map'] is not None:
# Use precomputed map (must match resolution)
if h != data['height'] or w != data['width']:
raise ValueError("Image resolution does not match precomputed map.")
points = data['xyz_map'].copy()
colors = img.astype(np.float32) / 255.0
return points, colors
# Otherwise, perform raycasting on the mesh
vertices = data['vertices']
faces = data['faces']
K = data['K']
world_to_cam = data['world_to_cam']
points = np.full((h, w, 3), np.nan, dtype=np.float32)
colors = img.astype(np.float32) / 255.0
# For speed, we can use a bounding volume hierarchy (BVH).
# Here we implement a simple brute-force over all faces (slow for large meshes).
# In practice, precomputing the xyz_map is strongly recommended.
for y in range(h):
for x in range(w):
origin, direction = pixel_to_3d_ray(x, y, K, world_to_cam)
best_t = np.inf
best_hit = None
for face in faces:
v0 = vertices[face[0]]
v1 = vertices[face[1]]
v2 = vertices[face[2]]
t = ray_intersect_triangle(origin, direction, v0, v1, v2)
if t is not None and t < best_t:
best_t = t
best_hit = origin + direction * t
if best_hit is not None:
points[y, x] = best_hit
return points, colors
def save_colored_pointcloud(points, colors, output_path):
"""Save as PLY file (colored point cloud)."""
# Flatten valid points
mask = ~np.isnan(points[:,:,0])
valid_points = points[mask].reshape(-1, 3)
valid_colors = (colors[mask].reshape(-1, 3) * 255).astype(np.uint8)
with open(output_path, 'w') as f:
f.write("ply\n")
f.write("format ascii 1.0\n")
f.write(f"element vertex {len(valid_points)}\n")
f.write("property float x\n")
f.write("property float y\n")
f.write("property float z\n")
f.write("property uchar red\n")
f.write("property uchar green\n")
f.write("property uchar blue\n")
f.write("end_header\n")
for pt, col in zip(valid_points, valid_colors):
f.write(f"{pt[0]} {pt[1]} {pt[2]} {col[0]} {col[1]} {col[2]}\n")
print(f"Saved point cloud to {output_path}")
if __name__ == "__main__":
data_dir = "./exported_data" # Path to exported files
image_path = "captured_image.png" # Input camera image
output_ply = "mapped_points.ply"
data = load_exported_data(data_dir)
points, colors = map_image_to_3d(image_path, data, use_precomputed=True)
save_colored_pointcloud(points, colors, output_ply)
# Optional: Visualize using Open3D
import open3d as o3d
pcd = o3d.geometry.PointCloud()
mask = ~np.isnan(points[:,:,0])
pcd.points = o3d.utility.Vector3dVector(points[mask].reshape(-1,3))
pcd.colors = o3d.utility.Vector3dVector(colors[mask].reshape(-1,3))
o3d.visualization.draw_geometries([pcd])
tools/3d_patch_generation_blender/random_patch_dd02c.py
+528
View File
@@ -0,0 +1,528 @@
"""
Texture Generator for Arbitrary 3D Surfaces (Blender Script)
- Mode: "patches" (colored patches) or "dots" (circular markers)
- Mesh UVs are filled white so that non‑visible parts become white.
- Texture background (no‑mesh areas) stays black.
"""
import bpy
import numpy as np
import os
import time
from mathutils import Vector
# ----------------------------------------------------------------------
# Common Utilities
# ----------------------------------------------------------------------
def draw_cross(pixels, width, height, cx_uv, cy_uv, length_px, thickness_px, color):
"""
Draw a cross (plus sign) in pixel space.
cx_uv, cy_uv: centre of cross in UV coordinates (0‑1), origin at bottom‑left.
length_px: total length of each arm (pixels).
thickness_px: line thickness (odd number gives symmetric result).
"""
# Convert UV centre to pixel coordinates (y from bottom)
cx = int(round(cx_uv * (width - 1)))
cy = int(round(cy_uv * (height - 1)))
half_len = length_px // 2
half_thick = thickness_px // 2
# Horizontal arm
x0 = max(0, cx - half_len)
x1 = min(width - 1, cx + half_len)
y0 = max(0, cy - half_thick)
y1 = min(height - 1, cy + half_thick)
for y in range(y0, y1 + 1):
for x in range(x0, x1 + 1):
idx = (y * width + x) * 4
pixels[idx:idx+4] = color
# Vertical arm
x0 = max(0, cx - half_thick)
x1 = min(width - 1, cx + half_thick)
y0 = max(0, cy - half_len)
y1 = min(height - 1, cy + half_len)
for y in range(y0, y1 + 1):
for x in range(x0, x1 + 1):
idx = (y * width + x) * 4
pixels[idx:idx+4] = color
def fill_background(pixels, width, height, color):
"""Fill the entire image with a solid RGBA color (0-1 floats)."""
r, g, b, a = color
for y in range(height):
for x in range(width):
idx = (y * width + x) * 4
pixels[idx] = r
pixels[idx+1] = g
pixels[idx+2] = b
pixels[idx+3] = a
def fill_polygon(image, pixels, width, height, polygon, color):
"""Fill a polygon in the pixel array (flattened, RGBA)."""
poly = [(int(round(x)), int(round(y))) for x, y in polygon]
if len(poly) < 3:
return
min_x = max(0, min(p[0] for p in poly))
max_x = min(width-1, max(p[0] for p in poly))
min_y = max(0, min(p[1] for p in poly))
max_y = min(height-1, max(p[1] for p in poly))
for y in range(min_y, max_y + 1):
intersections = []
for i in range(len(poly)):
x1, y1 = poly[i]
x2, y2 = poly[(i+1) % len(poly)]
if y1 == y2:
continue
if (y1 <= y < y2) or (y2 <= y < y1):
t = (y - y1) / (y2 - y1)
x = x1 + t * (x2 - x1)
intersections.append(x)
intersections.sort()
for k in range(0, len(intersections), 2):
if k+1 >= len(intersections):
break
x_start = int(round(intersections[k]))
x_end = int(round(intersections[k+1]))
x_start = max(0, min(width-1, x_start))
x_end = max(0, min(width-1, x_end))
for x in range(x_start, x_end + 1):
idx = (y * width + x) * 4
pixels[idx:idx+4] = color
def draw_ellipse(pixels, width, height, cx, cy, a, b, angle, color):
"""Filled ellipse with semi‑axes a,b (pixels), rotation angle (rad)."""
cos_a = np.cos(angle)
sin_a = np.sin(angle)
rx = abs(a * cos_a) + abs(b * sin_a)
ry = abs(a * sin_a) + abs(b * cos_a)
x_min = int(max(0, cx - rx))
x_max = int(min(width-1, cx + rx))
y_min = int(max(0, cy - ry))
y_max = int(min(height-1, cy + ry))
for y in range(y_min, y_max + 1):
for x in range(x_min, x_max + 1):
dx = x - cx
dy = y - cy
u = dx * cos_a + dy * sin_a
v = -dx * sin_a + dy * cos_a
if (u/a)**2 + (v/b)**2 <= 1.0:
idx = (y * width + x) * 4
pixels[idx:idx+4] = color
def get_uv_from_face(obj, face_index, location):
"""World‑space location → UV coordinate on the mesh face."""
me = obj.data
local_loc = obj.matrix_world.inverted() @ location
poly = me.polygons[face_index]
vert_indices = poly.vertices
loop_indices = poly.loop_indices
if len(vert_indices) < 3:
return None
v0 = me.vertices[vert_indices[0]].co
v1 = me.vertices[vert_indices[1]].co
v2 = me.vertices[vert_indices[2]].co
uv_layer = me.uv_layers.active.data
uv0 = uv_layer[loop_indices[0]].uv
uv1 = uv_layer[loop_indices[1]].uv
uv2 = uv_layer[loop_indices[2]].uv
def tri_area(a, b, c):
return 0.5 * ((b - a).cross(c - a)).length
total_area = tri_area(v0, v1, v2)
if total_area == 0:
return None
w0 = tri_area(local_loc, v1, v2) / total_area
w1 = tri_area(local_loc, v2, v0) / total_area
w2 = 1.0 - w0 - w1
uv = w0 * uv0 + w1 * uv1 + w2 * uv2
return Vector((uv.x, uv.y))
def raycast_from_camera(scene, depsgraph, cam, obj, u_ndc, v_ndc, frame_width, frame_height):
"""Camera raycast through normalised sensor coordinates (0‑1)."""
sx = (u_ndc - 0.5) * frame_width
sy = (v_ndc - 0.5) * frame_height
direction_cam = Vector((sx, sy, -1.0))
direction_world = cam.matrix_world.to_quaternion() @ direction_cam
direction_world.normalize()
origin = cam.matrix_world.translation
hit, location, normal, face_index, hit_obj, matrix = scene.ray_cast(
depsgraph, origin, direction_world)
if hit and hit_obj == obj:
return location, face_index
return None, None
def rasterize_mesh_uvs(obj, pixels, tex_width, tex_height, color=(1,1,1,1)):
"""
Fill all triangles of the active UV map with 'color' (default white).
This marks every part of the mesh, visible or not, as white.
"""
me = obj.data
uv_layer = me.uv_layers.active
if not uv_layer:
print("Warning: No active UV layer – mesh rasterization skipped.")
return
uv_data = uv_layer.data
for poly in me.polygons:
# Get loop indices for this polygon
loops = poly.loop_indices
# Build UV polygon
uv_poly = []
for loop_idx in loops:
uv = uv_data[loop_idx].uv
uv_poly.append((uv.x, uv.y))
# Convert UV (0‑1) to pixel coordinates
pixel_poly = [(u * (tex_width - 1), v * (tex_height - 1)) for u, v in uv_poly]
fill_polygon(None, pixels, tex_width, tex_height, pixel_poly, color)
# ----------------------------------------------------------------------
# Color generation (for patch mode)
# ----------------------------------------------------------------------
def merge_intervals(intervals):
if not intervals: return []
intervals.sort(key=lambda x: x[0])
merged = [intervals[0]]
for current in intervals[1:]:
last = merged[-1]
if current[0] <= last[1]:
merged[-1] = (last[0], max(last[1], current[1]))
else:
merged.append(current)
return merged
def get_allowed_intervals(neighbor_values, delta):
forbidden = []
for c in neighbor_values:
low = max(0, c - delta)
high = min(1, c + delta)
forbidden.append((low, high))
merged = merge_intervals(forbidden)
allowed = []
if merged:
if merged[0][0] > 0:
allowed.append((0, merged[0][0]))
for i in range(1, len(merged)):
allowed.append((merged[i-1][1], merged[i][0]))
if merged[-1][1] < 1:
allowed.append((merged[-1][1], 1))
else:
allowed.append((0, 1))
return allowed
def sample_from_allowed(allowed):
lengths = [end - start for start, end in allowed]
total_length = sum(lengths)
if total_length <= 0: return None
point = np.random.rand() * total_length
cumulative = 0
for (start, end), length in zip(allowed, lengths):
cumulative += length
if point <= cumulative:
return np.random.uniform(start, end)
return allowed[-1][0]
def sample_channel(neighbor_values, delta):
allowed = get_allowed_intervals(neighbor_values, delta)
if not allowed: return None
return sample_from_allowed(allowed)
def generate_valid_color(neighbors, r):
delta = np.sqrt(r)
color = np.zeros(3)
for channel in range(3):
neighbor_channel = [n[channel] for n in neighbors]
sampled = sample_channel(neighbor_channel, delta)
if sampled is None:
sampled = np.random.rand()
color[channel] = sampled
return color
# ----------------------------------------------------------------------
# Mode 1: Colored Patches
# ----------------------------------------------------------------------
def generate_patch_texture(
obj, cam, scene, depsgraph,
nx=40, ny=30, r=0.05,
texture_width=2048, texture_height=2048,
output_path="//patch_texture.png",uv_centre=None, cross_length=40, cross_thickness=3, cross_color=(1,0,0,1)
):
cam_data = cam.data
render = scene.render
res_x = render.resolution_x
res_y = render.resolution_y
if cam_data.sensor_fit == 'HORIZONTAL':
sensor_width = cam_data.sensor_width
sensor_height = sensor_width / (res_x / res_y)
else:
sensor_height = cam_data.sensor_height
sensor_width = sensor_height * (res_x / res_y)
focal = cam_data.lens
frame_width = sensor_width / focal
frame_height = sensor_height / focal
patch_uv_polygons = []
patch_colors = np.zeros((ny, nx, 3), dtype=np.float32)
print(f"Raycasting {nx*ny} patches...")
for j in range(ny):
v_start = j / ny
v_end = (j+1) / ny
for i in range(nx):
u_start = i / nx
u_end = (i+1) / nx
corners_ndc = [(u_start, v_start), (u_end, v_start),
(u_end, v_end), (u_start, v_end)]
uv_poly = []
for (u_ndc, v_ndc) in corners_ndc:
loc, face_idx = raycast_from_camera(scene, depsgraph, cam, obj,
u_ndc, v_ndc, frame_width, frame_height)
if loc is not None:
uv = get_uv_from_face(obj, face_idx, loc)
if uv is not None:
uv_poly.append(uv)
if len(uv_poly) < 3:
patch_uv_polygons.append(None)
continue
patch_uv_polygons.append(uv_poly)
neighbors = []
if j > 0: neighbors.append(patch_colors[j-1, i])
if i > 0: neighbors.append(patch_colors[j, i-1])
color = np.random.rand(3) if len(neighbors)==0 else generate_valid_color(neighbors, r)
patch_colors[j, i] = color
# Create image – start with black background
img_name = "PatchTexture"
if img_name in bpy.data.images:
bpy.data.images.remove(bpy.data.images[img_name])
image = bpy.data.images.new(img_name, width=texture_width, height=texture_height, alpha=True)
pixels = list(image.pixels)
# Fill entire image with black (no mesh areas)
fill_background(pixels, texture_width, texture_height, (0,0,0,1))
# Rasterize all mesh UVs in white (visible + non‑visible parts)
rasterize_mesh_uvs(obj, pixels, texture_width, texture_height, color=(1,1,1,1))
# Now draw the visible patches on top
idx = 0
for j in range(ny):
for i in range(nx):
uv_poly = patch_uv_polygons[idx]; idx += 1
if uv_poly is None: continue
pixel_poly = [(u*(texture_width-1), v*(texture_height-1)) for u,v in uv_poly]
color = patch_colors[j,i]
rgba = (color[0], color[1], color[2], 1.0)
fill_polygon(image, pixels, texture_width, texture_height, pixel_poly, rgba)
# (after drawing all patches)
if uv_centre is not None:
draw_cross(pixels, texture_width, texture_height,
uv_centre[0], uv_centre[1],
cross_length, cross_thickness, cross_color)
image.pixels[:] = pixels
image.update()
abs_path = bpy.path.abspath(output_path)
image.filepath_raw = abs_path
image.file_format = 'PNG'
image.save()
print(f"Patch texture saved to {abs_path}")
return image
# ----------------------------------------------------------------------
# Mode 2: Circular Dot Markers
# ----------------------------------------------------------------------
def generate_dot_texture(
obj, cam, scene, depsgraph,
dot_diameter_pixels=12,
spacing_pixels=40,
dot_color=(0,0,0,1), # RGBA 0-1
texture_width=2048, texture_height=2048,
output_path="//dot_texture.png",uv_centre=None, cross_length=40, cross_thickness=3, cross_color=(1,0,0,1)
):
render = scene.render
res_x = render.resolution_x
res_y = render.resolution_y
cam_data = cam.data
if cam_data.sensor_fit == 'HORIZONTAL':
sensor_width = cam_data.sensor_width
sensor_height = sensor_width / (res_x / res_y)
else:
sensor_height = cam_data.sensor_height
sensor_width = sensor_height * (res_x / res_y)
focal = cam_data.lens
frame_width = sensor_width / focal
frame_height = sensor_height / focal
margin = spacing_pixels // 2
x_centers = np.arange(margin, res_x, spacing_pixels)
y_centers = np.arange(margin, res_y, spacing_pixels)
offset_pix = max(2, dot_diameter_pixels // 4)
du_ndc = offset_pix / res_x
dv_ndc = offset_pix / res_y
dots = [] # (cx_uv, cy_uv, a_uv, b_uv, angle_uv)
print(f"Raycasting {len(x_centers)*len(y_centers)} dot centers...")
for y_pix in y_centers:
v_ndc = y_pix / res_y
for x_pix in x_centers:
u_ndc = x_pix / res_x
center_loc, face_idx = raycast_from_camera(scene, depsgraph, cam, obj,
u_ndc, v_ndc, frame_width, frame_height)
if center_loc is None: continue
center_uv = get_uv_from_face(obj, face_idx, center_loc)
if center_uv is None: continue
offsets = [
( u_ndc + du_ndc, v_ndc),
( u_ndc - du_ndc, v_ndc),
( u_ndc, v_ndc + dv_ndc),
( u_ndc, v_ndc - dv_ndc)
]
uv_offsets = []
valid = True
for (u_off, v_off) in offsets:
loc_off, f_idx_off = raycast_from_camera(scene, depsgraph, cam, obj,
u_off, v_off, frame_width, frame_height)
if loc_off is None:
valid = False; break
uv_off = get_uv_from_face(obj, f_idx_off, loc_off)
if uv_off is None:
valid = False; break
uv_offsets.append(uv_off)
if not valid: continue
scale_u = 1.0 / res_x
scale_v = 1.0 / res_y
du_dx = (uv_offsets[0] - uv_offsets[1]) / (2 * offset_pix * scale_u)
dv_dx = (uv_offsets[2] - uv_offsets[3]) / (2 * offset_pix * scale_v)
r_pix = dot_diameter_pixels / 2.0
a_uv = r_pix * du_dx.length * scale_u
b_uv = r_pix * dv_dx.length * scale_v
angle = np.arctan2(du_dx.y, du_dx.x) if du_dx.length > 0 else 0.0
dots.append((center_uv.x, center_uv.y, a_uv, b_uv, angle))
print(f"Found {len(dots)} valid dots.")
if uv_centre is not None:
draw_cross(pixels, texture_width, texture_height,
uv_centre[0], uv_centre[1],
cross_length, cross_thickness, cross_color)
# Create image
img_name = "DotTexture"
if img_name in bpy.data.images:
bpy.data.images.remove(bpy.data.images[img_name])
image = bpy.data.images.new(img_name, width=texture_width, height=texture_height, alpha=True)
pixels = list(image.pixels)
# Black background (no mesh)
fill_background(pixels, texture_width, texture_height, (0,0,0,1))
# White for all mesh areas
rasterize_mesh_uvs(obj, pixels, texture_width, texture_height, color=(1,1,1,1))
# Draw dots
for (cx_uv, cy_uv, a_uv, b_uv, angle_uv) in dots:
cx_px = cx_uv * (texture_width - 1)
cy_px = cy_uv * (texture_height - 1)
a_px = a_uv * texture_width
b_px = b_uv * texture_height
draw_ellipse(pixels, texture_width, texture_height,
cx_px, cy_px, a_px, b_px, angle_uv, dot_color)
image.pixels[:] = pixels
image.update()
abs_path = bpy.path.abspath(output_path)
image.filepath_raw = abs_path
image.file_format = 'PNG'
image.save()
print(f"Dot texture saved to {abs_path}")
return image
# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------
def main():
# ==================== CONFIGURATION ====================
MODE = "patches"
OBJ_NAME = "FingerMesh"
CAM_NAME = "Camera"
TEXTURE_WIDTH = 4096
TEXTURE_HEIGHT = 4096
OUTPUT_PATH = "//" + ("patch_texture.png" if MODE=="patches" else "dot_texture.png")
# Patch mode
PATCH_NX = 60; PATCH_NY = 50; COLOR_R = 0.05
# Dot mode
DOT_DIAMETER_PIX = 12; DOT_SPACING_PIX = 40; DOT_COLOR = (0,0,0,1)
# Cross
CROSS_LENGTH_PX = 40
CROSS_THICKNESS_PX = 3
CROSS_COLOR = (1, 0, 0, 1) # red cross
# ======================================================
start_time = time.time()
scene = bpy.context.scene
depsgraph = bpy.context.evaluated_depsgraph_get()
if OBJ_NAME:
obj = bpy.data.objects.get(OBJ_NAME)
else:
obj = bpy.context.active_object
if obj is None or obj.type != 'MESH':
raise ValueError("No valid mesh object found.")
cam = bpy.data.objects.get(CAM_NAME) if CAM_NAME else scene.camera
if cam is None or cam.type != 'CAMERA':
raise ValueError("No camera found.")
if not obj.data.uv_layers.active:
obj.data.uv_layers.new(name="UVMap")
# Compute UV centre
uv_layer = obj.data.uv_layers.active
uv_coords = np.array([uv_layer.data[i].uv.to_tuple() for i in range(len(uv_layer.data))])
uv_min = uv_coords.min(axis=0)
uv_max = uv_coords.max(axis=0)
uv_centre = (uv_min + uv_max) / 2.0
print(f"UV centre: ({uv_centre[0]:.4f}, {uv_centre[1]:.4f})")
# Compute UV bounding box centre
me = obj.data
uv_layer = me.uv_layers.active
uv_coords = np.array([uv_layer.data[i].uv.to_tuple() for i in range(len(uv_layer.data))])
uv_min = uv_coords.min(axis=0)
uv_max = uv_coords.max(axis=0)
uv_centre = (uv_min + uv_max) / 2.0
print(f"UV centre: {uv_centre[0]:.4f}, {uv_centre[1]:.4f}")
if MODE == "patches":
generate_patch_texture(
obj, cam, scene, depsgraph,
nx=PATCH_NX, ny=PATCH_NY, r=COLOR_R,
texture_width=TEXTURE_WIDTH, texture_height=TEXTURE_HEIGHT,
output_path=OUTPUT_PATH,
uv_centre=uv_centre,
cross_length=CROSS_LENGTH_PX,
cross_thickness=CROSS_THICKNESS_PX,
cross_color=CROSS_COLOR
)
elif MODE == "dots":
generate_dot_texture(
obj, cam, scene, depsgraph,
dot_diameter_pixels=DOT_DIAMETER_PIX,
spacing_pixels=DOT_SPACING_PIX,
dot_color=DOT_COLOR,
texture_width=TEXTURE_WIDTH, texture_height=TEXTURE_HEIGHT,
output_path=OUTPUT_PATH,
uv_centre=uv_centre,
cross_length=CROSS_LENGTH_PX,
cross_thickness=CROSS_THICKNESS_PX,
cross_color=CROSS_COLOR
)
print(f"Completed in {time.time()-start_time:.2f} seconds.")
if __name__ == "__main__":
main()
tools/3d_patch_generation_blender/uv_to_printer.py
+180
View File
@@ -0,0 +1,180 @@
"""
Convert a UV texture map to a physically scaled image for UV printing.
Usage: python uv_to_printer.py --model model.obj --texture texture.png --resolution 100 --output print_ready.png
"""
import argparse
import numpy as np
import cv2
import trimesh
import os
def compute_uv_physical_scale(mesh, uv_layer_name=None):
"""
Compute the physical dimensions (mm) of the UV space.
Assumes the UV map is a reasonably isometric unwrapping (e.g., cylindrical).
Returns: (width_mm, height_mm, scale_factor_u, scale_factor_v)
where scale_factor = physical length per UV unit.
"""
# Get UV coordinates and corresponding 3D positions
if not hasattr(mesh.visual, 'uv') or mesh.visual.uv is None:
raise ValueError("Mesh has no UV coordinates.")
uvs = mesh.visual.uv
faces = mesh.faces
vertices = mesh.vertices
# For each face, we have 3 UVs and 3 vertex positions
# We'll compute the physical edge lengths for all edges in UV space
# and compare to 3D edge lengths to estimate scaling.
total_uv_length = 0.0
total_3d_length = 0.0
edge_count = 0
# Also track UV bounding box for overall dimensions
uv_min = uvs.min(axis=0)
uv_max = uvs.max(axis=0)
uv_range = uv_max - uv_min
# Estimate scale factor by sampling edges
for face in faces:
for i in range(3):
j = (i + 1) % 3
# UV edge
uv0 = uvs[face[i]]
uv1 = uvs[face[j]]
uv_len = np.linalg.norm(uv1 - uv0)
# 3D edge
v0 = vertices[face[i]]
v1 = vertices[face[j]]
len3d = np.linalg.norm(v1 - v0)
if uv_len > 1e-6:
total_uv_length += uv_len
total_3d_length += len3d
edge_count += 1
if edge_count == 0:
raise ValueError("No valid edges found for scale estimation.")
# Average scale factor (physical mm per UV unit)
scale_factor = total_3d_length / total_uv_length
# Physical dimensions of the UV layout
width_mm = uv_range[0] * scale_factor
height_mm = uv_range[1] * scale_factor
print(f"UV bounding box: {uv_range[0]:.3f} x {uv_range[1]:.3f} UV units")
print(f"Scale factor: {scale_factor:.3f} mm per UV unit")
print(f"Physical size: {width_mm:.2f} mm x {height_mm:.2f} mm")
return width_mm, height_mm, scale_factor
def generate_print_image(texture_path, width_mm, height_mm, pixels_per_mm, output_path):
"""
Load the texture image, scale it to match the physical dimensions at the
specified resolution, and save the result.
"""
# Read texture (with alpha if present)
texture = cv2.imread(texture_path, cv2.IMREAD_UNCHANGED)
if texture is None:
raise FileNotFoundError(f"Cannot load texture: {texture_path}")
# Compute target pixel dimensions
target_width_px = int(round(width_mm * pixels_per_mm))
target_height_px = int(round(height_mm * pixels_per_mm))
print(f"Target image size: {target_width_px} x {target_height_px} pixels")
# Resize texture (using INTER_LINEAR for up/down scaling)
resized = cv2.resize(texture, (target_width_px, target_height_px),
interpolation=cv2.INTER_LINEAR)
# Save output (preserve alpha if present)
cv2.imwrite(output_path, resized)
print(f"Print-ready image saved to: {output_path}")
def process_exported_npz(npz_path, texture_path, pixels_per_mm, output_path):
"""
Alternative: load data exported from Blender script (vertices, faces, UVs).
"""
data = np.load(npz_path)
vertices = data['vertices']
faces = data['faces']
uv_data = data['uv_data']
loop_to_vertex = data['loop_to_vertex']
# Reconstruct per-vertex UVs (average)
# For simplicity, we'll compute physical scale using face edges as in trimesh method.
# We need to build a temporary mesh-like structure.
# Count number of faces
num_faces = len(faces)
# Build face UVs array (3 per face)
# The uv_data is per-loop, we need to group by face.
# Assuming faces are triangles.
face_uvs = []
face_vertices = []
uv_idx = 0
for face in faces:
f_uvs = []
f_verts = []
for _ in range(len(face)):
f_uvs.append(uv_data[uv_idx])
v_idx = loop_to_vertex[uv_idx]
f_verts.append(vertices[v_idx])
uv_idx += 1
face_uvs.append(np.array(f_uvs))
face_vertices.append(np.array(f_verts))
total_uv_len = 0.0
total_3d_len = 0.0
count = 0
uv_min = np.min(uv_data, axis=0)
uv_max = np.max(uv_data, axis=0)
uv_range = uv_max - uv_min
for f_uv, f_v in zip(face_uvs, face_vertices):
for i in range(3):
j = (i+1)%3
uv_len = np.linalg.norm(f_uv[i] - f_uv[j])
v_len = np.linalg.norm(f_v[i] - f_v[j])
if uv_len > 1e-6:
total_uv_len += uv_len
total_3d_len += v_len
count += 1
scale = total_3d_len / total_uv_len if count>0 else 1.0
width_mm = uv_range[0] * scale
height_mm = uv_range[1] * scale
print(f"Physical size from .npz: {width_mm:.2f} x {height_mm:.2f} mm")
generate_print_image(texture_path, width_mm, height_mm, pixels_per_mm, output_path)
def main():
parser = argparse.ArgumentParser(description="Convert UV texture to physically scaled print image.")
parser.add_argument("--model", help="3D model file (OBJ, PLY, STL) with UVs")
parser.add_argument("--npz", help="Alternatively, use exported .npz from Blender script")
parser.add_argument("--texture", required=True, help="Input texture image (PNG, JPG)")
parser.add_argument("--resolution", type=float, default=100.0,
help="Pixels per millimeter (default: 100)")
parser.add_argument("--output", default="print_ready.png", help="Output image path")
args = parser.parse_args()
if not args.model and not args.npz:
parser.error("Either --model or --npz must be provided.")
if args.npz:
process_exported_npz(args.npz, args.texture, args.resolution, args.output)
else:
# Load mesh with trimesh
mesh = trimesh.load(args.model, force='mesh')
if not isinstance(mesh, trimesh.Trimesh):
raise ValueError("Loaded object is not a single mesh.")
width_mm, height_mm, scale = compute_uv_physical_scale(mesh)
generate_print_image(args.texture, width_mm, height_mm, args.resolution, args.output)
if __name__ == "__main__":
main()
patch_generation/3.6mm.svg → tools/patch_generation/3.6mm.svg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 74 KiB

After

Width:  |  Height:  |  Size: 74 KiB

patch_generation/cmyk_generated_rgb.png → tools/patch_generation/cmyk_generated_rgb.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 929 KiB

After

Width:  |  Height:  |  Size: 929 KiB

patch_generation/cmyk_pattern.tiff → tools/patch_generation/cmyk_pattern.tiff
View File
patch_generation/dd01.svg → tools/patch_generation/dd01.svg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 2.6 MiB

After

Width:  |  Height:  |  Size: 2.6 MiB

patch_generation/patch_generation.py → tools/patch_generation/patch_generation.py
View File
patch_generation/patch_generation_3.py → tools/patch_generation/patch_generation_3.py
View File
tools/patch_generation/patch_generation_3_org.py
+140
View File
@@ -0,0 +1,140 @@
#%%
import numpy as np
import cv2
def merge_intervals(intervals):
# Merge overlapping intervals.
if not intervals:
return []
intervals.sort(key=lambda x: x[0])
merged = [intervals[0]]
for current in intervals[1:]:
last = merged[-1]
if current[0] <= last[1]:
merged[-1] = (last[0], max(last[1], current[1]))
else:
merged.append(current)
return merged
def get_allowed_intervals(neighbor_values, delta):
# Compute forbidden intervals for this channel from neighbor values.
forbidden = []
for c in neighbor_values:
low = max(0, c - delta)
high = min(1, c + delta)
forbidden.append((low, high))
# Merge overlapping forbidden intervals.
merged = merge_intervals(forbidden)
allowed = []
# Allowed region before the first forbidden interval.
if merged:
if merged[0][0] > 0:
allowed.append((0, merged[0][0]))
# Allowed regions between forbidden intervals.
for i in range(1, len(merged)):
allowed.append((merged[i-1][1], merged[i][0]))
# Allowed region after the last forbidden interval.
if merged[-1][1] < 1:
allowed.append((merged[-1][1], 1))
else:
allowed.append((0, 1))
return allowed
def sample_from_allowed(allowed):
# Choose a random value from the allowed intervals with probability proportional to interval length.
lengths = [end - start for start, end in allowed]
total_length = sum(lengths)
if total_length <= 0:
return None
point = np.random.rand() * total_length
cumulative = 0
for (start, end), length in zip(allowed, lengths):
cumulative += length
if point <= cumulative:
return np.random.uniform(start, end)
return allowed[-1][0]
def sample_channel(neighbor_values, delta):
# Compute allowed intervals and sample a value.
allowed = get_allowed_intervals(neighbor_values, delta)
if not allowed:
return None
return sample_from_allowed(allowed)
def generate_valid_color(neighbors, r):
"""
Generate a random color that satisfies, for each neighbor,
|R - R_neighbor| >= sqrt(r), and similarly for G and B.
neighbors: list of neighbor colors (each is an array of shape (3,))
r: threshold for squared difference.
"""
delta = np.sqrt(r)
color = np.zeros(3)
for channel in range(3):
neighbor_channel = [n[channel] for n in neighbors]
sampled = sample_channel(neighbor_channel, delta)
# If the allowed intervals are empty, use a fallback.
if sampled is None:
sampled = np.random.rand()
color[channel] = sampled
return color
def generate_color_pattern(h, w, d, r):
"""
Generate a dense color pattern.
h, w: height and width of the output image in pixels.
d: patch size in pixels (each patch is d x d).
r: randomness regulation factor. For each neighbor the condition is:
(|channel_new - channel_neighbor|)^2 >= r.
"""
pattern = np.zeros((h, w, 3), dtype=np.float32)
num_patches_x = h // d
num_patches_y = w // d
for i in range(num_patches_x):
for j in range(num_patches_y):
# Collect neighbor patch colors (using mean because patches have uniform color).
neighbors = []
if i > 0:
patch_above = pattern[(i-1)*d:i*d, j*d:(j+1)*d]
neighbors.append(np.mean(patch_above, axis=(0, 1)))
if j > 0:
patch_left = pattern[i*d:(i+1)*d, (j-1)*d:j*d]
neighbors.append(np.mean(patch_left, axis=(0, 1)))
if len(neighbors) == 0:
# For the first patch use a random color.
new_color = np.random.rand(3)
else:
new_color = generate_valid_color(neighbors, r)
pattern[i*d:(i+1)*d, j*d:(j+1)*d] = new_color
print('Patch ({}, {})'.format(i, j))
return pattern
def main():
# Parameters for the pattern.
# For example, using 100x100 pixels image with patches of 3x3 pixels.
# 11111111
height = 1200
width = 1000
patch_size = 12
# Set r to a value that gives sufficient difference.
# Note: r is the threshold for squared difference per channel, and the condition is
# |c - c_neighbor| >= sqrt(r). Very high r may make it impossible to satisfy.
r = 1000
pattern = generate_color_pattern(height, width, patch_size, r)
# Scale to 0-255 for display and saving.
pattern_display = (pattern * 255).astype(np.uint8)
# Display the pattern with OpenCV.
cv2.imshow('Pattern', pattern_display)
cv2.waitKey(0)
cv2.destroyAllWindows()
# Save the pattern as an image file.
cv2.imwrite('color_pattern.png', pattern_display)
if __name__ == '__main__':
main()
tools/patch_generation/patchgen_DD01.py
+139
View File
@@ -0,0 +1,139 @@
#%%
import numpy as np
import cv2
def merge_intervals(intervals):
# Merge overlapping intervals.
if not intervals:
return []
intervals.sort(key=lambda x: x[0])
merged = [intervals[0]]
for current in intervals[1:]:
last = merged[-1]
if current[0] <= last[1]:
merged[-1] = (last[0], max(last[1], current[1]))
else:
merged.append(current)
return merged
def get_allowed_intervals(neighbor_values, delta):
# Compute forbidden intervals for this channel from neighbor values.
forbidden = []
for c in neighbor_values:
low = max(0, c - delta)
high = min(1, c + delta)
forbidden.append((low, high))
# Merge overlapping forbidden intervals.
merged = merge_intervals(forbidden)
allowed = []
# Allowed region before the first forbidden interval.
if merged:
if merged[0][0] > 0:
allowed.append((0, merged[0][0]))
# Allowed regions between forbidden intervals.
for i in range(1, len(merged)):
allowed.append((merged[i-1][1], merged[i][0]))
# Allowed region after the last forbidden interval.
if merged[-1][1] < 1:
allowed.append((merged[-1][1], 1))
else:
allowed.append((0, 1))
return allowed
def sample_from_allowed(allowed):
# Choose a random value from the allowed intervals with probability proportional to interval length.
lengths = [end - start for start, end in allowed]
total_length = sum(lengths)
if total_length <= 0:
return None
point = np.random.rand() * total_length
cumulative = 0
for (start, end), length in zip(allowed, lengths):
cumulative += length
if point <= cumulative:
return np.random.uniform(start, end)
return allowed[-1][0]
def sample_channel(neighbor_values, delta):
# Compute allowed intervals and sample a value.
allowed = get_allowed_intervals(neighbor_values, delta)
if not allowed:
return None
return sample_from_allowed(allowed)
def generate_valid_color(neighbors, r):
"""
Generate a random color that satisfies, for each neighbor,
|R - R_neighbor| >= sqrt(r), and similarly for G and B.
neighbors: list of neighbor colors (each is an array of shape (3,))
r: threshold for squared difference.
"""
delta = np.sqrt(r)
color = np.zeros(3)
for channel in range(3):
neighbor_channel = [n[channel] for n in neighbors]
sampled = sample_channel(neighbor_channel, delta)
# If the allowed intervals are empty, use a fallback.
if sampled is None:
sampled = np.random.rand()
color[channel] = sampled
return color
def generate_color_pattern(h, w, d, r):
"""
Generate a dense color pattern.
h, w: height and width of the output image in pixels.
d: patch size in pixels (each patch is d x d).
r: randomness regulation factor. For each neighbor the condition is:
(|channel_new - channel_neighbor|)^2 >= r.
"""
pattern = np.zeros((h, w, 3), dtype=np.float32)
num_patches_x = h // d
num_patches_y = w // d
for i in range(num_patches_x):
for j in range(num_patches_y):
# Collect neighbor patch colors (using mean because patches have uniform color).
neighbors = []
if i > 0:
patch_above = pattern[(i-1)*d:i*d, j*d:(j+1)*d]
neighbors.append(np.mean(patch_above, axis=(0, 1)))
if j > 0:
patch_left = pattern[i*d:(i+1)*d, (j-1)*d:j*d]
neighbors.append(np.mean(patch_left, axis=(0, 1)))
if len(neighbors) == 0:
# For the first patch use a random color.
new_color = np.random.rand(3)
else:
new_color = generate_valid_color(neighbors, r)
pattern[i*d:(i+1)*d, j*d:(j+1)*d] = new_color
print('Patch ({}, {})'.format(i, j))
return pattern
def main():
# Parameters for the pattern.
# For example, using 100x100 pixels image with patches of 3x3 pixels.
height = 1200
width = 1000
patch_size = 12
# Set r to a value that gives sufficient difference.
# Note: r is the threshold for squared difference per channel, and the condition is
# |c - c_neighbor| >= sqrt(r). Very high r may make it impossible to satisfy.
r = 1000
pattern = generate_color_pattern(height, width, patch_size, r)
# Scale to 0-255 for display and saving.
pattern_display = (pattern * 255).astype(np.uint8)
# Display the pattern with OpenCV.
cv2.imshow('Pattern', pattern_display)
cv2.waitKey(0)
cv2.destroyAllWindows()
# Save the pattern as an image file.
cv2.imwrite('color_pattern.png', pattern_display)
if __name__ == '__main__':
main()
tools/patch_generation/patchgen_DD02C_dot.py
+44
View File
@@ -0,0 +1,44 @@
import numpy as np
from PIL import Image, ImageDraw
# Parameters
R = 9.0 # mm
H = 20.0 # mm
fov_half = np.radians(75) # 75 degrees half-angle
pixels_per_mm = 100
# Texture dimensions
width = int(np.pi * 2 * R * pixels_per_mm) # circumference
height = int(H * pixels_per_mm)
# Desired dot pattern in image space (normalized)
delta_theta = np.radians(10) # 10° azimuth spacing
delta_r = 0.1 # radial spacing in normalized units
rho = 0.08 # normalized dot radius (from reference)
# Generate dot centers in image space
r_vals = np.arange(R/H, np.tan(fov_half), delta_r)
theta_vals = np.arange(0, 2*np.pi, delta_theta)
# Create empty texture (white background)
img = Image.new('L', (width, height), color=255)
draw = ImageDraw.Draw(img)
for r_img in r_vals:
z = R / r_img
v = z / H
# Ellipse axes in texture coordinates
a_u = (rho * z) / (2 * np.pi * R)
a_v = (rho * z**2) / (R * H)
# Convert to pixels
a_u_px = a_u * width
a_v_px = a_v * height
for theta in theta_vals:
u = theta / (2 * np.pi)
cx = u * width
cy = v * height
# Draw ellipse (bounding box)
bbox = [cx - a_u_px, cy - a_v_px, cx + a_u_px, cy + a_v_px]
draw.ellipse(bbox, fill=0) # black dot
img.save('cylinder_marker_texture.png')
patch_generation/patch_generation_3_org.py → tools/patch_generation/patchgen_DD02M.py
Executable → Regular
View File
verify_fft_equivalence.py
-102
View File
@@ -1,102 +0,0 @@
"""
Verify that the OpenCL FFT convolution in solve_opencl / solve_opencl_batch2
produces the same result as scipy.signal.fftconvolve(f, G, mode='same').
This script simulates the OpenCL logic with numpy FFTs so we can compare
directly without needing a GPU.
"""
import numpy as np
from scipy.signal import fftconvolve
def make_green_2d(H, W, dx=1.0):
"""Reproduce PoissonSolver._create_radialGrid for 2D."""
rdims = (2*H - 1, 2*W - 1)
R = np.indices(rdims, dtype=np.float64)
R[0] += 1 - H
R[1] += 1 - W
R[0] *= dx
R[1] *= dx
np.hypot(R[1], R[0], R[0])
R = R[0]
zval = 0.5 * (dx + dx) / 2.0
R[H-1, W-1] = zval
# Apply 2D Green's function: log(r)/(2*pi)
G = np.log(R) * (0.5 / np.pi)
return G
def test_equivalence(H, W, dx=1.0, seed=42):
np.random.seed(seed)
f = np.random.randn(H, W).astype(np.float64)
G = make_green_2d(H, W, dx)
gdx_prod = dx ** 2
# --- CPU reference: scipy fftconvolve ---
p_cpu = fftconvolve(f, G, mode='same') * gdx_prod
# --- Simulate what OpenCL solve_opencl does ---
f32 = f.astype(np.float32)
G32 = G.astype(np.float32)
full_shape = (2*H - 1, 2*W - 1)
# G already has shape full_shape, no padding needed
# f is zero-padded to full_shape at [0:H, 0:W]
f_pad = np.zeros(full_shape, dtype=np.float32)
f_pad[:H, :W] = f32
# FFT at full_shape, multiply, iFFT
F_f = np.fft.rfftn(f_pad)
F_g = np.fft.rfftn(G32)
conv_circ = np.fft.irfftn(F_f * F_g, s=full_shape)
# OLD OpenCL extraction (BUGGY): same_starts = [(full_shape[i] - fshape[i]) // 2]
same_starts_old = [(full_shape[0] - H) // 2, (full_shape[1] - W) // 2]
p_old = conv_circ[
same_starts_old[0]:same_starts_old[0]+H,
same_starts_old[1]:same_starts_old[1]+W
].astype(np.float32) * np.float32(gdx_prod)
# FIXED OpenCL extraction: same_starts = [s - 1 for s in fshape]
same_starts_fixed = [H - 1, W - 1]
p_fixed = conv_circ[
same_starts_fixed[0]:same_starts_fixed[0]+H,
same_starts_fixed[1]:same_starts_fixed[1]+W
].astype(np.float32) * np.float32(gdx_prod)
# --- Also test with correct FFT size (3H-2) ---
correct_fft_shape = (H + (2*H-1) - 1, W + (2*W-1) - 1) # (3H-2, 3W-2)
F_f_full = np.fft.rfftn(f32, s=correct_fft_shape)
F_g_full = np.fft.rfftn(G32, s=correct_fft_shape)
conv_full = np.fft.irfftn(F_f_full * F_g_full, s=correct_fft_shape)
# scipy-style 'same' extraction from full linear conv
start_full = [(correct_fft_shape[0] - H) // 2, (correct_fft_shape[1] - W) // 2]
p_full_linear = conv_full[
start_full[0]:start_full[0]+H,
start_full[1]:start_full[1]+W
].astype(np.float32) * np.float32(gdx_prod)
# --- Compare ---
p_cpu32 = p_cpu.astype(np.float32)
scale = max(np.abs(p_cpu32).max(), 1e-10)
diff_old = np.abs(p_cpu32 - p_old)
diff_fixed = np.abs(p_cpu32 - p_fixed)
diff_full = np.abs(p_cpu32 - p_full_linear)
print(f"=== Test: f({H},{W}), G({2*H-1},{2*W-1}), full_shape={full_shape} ===")
print(f" scipy fftconvolve full linear conv shape: {correct_fft_shape}")
print(f" p_cpu range: [{p_cpu32.min():.4f}, {p_cpu32.max():.4f}]")
print()
print(f" [OLD BUG] same_starts={same_starts_old}:")
print(f" max_abs_diff = {diff_old.max():.6e} rel_max = {diff_old.max()/scale:.6e}")
print(f" [FIXED] same_starts={same_starts_fixed}:")
print(f" max_abs_diff = {diff_fixed.max():.6e} rel_max = {diff_fixed.max()/scale:.6e}")
print(f" [REF] Full linear conv (FFT size {correct_fft_shape}):")
print(f" max_abs_diff = {diff_full.max():.6e} rel_max = {diff_full.max()/scale:.6e}")
print()
return diff_old.max(), diff_fixed.max(), diff_full.max()
if __name__ == '__main__':
for size in [16, 32, 64, 128, 256, 384]:
d_ocl, d_cor, d_full = test_equivalence(size, size)
print()