202 lines
7.2 KiB
Python
202 lines
7.2 KiB
Python
"""
|
|
Pipeline runner for libSGM stereo disparity.
|
|
|
|
Resolves all paths from the project folder structure and drives
|
|
the stereosgm_new binary for each scan in a session (or all sessions on a date).
|
|
|
|
For each scan it takes the LAST matched lc_/rc_ image pair from 02_rect_images/
|
|
(images sorted by timestamp — highest timestamp = last acquired image).
|
|
|
|
Output layout per scan:
|
|
<processing_dir>/<project>/<date>/<session>/<scan>/
|
|
02_rect_images/ <- input (lc_ts<last>.png + rc_ts<same>.png)
|
|
03_sgm_disp_map/ <- disparity.xml + disparity_color.png (created here)
|
|
05_sgm_pcl/ <- untouched
|
|
|
|
Binary:
|
|
~/Speckle-Scanner/05_disparity/libsgm/build/sample/stereosgm_new
|
|
"""
|
|
|
|
import sys
|
|
import re
|
|
import argparse
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
# Resolve config.py from ~/Speckle-Scanner regardless of CWD
|
|
sys.path.insert(0, str(Path.home() / "Speckle-Scanner"))
|
|
import config # noqa: E402
|
|
|
|
BINARY = Path(__file__).parent / "build" / "sample" / "stereosgm_new"
|
|
|
|
|
|
def extract_ts_token(filename, prefix="lc_"):
|
|
"""Extract ts token from lc_ts1634840093.png or lc_ts1634840093_ck....png."""
|
|
m = re.search(rf"^{re.escape(prefix)}(ts\d+)", filename, 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 find_rc_for_ts(rect_dir, ts_token):
|
|
"""Match rc image by shared ts token (ck suffix optional)."""
|
|
rc_matches = sorted(rect_dir.glob(f"rc_{ts_token}_*.png"))
|
|
if not rc_matches:
|
|
rc_matches = sorted(rect_dir.glob(f"rc_{ts_token}*.png"))
|
|
return rc_matches[0] if rc_matches else None
|
|
|
|
|
|
def find_last_lc_rc_pair(rect_dir):
|
|
"""Return (lc_path, rc_path) for the highest-timestamp matched pair in rect_dir."""
|
|
rect_dir = Path(rect_dir)
|
|
pairs = []
|
|
for lc in rect_dir.glob("lc_ts*.png"):
|
|
ts_token, ts_int = extract_ts_token(lc.name, "lc_")
|
|
if ts_token is None:
|
|
continue
|
|
rc = find_rc_for_ts(rect_dir, ts_token)
|
|
if rc is None:
|
|
continue
|
|
pairs.append((ts_int, lc, rc))
|
|
|
|
if not pairs:
|
|
return None, None
|
|
|
|
pairs.sort(key=lambda item: item[0])
|
|
_, lc, rc = pairs[-1]
|
|
return lc, rc
|
|
|
|
|
|
def build_cmd(lc, rc, output_dir, sgm_args):
|
|
cmd = [
|
|
str(BINARY),
|
|
str(lc),
|
|
str(rc),
|
|
f"--output_dir={output_dir}",
|
|
"--no_display=1",
|
|
]
|
|
for key, val in sgm_args.items():
|
|
if val is not None:
|
|
cmd.append(f"--{key}={val}")
|
|
return cmd
|
|
|
|
|
|
def run_scan(project, date, session, scan, sgm_args):
|
|
rect_dir = config.PROCESSING_DIR / project / date / session / scan / "02_rect_images"
|
|
if not rect_dir.exists():
|
|
print(f"[SKIP] {session}/{scan}: 02_rect_images not found at {rect_dir}")
|
|
return False
|
|
|
|
lc, rc = find_last_lc_rc_pair(rect_dir)
|
|
if lc is None:
|
|
print(f"[SKIP] {session}/{scan}: no lc_ts*.png images found in {rect_dir}")
|
|
return False
|
|
if rc is None:
|
|
print(f"[SKIP] {session}/{scan}: no matching rc image for {lc.name}")
|
|
return False
|
|
|
|
output_dir = config.get_processing_step_dir(project, date, session, scan, "03_sgm_disp_map")
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"[SCAN] {session}/{scan}")
|
|
print(f" lc : {lc.name}")
|
|
print(f" rc : {rc.name}")
|
|
print(f" output : {output_dir}")
|
|
print(f"{'='*60}")
|
|
|
|
cmd = build_cmd(lc, rc, output_dir, sgm_args)
|
|
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, sgm_args):
|
|
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, sgm_args)
|
|
if not ok:
|
|
failed.append(f"{session}/{scan}")
|
|
return scans, failed
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="libSGM 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 (e.g. Scan000001); omit to process all scans in the session")
|
|
|
|
# SGM parameters — all optional, forwarded to stereosgm_new
|
|
parser.add_argument("--disp_size", type=int, default=None, help="Maximum disparity value (64, 128, or 256; default 256)")
|
|
parser.add_argument("--P1", type=int, default=None, help="SGM penalty for disparity change of ±1 (default 10)")
|
|
parser.add_argument("--P2", type=int, default=None, help="SGM penalty for disparity change >1 (default 120)")
|
|
parser.add_argument("--uniqueness", type=float, default=None, help="Uniqueness ratio threshold (default 0.80)")
|
|
parser.add_argument("--num_paths", type=int, default=None, choices=[4, 8], help="Scanlines for cost aggregation: 4 or 8 (default 8)")
|
|
parser.add_argument("--min_disp", type=int, default=None, help="Minimum disparity value (default -160)")
|
|
parser.add_argument("--LR_max_diff", type=int, default=None, help="Max left-right disparity difference (default 1)")
|
|
parser.add_argument("--census_type", type=int, default=None, choices=[0, 1], help="Census transform type: 0=CENSUS_9x7, 1=SYMMETRIC_CENSUS_9x7 (default 1)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not BINARY.exists():
|
|
print(f"ERROR: stereosgm_new binary not found at {BINARY}")
|
|
print("Build it first: cd ~/Speckle-Scanner/05_disparity/libsgm/build && make stereosgm_new")
|
|
sys.exit(1)
|
|
|
|
sgm_args = {
|
|
"disp_size": args.disp_size,
|
|
"P1": args.P1,
|
|
"P2": args.P2,
|
|
"uniqueness": args.uniqueness,
|
|
"num_paths": args.num_paths,
|
|
"min_disp": args.min_disp,
|
|
"LR_max_diff": args.LR_max_diff,
|
|
"census_type": args.census_type,
|
|
}
|
|
|
|
# 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, sgm_args
|
|
)
|
|
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()
|