rectify.py

import numpy as np
import cv2
import os
import glob
import argparse
import sys

# Define default paths
cam0_path = '/run/user/1000/gvfs/smb-share:server=unas-pro.local,share=dwe_nas/EuRoC/seq1/cam0'
cam1_path = '/run/user/1000/gvfs/smb-share:server=unas-pro.local,share=dwe_nas/EuRoC/seq1/cam1'
cam0_output = '/run/user/1000/gvfs/smb-share:server=unas-pro.local,share=dwe_nas/EuRoC/seq1/cam0_rec'
cam1_output = '/run/user/1000/gvfs/smb-share:server=unas-pro.local,share=dwe_nas/EuRoC/seq1/cam1_rec'
disp_output = '/run/user/1000/gvfs/smb-share:server=unas-pro.local,share=dwe_nas/EuRoC/seq1/disparity'
depth_output = '/run/user/1000/gvfs/smb-share:server=unas-pro.local,share=dwe_nas/EuRoC/seq1/depth_npy'

def main():
    # --- 0. PARSE ARGUMENTS ---
    parser = argparse.ArgumentParser(description="Rectify EuRoC stereo images, generate disparity, and save depth npy.")
    
    # Inputs
    parser.add_argument("--input_cam0", type=str, default=cam0_path, help="Path to cam0 (left) source.")
    parser.add_argument("--input_cam1", type=str, default=cam1_path, help="Path to cam1 (right) source.")
    
    # Outputs
    parser.add_argument("--output_cam0", type=str, default=cam0_output, help="Path to save rectified left.")
    parser.add_argument("--output_cam1", type=str, default=cam1_output, help="Path to save rectified right.")
    parser.add_argument("--output_disp", type=str, default=disp_output, help="Path to save color disparity maps.")
    parser.add_argument("--output_depth", type=str, default=depth_output, help="Path to save raw depth npy files.")
    
    parser.add_argument("--extension", type=str, default="png", help="Image extension (png, jpg).")
    
    args = parser.parse_args()

    if not os.path.isdir(args.input_cam0) or not os.path.isdir(args.input_cam1):
        print("Error: Input folders not found.")
        sys.exit(1)

    # --- 1. SETUP PARAMETERS (EuRoC) ---
    K0 = np.array([[458.654, 0.0, 367.215], [0.0, 457.296, 248.375], [0.0, 0.0, 1.0]])
    D0 = np.array([-0.28340811, 0.07395907, 0.00019359, 1.76187114e-05])
    T_IC0 = np.array([
        [0.0148655429818, -0.999880929698, 0.00414029679422, -0.0216401454975],
        [0.999557249008, 0.0149672133247, 0.025715529948, -0.064676986768],
        [-0.0257744366974, 0.00375618835797, 0.999660727178, 0.00981073058949],
        [0.0, 0.0, 0.0, 1.0]
    ])

    K1 = np.array([[457.587, 0.0, 379.999], [0.0, 456.134, 255.238], [0.0, 0.0, 1.0]])
    D1 = np.array([-0.28368365, 0.07451284, -0.00010473, -3.55590700e-05])
    T_IC1 = np.array([
        [0.0125552670891, -0.999755099723, 0.0182237714554, -0.0198435579556],
        [0.999598781151, 0.0130119051815, 0.0251588363115, 0.0453689425024],
        [-0.0253898008918, 0.0179005838253, 0.999517347078, 0.00786212447038],
        [0.0, 0.0, 0.0, 1.0]
    ])

    image_size = (752, 480)

    # --- 2. PRE-COMPUTE MAPS ---
    print("Computing rectification maps...")
    T_C1_C0 = np.linalg.inv(T_IC1) @ T_IC0
    R = T_C1_C0[0:3, 0:3]
    T = T_C1_C0[0:3, 3]

    # alpha=0: crop black borders
    # Q matrix is essential for Disparity -> Depth conversion
    R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(
        K0, D0, K1, D1, image_size, R, T, flags=cv2.CALIB_ZERO_DISPARITY, alpha=0
    )

    map0_x, map0_y = cv2.initUndistortRectifyMap(K0, D0, R1, P1, image_size, cv2.CV_32FC1)
    map1_x, map1_y = cv2.initUndistortRectifyMap(K1, D1, R2, P2, image_size, cv2.CV_32FC1)

    # --- 3. SETUP STEREO MATCHER (SGBM) ---
    min_disp = 0
    num_disp = 80 
    block_size = 11
    
    stereo = cv2.StereoSGBM_create(
        minDisparity=min_disp,
        numDisparities=num_disp,
        blockSize=block_size,
        P1=8 * 1 * block_size**2,
        P2=32 * 1 * block_size**2,
        disp12MaxDiff=1,
        uniquenessRatio=10,
        speckleWindowSize=100,
        speckleRange=32
    )

    # --- 4. PROCESSING LOOP ---
    process_folder(
        args.input_cam0, args.input_cam1,
        args.output_cam0, args.output_cam1, 
        args.output_disp, args.output_depth,
        map0_x, map0_y, map1_x, map1_y, 
        stereo, num_disp, Q, args.extension
    )


