179 lines
6.8 KiB
Python
179 lines
6.8 KiB
Python
"""
|
|
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()
|