Spaces:
Running
on
L4
Running
on
L4
File size: 7,173 Bytes
2df809d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 |
#!/usr/bin/env python3
"""
Preprocess the IRS dataset.
This script converts disparity EXR files into depth maps, copies corresponding RGB images,
and saves camera intrinsics computed from a given focal length and baseline. Processing is
done per sequence directory using parallel processing.
Usage:
python preprocess_irs.py
--root_dir /path/to/data_irs
--out_dir /path/to/processed_irs
"""
import os
import shutil
import re
import glob
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
import numpy as np
import OpenEXR
import Imath
import imageio
from PIL import Image
from tqdm import tqdm
import argparse
# Ensure OpenEXR support in OpenCV if needed.
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
def exr2hdr(exrpath):
"""
Read an OpenEXR file and return an HDR image as a NumPy array.
"""
file = OpenEXR.InputFile(exrpath)
pixType = Imath.PixelType(Imath.PixelType.FLOAT)
dw = file.header()["dataWindow"]
num_channels = len(file.header()["channels"].keys())
if num_channels > 1:
channels = ["R", "G", "B"]
num_channels = 3
else:
channels = ["G"]
size = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1)
pixels = [
np.fromstring(file.channel(c, pixType), dtype=np.float32) for c in channels
]
hdr = np.zeros((size[1], size[0], num_channels), dtype=np.float32)
if num_channels == 1:
hdr[:, :, 0] = np.reshape(pixels[0], (size[1], size[0]))
else:
hdr[:, :, 0] = np.reshape(pixels[0], (size[1], size[0]))
hdr[:, :, 1] = np.reshape(pixels[1], (size[1], size[0]))
hdr[:, :, 2] = np.reshape(pixels[2], (size[1], size[0]))
return hdr
def writehdr(hdrpath, hdr):
"""
Write an HDR image to a file using the HDR format.
If the input has one channel, duplicate it across R, G, and B.
"""
h, w, c = hdr.shape
if c == 1:
hdr = np.pad(hdr, ((0, 0), (0, 0), (0, 2)), "constant")
hdr[:, :, 1] = hdr[:, :, 0]
hdr[:, :, 2] = hdr[:, :, 0]
imageio.imwrite(hdrpath, hdr, format="hdr")
def load_exr(filename):
"""
Load an EXR file and return the HDR image as a NumPy array.
"""
hdr = exr2hdr(filename)
h, w, c = hdr.shape
if c == 1:
hdr = np.squeeze(hdr)
return hdr
def process_basename(args):
"""
Process a single basename:
- Load an RGB image and disparity (EXR) file.
- Compute a depth map from disparity using: depth = (baseline * f) / disparity.
- Copy the RGB image and save the computed depth and camera intrinsics.
Parameters:
args: tuple containing
(basename, seq_dir, out_rgb_dir, out_depth_dir, out_cam_dir, f, baseline)
Returns:
None on success or an error string on failure.
"""
basename, seq_dir, out_rgb_dir, out_depth_dir, out_cam_dir, f, baseline = args
out_img_path = os.path.join(out_rgb_dir, f"{basename}.png")
out_depth_path = os.path.join(out_depth_dir, f"{basename}.npy")
out_cam_path = os.path.join(out_cam_dir, f"{basename}.npz")
if os.path.exists(out_cam_path):
return
try:
img_file = os.path.join(seq_dir, f"l_{basename}.png")
disp_file = os.path.join(seq_dir, f"d_{basename}.exr")
# Load image using PIL.
img = Image.open(img_file)
# Load disparity using the custom load_exr function.
disp = load_exr(disp_file).astype(np.float32)
H, W = disp.shape
# Verify that the image size matches the disparity map.
if img.size != (W, H):
return f"Size mismatch for {basename}: Image size {img.size}, Disparity size {(W, H)}"
# Create a simple camera intrinsics matrix.
K = np.eye(3, dtype=np.float32)
K[0, 0] = f
K[1, 1] = f
K[0, 2] = W // 2
K[1, 2] = H // 2
# Compute depth from disparity.
depth = baseline * f / disp
# Copy the RGB image.
shutil.copyfile(img_file, out_img_path)
# Save the depth map.
np.save(out_depth_path, depth)
# Save the camera intrinsics.
np.savez(out_cam_path, intrinsics=K)
except Exception as e:
return f"Error processing {basename}: {e}"
return None
def main():
parser = argparse.ArgumentParser(
description="Preprocess IRS dataset: convert EXR disparity to depth, "
"copy RGB images, and save camera intrinsics."
)
parser.add_argument(
"--root_dir",
type=str,
default="/path/to/data_raw_videos/data_irs",
help="Root directory of the raw IRS data.",
)
parser.add_argument(
"--out_dir",
type=str,
default="/path/to/data_raw_videos/processed_irs",
help="Output directory for processed IRS data.",
)
args = parser.parse_args()
# Example parameters (adjust as needed)
baseline = 0.1
f = 480
root = args.root_dir
out_dir = args.out_dir
# Gather sequence directories.
seq_dirs = []
for d in os.listdir(root):
if os.path.isdir(os.path.join(root, d)):
if d == "Store":
for sub in os.listdir(os.path.join(root, d)):
if os.path.isdir(os.path.join(root, d, sub)):
seq_dirs.append(os.path.join(d, sub))
elif d == "IRS_small":
for sub in os.listdir(os.path.join(root, d)):
if os.path.isdir(os.path.join(root, d, sub)):
for subsub in os.listdir(os.path.join(root, d, sub)):
if os.path.isdir(os.path.join(root, d, sub, subsub)):
seq_dirs.append(os.path.join(d, sub, subsub))
else:
seq_dirs.append(d)
seq_dirs.sort()
# Process each sequence.
for seq in seq_dirs:
seq_dir = os.path.join(root, seq)
out_rgb_dir = os.path.join(out_dir, seq, "rgb")
out_depth_dir = os.path.join(out_dir, seq, "depth")
out_cam_dir = os.path.join(out_dir, seq, "cam")
os.makedirs(out_rgb_dir, exist_ok=True)
os.makedirs(out_depth_dir, exist_ok=True)
os.makedirs(out_cam_dir, exist_ok=True)
# Get basenames from disparity files.
basenames = sorted([d[2:-4] for d in os.listdir(seq_dir) if d.endswith(".exr")])
tasks = []
for basename in basenames:
task = (
basename,
seq_dir,
out_rgb_dir,
out_depth_dir,
out_cam_dir,
f,
baseline,
)
tasks.append(task)
num_workers = os.cpu_count() // 2
with ProcessPoolExecutor(max_workers=num_workers) as executor:
futures = {
executor.submit(process_basename, task): task[0] for task in tasks
}
for future in tqdm(
as_completed(futures), total=len(futures), desc=f"Processing {seq}"
):
error = future.result()
if error:
print(error)
if __name__ == "__main__":
main()
|