diff --git a/.gitignore b/.gitignore index b18cd66..e986e81 100644 --- a/.gitignore +++ b/.gitignore @@ -1,17 +1,37 @@ +# Python bytecode / caches +__pycache__/ +*.py[cod] +.pytest_cache/ + +# Packaging / build outputs /build/ /dist/ /dist_obfuscated/ -/release/ -__pycache__/ -orisys.egg-info/ -/.cursor/ /dist_pyd/ /dist_single_pyd/ -/old_orisys/ /dist_nuitka/ -/build_nuitka -/data/ +/build_nuitka/ +/release/ +/SDK/ +*.egg-info/ + +# Local tools / editor state +/.cursor/ +/.vscode/ + +# Local data / models /models/ -/doc/.quarto/ -*.npy -*.png \ No newline at end of file +/data/* +!/data/sample/ +!/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/ \ No newline at end of file diff --git a/API文档-0.3.0.pdf b/API文档-0.3.0.pdf deleted file mode 100644 index acf4348..0000000 Binary files a/API文档-0.3.0.pdf and /dev/null differ diff --git a/RAFT_test.py b/RAFT_test.py deleted file mode 100644 index 416a4cf..0000000 --- a/RAFT_test.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..a9ac9ce --- /dev/null +++ b/README.md @@ -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. diff --git a/Readme.docx b/Readme.docx deleted file mode 100644 index c2b5f4e..0000000 Binary files a/Readme.docx and /dev/null differ diff --git a/build_nuitka.py b/build_nuitka.py deleted file mode 100644 index 90f5cab..0000000 --- a/build_nuitka.py +++ /dev/null @@ -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() - diff --git a/calibration_data_collection.py b/calibration_data_collection.py deleted file mode 100644 index 0ed3600..0000000 --- a/calibration_data_collection.py +++ /dev/null @@ -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() diff --git a/config/compound_eye.png b/config/compound_eye.png deleted file mode 100644 index 0fa802b..0000000 Binary files a/config/compound_eye.png and /dev/null differ diff --git a/config/ddjx01.json b/config/ddjx01.json index 988ab1e..21b5e96 100644 --- a/config/ddjx01.json +++ b/config/ddjx01.json @@ -11,7 +11,7 @@ }, "stitcher": { "min_area_frac": 0.01, - "max_area_frac": 0.04, + "max_area_frac": 0.05, "max_aspect_ratio": 50, "index": [1,2,3,6,7,8,11,12,13], "margin": 10, diff --git a/data/RGB/c0.5mm.png b/data/RGB/c0.5mm.png new file mode 100644 index 0000000..9627963 Binary files /dev/null and b/data/RGB/c0.5mm.png differ diff --git a/data/RGB/init_2.png b/data/RGB/init_2.png new file mode 100644 index 0000000..3b84385 Binary files /dev/null and b/data/RGB/init_2.png differ diff --git a/data/jx01/jx01_hex1.3_arrows.mp4 b/data/jx01/jx01_hex1.3_arrows.mp4 new file mode 100644 index 0000000..4e3969e Binary files /dev/null and b/data/jx01/jx01_hex1.3_arrows.mp4 differ diff --git a/data/jx01/jx01_hex1.3_depth.mp4 b/data/jx01/jx01_hex1.3_depth.mp4 new file mode 100644 index 0000000..4fd91e7 Binary files /dev/null and b/data/jx01/jx01_hex1.3_depth.mp4 differ diff --git a/data/jx01/jx01_hex1.3_pyqt.mp4 b/data/jx01/jx01_hex1.3_pyqt.mp4 new file mode 100644 index 0000000..4d4fea9 Binary files /dev/null and b/data/jx01/jx01_hex1.3_pyqt.mp4 differ diff --git a/data/jx01/jx01_hex1.5_2_arrows.mp4 b/data/jx01/jx01_hex1.5_2_arrows.mp4 new file mode 100644 index 0000000..bf92053 Binary files /dev/null and b/data/jx01/jx01_hex1.5_2_arrows.mp4 differ diff --git a/data/jx01/jx01_hex1.5_2_depth.mp4 b/data/jx01/jx01_hex1.5_2_depth.mp4 new file mode 100644 index 0000000..a09711a Binary files /dev/null and b/data/jx01/jx01_hex1.5_2_depth.mp4 differ diff --git a/data/jx01/jx01_hex1.5_2_pyqt.mp4 b/data/jx01/jx01_hex1.5_2_pyqt.mp4 new file mode 100644 index 0000000..daf398a Binary files /dev/null and b/data/jx01/jx01_hex1.5_2_pyqt.mp4 differ diff --git a/patch_generation/s1.jpg b/data/sample/s1.jpg similarity index 100% rename from patch_generation/s1.jpg rename to data/sample/s1.jpg diff --git a/s1_480p.mp4 b/data/sample/s1_480p.mp4 similarity index 100% rename from s1_480p.mp4 rename to data/sample/s1_480p.mp4 diff --git a/s1_720p.mp4 b/data/sample/s1_720p.mp4 similarity index 100% rename from s1_720p.mp4 rename to data/sample/s1_720p.mp4 diff --git a/doc/get_started.docx b/doc/get_started.docx deleted file mode 100644 index 5426117..0000000 Binary files a/doc/get_started.docx and /dev/null differ diff --git a/doc/软件使用说明-0.3.0.pdf b/doc/软件使用说明-0.3.0.pdf deleted file mode 100644 index 620aa54..0000000 Binary files a/doc/软件使用说明-0.3.0.pdf and /dev/null differ diff --git a/Readme.md b/docs/api.md similarity index 93% rename from Readme.md rename to docs/api.md index e647c07..ed310cf 100644 --- a/Readme.md +++ b/docs/api.md @@ -72,8 +72,8 @@ conda install conda-forge::pyvkfft ### 输入参数 * **vid_src** (`int | str`, 必填): 传感器 ID、相机编号或视频路径。默认为 `0`。 -* **config_name** (`str`，可选): 传感器配置`json`文件路径，用于加载对应的参数配置，默认为 `./config/ddjx01`。 -* **cal_path** (`str`，可选): 复眼拼接校准文件路径，默认为 `./config/cal.npy`。有多个传感器同时运行时建议传感器类创建时建议填写，避免校准文件覆写。在首次运行时程序会在找不到文件时自动创建。 +* **config_name** (`str`，可选): 传感器配置文件名称或 `json` 文件路径，用于加载对应的参数配置；例如 `ddjx01` 或 `./config/ddjx01.json`。 +* **cal_path** (`str`，可选，已弃用): 该参数仅为兼容旧代码保留，当前版本中不再使用，传入后也会被忽略。 * **backend** (`str`，可选): 光流场解耦计算后端，默认为 `auto`，可选`cupy`,`opencl`,`cpu`。 * **verbose** (`bool`，可选): 是否输出初始化及运行过程中的详细日志信息，默认为 `False`。 @@ -99,12 +99,12 @@ sensor = Sensor("/data/data.mp4") 采集并预处理一帧图像并存到`Sensor` 类别中。 + ### 返回 -* `np.ndarray`：处理后的BGR格式图像(`uint8`，尺寸为json文件中`resolution_stitch`对应分辨率)。 -* `None`：当输入源读取失败（例如视频结束、摄像头断开）时返回。 +* `bool`：是否成功读取图像。 ### 状态更新 -* 成功时更新`sensor`类里面的图像。 +* 成功时更新`sensor`类里面的`sensor.img`。处理后的BGR格式图像(`uint8`，尺寸为json文件中`resolution_stitch`对应分辨率)。 ### 示例 ```python @@ -123,6 +123,8 @@ sensor.get_img() ### 输入参数 * **check_motion**(`bool`，可选): 为 `True` 时先根据光流判断是否有运动，有运动才做 HHD 分解；为 `False` 时每帧都做 HHD。默认 `True`。 * **threshold**(`float`，可选): **check_motion** 功能的阈值，默认为`6`。 +* **tail_frames** (`int`，可选): 光流变化低于阈值后继续解算多少帧，确保传感器回复，，默认为 `10`. + ### 返回 * `bool`：是否接触到物体，如果**check_motion**为`False`时返回`None`。 diff --git a/doc/get_started.md b/docs/get_started.md similarity index 98% rename from doc/get_started.md rename to docs/get_started.md index 60916cb..fbc124e 100644 --- a/doc/get_started.md +++ b/docs/get_started.md @@ -1,6 +1,6 @@ 叠动科技软件使用手册 版本： 0.3.2 -更新日期： 2026年3月6日 +更新日期： 2026年4月16日 # 简介 本文件介绍如何上手和运行叠动视触觉传感器配套的开发套件。本套件是配合叠动视触觉传感器进行图像处理，计算形变、力以及接触信息。它可以从摄像头或视频文件中读取图像，处理后打印和展示处理后信息。 @@ -120,7 +120,7 @@ print(orisys.__version__) 3. 如果第一步没有输出，传感器的ID就是“0”，如果第一步有一个设备，传感器的ID就是“1”，如此类推。 4. 在确定了传感器的ID（下方例子假设为“0”）后，打开Anaconda Prompt, 切换到工作路径，然后在命令行中输入： ``` bash - python example1_basic_pipeline.py -v 0 + python examples/basic_pipeline.py -v 0 ``` 5. 程序会运行，并自动打开可视化窗口和打印资料，在窗口上按键盘”q”键可以停止运行，按”r”键可以重设数据 @@ -130,7 +130,7 @@ print(orisys.__version__) 3. 根据新增的输出，例如”/dev/video0”，确定传感器的ID为“0”, 如果是”/dev/video3”的话ID为“3” ，如此类推。 4. 在确定了传感器的ID（下方例子假设为“0”）后，打开终端, 激活相应的python环境，切换到工作路径，然后在命令行中输入： ``` bash - python example1_basic_pipeline.py -v 0 + python examples/basic_pipeline.py -v 0 ``` 5. 程序会运行，并自动打开可视化窗口和打印资料，在窗口上按键盘”q”键可以停止运行，按”r”键可以重设数据 diff --git a/res/orisys_logo.png b/docs/images/orisys_logo.png similarity index 100% rename from res/orisys_logo.png rename to docs/images/orisys_logo.png diff --git a/res/orisys_logo_notext.png b/docs/images/orisys_logo_notext.png similarity index 100% rename from res/orisys_logo_notext.png rename to docs/images/orisys_logo_notext.png diff --git a/doc/render_command.txt b/docs/render_command.txt similarity index 100% rename from doc/render_command.txt rename to docs/render_command.txt diff --git a/example1_basic_pipeline.py b/example1_basic_pipeline.py deleted file mode 100644 index d8763c6..0000000 --- a/example1_basic_pipeline.py +++ /dev/null @@ -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() diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..b05b54e --- /dev/null +++ b/examples/README.md @@ -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. diff --git a/examples/basic_pipeline.py b/examples/basic_pipeline.py new file mode 100644 index 0000000..7404637 --- /dev/null +++ b/examples/basic_pipeline.py @@ -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() diff --git a/example3_plot_data.py b/examples/plot_data.py similarity index 53% rename from example3_plot_data.py rename to examples/plot_data.py index 1463545..65aa22c 100644 --- a/example3_plot_data.py +++ b/examples/plot_data.py @@ -1,3 +1,15 @@ +""" +示例 3：实时曲线与幅值图显示 + +本示例演示如何在读取传感器结果的同时： +1. 显示形变矢量场 +2. 显示流场幅值热力图 +3. 使用 PyQtGraph 实时绘制法向力与切向力曲线 +4. 按键控制录像开始与结束 + +SDK version: 0.3.1 +""" + import os os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" import cv2 @@ -5,10 +17,45 @@ import orisys import numpy as np import time 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(): - p1 = win.addPlot(title="法相力") + p1 = win.addPlot(title="法向力") p1.setLabel('bottom', '时间', units='s') win.nextRow() 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: """ - Draw color map of pointwise magnitude of velocity field. + 绘制流场幅值伪彩色图。 Args: - vfield: Velocity field to visualize (2D vector field) - threshold: Minimum magnitude to display (below threshold shows as zero/blue) - colormap: OpenCV colormap to use (default: COLORMAP_JET) + vfield: 需要可视化的二维矢量场 + threshold: 最小显示阈值；低于阈值的区域显示为低值颜色 + colormap: OpenCV 颜色映射类型（默认 COLORMAP_JET） Returns: - np.ndarray: 400x400 BGR image with colored magnitude map + np.ndarray: 400x400 的 BGR 伪彩图像 """ flow_field = vfield - # Compute magnitude at each point + # 计算每个像素点的幅值 magnitude = np.sqrt(flow_field[:, :, 0]**2 + flow_field[:, :, 1]**2) - # Smooth by downsampling and interpolating back + # 先降采样再插值回原尺寸，使显示更平滑 downsample_factor = 10 h, w = magnitude.shape magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor] 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 < threshold] = 0 - # Normalize to 0-255 range + # 归一化到 0-255 mag_min, mag_max = magnitude_thresholded.min(), magnitude_thresholded.max() if mag_max > mag_min: magnitude_norm = ((magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255).astype(np.uint8) else: magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8) - # Apply colormap + # 应用伪彩色映射 magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap) - # Resize to output size + # 缩放到固定输出尺寸 out_img = cv2.resize(magnitude_colored, (400, 400), interpolation=cv2.INTER_CUBIC) return out_img # 实时图表 -npoints = 100 # 图表数据长度 +npoints = 100 # 曲线缓冲区长度 pg.setConfigOption('background', 'w') pg.setConfigOption('foreground', 'k') -app = pg.mkQApp("Plotting Example") -win = pg.GraphicsLayoutWidget(show=True, title="Basic plotting examples") -win.setWindowTitle('pyqtgraph example: Plotting') +app = pg.mkQApp("Orisys Plot Example") +win = pg.GraphicsLayoutWidget(show=True, title="Orisys 实时力曲线") +win.setWindowTitle('Orisys 实时力曲线') p1, p2, p3, plot_af, plot_tfx, plot_tfy = create_plot() data = np.zeros([npoints,2,3]) start_time = None # Will be set when sensor starts for i in range(3): - data[:,0,i] = np.linspace(-npoints,0,npoints) + data[:,0,i] = np.linspace(-npoints,0,npoints)/30 def main(): global start_time 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("--video","-v", type=str, default="0", help="视频源：摄像头编号或视频文件路径") + parser.add_argument("--config","-c", type=str, default="./config/ddjx01.json", help="配置文件名称或路径") + parser.add_argument("--cal","-cal", type=str, default="./config/ddjx01.npy", help="标定文件路径") + parser.add_argument("--verbose","-verbose", type=bool, default=True, help="是否输出详细日志") args = parser.parse_args() - # 步骤1：创建传感器对象 (必要步骤)， - # config_name为配置文件名（内置配置文件名）或路径（自定义配置文件路径） - # cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件 + # 步骤 1：创建传感器对象（必需） + # config_name 可以是内置配置名称，也可以是自定义配置文件路径 + # cal_path 为标定文件路径；每台设备建议使用独立的标定文件 - # 方式1: 从视频文件读取 + # 输入既可以是摄像头编号，也可以是视频文件路径 try: input = int(args.video) is_camera = True @@ -109,23 +156,23 @@ def main(): except: is_camera = False 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) start_time = time.time() is_save = False frame_count = 0 writer = None + save_path = None + print("\n按键说明：'q' 退出程序，'s' 开始录像，'e' 结束录像。") while True: - sensor.get_img() # 提取并拼接复眼图像 + sensor.get_img() # 获取并拼接复眼图像 if is_save and writer is not None and sensor.frame is not None: frame_count += 1 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) == 3: save_img = save_img.squeeze(2) @@ -133,44 +180,48 @@ def main(): writer.write(save_img) - sensor.compute_deformation(check_motion=False) # 计算形变和分解 - # 提取信息： - # 向量场: VRAW, VNORMAL, VSHEAR: 原始，法向，切向 - # 帧率，时间戳: FPS, TIMESTAMP - # 变形量总和: FNORMAL , FSHEARX, FSHEARY : 法向，切向（X)，切向（Y) + sensor.compute_deformation(check_motion=True, threshold=0) # 计算形变与分解结果 + # 读取信息： + # 向量场：VRAW、VNORMAL、VSHEAR -> 原始 / 法向 / 切向 + # 标量：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) - 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) - cv2.imshow("Deformation", arrows) - # 显示深度图 + cv2.imshow("形变矢量场", arrows) + # 显示流场幅值图 mag = draw_magnitude_map(flow) - cv2.imshow("Magnitude", mag) + cv2.imshow("流场幅值图", mag) - # 显示图表 - 更新时间轴 + # 更新实时曲线 current_time = time.time() - start_time + data_filtered = data.copy() for i in range(3): - data[:-1, 0, i] = data[1:, 0, i] # shift time data left - data[:-1, 1, i] = data[1:, 1, i] # shift force data left - data[-1, 0, :] = current_time # set current time for all plots + data[:-1, 0, i] = data[1:, 0, i] # 时间轴左移 + data[:-1, 1, i] = data[1:, 1, i] # 力数据左移 + data[-1, 0, :] = current_time # 写入当前时间 data[-1, 1, 0] = fn data[-1, 1, 1] = fx 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) if key & 0xFF == ord("q"): - print("\n 停止程序") + print("\n正在退出示例程序...") break if key == ord("s"): save_path = f"./data/pulse_test.mp4" - # Ensure directory exists + # 确保输出目录存在 os.makedirs(os.path.dirname(save_path), exist_ok=True) - # Get frame dimensions from the actual frame + # 从当前帧获取视频尺寸 save_img = sensor.frame 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() if len(temp_img.shape) == 2 or (len(temp_img.shape) == 3 and temp_img.shape[2] == 1): 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)) if writer.isOpened(): is_save = True - print(f"Start writing to {save_path}") + print(f"开始录像：{save_path}") else: - print(f"Failed to open video writer for {save_path}") + print(f"无法创建录像文件：{save_path}") writer = None else: - print("No frame available to determine video dimensions") + print("当前没有可用帧，无法确定录像尺寸。") 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}") + print(f"录像结束：{save_path}，总帧数：{frame_count}") else: - print(f"Stop writing, total frames: {frame_count}") + print(f"录像结束，总帧数：{frame_count}") frame_count = 0 - sensor.disconnect() # 断开传感器连接 + sensor.disconnect() # 断开传感器连接并释放资源 diff --git a/example2_two_sensor_pipeline.py b/examples/two_sensor_pipeline.py similarity index 66% rename from example2_two_sensor_pipeline.py rename to examples/two_sensor_pipeline.py index 60e5fd4..0bf7c5a 100644 --- a/example2_two_sensor_pipeline.py +++ b/examples/two_sensor_pipeline.py @@ -1,22 +1,24 @@ """ -示例 4: 多传感器并行处理（多进程模式 - 不使用 MultiSensor） +示例 2：双传感器并行处理（多进程） 本示例演示如何使用 multiprocessing 同时运行两个传感器： -- 每个传感器运行在独立的进程中 +- 每个传感器运行在独立进程中 - 每个进程拥有独立的 CUDA 上下文 -- 使用 multiprocessing.Process 和 Queue 进行进程间通信 +- 通过 Process 和 Queue 在主进程中汇总结果并显示 架构说明： - 进程 1 (传感器 0): 采集 → 光流 → HHD → 接触检测 → 队列 - 进程 2 (传感器 1): 采集 → 光流 → HHD → 接触检测 → 队列 + 进程 1（传感器 0）：采集 → 光流 → HHD → 接触检测 → 队列 + 进程 2（传感器 1）：采集 → 光流 → HHD → 接触检测 → 队列 ↓ - 主进程: 显示 ← ───────────────┘ + 主进程： 显示 ← ───────────────┘ -适用场景：多指触觉传感器、多视角采集、需要 GPU 加速的多路处理 +适用场景：多指触觉传感器、多视角采集，以及需要 GPU 加速的多路处理。 注意事项： - Windows 下必须在 if __name__ == '__main__' 中调用 mp.freeze_support() -- 每个传感器进程会占用独立的 GPU 显存 +- 每个传感器进程都会占用独立的 GPU 显存 + +SDK version: 0.3.1 """ import os os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" @@ -26,20 +28,20 @@ import numpy as np import multiprocessing as mp from multiprocessing import Queue, Process 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: sensor_id: 传感器编号 vid_src: 视频源（文件路径或摄像头索引） config_name: 配置名称 - cal_path: 标定路径 result_queue: 结果队列，用于向主进程发送数据 stop_event: 停止事件，用于通知工作进程退出 - target_fps: 目标帧率；若设置则用定时器限速，使多传感器 FPS 接近一致 + target_fps: 目标帧率；设置后将对处理循环限速，使多路传感器 FPS 更接近 """ try: # 在每个进程中创建独立的传感器实例 @@ -48,27 +50,26 @@ def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_ isstitch=True, cuda=True, verbose=False, # 减少子进程输出 - config_name=config_name, - cal_path=cal_path + config_name=config_name ) - 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 while not stop_event.is_set(): t_start = time.perf_counter() try: - # 步骤1: 获取图像 + # 步骤 1：获取图像 sensor.get_img() - # 步骤2: 计算形变场 + # 步骤 2：计算形变场 sensor.compute_deformation() - # 步骤3: 计算触碰点 + # 步骤 3：计算接触区域（如需要可取消注释） # sensor.compute_contact() - # 步骤4: 读取信息 + # 步骤 4：读取结果 fps, flow, vnormal, img, contour, centroid, depth_map = sensor.read_info( sensor.info.FPS, sensor.info.VRAW, @@ -79,7 +80,7 @@ def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_ sensor.info.DEPTH ) - # 步骤5: 准备数据字典（确保所有数据都是可序列化的） + # 步骤 5：整理结果字典（确保数据可序列化） data = { 'sensor_id': sensor_id, 'fps': fps, @@ -92,100 +93,104 @@ def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_ 'timestamp': time.time() } - # 步骤6: 发送数据到主进程（非阻塞） + # 步骤 6：将结果发送到主进程（非阻塞） try: result_queue.put_nowait(data) except Exception: - # 如果队列已满，跳过这一帧 + # 如果队列已满，则丢弃当前帧，避免阻塞工作进程 pass except Exception as e: - print(f"传感器 {sensor_id} 处理错误: {e}") + print(f"传感器 {sensor_id} 处理失败：{e}") time.sleep(0.1) - # 步骤7: 按 target_fps 限速，使多传感器 FPS 接近 + # 步骤 7：按 target_fps 限速，使多路传感器 FPS 更接近 if frame_interval > 0: elapsed = time.perf_counter() - t_start sleep_time = frame_interval - elapsed if sleep_time > 0.001: # 至少 1ms 才 sleep，避免无意义调度 time.sleep(sleep_time) - # 清理资源 + # 释放资源 sensor.disconnect() - print(f"传感器 {sensor_id} 已停止") + print(f"传感器 {sensor_id} 已停止。") except Exception as e: - print(f"传感器 {sensor_id} 初始化失败: {e}") + print(f"传感器 {sensor_id} 初始化失败：{e}") result_queue.put({'sensor_id': sensor_id, 'error': str(e)}) 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 num_sensors = 2 sensor_configs = [ { 'sensor_id': 0, - 'vid_src': 2, - 'config_name': './config/prototype_of3_720p.json', - 'cal_path': './config/ddjx01.npy' + 'vid_src': args.video1, + 'config_name': './config/ddjx01.json' }, { 'sensor_id': 1, - 'vid_src': 4, # 可以改为不同的视频源或摄像头索引 - 'config_name': './config/prototype_of3_720p.json', - 'cal_path': './config/ddjx01.npy' + 'vid_src': args.video2, # 可替换为其他视频源或摄像头索引 + 'config_name': './config/ddjx01.json' }, ] # ========================================================================= # 2. 创建进程间通信对象 # ========================================================================= - # 增大队列大小以缓冲更多帧，避免因主进程处理延迟导致的丢帧 + # 适当增大队列容量，以缓冲更多帧并减轻主进程偶发延迟带来的丢帧 result_queues = [Queue(maxsize=30) for _ in sensor_configs] stop_events = [mp.Event() for _ in sensor_configs] # ========================================================================= - # 3. 启动传感器工作进程 + # 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'], + args=(i, config['vid_src'], config['config_name'], result_queues[i], stop_events[i], target_fps) ) p.start() processes.append(p) - print(f"\n已启动 {len(processes)} 个传感器进程") - print("按 'q' 退出程序") + print(f"\n已启动 {len(processes)} 个传感器进程。") + print("按键说明：'q' 退出程序。") - # 用于存储每个传感器的最新数据 + # 保存每个传感器的最新一帧数据 latest_data = [None] * len(sensor_configs) - # 用于控制打印频率，减少I/O开销 + # 控制打印频率，减少 I/O 开销 print_counter = 0 - print_interval = 30 # 每30帧打印一次，进一步减少I/O + print_interval = 30 # 每 30 帧打印一次 - # 用于控制按键检查频率 + # 控制按键检查频率 key_check_counter = 0 - key_check_interval = 5 # 每5帧检查一次按键 + key_check_interval = 5 # 每 5 帧检查一次按键 try: # ===================================================================== - # 4. 主循环：从队列获取数据并显示 + # 4. 主循环：从队列获取结果并显示 # ===================================================================== while True: - # 从所有队列中快速获取最新数据（高效清空队列，只保留最新帧） + # 从所有队列中快速获取最新数据，只保留最后一帧 for i, queue in enumerate(result_queues): - # 快速清空队列，只保留最后一帧 + # 清空队列，仅保留最后一帧 data = None - # 一次性获取所有可用数据，只保留最后一个 + # 读取所有可用数据，只保留最新结果 while True: try: data = queue.get_nowait() @@ -195,14 +200,14 @@ def main(): 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'], @@ -211,7 +216,7 @@ def main(): grid_spacing=10, arrow_scale=1.0 ) - cv2.imshow(f"Deformation {sid}", arrows) + cv2.imshow(f"传感器 {sid} - 形变矢量场", arrows) # # 显示触碰点和区域 # if data['img'] is not None: @@ -233,33 +238,33 @@ def main(): # cv2.imshow(f"Depth Map {sid}", # cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET)) - # 降低打印频率以减少I/O开销（在主循环级别控制） + # 降低打印频率，减少主循环中的 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(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退出程序...") + print("\n正在退出示例程序...") break key_check_counter = 0 finally: # ===================================================================== - # 5. 清理资源 + # 5. 释放资源 # ===================================================================== # 通知所有工作进程停止 for event in stop_events: event.set() - # 等待所有进程结束 + # 等待所有工作进程结束 for p in processes: p.join(timeout=2) if p.is_alive(): @@ -267,10 +272,10 @@ def main(): p.join() cv2.destroyAllWindows() - print("所有传感器进程已停止") + print("所有传感器进程已停止。") if __name__ == '__main__': - # Windows 下多进程必需的初始化 + # Windows 下运行多进程示例所需的初始化 mp.freeze_support() main() diff --git a/five_sensors.py b/five_sensors.py deleted file mode 100644 index 6277747..0000000 --- a/five_sensors.py +++ /dev/null @@ -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() diff --git a/hand_test.py b/hand_test.py deleted file mode 100644 index 3f5d3de..0000000 --- a/hand_test.py +++ /dev/null @@ -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() diff --git a/patch_generation/color_pattern.png b/patch_generation/color_pattern.png deleted file mode 100644 index 221c4aa..0000000 Binary files a/patch_generation/color_pattern.png and /dev/null differ diff --git a/patch_generation/raw_1.jpg b/patch_generation/raw_1.jpg deleted file mode 100644 index 24bfbdd..0000000 Binary files a/patch_generation/raw_1.jpg and /dev/null differ diff --git a/patch_generation/red_pattern.png b/patch_generation/red_pattern.png deleted file mode 100644 index eaeab93..0000000 Binary files a/patch_generation/red_pattern.png and /dev/null differ diff --git a/patch_generation/test_hue_gray.py b/patch_generation/test_hue_gray.py deleted file mode 100644 index f96defa..0000000 --- a/patch_generation/test_hue_gray.py +++ /dev/null @@ -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() diff --git a/patch_generation/yellow_pattern.png b/patch_generation/yellow_pattern.png deleted file mode 100644 index 6dc0e91..0000000 Binary files a/patch_generation/yellow_pattern.png and /dev/null differ diff --git a/benchmark.py b/scripts/benchmark.py similarity index 100% rename from benchmark.py rename to scripts/benchmark.py diff --git a/scripts/build_nuitka.py b/scripts/build_nuitka.py new file mode 100644 index 0000000..99b82b2 --- /dev/null +++ b/scripts/build_nuitka.py @@ -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() + diff --git a/calibration_script.py b/scripts/calibration_script.py similarity index 100% rename from calibration_script.py rename to scripts/calibration_script.py diff --git a/scripts/demo_wuxi.py b/scripts/demo_wuxi.py new file mode 100644 index 0000000..d6cfe1c --- /dev/null +++ b/scripts/demo_wuxi.py @@ -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() diff --git a/promote_video.py b/scripts/promote_video.py similarity index 94% rename from promote_video.py rename to scripts/promote_video.py index 2d301b2..56a80b9 100644 --- a/promote_video.py +++ b/scripts/promote_video.py @@ -153,7 +153,9 @@ def main(): frame_count = 0 writer = None while True: - sensor1.get_img() # 提取并拼接复眼图像 + isimg = sensor1.get_img() # 提取并拼接复眼图像 + if isimg is None: + break 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) print(fps,fn,fx,fy) # 显示向量场 - black_img = np.zeros_like(sensor1.img) - arrows = orisys.util.draw_arrows(sensor1.img, flow, threshold=2, grid_spacing=20, arrow_scale=0.5) + _bg_bgr = (22, 16, 11) + 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) # 显示深度图 diff --git a/scripts/test_ffmpeg.py b/scripts/test_ffmpeg.py new file mode 100644 index 0000000..1c3396f --- /dev/null +++ b/scripts/test_ffmpeg.py @@ -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() diff --git a/setup.py b/setup.py index b229e47..91fca3c 100644 --- a/setup.py +++ b/setup.py @@ -1,11 +1,12 @@ from setuptools import setup, find_packages from pathlib import Path -# Find all JSON files in the configs directory -configs_dir = Path("src/orisys/configs") -json_files = [] -if configs_dir.exists(): - json_files = [str(p.relative_to("src")) for p in configs_dir.glob("*.json")] +config_dir = Path("config") +data_files = [] +if config_dir.exists(): + json_files = [str(p) for p in sorted(config_dir.glob("*.json"))] + if json_files: + data_files.append(("orisys/configs", json_files)) setup( name="orisys", @@ -14,10 +15,8 @@ setup( author="Your Name", package_dir={"": "src"}, packages=find_packages(where="src"), - package_data={ - "orisys": ["configs/*.json"], - }, include_package_data=True, + data_files=data_files, install_requires=[ "numpy>=1.24.4", "opencv-python>=4.10.0.84", diff --git a/src/orisys/HHD.py b/src/orisys/HHD.py deleted file mode 100644 index 0a2988e..0000000 --- a/src/orisys/HHD.py +++ /dev/null @@ -1,2187 +0,0 @@ -''' -Copyright (c) 2015, Harsh Bhatia (bhatia4@llnl.gov) -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: -1. Redistributions of source code must retain the above copyright notice, this -list of conditions and the following disclaimer. -2. Redistributions in binary form must reproduce the above copyright notice, -this list of conditions and the following disclaimer in the documentation -and/or other materials provided with the distribution. -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED -WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -''' - -import numpy -import timeit -import sys -# import cv2 -from scipy import signal, spatial -from numba import njit - -from .backend_availability import ( - CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE, - cp, cp_signal, _torque_kernel, - cl, cla, cl_rfftn, cl_irfftn -) - - -# -- Numba-accelerated finite-difference kernels for 2D HHD ------------------ -# Matches numpy.gradient's second-order boundary stencil exactly. - -@njit(cache=True) -def _divcurl_2d_nb(vfield, dy, dx, div_out, curl_out): - """Fused divergence + curl of a 2D vector field via finite differences.""" - H = vfield.shape[0] - W = vfield.shape[1] - idy2 = 1.0 / (2.0 * dy) - idx2 = 1.0 / (2.0 * dx) - idy_edge = 1.0 / (2.0 * dy) - idx_edge = 1.0 / (2.0 * dx) - for y in range(H): - for x in range(W): - # du/dx - if x == 0: - dudx = (-3.0*vfield[y,0,0] + 4.0*vfield[y,1,0] - vfield[y,2,0]) * idx_edge - elif x == W - 1: - dudx = (vfield[y,W-3,0] - 4.0*vfield[y,W-2,0] + 3.0*vfield[y,W-1,0]) * idx_edge - else: - dudx = (vfield[y,x+1,0] - vfield[y,x-1,0]) * idx2 - # dv/dy - if y == 0: - dvdy = (-3.0*vfield[0,x,1] + 4.0*vfield[1,x,1] - vfield[2,x,1]) * idy_edge - elif y == H - 1: - dvdy = (vfield[H-3,x,1] - 4.0*vfield[H-2,x,1] + 3.0*vfield[H-1,x,1]) * idy_edge - else: - dvdy = (vfield[y+1,x,1] - vfield[y-1,x,1]) * idy2 - # dv/dx - if x == 0: - dvdx = (-3.0*vfield[y,0,1] + 4.0*vfield[y,1,1] - vfield[y,2,1]) * idx_edge - elif x == W - 1: - dvdx = (vfield[y,W-3,1] - 4.0*vfield[y,W-2,1] + 3.0*vfield[y,W-1,1]) * idx_edge - else: - dvdx = (vfield[y,x+1,1] - vfield[y,x-1,1]) * idx2 - # du/dy - if y == 0: - dudy = (-3.0*vfield[0,x,0] + 4.0*vfield[1,x,0] - vfield[2,x,0]) * idy_edge - elif y == H - 1: - dudy = (vfield[H-3,x,0] - 4.0*vfield[H-2,x,0] + 3.0*vfield[H-1,x,0]) * idy_edge - else: - dudy = (vfield[y+1,x,0] - vfield[y-1,x,0]) * idy2 - - div_out[y, x] = dudx + dvdy - curl_out[y, x] = dvdx - dudy - - -@njit(cache=True) -def _grad_rotgrad_h_2d_nb(nD, nRu, vfield, dy, dx, d_out, r_out, h_out): - """Fused gradient(nD) + rotated_gradient(nRu) + h = vfield - d - r.""" - H = nD.shape[0] - W = nD.shape[1] - idy2 = 1.0 / (2.0 * dy) - idx2 = 1.0 / (2.0 * dx) - idy_edge = 1.0 / (2.0 * dy) - idx_edge = 1.0 / (2.0 * dx) - for y in range(H): - for x in range(W): - # gradient(nD) → d = (ddx_nD, ddy_nD) - if x == 0: - ddx_nD = (-3.0*nD[y,0] + 4.0*nD[y,1] - nD[y,2]) * idx_edge - ddx_nRu = (-3.0*nRu[y,0] + 4.0*nRu[y,1] - nRu[y,2]) * idx_edge - elif x == W - 1: - ddx_nD = (nD[y,W-3] - 4.0*nD[y,W-2] + 3.0*nD[y,W-1]) * idx_edge - ddx_nRu = (nRu[y,W-3] - 4.0*nRu[y,W-2] + 3.0*nRu[y,W-1]) * idx_edge - else: - ddx_nD = (nD[y,x+1] - nD[y,x-1]) * idx2 - ddx_nRu = (nRu[y,x+1] - nRu[y,x-1]) * idx2 - - if y == 0: - ddy_nD = (-3.0*nD[0,x] + 4.0*nD[1,x] - nD[2,x]) * idy_edge - ddy_nRu = (-3.0*nRu[0,x] + 4.0*nRu[1,x] - nRu[2,x]) * idy_edge - elif y == H - 1: - ddy_nD = (nD[H-3,x] - 4.0*nD[H-2,x] + 3.0*nD[H-1,x]) * idy_edge - ddy_nRu = (nRu[H-3,x] - 4.0*nRu[H-2,x] + 3.0*nRu[H-1,x]) * idy_edge - else: - ddy_nD = (nD[y+1,x] - nD[y-1,x]) * idy2 - ddy_nRu = (nRu[y+1,x] - nRu[y-1,x]) * idy2 - - d0 = ddx_nD - d1 = ddy_nD - r0 = -ddy_nRu - r1 = ddx_nRu - - d_out[y, x, 0] = d0 - d_out[y, x, 1] = d1 - r_out[y, x, 0] = r0 - r_out[y, x, 1] = r1 - h_out[y, x, 0] = vfield[y, x, 0] - d0 - r0 - h_out[y, x, 1] = vfield[y, x, 1] - d1 - r1 - - -# --------------------------------------------------------------------------- -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._decompose_bufs = None - self.psolver.prepare(False) - - def _ensure_decompose_bufs(self, shape2d): - """Lazily allocate reusable buffers for the 2D decompose fast path.""" - if self._decompose_bufs is not None and self._decompose_bufs['shape'] == shape2d: - return self._decompose_bufs - self._decompose_bufs = { - 'shape': shape2d, - 'div': numpy.empty(shape2d, dtype=numpy.float32), - 'curl': numpy.empty(shape2d, dtype=numpy.float32), - 'd': numpy.empty(shape2d + (2,), dtype=numpy.float32), - 'r': numpy.empty(shape2d + (2,), dtype=numpy.float32), - 'h': numpy.empty(shape2d + (2,), dtype=numpy.float32), - } - return self._decompose_bufs - - 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: - shape2d = (vfield.shape[0], vfield.shape[1]) - bufs = self._ensure_decompose_bufs(shape2d) - dy, dx = self.mesh.dx[0], self.mesh.dx[1] - - vf32 = vfield if vfield.dtype == numpy.float32 else vfield.astype(numpy.float32) - _divcurl_2d_nb(vf32, dy, dx, bufs['div'], bufs['curl']) - self.div = bufs['div'] - self.curlw = bufs['curl'] - - 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()) - - 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()) - - nD_f32 = self.nD if self.nD.dtype == numpy.float32 else self.nD.astype(numpy.float32) - nRu_f32 = self.nRu if self.nRu.dtype == numpy.float32 else self.nRu.astype(numpy.float32) - _grad_rotgrad_h_2d_nb(nD_f32, nRu_f32, vf32, dy, dx, - bufs['d'], bufs['r'], bufs['h']) - self.d = bufs['d'] - self.r = bufs['r'] - self.h = bufs['h'] - - if verbose > 0: - print( 'd =', self.d.shape, self.d.min(), self.d.max()) - print( 'r =', self.r.shape, self.r.min(), self.r.max()) - print( 'h =', self.h.shape, self.h.min(), self.h.max()) - return - 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 - - # Also store CPU copies for compatibility - # self.d = cp.asnumpy(d_gpu) - # self.h = cp.asnumpy(h_gpu) - # self.r = cp.asnumpy(r_gpu) - # self.nRu = cp.asnumpy(nRu_gpu) - # self.nD = cp.asnumpy(nD_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 - # return self.d, self.h, numpy.cross(r_p, self.r) , numpy.cross(r_n, self.r) - - 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) - - # Keep data on GPU - avoid transfers by computing directly - # d_gpu = cp.asarray(self.d, dtype=cp.float32) - # h_gpu = cp.asarray(self.h, dtype=cp.float32) - # r_gpu = cp.asarray(self.r, dtype=cp.float32) - # nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32) - - 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] - - # Compute torque with fused sum - # 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 - # )) - # Replace lines 413-416: - 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)) - - 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() - - - -# --------------------------------------------------------------------------- - -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() - -# --------------------------------------------------------------------------- - -# Try to import CuPy for GPU acceleration -# try: -# import cupy as cp -# CUPY_AVAILABLE = True -# except ImportError: -# cp = None -# CUPY_AVAILABLE = False - -# ------------------------------------------------------------------------------ -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() - - # -------------------------------------------------------------------------- - # -------------------------------------------------------------------------- - - -# import numpy -# from scipy import signal, spatial -# # from .timer import Timer - -# # Try to import CuPy for GPU acceleration -# try: -# import cupy as cp -# import cupyx.scipy.signal as cp_signal -# CUPY_AVAILABLE = True -# except ImportError: -# cp = None -# cp_signal = None -# CUPY_AVAILABLE = False - -# Try to import OpenCL + pyvkfft for FFT-based convolution - - - -''' -# ------------------------------------------------------------------------------ - 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 - - # def solve_opencl(self, f, verbose=False, profile=True): - # """ - # OpenCL Poisson solver using pyvkfft FFT convolution. - # Replaces scipy.signal.fftconvolve(..., mode='same') for structured grids. - # 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(f, verbose) - - # if not self.ready: - # self.prepare(verbose) - - # fshape = f.shape - # gtimer = Timer() - # if verbose: - # print('\nSolving Poisson Eq. with OpenCL FFT') - - # if self.ptype == 'G': - # if self.gdims != fshape: - # print(self.gdims, fshape) - # raise ValueError("Shape of function should match shape of grid") - - # if self.stype == 'F': - # 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() - # f_host = numpy.ascontiguousarray(f, dtype=numpy.float32) - # full_shape = tuple(2 * s - 1 for s in f_host.shape) - # timing['host_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # cache_valid = ( - # self._opencl_cache is not None and - # self._opencl_cache['fshape'] == f_host.shape and - # self._opencl_cache['full_shape'] == full_shape and - # 'f_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 = [f_host.shape[i] - 1 for i in range(self.dim)] - # same_slices = tuple( - # slice(same_starts[i], same_starts[i] + f_host.shape[i]) - # for i in range(self.dim) - # ) - # fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1]) - # self._opencl_cache = { - # 'fshape': f_host.shape, - # 'full_shape': full_shape, - # 'fft_shape': fft_shape, - # 'g_fft_dev': g_fft_dev, - # 'same_slices': same_slices, - # 'same_starts': same_starts, - # 'gdx_prod': numpy.prod(self.gdx), - # # Zero-init so padding stays zero across frames - # 'f_dev': cla.zeros(queue, full_shape, dtype=numpy.float32), - # # Small staging buffer for input-sized H2D - # 'f_stage_dev': cla.empty(queue, f_host.shape, dtype=numpy.float32), - # 'f_fft_dev': cla.empty(queue, fft_shape, dtype=numpy.complex64), - # 'conv_full_dev': cla.empty(queue, full_shape, dtype=numpy.float32), - # 'p_dev': cla.empty(queue, f_host.shape, dtype=numpy.float32), - # 'p_host': numpy.empty(f_host.shape, dtype=numpy.float32), - # } - # timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # # H2D: upload only input-sized data, scatter into padded device buffer - # t0 = timeit.default_timer() - # self._opencl_cache['f_stage_dev'].set(f_host, queue=queue) - # if self.dim == 2: - # scatter_kern = self._get_opencl_scatter_kernel(queue) - # fshape_c = self._opencl_cache['fshape'] - # full_shape_c = self._opencl_cache['full_shape'] - # scatter_kern( - # queue, (fshape_c[1], fshape_c[0]), None, - # self._opencl_cache['f_stage_dev'].data, - # self._opencl_cache['f_dev'].data, - # numpy.int32(fshape_c[1]), - # numpy.int32(fshape_c[0]), - # numpy.int32(full_shape_c[1]) - # ) - # else: - # # 3D fallback: host-side pad + full upload - # f_pad = numpy.zeros(full_shape, dtype=numpy.float32) - # f_slices = tuple(slice(0, s) for s in f_host.shape) - # f_pad[f_slices] = f_host - # self._opencl_cache['f_dev'].set(f_pad, queue=queue) - # if profile: - # queue.finish() - # timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # t0 = timeit.default_timer() - # f_fft_dev = cl_rfftn(self._opencl_cache['f_dev'], self._opencl_cache['f_fft_dev']) - # if profile: - # queue.finish() - # timing['fft_f_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # t0 = timeit.default_timer() - # f_fft_dev *= self._opencl_cache['g_fft_dev'] - # if profile: - # queue.finish() - # timing['mul_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # t0 = timeit.default_timer() - # conv_full_dev = cl_irfftn(f_fft_dev, self._opencl_cache['conv_full_dev']) - # if profile: - # queue.finish() - # timing['ifft_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # # Optimized path: use OpenCL kernel to crop and scale on device, then transfer only cropped result. - # t0 = timeit.default_timer() - # p_dev = self._opencl_cache['p_dev'] - # same_starts = self._opencl_cache['same_starts'] - # fshape = self._opencl_cache['fshape'] - # full_shape_cache = self._opencl_cache['full_shape'] - - # if self.dim == 2: - # # Use custom kernel for 2D cropping and scaling - # crop_kernel = self._get_opencl_crop_kernel(queue) - # crop_kernel( - # queue, (fshape[1], fshape[0]), None, - # conv_full_dev.data, - # p_dev.data, - # numpy.int32(full_shape_cache[1]), # src_width - # numpy.int32(fshape[1]), # dst_width - # numpy.int32(fshape[0]), # dst_height - # numpy.int32(same_starts[1]), # start_x - # numpy.int32(same_starts[0]), # start_y - # numpy.float32(self._opencl_cache['gdx_prod']) # scale - # ) - # if profile: - # queue.finish() - # timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # t0 = timeit.default_timer() - # p = p_dev.get(ary=self._opencl_cache['p_host']) - # timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0 - # else: - # # Fallback for 3D: transfer full, crop on CPU - # conv_full_dev *= self._opencl_cache['gdx_prod'] - # queue.finish() - # timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # t0 = timeit.default_timer() - # conv_full = conv_full_dev.get() - # timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0 - - # t0 = timeit.default_timer() - # p = conv_full[self._opencl_cache['same_slices']].copy() - # timing['crop_scale_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) finished') - # if profile: - # print( - # "OpenCL 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}" - # ) - # gtimer.end() - # return p - - # elif self.stype == 'S': - # raise SyntaxError('Spatial solver not suggested for structured grids') - - # elif self.ptype == 'P': - # if verbose: - # print('OpenCL solver not optimized for unstructured points, using CPU') - # return self.solve(f, verbose) - - diff --git a/src/orisys/configs.py b/src/orisys/configs.py new file mode 100644 index 0000000..ebdfad8 --- /dev/null +++ b/src/orisys/configs.py @@ -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) diff --git a/src/orisys/configs/__init__.py b/src/orisys/configs/__init__.py deleted file mode 100644 index 97d8b65..0000000 --- a/src/orisys/configs/__init__.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/src/orisys/configs/ddjx01.json b/src/orisys/configs/ddjx01.json deleted file mode 100644 index dc1e174..0000000 --- a/src/orisys/configs/ddjx01.json +++ /dev/null @@ -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 - } -} - - diff --git a/src/orisys/configs/prototype.json b/src/orisys/configs/prototype.json deleted file mode 100644 index b4aecb7..0000000 --- a/src/orisys/configs/prototype.json +++ /dev/null @@ -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 - } -} - - diff --git a/src/orisys/hhd/__init__.py b/src/orisys/hhd/__init__.py new file mode 100644 index 0000000..5a2f11c --- /dev/null +++ b/src/orisys/hhd/__init__.py @@ -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"] diff --git a/src/orisys/hhd/api.py b/src/orisys/hhd/api.py new file mode 100644 index 0000000..e03c2c4 --- /dev/null +++ b/src/orisys/hhd/api.py @@ -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() diff --git a/src/orisys/hhd/grid.py b/src/orisys/hhd/grid.py new file mode 100644 index 0000000..cdcf20b --- /dev/null +++ b/src/orisys/hhd/grid.py @@ -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() + + # -------------------------------------------------------------------------- + # -------------------------------------------------------------------------- diff --git a/src/orisys/hhd/kernels.py b/src/orisys/hhd/kernels.py new file mode 100644 index 0000000..7261f69 --- /dev/null +++ b/src/orisys/hhd/kernels.py @@ -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__ = [] diff --git a/src/orisys/hhd/poisson.py b/src/orisys/hhd/poisson.py new file mode 100644 index 0000000..a1a6a17 --- /dev/null +++ b/src/orisys/hhd/poisson.py @@ -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 + \ No newline at end of file diff --git a/src/orisys/hhd/timing.py b/src/orisys/hhd/timing.py new file mode 100644 index 0000000..9253c5a --- /dev/null +++ b/src/orisys/hhd/timing.py @@ -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() \ No newline at end of file diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index 3bb3382..dce081d 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -19,7 +19,7 @@ from typing import Tuple from .Stitcher import Stitcher from .tracker import Tracker -from .HHD import nHHD +from .hhd import nHHD from . import configs from .get_depth import Contact3D from .backend_availability import (check_cupy_available, check_opencl_available) @@ -202,7 +202,7 @@ class Sensor: Capture and preprocess a frame. 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 # Timing measurement helps identify if capture or processing is the bottleneck @@ -265,39 +265,66 @@ class Sensor: def detect_motion(self, threshold=6): # 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()}") 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. - - Returns: - np.ndarray: Optical flow field + Compute optical flow and optionally decompose it into force components. + + Parameters + ---------- + 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: raise Exception("Warning: Get image first") - + self.compute_flow() - # self.compute_hhd() - if check_motion: - 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 not check_motion: + self.tail_frames_count = 0 if decompose: self.compute_hhd() - 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): if self.img is None: diff --git a/test_3Dflow.py b/test_3Dflow.py deleted file mode 100644 index 53bc26f..0000000 --- a/test_3Dflow.py +++ /dev/null @@ -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() diff --git a/test_cam.py b/test_cam.py deleted file mode 100644 index 6f5e2c7..0000000 --- a/test_cam.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/test_downsample.py b/test_downsample.py deleted file mode 100644 index d9e1f77..0000000 --- a/test_downsample.py +++ /dev/null @@ -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() diff --git a/test_ffmpeg.py b/test_ffmpeg.py deleted file mode 100644 index 2ed320f..0000000 --- a/test_ffmpeg.py +++ /dev/null @@ -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() diff --git a/test_opencl.py b/test_opencl.py deleted file mode 100644 index 7a0e9a6..0000000 --- a/test_opencl.py +++ /dev/null @@ -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!") \ No newline at end of file diff --git a/test_stitch.py b/test_stitch.py deleted file mode 100644 index 35c4689..0000000 --- a/test_stitch.py +++ /dev/null @@ -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() diff --git a/test_usb_fps.py b/test_usb_fps.py deleted file mode 100644 index 4560170..0000000 --- a/test_usb_fps.py +++ /dev/null @@ -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) - diff --git a/tools/3d_patch_generation_blender/blender_export.py b/tools/3d_patch_generation_blender/blender_export.py new file mode 100644 index 0000000..2f4d2c2 --- /dev/null +++ b/tools/3d_patch_generation_blender/blender_export.py @@ -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() \ No newline at end of file diff --git a/tools/3d_patch_generation_blender/checker_patch_dd02c.py b/tools/3d_patch_generation_blender/checker_patch_dd02c.py new file mode 100644 index 0000000..f8d0aa8 --- /dev/null +++ b/tools/3d_patch_generation_blender/checker_patch_dd02c.py @@ -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() \ No newline at end of file diff --git a/tools/3d_patch_generation_blender/dd02c.blend b/tools/3d_patch_generation_blender/dd02c.blend new file mode 100644 index 0000000..119edbd Binary files /dev/null and b/tools/3d_patch_generation_blender/dd02c.blend differ diff --git a/tools/3d_patch_generation_blender/mapto3d.py b/tools/3d_patch_generation_blender/mapto3d.py new file mode 100644 index 0000000..d2d936f --- /dev/null +++ b/tools/3d_patch_generation_blender/mapto3d.py @@ -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]) \ No newline at end of file diff --git a/tools/3d_patch_generation_blender/random_patch_dd02c.py b/tools/3d_patch_generation_blender/random_patch_dd02c.py new file mode 100644 index 0000000..a010717 --- /dev/null +++ b/tools/3d_patch_generation_blender/random_patch_dd02c.py @@ -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() \ No newline at end of file diff --git a/tools/3d_patch_generation_blender/uv_to_printer.py b/tools/3d_patch_generation_blender/uv_to_printer.py new file mode 100644 index 0000000..7443ef4 --- /dev/null +++ b/tools/3d_patch_generation_blender/uv_to_printer.py @@ -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() \ No newline at end of file diff --git a/patch_generation/3.6mm.svg b/tools/patch_generation/3.6mm.svg similarity index 100% rename from patch_generation/3.6mm.svg rename to tools/patch_generation/3.6mm.svg diff --git a/patch_generation/cmyk_generated_rgb.png b/tools/patch_generation/cmyk_generated_rgb.png similarity index 100% rename from patch_generation/cmyk_generated_rgb.png rename to tools/patch_generation/cmyk_generated_rgb.png diff --git a/patch_generation/cmyk_pattern.tiff b/tools/patch_generation/cmyk_pattern.tiff similarity index 100% rename from patch_generation/cmyk_pattern.tiff rename to tools/patch_generation/cmyk_pattern.tiff diff --git a/patch_generation/dd01.svg b/tools/patch_generation/dd01.svg similarity index 100% rename from patch_generation/dd01.svg rename to tools/patch_generation/dd01.svg diff --git a/patch_generation/patch_generation.py b/tools/patch_generation/patch_generation.py similarity index 100% rename from patch_generation/patch_generation.py rename to tools/patch_generation/patch_generation.py diff --git a/patch_generation/patch_generation_3.py b/tools/patch_generation/patch_generation_3.py similarity index 100% rename from patch_generation/patch_generation_3.py rename to tools/patch_generation/patch_generation_3.py diff --git a/tools/patch_generation/patch_generation_3_org.py b/tools/patch_generation/patch_generation_3_org.py new file mode 100644 index 0000000..0a09359 --- /dev/null +++ b/tools/patch_generation/patch_generation_3_org.py @@ -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() diff --git a/tools/patch_generation/patchgen_DD01.py b/tools/patch_generation/patchgen_DD01.py new file mode 100644 index 0000000..ca60363 --- /dev/null +++ b/tools/patch_generation/patchgen_DD01.py @@ -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() diff --git a/tools/patch_generation/patchgen_DD02C_dot.py b/tools/patch_generation/patchgen_DD02C_dot.py new file mode 100644 index 0000000..948d6b3 --- /dev/null +++ b/tools/patch_generation/patchgen_DD02C_dot.py @@ -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') diff --git a/patch_generation/patch_generation_3_org.py b/tools/patch_generation/patchgen_DD02M.py old mode 100755 new mode 100644 similarity index 100% rename from patch_generation/patch_generation_3_org.py rename to tools/patch_generation/patchgen_DD02M.py diff --git a/verify_fft_equivalence.py b/verify_fft_equivalence.py deleted file mode 100644 index a5a17ab..0000000 --- a/verify_fft_equivalence.py +++ /dev/null @@ -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()