def process_folder(in_cam0, in_cam1, out_cam0, out_cam1, out_disp, out_depth,
                   map0_x, map0_y, map1_x, map1_y, stereo, num_disp, Q, ext):
    
    # Create directories
    os.makedirs(out_cam0, exist_ok=True)
    os.makedirs(out_cam1, exist_ok=True)
    os.makedirs(out_disp, exist_ok=True)
    os.makedirs(out_depth, exist_ok=True)

    pattern = f"*.{ext}"
    images0 = sorted(glob.glob(os.path.join(in_cam0, pattern)))
    images1 = sorted(glob.glob(os.path.join(in_cam1, pattern)))

    if not images0:
        print("No images found.")
        return

    count = min(len(images0), len(images1))
    print(f"Processing {count} image pairs...")
    print(f"Saving depth .npy to {out_depth}")

    for i, (img0_file, img1_file) in enumerate(zip(images0, images1)):
        filename = os.path.basename(img0_file)
        # Change extension for npy file (e.g. image.png -> image.npy)
        filename_npy = os.path.splitext(filename)[0] + '.npy'
        
        img0 = cv2.imread(img0_file, cv2.IMREAD_GRAYSCALE)
        img1 = cv2.imread(img1_file, cv2.IMREAD_GRAYSCALE)

        if img0 is None or img1 is None:
            continue

        # 1. Rectify
        rect0 = cv2.remap(img0, map0_x, map0_y, cv2.INTER_LINEAR)
        rect1 = cv2.remap(img1, map1_x, map1_y, cv2.INTER_LINEAR)

        # 2. Compute Disparity
        disp_16s = stereo.compute(rect0, rect1)
        
        # Convert to float (pixels)
        disp_float = disp_16s.astype(np.float32) / 16.0
        
        # 3. Compute Depth Map
        # cv2.reprojectImageTo3D uses Q to map (x, y, disparity) -> (X, Y, Z)
        # We only need Z.
        points_3d = cv2.reprojectImageTo3D(disp_float, Q)
        depth_map = points_3d[:, :, 2] # Extract Z channel
        
        # Handle invalid values:
        # SGBM produces negative/zero disparity for invalid matches.
        # Reprojecting 0 disparity results in infinite depth.
        # Mask out low disparity values to keep depth finite and clean.
        min_valid_disp = 1.0 # Minimum disparity to consider valid
        mask = disp_float < min_valid_disp
        depth_map[mask] = 0.0 # Set invalid depth to 0.0 (standard convention)

        # 4. Color Map for Visualization (Optional: using same disp_float)
        disp_norm = (disp_float / num_disp) * 255.0
        disp_norm = np.clip(disp_norm, 0, 255).astype(np.uint8)
        disp_color = cv2.applyColorMap(disp_norm, cv2.COLORMAP_JET)

        # 5. Save all
        cv2.imwrite(os.path.join(out_cam0, filename), rect0)
        cv2.imwrite(os.path.join(out_cam1, filename), rect1)
        cv2.imwrite(os.path.join(out_disp, filename), disp_color)
        np.save(os.path.join(out_depth, filename_npy), depth_map)
        
        if i % 50 == 0:
            print(f"Processed {i}/{count}...")
            
    print(f"Done! All outputs saved.")

if __name__ == "__main__":
    main()