convert_data.py

import os
import re
import cv2
import numpy as np
import argparse

# --- Parse command-line arguments ---
parser = argparse.ArgumentParser(description="Rename left images and stack depth images into depth.npy")
parser.add_argument('--directory', type=str, default="sbin_new", help='Directory containing left/depth images')
args = parser.parse_args()

directory = args.directory

# --- Regex Patterns ---
left_pattern = re.compile(r"left(\d+)\.png")
depth_pattern = re.compile(r"depth(\d+)\.png")

depth_frames = []
depth_files = []

# --- Process all files in directory ---
for filename in os.listdir(directory):
    # --- Rename leftXXXXXX.png to 4-digit.jpg ---
    left_match = left_pattern.match(filename)
    if left_match:
        index = int(left_match.group(1))
        new_filename = f"{index:04d}.jpg"
        os.rename(
            os.path.join(directory, filename),
            os.path.join(directory, new_filename)
        )
        print(f"Renamed: {filename} -> {new_filename}")

    # --- Collect depthXXXXXX.png for stacking ---
    depth_match = depth_pattern.match(filename)
    if depth_match:
        index = int(depth_match.group(1))
        depth_files.append((index, filename))

# --- Sort and Stack Depth Images ---
depth_files.sort()
for index, filename in depth_files:
    path = os.path.join(directory, filename)
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if img is None:
        print(f"Warning: Could not read {filename}")
        continue
    depth_frames.append(img)

# --- Save as depth.npy and delete originals ---
if depth_frames:
    depth_array = np.stack(depth_frames)
    np.save(os.path.join(directory, "depth.npy"), depth_array)
    print(f"Saved depth.npy with shape {depth_array.shape} and dtype {depth_array.dtype}")
    for _, filename in depth_files:
        path = os.path.join(directory, filename)
        try:
            os.remove(path)
            print(f"Deleted: {filename}")
        except Exception as e:
            print(f"Failed to delete {filename}: {e}")
else:
    print("No depth frames loaded.")