Spaces:
Runtime error
Runtime error
File size: 6,040 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 |
#!/usr/bin/env python3
"""
Preprocess Script for SmartPortraits Dataset
This script processes each sequence in a specified input directory. Each sequence must contain:
- An "association.txt" file listing (timestamp_rgb, rgb_filename, timestamp_depth, depth_filename)
- Pairs of .png files (one for RGB and one for depth)
The script copies each RGB .png file to an output "rgb" folder and converts each 16-bit depth
image to a float32 .npy file in an output "depth" folder. It runs in parallel using
ProcessPoolExecutor for faster performance on multi-core systems.
Usage:
python preprocess_smartportraits.py \
--input_dir /path/to/processed_smartportraits1 \
--output_dir /path/to/processed_smartportraits
"""
import os
import shutil
import argparse
import numpy as np
import cv2
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor, as_completed
def process_pair(args):
"""
Process a single (RGB, depth) pair by:
- Reading the depth .png file and converting it to float32 (depth_in_meters = depth_val / 5000).
- Copying the RGB file to the output directory.
- Saving the converted depth to a .npy file.
Args:
args (tuple): A tuple containing:
- seq_dir (str): Path to the sequence directory.
- seq (str): The name of the current sequence.
- pair_index (int): Index of the pair in the association file (for naming outputs).
- pair (tuple): (rgb_filename, depth_filename).
- out_rgb_dir (str): Output directory for RGB images.
- out_depth_dir (str): Output directory for depth .npy files.
Returns:
None or str:
- Returns None upon successful processing.
- Returns an error message (str) if something fails.
"""
seq_dir, seq, pair_index, pair, out_rgb_dir, out_depth_dir = args
out_rgb_path = os.path.join(out_rgb_dir, f"{pair_index:06d}.png")
out_depth_path = os.path.join(out_depth_dir, f"{pair_index:06d}.npy")
# Skip if both output files already exist
if os.path.exists(out_rgb_path) and os.path.exists(out_depth_path):
return None
try:
rgb_path = os.path.join(seq_dir, pair[0])
depth_path = os.path.join(seq_dir, pair[1])
if not os.path.isfile(rgb_path):
return f"RGB image not found: {rgb_path}"
if not os.path.isfile(depth_path):
return f"Depth image not found: {depth_path}"
# Read the 16-bit depth file
depth = cv2.imread(depth_path, cv2.IMREAD_ANYDEPTH)
if depth is None:
return f"Failed to read depth image: {depth_path}"
# Convert depth values to float32, scale by 1/5000
depth = depth.astype(np.float32) / 5000.0
# Copy the RGB image
shutil.copyfile(rgb_path, out_rgb_path)
# Save depth as a .npy file
np.save(out_depth_path, depth)
except Exception as e:
return f"Error processing pair {pair_index} in sequence '{seq}': {e}"
return None
def process_sequence(seq, input_dir, output_dir):
"""
Process all (RGB, depth) pairs within a single sequence directory.
Args:
seq (str): Name of the sequence (subdirectory).
input_dir (str): Base input directory containing all sequences.
output_dir (str): Base output directory where processed data will be stored.
"""
seq_dir = os.path.join(input_dir, seq)
assoc_file = os.path.join(seq_dir, "association.txt")
# If the association file does not exist, skip this sequence
if not os.path.isfile(assoc_file):
tqdm.write(f"No association.txt found for sequence {seq}. Skipping.")
return
# Prepare output directories
out_rgb_dir = os.path.join(output_dir, seq, "rgb")
out_depth_dir = os.path.join(output_dir, seq, "depth")
os.makedirs(out_rgb_dir, exist_ok=True)
os.makedirs(out_depth_dir, exist_ok=True)
# Read the association file
pairs = []
with open(assoc_file, "r") as f:
for line in f:
items = line.strip().split()
# Format: <timestamp_rgb> <rgb_filename> <timestamp_depth> <depth_filename>
if len(items) < 4:
continue
rgb_file = items[1]
depth_file = items[3]
pairs.append((rgb_file, depth_file))
# Build a list of tasks for parallel processing
tasks = []
for i, pair in enumerate(pairs):
task_args = (seq_dir, seq, i, pair, out_rgb_dir, out_depth_dir)
tasks.append(task_args)
# Process pairs in parallel
num_workers = max(1, os.cpu_count() or 1)
with ProcessPoolExecutor(max_workers=num_workers) as executor:
futures = {executor.submit(process_pair, t): t for t in tasks}
for future in tqdm(
as_completed(futures), total=len(futures), desc=f"Processing sequence {seq}"
):
error = future.result()
if error:
tqdm.write(error)
def main():
parser = argparse.ArgumentParser(description="Preprocess SmartPortraits dataset.")
parser.add_argument(
"--input_dir",
required=True,
help="Path to the directory containing all sequences with association.txt files.",
)
parser.add_argument(
"--output_dir",
required=True,
help="Path to the directory where processed results will be saved.",
)
args = parser.parse_args()
# Gather sequences
if not os.path.isdir(args.input_dir):
raise ValueError(f"Input directory not found: {args.input_dir}")
seqs = sorted(
[
d
for d in os.listdir(args.input_dir)
if os.path.isdir(os.path.join(args.input_dir, d))
]
)
if not seqs:
raise ValueError(f"No valid subdirectories found in {args.input_dir}")
# Process each sequence
for seq in tqdm(seqs, desc="Sequences"):
process_sequence(seq, args.input_dir, args.output_dir)
if __name__ == "__main__":
main()
|