Initial commit: Speckle-Scanner 3D pipeline with setup README
This commit is contained in:
@@ -0,0 +1,173 @@
|
||||
import argparse
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
# Default paths: repo layout has input/ and rectified/ next to this package folder.
|
||||
_DEFAULT_DISPARITY = os.path.normpath(
|
||||
os.path.join(_SCRIPT_DIR, "..", "input", "disparity.xml")
|
||||
)
|
||||
_DEFAULT_Q = os.path.normpath(os.path.join(_SCRIPT_DIR, "..", "rectified", "Q.cvstore"))
|
||||
|
||||
|
||||
def load_disparity_map(disparity_file):
|
||||
"""Load disparity from .xml (OpenCV FileStorage), .png, or .npy."""
|
||||
disparity_file = os.fspath(disparity_file)
|
||||
if disparity_file.endswith(".xml"):
|
||||
fs = cv2.FileStorage(disparity_file, cv2.FILE_STORAGE_READ)
|
||||
disp_map = fs.getNode("disparity").mat().astype(np.float32)
|
||||
fs.release()
|
||||
elif disparity_file.endswith(".png"):
|
||||
disp_map = cv2.imread(disparity_file, cv2.IMREAD_UNCHANGED).astype(np.float32)
|
||||
elif disparity_file.endswith(".npy"):
|
||||
disp_map = np.load(disparity_file).astype(np.float32)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported disparity format: {disparity_file}. Use .xml, .png, or .npy."
|
||||
)
|
||||
return disp_map
|
||||
|
||||
|
||||
def load_q_matrix(Q):
|
||||
"""
|
||||
Accept a 4x4 ndarray or a file path:
|
||||
- .npy: array saved with numpy
|
||||
- OpenCV FileStorage (.cvstore, .xml, .yaml, .yml): matrix node named 'Q'
|
||||
"""
|
||||
if isinstance(Q, (str, os.PathLike)):
|
||||
path = os.fspath(Q)
|
||||
lower = path.lower()
|
||||
if lower.endswith(".npy"):
|
||||
Q = np.load(path)
|
||||
else:
|
||||
fs = cv2.FileStorage(path, cv2.FILE_STORAGE_READ)
|
||||
if not fs.isOpened():
|
||||
raise ValueError(f"Could not open Q file: {path}")
|
||||
q_node = fs.getNode("Q")
|
||||
if q_node.empty():
|
||||
fs.release()
|
||||
raise ValueError(
|
||||
f"No matrix 'Q' found in {path}. Expected OpenCV FileStorage node 'Q'."
|
||||
)
|
||||
Q = q_node.mat()
|
||||
fs.release()
|
||||
Q = np.asarray(Q, dtype=np.float64)
|
||||
if Q.shape != (4, 4):
|
||||
raise ValueError(f"Q must be 4x4, got shape {Q.shape}.")
|
||||
return Q
|
||||
|
||||
|
||||
def disparity_to_pointcloud(disparity_file, Q, valid_mask=None):
|
||||
"""
|
||||
Build a 3D point cloud from a disparity map and stereo rectification Q matrix.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
disparity_file : str or os.PathLike
|
||||
Path to disparity (.xml with 'disparity' node, .png, or .npy).
|
||||
Q : ndarray (4, 4) or str path
|
||||
4x4 reprojection matrix; .npy or OpenCV store (.cvstore / .xml / .yaml) with node 'Q'.
|
||||
valid_mask : ndarray of bool, optional
|
||||
If None, uses pixels where disparity > min(disparity) (same as original script).
|
||||
|
||||
Returns
|
||||
-------
|
||||
points_xyz : ndarray, shape (N, 3)
|
||||
3D points in rectified camera coordinates.
|
||||
"""
|
||||
disp_map = load_disparity_map(disparity_file)
|
||||
Qm = load_q_matrix(Q)
|
||||
point_cloud = cv2.reprojectImageTo3D(disp_map, Qm)
|
||||
if valid_mask is None:
|
||||
valid_mask = disp_map > disp_map.min()
|
||||
else:
|
||||
valid_mask = np.asarray(valid_mask, dtype=bool)
|
||||
points_xyz = point_cloud[valid_mask].reshape(-1, 3)
|
||||
return points_xyz
|
||||
|
||||
|
||||
def pointcloud_with_depth_colors(points_xyz):
|
||||
"""Attach jet colormap RGB in [0,1] from Z depth (same visualization as before)."""
|
||||
depth = points_xyz[:, 2]
|
||||
dmin, dmax = depth.min(), depth.max()
|
||||
if dmax > dmin:
|
||||
depth_norm = (depth - dmin) / (dmax - dmin)
|
||||
else:
|
||||
depth_norm = np.zeros_like(depth)
|
||||
colors = plt.get_cmap("jet")(depth_norm)[:, :3]
|
||||
return np.hstack((points_xyz, colors))
|
||||
|
||||
|
||||
def save_ply(filename, points):
|
||||
with open(filename, "w") as f:
|
||||
f.write("ply\nformat ascii 1.0\n")
|
||||
f.write(f"element vertex {points.shape[0]}\n")
|
||||
f.write("property float x\nproperty float y\nproperty float z\n")
|
||||
f.write("property uchar red\nproperty uchar green\nproperty uchar blue\nend_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 main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert disparity map + Q matrix to a point cloud (PLY/TXT)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--disparity",
|
||||
"-d",
|
||||
default=_DEFAULT_DISPARITY,
|
||||
help=f"Disparity file (.xml, .png, .npy). Default: {_DEFAULT_DISPARITY}",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--q",
|
||||
"-q",
|
||||
default=_DEFAULT_Q,
|
||||
help=f"Q matrix: .npy or OpenCV FileStorage (.cvstore, .xml, .yaml) with node 'Q'. Default: {_DEFAULT_Q}",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--out-dir",
|
||||
default="./output",
|
||||
help="Directory for PLY and TXT exports (default: ./output).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-save",
|
||||
action="store_true",
|
||||
help="Only print point count; do not write PLY/TXT.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--troubleshooting",
|
||||
action="store_true",
|
||||
help="Also save ASCII .txt point cloud (default: PLY only)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
points_xyz = disparity_to_pointcloud(args.disparity, args.q)
|
||||
print(f"Point cloud shape: {points_xyz.shape} (N x 3)")
|
||||
|
||||
if args.no_save:
|
||||
return
|
||||
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
base = os.path.splitext(os.path.basename(os.fspath(args.disparity)))[0]
|
||||
point_cloud_color = pointcloud_with_depth_colors(points_xyz)
|
||||
out_ply = os.path.join(args.out_dir, f"Pointcloud_{base}.ply")
|
||||
save_ply(out_ply, point_cloud_color)
|
||||
print(f"Saved: {out_ply}")
|
||||
if args.troubleshooting:
|
||||
out_txt = os.path.join(args.out_dir, f"Pointcloud_{base}.txt")
|
||||
save_ascii_point_cloud(out_txt, point_cloud_color)
|
||||
print(f"Saved: {out_txt}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user