Initial commit: Speckle-Scanner 3D pipeline with setup README
This commit is contained in:
@@ -0,0 +1,323 @@
|
||||
# **Stereo Disparity Map and Point Cloud Generation**
|
||||
|
||||
This project implements a dense stereo disparity map and point cloud generation pipeline using spatio-temporal Zero-Normalized Cross-Correlation (ZNCC) with subpixel interpolation methods and other advanced features.
|
||||
|
||||
## **Features**
|
||||
|
||||
- **Dense Disparity Map Calculation**: Generates Disparity and vertical shift maps.
|
||||
- **Subpixel Interpolation**: Supports parabolic, Gaussian, and equiangular interpolation methods for subpixel accuracy.
|
||||
- **Correlation Map**: Computes ZNCC correlation for each pixel.
|
||||
- **Point Cloud Generation**: Reprojects disparity maps into 3D space using a Q matrix.
|
||||
- **Color Mapping**: Normalizes depth values and applies a colormap for visualization.
|
||||
- **Optional Median Filtering**: Smoothens the disparity map.
|
||||
- **Visualization**: Saves disparity maps, vertical shift maps, and correlation maps with colorbars.
|
||||
|
||||
## **Requirements**
|
||||
|
||||
- Python 3.9 or later
|
||||
|
||||
```bash
|
||||
# This step only
|
||||
pip install -r ~/Speckle-Scanner/05_disparity/zncc/requirements.txt
|
||||
|
||||
# Or install everything for the full pipeline
|
||||
pip install -r ~/Speckle-Scanner/requirements.txt
|
||||
|
||||
# Optional GPU acceleration (pick one matching your CUDA version)
|
||||
pip install cupy-cuda12x # CUDA 12.x
|
||||
pip install cupy-cuda11x # CUDA 11.2–11.8
|
||||
```
|
||||
|
||||
Packages: `numpy`, `opencv-python`, `matplotlib`, `scipy`, `numba`, `open3d`. Without `cupy`, the pipeline falls back to CPU (`numba`).
|
||||
|
||||
## **Input Requirements**
|
||||
|
||||
- **Stereo Rectified Images**:
|
||||
- Images must be stored in `.png` or `.bmp` format.
|
||||
- **Pipeline pairing** (in `02_rect_images/`): `lc_ts1634840093_ck….png` is matched with `rc_ts1634840093_ck….png` on the shared `ts` token; the `ck…` suffix is ignored.
|
||||
- Pairs are sorted by timestamp; `--num_images` defaults to **all** pairs; if limited, the **last N** (highest timestamps) are used.
|
||||
- Standalone runs may still use separate `left/` and `right/` folders with any naming; pipeline mode uses one folder with `lc_` / `rc_` prefixes.
|
||||
- PNG is preferred when both PNG and BMP pairs exist.
|
||||
- **Q Matrix File**:
|
||||
- A `.cvstore` file containing the Q matrix for 3D reprojection.
|
||||
|
||||
## **Outputs**
|
||||
|
||||
- **Disparity Map**: Saved as `Disparity.npy`.
|
||||
- **Vertical Shift Map**: Saved as `vertical_shift_map.npy`.
|
||||
- **Correlation Map**: Saved as `correlation_map.npy`.
|
||||
- **Disparity and Shift Maps with Colorbars**:
|
||||
- `Disparity_map_colorbar.png`
|
||||
- `vertical_shift_map_colorbar.png`
|
||||
- `correlation_map_colorbar.png`
|
||||
- **Point Cloud**:
|
||||
- Saved as `Point_cloud.ply` in PLY format with depth-based color information.
|
||||
|
||||
## **Usage**
|
||||
|
||||
1. Clone the repository and navigate to the project directory.
|
||||
2. Prepare your stereo images and paste them in data folder
|
||||
Example directory structure:
|
||||
```
|
||||
|
||||
project_directory/
|
||||
|
||||
.data/
|
||||
├── left/
|
||||
│ ├── 0.png
|
||||
│ ├── 1.png
|
||||
│ └── ...
|
||||
└── right/
|
||||
├── 0.png
|
||||
├── 1.png
|
||||
└── ...
|
||||
|
||||
Or
|
||||
|
||||
.data/
|
||||
├── left/
|
||||
│ ├── lc00000.bmp
|
||||
│ └── lc00001.bmp
|
||||
└── right/
|
||||
├── rc00000.bmp
|
||||
└── rc00001.bmp
|
||||
|
||||
```
|
||||
3. The Q matrix must be provided in a YAML file (e.g., `Q.cvstore`).
|
||||
Format example:
|
||||
```
|
||||
|
||||
%YAML:1.0
|
||||
---
|
||||
Q: !!opencv-matrix
|
||||
rows: 4
|
||||
cols: 4
|
||||
dt: f
|
||||
data: [ 1.0, 0.0, 0.0, -5.5863211059570312e+02,
|
||||
0.0, 1.0, 0.0, -7.3320780181884766e+02,
|
||||
0.0, 0.0, 0.0, 3.3665836719173171e+03,
|
||||
0.0, 0.0, 1.4262549294100269e+00, 1.0682440279589944e+03 ]
|
||||
|
||||
```
|
||||
This matrix is essential for converting disparity maps into a 3D point cloud.
|
||||
3. Run the script with desired parameters (see below).
|
||||
|
||||
### **Basic Command**
|
||||
|
||||
```bash
|
||||
# stereo_disparity_main.py auto-selects GPU (CUDA) or CPU.
|
||||
|
||||
cd ~/Speckle-Scanner/05_disparity/zncc
|
||||
|
||||
# Default parameters (saves disparity map only)
|
||||
python stereo_disparity_main.py \
|
||||
--left_dir <path_to_rect_images> \
|
||||
--right_dir <path_to_rect_images> \
|
||||
--left_prefix lc_ \
|
||||
--right_prefix rc_ \
|
||||
--q_file <path_to_lc-rc_Q.cvstore> \
|
||||
--disp_output_dir <path_to_04_zncc_disp_map> \
|
||||
--pcl_output_dir <path_to_06_zncc_pcl>
|
||||
|
||||
# Troubleshooting mode — saves all outputs (vertical shift, correlation, point cloud)
|
||||
python stereo_disparity_main.py \
|
||||
--left_dir <path_to_rect_images> \
|
||||
--right_dir <path_to_rect_images> \
|
||||
--left_prefix lc_ \
|
||||
--right_prefix rc_ \
|
||||
--q_file <path_to_lc-rc_Q.cvstore> \
|
||||
--disp_output_dir <path_to_04_zncc_disp_map> \
|
||||
--pcl_output_dir <path_to_06_zncc_pcl> \
|
||||
--troubleshooting
|
||||
```
|
||||
This runs the program with default parameters.
|
||||
Same commands for stereo_disparity_cpu.py and stereo_disparity_gpu.py
|
||||
|
||||
### **Custom Parameters**
|
||||
|
||||
|
||||
You can override any parameter through command-line arguments. Below is an example with some customized parameters:
|
||||
|
||||
```bash
|
||||
python stereo_disparity_main.py \
|
||||
--window_size 5 \
|
||||
--H_neg_range -165 \
|
||||
--H_pos_range 65 \
|
||||
--v_neg_range -4 \
|
||||
--v_pos_range 0 \
|
||||
--zncc_threshold 0.4 \
|
||||
--num_images 10 \ # optional: last 10 pairs by timestamp; omit for all pairs
|
||||
--noise_filter "open3d" \
|
||||
--interpolation True \
|
||||
--noise_remove False \
|
||||
--method "gaussian" \
|
||||
--q_file "./data/Q_matrix.cvstore" \
|
||||
--left_dir "./data/left" \
|
||||
--right_dir "./data/right" \
|
||||
--disp_output_dir "./output/disp" \
|
||||
--pcl_output_dir "./output/pcl" \
|
||||
--troubleshooting
|
||||
```
|
||||
|
||||
|
||||
## **Available Parameters**
|
||||
|
||||
| Parameter | Default Value | Description |
|
||||
| ------------------ | ------------------- | --------------------------------------------------------------------------------------|
|
||||
| `--window_size` | `5` | Size of the window for block matching. |
|
||||
| `--H_neg_range` | `0` | Horizontal negative disparity range. |
|
||||
| `--H_pos_range` | `65` | Horizontal positive disparity range. |
|
||||
| `--v_neg_range` | `-4` | Vertical negative disparity range. |
|
||||
| `--v_pos_range` | `2` | Vertical positive disparity range. |
|
||||
| `--zncc_threshold` | `0.4` | ZNCC threshold for valid matches. |
|
||||
| `--num_images` | all pairs | Stereo pairs to use per scan. Default: all `lc_ts*`/`rc_ts*` pairs matched on `ts`. If fewer than available, the **last N** pairs (highest timestamps) are used; `ck*` suffix is ignored. |
|
||||
| `--noise_filter` | `open3d` | Post-processing noise filter method (`median`[Disparity Map] or `open3d`[Pointcloud]).|
|
||||
| `--interpolation` | `True` | Enable subpixel interpolation (`True` or `False`). |
|
||||
| `--noise_remove` | `True` | noise removal using bidirectional disparity estimation (`True` or `False`). |
|
||||
| `--method` | `parabolic` | Subpixel interpolation method (`parabolic`, `gaussian`, or `equiangular`). |
|
||||
| `--q_file` | `./data1/Q.cvstore` | Path to the Q matrix file. |
|
||||
| `--left_dir` | `./data1/left` | Path to the directory containing left images. |
|
||||
| `--right_dir` | `./data1/right` | Path to the directory containing right images. |
|
||||
| `--left_prefix` | `""` | Filename prefix to filter left images (e.g. `lc_`) when both cameras share a folder. |
|
||||
| `--right_prefix` | `""` | Filename prefix to filter right images (e.g. `rc_`) when both cameras share a folder. |
|
||||
| `--output_dir` | `./results` | Fallback output directory (used when `--disp_output_dir`/`--pcl_output_dir` are not set). |
|
||||
| `--disp_output_dir` | same as `--output_dir` | Directory for disparity map outputs (npy + png). |
|
||||
| `--pcl_output_dir` | same as `--output_dir` | Directory for point cloud outputs (PLY + TXT). |
|
||||
| `--troubleshooting` | `False` | When set, saves all outputs (vertical shift map, correlation map, point cloud); otherwise only disparity map is saved. |
|
||||
|
||||
|
||||
## **Details of Post-processing noise filter method**
|
||||
|
||||
- **open3d**: used for denoising point clouds.
|
||||
- **median**: used to reduce noise in the disparity map.
|
||||
|
||||
## **Details of Subpixel Interpolation Methods**
|
||||
|
||||
- **Parabolic**: Fits a parabola to the ZNCC values and finds the peak.
|
||||
- **Gaussian**: Uses the logarithm of ZNCC values to model a Gaussian distribution and find the peak.
|
||||
- **Equiangular**: Uses angular interpolation for smoother and more robust subpixel shifts.
|
||||
|
||||
## **Notes**
|
||||
|
||||
- Ensure stereo images are rectified before processing.
|
||||
- Adjust the disparity ranges (`H_neg_range`, `H_pos_range`, etc.) based on your dataset.
|
||||
- If disparity maps appear noisy, enable median filtering or refine the ZNCC threshold.
|
||||
|
||||
## **Execution Time**
|
||||
|
||||
- Execution time depends on the number of images, window size, and disparity ranges.
|
||||
- Subpixel interpolation adds computational overhead but improves accuracy.
|
||||
|
||||
## **Point Cloud File**
|
||||
|
||||
- The generated `.ply` file can be visualized using point cloud tools like MeshLab or CloudCompare.
|
||||
- Depth values are normalized and color-coded using the `jet` colormap.
|
||||
|
||||
## **Troubleshooting**
|
||||
|
||||
- Ensure the `Q.cvstore` file exists in the specified path.
|
||||
- Check file permissions.
|
||||
- Ensure the left and right images are rectified and have the same resolution.
|
||||
- Reduce the number of images or use a smaller window size.
|
||||
- Run on a machine with higher processing power.
|
||||
|
||||
---
|
||||
|
||||
## **Pipeline Usage (Automated Path Resolution)**
|
||||
|
||||
Use `run_zncc_pipeline.py` instead of calling `stereo_disparity_main.py` directly.
|
||||
It resolves all paths automatically from the project folder structure and processes
|
||||
every scan in a session (or a single scan you name).
|
||||
|
||||
### **Folder structure assumed**
|
||||
|
||||
```
|
||||
~/Speckle-Scanner_Processing_data/
|
||||
└── <project>/
|
||||
└── <date>/
|
||||
└── <session>/
|
||||
├── params_link/
|
||||
│ └── lc-rc_Q.cvstore ← Q matrix (input)
|
||||
└── <ScanXXXXXX>/
|
||||
├── 02_rect_images/ ← lc_ts*.png + rc_ts*.png (input)
|
||||
├── 04_zncc_disp_map/ ← disparity .npy + colorbar .png (created)
|
||||
└── 06_zncc_pcl/ ← Point_cloud.ply + .txt (created)
|
||||
```
|
||||
|
||||
### **Commands**
|
||||
|
||||
```bash
|
||||
cd ~/Speckle-Scanner/05_disparity/zncc
|
||||
|
||||
# Process ALL scans in a session (all matched lc/rc pairs per scan)
|
||||
python run_zncc_pipeline.py \
|
||||
--project Olsen_wings \
|
||||
--date 2026-05-12 \
|
||||
--session session1
|
||||
|
||||
# Process ALL sessions on a date (omit --session)
|
||||
python run_zncc_pipeline.py \
|
||||
--project Olsen_wings \
|
||||
--date 2026-05-12
|
||||
|
||||
# Process a SINGLE scan — only the last 3 pairs (highest timestamps)
|
||||
python run_zncc_pipeline.py \
|
||||
--project Olsen_wings \
|
||||
--date 2026-05-12 \
|
||||
--session session1 \
|
||||
--scan Scan000001 \
|
||||
--num_images 3
|
||||
|
||||
# Full outputs including point cloud and debug maps (troubleshooting mode)
|
||||
python run_zncc_pipeline.py \
|
||||
--project Olsen_wings \
|
||||
--date 2026-05-12 \
|
||||
--session session1 \
|
||||
--scan Scan000001 \
|
||||
--troubleshooting
|
||||
|
||||
# Custom ZNCC parameters
|
||||
python run_zncc_pipeline.py \
|
||||
--project Olsen_wings \
|
||||
--date 2026-05-12 \
|
||||
--session session1 \
|
||||
--window_size 7 \
|
||||
--H_neg_range 0 \
|
||||
--H_pos_range 80 \
|
||||
--v_neg_range -4 \
|
||||
--v_pos_range 2 \
|
||||
--zncc_threshold 0.5 \
|
||||
--noise_filter open3d \
|
||||
--method gaussian
|
||||
```
|
||||
|
||||
### **Pipeline parameters**
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|---------------------|---------|------------------------------------------------------------------------------------|
|
||||
| `--project` | — | Project name (e.g. `Olsen_wings`) |
|
||||
| `--date` | — | Date string (e.g. `2026-05-12`) |
|
||||
| `--session` | all | Session name (e.g. `session1`); omit to process **all sessions** on that date |
|
||||
| `--scan` | all | Single scan to process (e.g. `Scan000001`); omit to process all scans in session |
|
||||
| `--troubleshooting` | off | When set, saves vertical shift map, correlation map, and point cloud PLY/TXT; by default only `disparity.npy` and `Disparity_map_colorbar.png` are saved |
|
||||
| `--num_images` | all | Image pairs per scan (matched on `ts`; last N if limited) |
|
||||
| All other ZNCC params | see above table | Forwarded directly to `stereo_disparity_main.py` |
|
||||
|
||||
### **What gets saved**
|
||||
|
||||
| File | Default | `--troubleshooting` |
|
||||
|------|:-------:|:-------------------:|
|
||||
| `04_zncc_disp_map/disparity.npy` | ✓ | ✓ |
|
||||
| `04_zncc_disp_map/Disparity_map_colorbar.png` | ✓ | ✓ |
|
||||
| `04_zncc_disp_map/vertical_shift_map.npy` | | ✓ |
|
||||
| `04_zncc_disp_map/vertical_shift_map_colorbar.png` | | ✓ |
|
||||
| `04_zncc_disp_map/correlation_map.npy` | | ✓ |
|
||||
| `04_zncc_disp_map/correlation_map_colorbar.png` | | ✓ |
|
||||
| `06_zncc_pcl/Point_cloud.ply` | | ✓ |
|
||||
| `06_zncc_pcl/Point_cloud.txt` | | ✓ |
|
||||
|
||||
---
|
||||
|
||||
Contributions and improvements are welcome! Feel free to open an issue or submit a pull request on GitHub.
|
||||
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
# 05_disparity/zncc — Python dependencies
|
||||
# Install: pip install -r requirements.txt
|
||||
# Full pipeline (all steps): pip install -r ~/Speckle-Scanner/requirements.txt
|
||||
#
|
||||
# GPU acceleration (optional — falls back to CPU numba if absent):
|
||||
# CUDA 11.x: pip install cupy-cuda11x
|
||||
# CUDA 12.x: pip install cupy-cuda12x
|
||||
|
||||
numpy>=1.21
|
||||
opencv-python>=4.8
|
||||
matplotlib>=3.5
|
||||
scipy>=1.7
|
||||
numba>=0.56
|
||||
open3d>=0.16
|
||||
@@ -0,0 +1,178 @@
|
||||
"""
|
||||
Pipeline runner for ZNCC stereo disparity.
|
||||
|
||||
Resolves all paths from the project folder structure and drives
|
||||
stereo_disparity_main.py for each scan in a session (or all sessions in a date).
|
||||
|
||||
Output layout per scan:
|
||||
<processing_dir>/<project>/<date>/<session>/<scan>/
|
||||
02_rect_images/ <- input (must exist)
|
||||
04_zncc_disp_map/ <- disparity.npy + Disparity_map_colorbar.png (always)
|
||||
+ vertical shift & correlation maps (troubleshooting only)
|
||||
06_zncc_pcl/ <- Point_cloud.ply + .txt (troubleshooting only)
|
||||
|
||||
Q matrix is read from:
|
||||
<processing_dir>/<project>/<date>/<session>/params_link/lc-rc_Q.cvstore
|
||||
"""
|
||||
|
||||
import sys
|
||||
import argparse
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
# Resolve config.py relative to ~/Speckle-Scanner regardless of CWD
|
||||
sys.path.insert(0, str(Path.home() / "Speckle-Scanner"))
|
||||
import config # noqa: E402
|
||||
|
||||
SCRIPT = Path(__file__).parent / "stereo_disparity_main.py"
|
||||
|
||||
|
||||
def build_cmd(rect_dir, q_file, disp_out, pcl_out, zncc_args, troubleshooting):
|
||||
cmd = [
|
||||
sys.executable, str(SCRIPT),
|
||||
"--left_dir", str(rect_dir),
|
||||
"--right_dir", str(rect_dir),
|
||||
"--left_prefix", "lc_",
|
||||
"--right_prefix", "rc_",
|
||||
"--q_file", str(q_file),
|
||||
"--disp_output_dir", str(disp_out),
|
||||
"--pcl_output_dir", str(pcl_out),
|
||||
]
|
||||
if troubleshooting:
|
||||
cmd.append("--troubleshooting")
|
||||
# Forward ZNCC tuning params
|
||||
for key, val in zncc_args.items():
|
||||
if val is not None:
|
||||
cmd += [f"--{key}", str(val)]
|
||||
return cmd
|
||||
|
||||
|
||||
def run_scan(project, date, session, scan, zncc_args, troubleshooting):
|
||||
rect_dir = config.PROCESSING_DIR / project / date / session / scan / "02_rect_images"
|
||||
q_file = config.get_params_link_dir(project, date, session) / "lc-rc_Q.cvstore"
|
||||
disp_out = config.get_processing_step_dir(project, date, session, scan, "04_zncc_disp_map")
|
||||
pcl_out = config.get_processing_step_dir(project, date, session, scan, "06_zncc_pcl")
|
||||
|
||||
if not rect_dir.exists():
|
||||
print(f"[SKIP] {session}/{scan}: 02_rect_images not found at {rect_dir}")
|
||||
return False
|
||||
if not q_file.exists():
|
||||
print(f"[SKIP] {session}/{scan}: Q matrix not found at {q_file}")
|
||||
return False
|
||||
|
||||
mode = "troubleshooting" if troubleshooting else "disparity-only"
|
||||
print(f"\n{'='*60}")
|
||||
print(f"[SCAN] {session}/{scan} [{mode}]")
|
||||
print(f" rect : {rect_dir}")
|
||||
print(f" Q : {q_file}")
|
||||
print(f" disp : {disp_out}")
|
||||
if troubleshooting:
|
||||
print(f" pcl : {pcl_out}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
cmd = build_cmd(rect_dir, q_file, disp_out, pcl_out, zncc_args, troubleshooting)
|
||||
result = subprocess.run(cmd)
|
||||
if result.returncode != 0:
|
||||
print(f"[FAIL] {session}/{scan} exited with code {result.returncode}")
|
||||
return False
|
||||
print(f"[DONE] {session}/{scan}")
|
||||
return True
|
||||
|
||||
|
||||
def run_session(project, date, session, scan_arg, zncc_args, troubleshooting):
|
||||
if scan_arg:
|
||||
scans = [scan_arg]
|
||||
else:
|
||||
scans = config.list_scan_dirs(project, date, session)
|
||||
if not scans:
|
||||
print(f"[WARN] No scan folders found in {project}/{date}/{session}")
|
||||
return [], []
|
||||
print(f"\n Session {session}: {len(scans)} scan(s) found")
|
||||
|
||||
failed = []
|
||||
for scan in scans:
|
||||
ok = run_scan(project, date, session, scan, zncc_args, troubleshooting)
|
||||
if not ok:
|
||||
failed.append(f"{session}/{scan}")
|
||||
return scans, failed
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="ZNCC disparity pipeline runner — resolves paths from project structure"
|
||||
)
|
||||
# Project location
|
||||
parser.add_argument("--project", required=True, help="Project name (e.g. Olsen_wings)")
|
||||
parser.add_argument("--date", required=True, help="Date string (e.g. 2026-05-12)")
|
||||
parser.add_argument("--session", default=None, help="Session name (e.g. session1); omit to process ALL sessions on that date")
|
||||
parser.add_argument("--scan", default=None, help="Single scan to process (e.g. Scan000001); omit to process all scans in the session")
|
||||
|
||||
# Output mode
|
||||
parser.add_argument("--troubleshooting", action="store_true",
|
||||
help="Save all outputs (vertical shift map, correlation map, point cloud PLY/TXT); "
|
||||
"by default only disparity map (npy + png) is saved")
|
||||
|
||||
# ZNCC tuning — all optional, forwarded to stereo_disparity_main.py
|
||||
parser.add_argument(
|
||||
"--num_images",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Stereo pairs per scan. Default: all matched pairs. "
|
||||
"If fewer than available, uses the last N (highest timestamps).",
|
||||
)
|
||||
parser.add_argument("--window_size", type=int, default=None)
|
||||
parser.add_argument("--H_neg_range", type=int, default=None)
|
||||
parser.add_argument("--H_pos_range", type=int, default=None)
|
||||
parser.add_argument("--v_neg_range", type=int, default=None)
|
||||
parser.add_argument("--v_pos_range", type=int, default=None)
|
||||
parser.add_argument("--zncc_threshold", type=float, default=None)
|
||||
parser.add_argument("--noise_filter", type=str, default=None, choices=["median", "open3d"])
|
||||
parser.add_argument("--interpolation", type=str, default=None)
|
||||
parser.add_argument("--noise_remove", type=str, default=None)
|
||||
parser.add_argument("--method", type=str, default=None, choices=["parabolic", "gaussian", "equiangular"])
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
zncc_args = {
|
||||
"num_images": args.num_images,
|
||||
"window_size": args.window_size,
|
||||
"H_neg_range": args.H_neg_range,
|
||||
"H_pos_range": args.H_pos_range,
|
||||
"v_neg_range": args.v_neg_range,
|
||||
"v_pos_range": args.v_pos_range,
|
||||
"zncc_threshold": args.zncc_threshold,
|
||||
"noise_filter": args.noise_filter,
|
||||
"interpolation": args.interpolation,
|
||||
"noise_remove": args.noise_remove,
|
||||
"method": args.method,
|
||||
}
|
||||
|
||||
# Determine sessions to process
|
||||
if args.session:
|
||||
sessions = [args.session]
|
||||
else:
|
||||
sessions = config.list_session_dirs(args.project, args.date)
|
||||
if not sessions:
|
||||
print(f"No session folders found under {args.project}/{args.date}")
|
||||
sys.exit(1)
|
||||
print(f"Found {len(sessions)} session(s): {sessions}")
|
||||
|
||||
total_scans = 0
|
||||
all_failed = []
|
||||
for session in sessions:
|
||||
scans, failed = run_session(
|
||||
args.project, args.date, session,
|
||||
args.scan, zncc_args, args.troubleshooting
|
||||
)
|
||||
total_scans += len(scans)
|
||||
all_failed.extend(failed)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Finished: {total_scans - len(all_failed)}/{total_scans} scans succeeded.")
|
||||
if all_failed:
|
||||
print(f"Failed: {all_failed}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,626 @@
|
||||
import os
|
||||
import argparse
|
||||
import numpy as np
|
||||
import cv2
|
||||
import subprocess
|
||||
import glob
|
||||
import re
|
||||
import open3d as o3d
|
||||
import matplotlib.pyplot as plt
|
||||
import time
|
||||
from scipy.ndimage import median_filter
|
||||
from numba import njit, prange
|
||||
|
||||
# Save the point cloud in PLY format.
|
||||
def save_ply(filename, points):
|
||||
with open(filename, 'w') as f:
|
||||
f.write("ply\n")
|
||||
f.write("format ascii 1.0\n")
|
||||
f.write(f"element vertex {points.shape[0]}\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 point in points:
|
||||
x, y, z, r, g, b = point
|
||||
f.write(f"{x} {y} {z} {int(r * 255)} {int(g * 255)} {int(b * 255)}\n")
|
||||
|
||||
def save_ascii_point_cloud(filename, points):
|
||||
with open(filename, 'w') as f:
|
||||
for point in points:
|
||||
x, y, z, r, g, b = point
|
||||
f.write(f"{x} {y} {z} {int(r * 255)} {int(g * 255)} {int(b * 255)}\n")
|
||||
|
||||
def process_disparity(disp_left, disp_right, tolerance=1.0, filter_size=10):
|
||||
disp_right_aligned = -disp_right # Align disparities
|
||||
height, width = disp_left.shape
|
||||
invalid_value = np.min(disp_left)
|
||||
valid_disp = np.full_like(disp_left, invalid_value)
|
||||
|
||||
for y in range(height):
|
||||
for x in range(width):
|
||||
d_left = disp_left[y, x]
|
||||
if d_left == invalid_value:
|
||||
continue
|
||||
x_right = int(x - d_left)
|
||||
if x_right < 0 or x_right >= width:
|
||||
continue
|
||||
d_right = disp_right_aligned[y, x_right]
|
||||
if abs(d_left - d_right) <= tolerance:
|
||||
valid_disp[y, x] = d_left
|
||||
|
||||
return median_filter(valid_disp, size=filter_size)
|
||||
|
||||
def save_with_colorbar(matrix, output_path, cmap="jet"):
|
||||
plt.figure(figsize=(10, 6))
|
||||
plt.imshow(matrix, cmap=cmap, aspect="auto")
|
||||
plt.colorbar()
|
||||
plt.title("Matrix Visualization with Colorbar")
|
||||
plt.savefig(output_path, bbox_inches="tight")
|
||||
plt.close()
|
||||
|
||||
def load_q_matrix(q_file):
|
||||
fs = cv2.FileStorage(q_file, cv2.FILE_STORAGE_READ)
|
||||
Q = fs.getNode("Q").mat()
|
||||
fs.release()
|
||||
if Q is None:
|
||||
raise ValueError("Failed to load Q matrix from the provided file.")
|
||||
return Q
|
||||
|
||||
def extract_ts_token(filename, prefix):
|
||||
"""Extract ts token from e.g. lc_ts1634840093_ck....png -> ('ts1634840093', 1634840093)."""
|
||||
name = os.path.basename(filename)
|
||||
m = re.search(rf"^{re.escape(prefix)}(ts\d+)", name, re.IGNORECASE)
|
||||
if not m:
|
||||
return None, None
|
||||
ts_token = m.group(1).lower()
|
||||
ts_int = int(re.search(r"\d+", ts_token).group())
|
||||
return ts_token, ts_int
|
||||
|
||||
|
||||
def discover_stereo_pairs(left_dir, right_dir, left_prefix="", right_prefix="", extensions=(".png", ".bmp")):
|
||||
"""
|
||||
Match lc/rc files by shared ts token (ck* suffix ignored).
|
||||
Returns list of (left_path, right_path, ts_int) sorted by timestamp ascending.
|
||||
"""
|
||||
best_pairs = []
|
||||
best_ext = None
|
||||
|
||||
for ext in extensions:
|
||||
pairs_by_ts = {}
|
||||
for lc_path in glob.glob(os.path.join(left_dir, f"{left_prefix}*{ext}")):
|
||||
ts_token, ts_int = extract_ts_token(lc_path, left_prefix)
|
||||
if ts_token is None:
|
||||
continue
|
||||
|
||||
rc_pattern = os.path.join(right_dir, f"{right_prefix}{ts_token}_*{ext}")
|
||||
rc_matches = sorted(glob.glob(rc_pattern))
|
||||
if not rc_matches:
|
||||
rc_pattern = os.path.join(right_dir, f"{right_prefix}{ts_token}*{ext}")
|
||||
rc_matches = sorted(glob.glob(rc_pattern))
|
||||
if not rc_matches:
|
||||
continue
|
||||
|
||||
pairs_by_ts[ts_int] = (lc_path, rc_matches[0], ts_int)
|
||||
|
||||
pairs = [pairs_by_ts[k] for k in sorted(pairs_by_ts)]
|
||||
if len(pairs) > len(best_pairs):
|
||||
best_pairs = pairs
|
||||
best_ext = ext.lstrip(".")
|
||||
|
||||
return best_pairs, best_ext
|
||||
|
||||
|
||||
def load_stereo_stacks(left_dir, right_dir, num_images=None, left_prefix="", right_prefix=""):
|
||||
"""
|
||||
Load stereo image stacks matched by ts token.
|
||||
|
||||
num_images:
|
||||
None or <= 0 -> use all matched pairs in the folder
|
||||
N < available -> use the last N pairs (highest timestamps)
|
||||
N >= available -> use all available pairs
|
||||
"""
|
||||
pairs, ext = discover_stereo_pairs(left_dir, right_dir, left_prefix, right_prefix)
|
||||
if not pairs:
|
||||
raise ValueError(
|
||||
"No matching stereo image pairs (.png or .bmp) found. "
|
||||
f"Expected {left_prefix}ts<timestamp>_*.{{png,bmp}} paired with "
|
||||
f"{right_prefix}ts<same_timestamp>_*.{{png,bmp}}."
|
||||
)
|
||||
|
||||
total = len(pairs)
|
||||
if num_images is None or num_images <= 0:
|
||||
selected = pairs
|
||||
elif num_images >= total:
|
||||
if num_images > total:
|
||||
print(f"Warning: Only {total} pair(s) available (requested {num_images}). Using all.")
|
||||
selected = pairs
|
||||
else:
|
||||
selected = pairs[-num_images:]
|
||||
|
||||
ts_first = selected[0][2]
|
||||
ts_last = selected[-1][2]
|
||||
print(
|
||||
f"Using {len(selected)}/{total} {ext.upper()} pair(s) "
|
||||
f"(timestamps {ts_first} … {ts_last}; matched on ts, ck suffix ignored)."
|
||||
)
|
||||
|
||||
left_images = []
|
||||
right_images = []
|
||||
for lc_path, rc_path, _ in selected:
|
||||
left_img = cv2.imread(lc_path, cv2.IMREAD_GRAYSCALE)
|
||||
right_img = cv2.imread(rc_path, cv2.IMREAD_GRAYSCALE)
|
||||
if left_img is None or right_img is None:
|
||||
raise ValueError(f"Failed to load pair: {lc_path} / {rc_path}")
|
||||
if left_img.shape != right_img.shape:
|
||||
raise ValueError(
|
||||
f"Shape mismatch for ts pair: {os.path.basename(lc_path)} "
|
||||
f"{left_img.shape} vs {os.path.basename(rc_path)} {right_img.shape}"
|
||||
)
|
||||
left_images.append(left_img)
|
||||
right_images.append(right_img)
|
||||
|
||||
return np.stack(left_images), np.stack(right_images)
|
||||
|
||||
def is_gpu_available():
|
||||
try:
|
||||
output = subprocess.check_output("nvidia-smi", stderr=subprocess.STDOUT, shell=True)
|
||||
return True
|
||||
except subprocess.CalledProcessError:
|
||||
return False
|
||||
|
||||
if is_gpu_available():
|
||||
print("GPU is available!")
|
||||
import cupy as cp
|
||||
############# GPU Code Starts #############
|
||||
|
||||
cuda_kernel = cp.RawKernel(r'''
|
||||
extern "C" __global__
|
||||
void disparity_kernel(
|
||||
const float* left_stack, const float* right_stack, float* disp_matrix, int* vert_disp_matrix, float* corr_matrix,
|
||||
int num_pairs, int height, int width, int window_size, int negative_range, int positive_range,
|
||||
int v_neg_range, int v_pos_range, float zncc_threshold, bool interpolation, int interpolation_method
|
||||
) {
|
||||
int row = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
int col1 = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
int half_window = window_size / 2;
|
||||
int patch_size = window_size * window_size;
|
||||
|
||||
// Check if the current pixel is within valid bounds
|
||||
if (row < half_window || row >= height - half_window || col1 < half_window || col1 >= width - half_window) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Local array to store correlation values for each disparity
|
||||
const int max_disparity = positive_range - negative_range + 1; // Total disparities: 54 - (-160) + 1 = 215
|
||||
float corrs[4000]; // Adjust size to match max_disparity
|
||||
int shifts[4000]; // Store vertical shifts for each disparity
|
||||
|
||||
// Initialize correlation values to -1
|
||||
for (int i = 0; i < max_disparity; i++) {
|
||||
corrs[i] = -1.0f;
|
||||
shifts[i] = 0;
|
||||
}
|
||||
|
||||
// Iterate over possible horizontal disparities
|
||||
for (int col2 = max(half_window, col1 - positive_range); col2 <= min(width - half_window - 1, col1 - negative_range); col2++) {
|
||||
int disparity = col1 - col2;
|
||||
|
||||
// Check if disparity is within the valid range [-160, 54]
|
||||
if (disparity < negative_range || disparity > positive_range) {
|
||||
continue; // Skip invalid disparities
|
||||
}
|
||||
|
||||
// Map disparity to index in corrs array
|
||||
int disparity_idx = disparity - negative_range;
|
||||
|
||||
// Iterate over possible vertical disparities
|
||||
for (int v_shift = v_neg_range; v_shift <= v_pos_range; v_shift++) {
|
||||
int row_shifted = row + v_shift;
|
||||
if (row_shifted < half_window || row_shifted >= height - half_window) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Compute mean of the patches in the left and right images
|
||||
float mean1 = 0.0f, mean2 = 0.0f;
|
||||
for (int pair = 0; pair < num_pairs; pair++) {
|
||||
for (int i = 0; i < patch_size; i++) {
|
||||
int patch_row = i / window_size;
|
||||
int patch_col = i % window_size;
|
||||
mean1 += left_stack[(pair * height + row - half_window + patch_row) * width + (col1 - half_window + patch_col)];
|
||||
mean2 += right_stack[(pair * height + row_shifted - half_window + patch_row) * width + (col2 - half_window + patch_col)];
|
||||
}
|
||||
}
|
||||
mean1 /= (num_pairs * patch_size);
|
||||
mean2 /= (num_pairs * patch_size);
|
||||
|
||||
// Compute standard deviation of the patches
|
||||
float std1 = 0.0f, std2 = 0.0f;
|
||||
for (int pair = 0; pair < num_pairs; pair++) {
|
||||
for (int i = 0; i < patch_size; i++) {
|
||||
int patch_row = i / window_size;
|
||||
int patch_col = i % window_size;
|
||||
float val1 = left_stack[(pair * height + row - half_window + patch_row) * width + (col1 - half_window + patch_col)] - mean1;
|
||||
float val2 = right_stack[(pair * height + row_shifted - half_window + patch_row) * width + (col2 - half_window + patch_col)] - mean2;
|
||||
std1 += val1 * val1;
|
||||
std2 += val2 * val2;
|
||||
}
|
||||
}
|
||||
std1 = sqrt(std1 / (num_pairs * patch_size));
|
||||
std2 = sqrt(std2 / (num_pairs * patch_size));
|
||||
|
||||
if (std1 == 0 || std2 == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Compute ZNCC (Zero-Normalized Cross-Correlation)
|
||||
float sum = 0.0f;
|
||||
for (int pair = 0; pair < num_pairs; pair++) {
|
||||
for (int i = 0; i < patch_size; i++) {
|
||||
int patch_row = i / window_size;
|
||||
int patch_col = i % window_size;
|
||||
float val1 = left_stack[(pair * height + row - half_window + patch_row) * width + (col1 - half_window + patch_col)] - mean1;
|
||||
float val2 = right_stack[(pair * height + row_shifted - half_window + patch_row) * width + (col2 - half_window + patch_col)] - mean2;
|
||||
sum += val1 * val2;
|
||||
}
|
||||
}
|
||||
float rho = sum / (std1 * std2 * num_pairs * patch_size);
|
||||
|
||||
// Store the maximum correlation for this disparity
|
||||
if (rho > corrs[disparity_idx]) {
|
||||
corrs[disparity_idx] = rho;
|
||||
shifts[disparity_idx] = v_shift;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Find the disparity with the maximum correlation
|
||||
int max_corr_idx = 0;
|
||||
for (int i = 1; i < max_disparity; i++) {
|
||||
if (corrs[i] > corrs[max_corr_idx]) {
|
||||
max_corr_idx = i;
|
||||
}
|
||||
}
|
||||
|
||||
// Perform interpolation if enabled
|
||||
if (corrs[max_corr_idx] > zncc_threshold) {
|
||||
float d_subpixel = (float)(max_corr_idx + negative_range);
|
||||
|
||||
if (interpolation && max_corr_idx > 0 && max_corr_idx < max_disparity - 1) {
|
||||
float corr_left = corrs[max_corr_idx - 1];
|
||||
float corr_best = corrs[max_corr_idx];
|
||||
float corr_right = corrs[max_corr_idx + 1];
|
||||
|
||||
if (interpolation_method == 1) {
|
||||
// Parabolic interpolation
|
||||
float denom = (corr_left - 2 * corr_best + corr_right);
|
||||
if (denom != 0.0f) {
|
||||
d_subpixel += (corr_left - corr_right) / (2 * denom);
|
||||
}
|
||||
} else if (interpolation_method == 2) {
|
||||
// Gaussian interpolation
|
||||
float log_minus_1 = logf(max(corr_left, 1e-10f));
|
||||
float log_0 = logf(max(corr_best, 1e-10f));
|
||||
float log_plus_1 = logf(max(corr_right, 1e-10f));
|
||||
|
||||
float a = (log_plus_1 - 2 * log_0 + log_minus_1) / 2.0f;
|
||||
float b = (log_plus_1 - log_minus_1) / 2.0f;
|
||||
|
||||
if (a != 0.0f) {
|
||||
d_subpixel += -b / (2 * a);
|
||||
}
|
||||
} else if (interpolation_method == 3) {
|
||||
// Angular interpolation
|
||||
float angle_minus_1 = atan2f(corr_left, 1.0f);
|
||||
float angle_0 = atan2f(corr_best, 1.0f);
|
||||
float angle_plus_1 = atan2f(corr_right, 1.0f);
|
||||
|
||||
float mean_angle = (angle_plus_1 - angle_minus_1) / 2.0f;
|
||||
d_subpixel -= tanf(mean_angle);
|
||||
}
|
||||
}
|
||||
|
||||
// Store the results
|
||||
int output_row = row - half_window;
|
||||
int output_col = col1 - half_window;
|
||||
disp_matrix[output_row * (width - window_size) + output_col] = d_subpixel;
|
||||
vert_disp_matrix[output_row * (width - window_size) + output_col] = shifts[max_corr_idx];
|
||||
corr_matrix[output_row * (width - window_size) + output_col] = corrs[max_corr_idx];
|
||||
}
|
||||
}
|
||||
''', 'disparity_kernel')
|
||||
|
||||
def compute_disparity_map_with_subpixel_GPU(
|
||||
left_stack, right_stack, window_size, negative_range, positive_range, v_neg_range, v_pos_range, zncc_threshold,
|
||||
interpolation, method
|
||||
):
|
||||
# Map method strings to numeric values
|
||||
methods = {
|
||||
"parabolic": 1,
|
||||
"gaussian": 2,
|
||||
"equiangular": 3
|
||||
}
|
||||
|
||||
# Validate method
|
||||
if method not in methods:
|
||||
raise ValueError(f"Invalid method: {method}. Must be one of {list(methods.keys())}")
|
||||
|
||||
# Get the numeric value for the method
|
||||
method_code = methods[method]
|
||||
|
||||
num_pairs, height, width = left_stack.shape
|
||||
disp_matrix = cp.full((height - window_size, width - window_size), negative_range, dtype=cp.float32)
|
||||
vert_disp_matrix = cp.zeros((height - window_size, width - window_size), dtype=cp.int32)
|
||||
corr_matrix = cp.zeros((height - window_size, width - window_size), dtype=cp.float32)
|
||||
|
||||
# Transfer data to GPU
|
||||
left_stack_gpu = cp.asarray(left_stack, dtype=cp.float32)
|
||||
right_stack_gpu = cp.asarray(right_stack, dtype=cp.float32)
|
||||
|
||||
# Define block and grid sizes
|
||||
block_size = (16, 16)
|
||||
grid_size = (
|
||||
(width - window_size + block_size[0] - 1) // block_size[0],
|
||||
(height - window_size + block_size[1] - 1) // block_size[1]
|
||||
)
|
||||
|
||||
# Launch the CUDA kernel
|
||||
cuda_kernel(
|
||||
grid_size, block_size,
|
||||
(left_stack_gpu, right_stack_gpu, disp_matrix, vert_disp_matrix, corr_matrix,
|
||||
num_pairs, height, width, window_size, negative_range, positive_range,
|
||||
v_neg_range, v_pos_range, zncc_threshold, interpolation, method_code)
|
||||
)
|
||||
|
||||
disp_matrix_np = disp_matrix.get()
|
||||
vert_disp_matrix_np = vert_disp_matrix.get()
|
||||
corr_matrix_np = corr_matrix.get()
|
||||
|
||||
return disp_matrix_np, vert_disp_matrix_np, corr_matrix_np
|
||||
|
||||
############# GPU Code Ends #############
|
||||
|
||||
|
||||
|
||||
############# CPU Code Starts #############
|
||||
|
||||
# Subpixel interpolation methods
|
||||
@njit
|
||||
def parabolic_interpolation(corrs, idx):
|
||||
zncc_d_minus_1 = corrs[idx - 1]
|
||||
zncc_d = corrs[idx]
|
||||
zncc_d_plus_1 = corrs[idx + 1]
|
||||
|
||||
a = (zncc_d_plus_1 - 2 * zncc_d + zncc_d_minus_1) / 2.0
|
||||
b = (zncc_d_plus_1 - zncc_d_minus_1) / 2.0
|
||||
|
||||
if a != 0.0:
|
||||
return -b / (2 * a)
|
||||
return 0.0
|
||||
|
||||
@njit
|
||||
def gaussian_interpolation(corrs, idx):
|
||||
zncc_d_minus_1 = corrs[idx - 1]
|
||||
zncc_d = corrs[idx]
|
||||
zncc_d_plus_1 = corrs[idx + 1]
|
||||
|
||||
log_minus_1 = np.log(max(zncc_d_minus_1, 1e-10))
|
||||
log_0 = np.log(max(zncc_d, 1e-10))
|
||||
log_plus_1 = np.log(max(zncc_d_plus_1, 1e-10))
|
||||
|
||||
a = (log_plus_1 - 2 * log_0 + log_minus_1) / 2.0
|
||||
b = (log_plus_1 - log_minus_1) / 2.0
|
||||
|
||||
if a != 0.0:
|
||||
return -b / (2 * a)
|
||||
return 0.0
|
||||
|
||||
@njit
|
||||
def equiangular_interpolation(corrs, idx):
|
||||
zncc_d_minus_1 = corrs[idx - 1]
|
||||
zncc_d = corrs[idx]
|
||||
zncc_d_plus_1 = corrs[idx + 1]
|
||||
|
||||
# Convert to angles (avoid division by zero)
|
||||
angle_minus_1 = np.arctan2(zncc_d_minus_1, 1.0)
|
||||
angle_0 = np.arctan2(zncc_d, 1.0)
|
||||
angle_plus_1 = np.arctan2(zncc_d_plus_1, 1.0)
|
||||
|
||||
mean_angle = (angle_plus_1 - angle_minus_1) / 2.0
|
||||
|
||||
return np.tan(mean_angle)
|
||||
@njit
|
||||
def subpixel_interpolation(corrs, max_corr_idx, method):
|
||||
if max_corr_idx <= 0 or max_corr_idx >= len(corrs) - 1:
|
||||
return 0.0 # Boundary condition: cannot interpolate
|
||||
|
||||
if method == "parabolic":
|
||||
return parabolic_interpolation(corrs, max_corr_idx)
|
||||
elif method == "gaussian":
|
||||
return gaussian_interpolation(corrs, max_corr_idx)
|
||||
elif method == "equiangular":
|
||||
return equiangular_interpolation(corrs, max_corr_idx)
|
||||
else:
|
||||
raise ValueError(f"Unknown interpolation method: {method}")
|
||||
|
||||
@njit
|
||||
def zncc_patch(patch1, patch2):
|
||||
mean1 = np.mean(patch1)
|
||||
mean2 = np.mean(patch2)
|
||||
std1 = np.std(patch1)
|
||||
std2 = np.std(patch2)
|
||||
|
||||
if std1 == 0 or std2 == 0:
|
||||
return -1.0
|
||||
|
||||
return np.sum((patch1 - mean1) * (patch2 - mean2)) / (std1 * std2 * patch1.size)
|
||||
|
||||
@njit(parallel=True)
|
||||
def compute_disparity_map_with_subpixel_CPU(
|
||||
left_stack, right_stack, window_size, negative_range, positive_range, v_neg_range, v_pos_range, zncc_threshold, interpolation, method
|
||||
):
|
||||
num_pairs, height, width = left_stack.shape
|
||||
disp_matrix = np.full((height - window_size, width - window_size), negative_range, dtype=np.float32)
|
||||
vert_disp_matrix = np.zeros_like(disp_matrix, dtype=np.int32)
|
||||
corr_matrix = np.zeros_like(disp_matrix, dtype=np.float32)
|
||||
|
||||
half_window = window_size // 2
|
||||
|
||||
for row in prange(half_window, height - half_window):
|
||||
for col1 in range(half_window, width - half_window):
|
||||
win1_stack = left_stack[:, row - half_window : row + half_window + 1, col1 - half_window : col1 + half_window + 1]
|
||||
corrs = np.full((positive_range - negative_range), -1.0, dtype=np.float32)
|
||||
shifts = np.full((positive_range - negative_range, 2), 0, dtype=np.int32)
|
||||
|
||||
for idx, col2 in enumerate(range(max(half_window, col1 - positive_range), min(width - half_window, col1 - negative_range))):
|
||||
max_corr = -1.0
|
||||
best_disp_v = 0
|
||||
for v_shift in range(v_neg_range, v_pos_range + 1):
|
||||
row_shifted = row + v_shift
|
||||
if not (half_window <= row_shifted < height - half_window):
|
||||
continue
|
||||
|
||||
win2_stack = right_stack[:, row_shifted - half_window : row_shifted + half_window + 1,
|
||||
col2 - half_window : col2 + half_window + 1]
|
||||
|
||||
rho = zncc_patch(win1_stack.ravel(), win2_stack.ravel())
|
||||
|
||||
if rho > max_corr:
|
||||
max_corr = rho
|
||||
best_disp_v = v_shift
|
||||
corrs[idx] = max_corr
|
||||
shifts[idx, :] = np.array([col1 - col2, best_disp_v])
|
||||
|
||||
max_corr_idx = np.argmax(corrs)
|
||||
max_corr_value = corrs[max_corr_idx]
|
||||
|
||||
if max_corr_value > zncc_threshold:
|
||||
disparity_h, disparity_v = shifts[max_corr_idx]
|
||||
|
||||
if interpolation:
|
||||
subpixel_offset = subpixel_interpolation(corrs, max_corr_idx, method)
|
||||
disparity_h = disparity_h - subpixel_offset
|
||||
disp_matrix[row - half_window, col1 - half_window] = disparity_h
|
||||
vert_disp_matrix[row - half_window, col1 - half_window] = disparity_v
|
||||
corr_matrix[row - half_window, col1 - half_window] = max_corr_value
|
||||
|
||||
return disp_matrix, vert_disp_matrix, corr_matrix
|
||||
|
||||
############# CPU Code Ends #############
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Stereo Disparity Map Generator")
|
||||
parser.add_argument("--window_size", type=int, default=5, help="Window size for block matching")
|
||||
parser.add_argument("--H_neg_range", type=int, default=0, help="Horizontal negative range for disparity")
|
||||
parser.add_argument("--H_pos_range", type=int, default=65, help="Horizontal positive range for disparity")
|
||||
parser.add_argument("--v_neg_range", type=int, default=-4, help="Vertical negative range")
|
||||
parser.add_argument("--v_pos_range", type=int, default=2, help="Vertical positive range")
|
||||
parser.add_argument("--zncc_threshold", type=float, default=0.4, help="ZNCC threshold")
|
||||
parser.add_argument(
|
||||
"--num_images",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Number of stereo pairs to use. Default: all matched pairs in the folder. "
|
||||
"If fewer than available, the last N pairs (highest timestamps) are used.",
|
||||
)
|
||||
parser.add_argument("--noise_filter", type=str, default="open3d", choices=["median", "open3d", None], help="Post-processing noise filter method: 'median' or 'open3d'")
|
||||
parser.add_argument("--interpolation", type=bool, default=True, help="Use sub-pixel interpolation")
|
||||
parser.add_argument("--method", type=str, default="parabolic", help="Interpolation method (parabolic, gaussian, equiangular)")
|
||||
parser.add_argument("--q_file", type=str, default="./data1/Q.cvstore", help="Path to Q matrix file")
|
||||
parser.add_argument("--left_dir", type=str, default="./data1/left", help="Path to left images")
|
||||
parser.add_argument("--right_dir", type=str, default="./data1/right", help="Path to right images")
|
||||
parser.add_argument("--output_dir", type=str, default="./results", help="Output directory (fallback when disp/pcl dirs not set)")
|
||||
parser.add_argument("--noise_remove", type=bool, default=True, help="Enable noise removal using left-right consistency check")
|
||||
parser.add_argument("--left_prefix", type=str, default="", help="Filename prefix for left images (e.g. 'lc_') when both cameras share a folder")
|
||||
parser.add_argument("--right_prefix", type=str, default="", help="Filename prefix for right images (e.g. 'rc_') when both cameras share a folder")
|
||||
parser.add_argument("--disp_output_dir", type=str, default=None, help="Output dir for disparity/map files; falls back to --output_dir")
|
||||
parser.add_argument("--pcl_output_dir", type=str, default=None, help="Output dir for point cloud files; falls back to --output_dir")
|
||||
parser.add_argument("--troubleshooting", action="store_true", help="Save all outputs (vertical shift, correlation maps, point cloud); by default only disparity map is saved")
|
||||
args = parser.parse_args()
|
||||
|
||||
start_time = time.time()
|
||||
disp_out = args.disp_output_dir or args.output_dir
|
||||
pcl_out = args.pcl_output_dir or args.output_dir
|
||||
os.makedirs(disp_out, exist_ok=True)
|
||||
os.makedirs(pcl_out, exist_ok=True)
|
||||
|
||||
left_stack, right_stack = load_stereo_stacks(
|
||||
args.left_dir, args.right_dir, args.num_images,
|
||||
args.left_prefix, args.right_prefix
|
||||
)
|
||||
left_shap = left_stack.shape
|
||||
print(f"{left_shap[0]} image pairs loaded and their shapes: {left_shap}, {right_stack.shape}")
|
||||
Q = load_q_matrix(args.q_file)
|
||||
print("Q matrix loaded successfully.")
|
||||
if is_gpu_available():
|
||||
print("Running on GPU...")
|
||||
disp_matrix_LR, vert_disp_matrix, corr_matrix = compute_disparity_map_with_subpixel_GPU(
|
||||
left_stack, right_stack, args.window_size, args.H_neg_range, args.H_pos_range, args.v_neg_range, args.v_pos_range, np.float32(args.zncc_threshold), args.interpolation, args.method
|
||||
)
|
||||
if args.noise_remove:
|
||||
disp_matrix_RL, _, _ = compute_disparity_map_with_subpixel_GPU(
|
||||
right_stack, left_stack, args.window_size, -args.H_pos_range, -args.H_neg_range, -args.v_pos_range, -args.v_neg_range, np.float32(args.zncc_threshold), args.interpolation, args.method
|
||||
)
|
||||
disp_matrix = process_disparity(disp_matrix_LR, disp_matrix_RL)
|
||||
else:
|
||||
disp_matrix = disp_matrix_LR
|
||||
|
||||
else:
|
||||
print("Running on CPU...")
|
||||
disp_matrix_LR, vert_disp_matrix, corr_matrix = compute_disparity_map_with_subpixel_CPU(
|
||||
left_stack, right_stack, args.window_size, args.H_neg_range, args.H_pos_range, args.v_neg_range, args.v_pos_range, args.zncc_threshold, args.interpolation, args.method
|
||||
)
|
||||
if args.noise_remove:
|
||||
disp_matrix_RL, _, _ = compute_disparity_map_with_subpixel_CPU(
|
||||
right_stack, left_stack, args.window_size, -args.H_pos_range, -args.H_neg_range, -args.v_pos_range, -args.v_neg_range, args.zncc_threshold, args.interpolation, args.method
|
||||
)
|
||||
disp_matrix = process_disparity(disp_matrix_LR, disp_matrix_RL)
|
||||
else:
|
||||
disp_matrix = disp_matrix_LR
|
||||
|
||||
# Always save disparity map
|
||||
np.save(f"{disp_out}/disparity.npy", disp_matrix)
|
||||
save_with_colorbar(disp_matrix, f"{disp_out}/Disparity_map_colorbar.png")
|
||||
print(f"Disparity map saved to {disp_out}")
|
||||
|
||||
if args.troubleshooting:
|
||||
np.save(f"{disp_out}/vertical_shift_map.npy", vert_disp_matrix)
|
||||
np.save(f"{disp_out}/correlation_map.npy", corr_matrix)
|
||||
save_with_colorbar(vert_disp_matrix, f"{disp_out}/vertical_shift_map_colorbar.png")
|
||||
save_with_colorbar(corr_matrix, f"{disp_out}/correlation_map_colorbar.png")
|
||||
|
||||
if args.noise_filter == "median":
|
||||
disp_matrix = median_filter(disp_matrix, size=15)
|
||||
|
||||
disp_map_loaded = disp_matrix.astype(np.float32)
|
||||
point_cloud = cv2.reprojectImageTo3D(disp_map_loaded, Q)
|
||||
mask_map = disp_map_loaded > disp_map_loaded.min()
|
||||
point_cloud = point_cloud[mask_map]
|
||||
|
||||
depth_values = point_cloud[:, 2]
|
||||
depth_normalized = (depth_values - depth_values.min()) / (depth_values.max() - depth_values.min())
|
||||
colormap = plt.get_cmap('jet')
|
||||
colors = colormap(depth_normalized)[:, :3]
|
||||
|
||||
if args.noise_filter == "open3d":
|
||||
pcd = o3d.geometry.PointCloud()
|
||||
pcd.points = o3d.utility.Vector3dVector(point_cloud)
|
||||
pcd.colors = o3d.utility.Vector3dVector(colors)
|
||||
pcd_clean, ind = pcd.remove_statistical_outlier(nb_neighbors=50, std_ratio=0.004)
|
||||
point_cloud = np.asarray(pcd_clean.points)
|
||||
colors = np.asarray(pcd_clean.colors)
|
||||
|
||||
point_cloud_with_colors = np.hstack((point_cloud, colors))
|
||||
output_file_ply = f"{pcl_out}/Point_cloud.ply"
|
||||
output_file_txt = f"{pcl_out}/Point_cloud.txt"
|
||||
save_ply(output_file_ply, point_cloud_with_colors)
|
||||
save_ascii_point_cloud(output_file_txt, point_cloud_with_colors)
|
||||
print(f"Vertical shift, correlation maps and point cloud saved (troubleshooting mode)")
|
||||
print(f"Point cloud saved to {pcl_out}")
|
||||
|
||||
print(f"Execution time: {time.time() - start_time:.2f} seconds")
|
||||
Reference in New Issue
Block a user