import os import json import shutil import numpy as np import cv2 as cv import imageio from tqdm import tqdm from concurrent.futures import ProcessPoolExecutor, as_completed import open3d as o3d import scipy.ndimage import pickle # Set environment variable to limit OpenBLAS threads os.environ["OPENBLAS_NUM_THREADS"] = "1" DEPTH_SCALE_FACTOR = 5000 # Point cloud from depth def pointcloudify_depth(depth, intrinsics, dist_coeff, undistort=True): shape = depth.shape[::-1] if undistort: undist_intrinsics, _ = cv.getOptimalNewCameraMatrix( intrinsics, dist_coeff, shape, 1, shape ) inv_undist_intrinsics = np.linalg.inv(undist_intrinsics) map_x, map_y = cv.initUndistortRectifyMap( intrinsics, dist_coeff, None, undist_intrinsics, shape, cv.CV_32FC1 ) undist_depth = cv.remap(depth, map_x, map_y, cv.INTER_NEAREST) else: inv_undist_intrinsics = np.linalg.inv(intrinsics) undist_depth = depth # Generate x,y grid for H x W image grid_x, grid_y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1])) grid = np.stack((grid_x, grid_y, np.ones_like(grid_x)), axis=-1) # Reshape and compute local grid grid_flat = grid.reshape(-1, 3).T local_grid = inv_undist_intrinsics @ grid_flat # Multiply by depth local_grid = local_grid.T * undist_depth.reshape(-1, 1) return local_grid.astype(np.float32) def project_pcd_to_depth(pcd, undist_intrinsics, img_size, config): h, w = img_size points = np.asarray(pcd.points) d = points[:, 2] normalized_points = points / points[:, 2][:, np.newaxis] proj_pcd = np.round((undist_intrinsics @ normalized_points.T).T).astype(np.int64) proj_mask = ( (proj_pcd[:, 0] >= 0) & (proj_pcd[:, 0] < w) & (proj_pcd[:, 1] >= 0) & (proj_pcd[:, 1] < h) ) proj_pcd = proj_pcd[proj_mask] d = d[proj_mask] pcd_image = np.zeros((config["res_h"], config["res_w"]), dtype=np.float32) pcd_image[proj_pcd[:, 1], proj_pcd[:, 0]] = d return pcd_image def smooth_depth(depth): MAX_DEPTH_VAL = 1e5 KERNEL_SIZE = 11 depth = depth.copy() depth[depth == 0] = MAX_DEPTH_VAL smoothed_depth = scipy.ndimage.minimum_filter(depth, KERNEL_SIZE) smoothed_depth[smoothed_depth == MAX_DEPTH_VAL] = 0 return smoothed_depth def align_rgb_depth(rgb, depth, roi, config, rgb_cnf, config_dict, T): # Undistort rgb image undist_rgb = cv.undistort( rgb, rgb_cnf["intrinsics"], rgb_cnf["dist_coeff"], None, rgb_cnf["undist_intrinsics"], ) # Create point cloud from depth pcd = o3d.geometry.PointCloud() points = pointcloudify_depth( depth, config_dict["depth"]["dist_mtx"], config_dict["depth"]["dist_coef"] ) pcd.points = o3d.utility.Vector3dVector(points) # Align point cloud with depth reference frame pcd.transform(T) # Project aligned point cloud to rgb aligned_depth = project_pcd_to_depth( pcd, rgb_cnf["undist_intrinsics"], rgb.shape[:2], config ) smoothed_aligned_depth = smooth_depth(aligned_depth) x, y, w, h = roi depth_res = smoothed_aligned_depth[y : y + h, x : x + w] rgb_res = undist_rgb[y : y + h, x : x + w] return rgb_res, depth_res, rgb_cnf["undist_intrinsics"] def process_pair(args): ( pair, smartphone_folder, azure_depth_folder, final_folder, config, rgb_cnf, config_dict, T, ) = args try: rgb_image = cv.imread(os.path.join(smartphone_folder, f"{pair[0]}.png")) depth_array = np.load( os.path.join(azure_depth_folder, f"{pair[1]}.npy"), allow_pickle=True ) rgb_image_aligned, depth_array_aligned, intrinsics = align_rgb_depth( rgb_image, depth_array, (0, 0, config["res_w"], config["res_h"]), config, rgb_cnf, config_dict, T, ) # Save rgb as 8-bit png cv.imwrite( os.path.join(final_folder, "rgb", f"{pair[0]}.png"), rgb_image_aligned ) # # Save depth as 16-bit unsigned int with scale factor # depth_array_aligned = (depth_array_aligned * # DEPTH_SCALE_FACTOR).astype(np.uint16) # imageio.imwrite(os.path.join(final_folder, 'depth', f"{pair[1]}.png"), depth_array_aligned) np.save( os.path.join(final_folder, "depth", f"{pair[0]}.npy"), depth_array_aligned ) np.savez( os.path.join(final_folder, "cam", f"{pair[0]}.npz"), intrinsics=intrinsics ) except Exception as e: return f"Error processing pair {pair}: {e}" return None def main(): DATA_DIR_ = "data_smartportraits/SmartPortraits" # REPLACE WITH YOUR OWN DATA PATH! DATA_DIR = DATA_DIR_.rstrip("/") print(f"{DATA_DIR_} {DATA_DIR}/") # Folder where the data in TUM format will be put curr_dir = os.path.dirname(os.path.abspath(__file__)) with open(os.path.join(curr_dir, "config.json")) as conf_f: config = json.load(conf_f) # Pre-load shared data with open(os.path.join(curr_dir, config["depth_conf"]), "rb") as config_f: config_dict = pickle.load(config_f) rgb_cnf = np.load( os.path.join(curr_dir, config["rgb_intristics"]), allow_pickle=True ).item() T = np.load(os.path.join(curr_dir, config["transform_intristics"])) final_root = "processed_smartportraits1" # REPLACE WITH YOUR OWN DATA PATH! seqs = [] for scene in os.listdir(DATA_DIR): scene_path = os.path.join(DATA_DIR, scene) if not os.path.isdir(scene_path): continue for s in os.listdir(scene_path): s_path = os.path.join(scene_path, s) if not os.path.isdir(s_path): continue for date in os.listdir(s_path): date_path = os.path.join(s_path, date) if os.path.isdir(date_path): seqs.append((scene, s, date)) for seq in tqdm(seqs): scene, s, date = seq dataset_path = os.path.join(DATA_DIR, scene, s, date) final_folder = os.path.join(final_root, "_".join([scene, s, date])) azure_depth_folder = os.path.join(dataset_path, "_azure_depth_image_raw") smartphone_folder = os.path.join(dataset_path, "smartphone_video_frames") depth_files = [ file for file in os.listdir(azure_depth_folder) if file.endswith(".npy") ] depth_ts = np.array([int(file.split(".")[0]) for file in depth_files]) depth_ts.sort() rgb_files = [ file for file in os.listdir(smartphone_folder) if file.endswith(".png") ] rgb_ts = np.array([int(file.split(".")[0]) for file in rgb_files]) rgb_ts.sort() print( f"Depth timestamps from {depth_ts[0]} to {depth_ts[-1]} (cnt {len(depth_ts)})" ) print(f"RGB timestamps from {rgb_ts[0]} to {rgb_ts[-1]} (cnt {len(rgb_ts)})") # Build correspondences between depth and rgb by nearest neighbour algorithm rgbd_pairs = [] for depth_t in depth_ts: idx = np.argmin(np.abs(rgb_ts - depth_t)) closest_rgb_t = rgb_ts[idx] rgbd_pairs.append((closest_rgb_t, depth_t)) # Prepare folder infrastructure if os.path.exists(final_folder): shutil.rmtree(final_folder) os.makedirs(os.path.join(final_folder, "depth"), exist_ok=True) os.makedirs(os.path.join(final_folder, "rgb"), exist_ok=True) os.makedirs(os.path.join(final_folder, "cam"), exist_ok=True) # Prepare arguments for processing tasks = [ ( pair, smartphone_folder, azure_depth_folder, final_folder, config, rgb_cnf, config_dict, T, ) for pair in rgbd_pairs ] num_workers = os.cpu_count() with ProcessPoolExecutor(max_workers=num_workers) as executor: futures = {executor.submit(process_pair, task): task[0] for task in tasks} for future in tqdm( as_completed(futures), total=len(futures), desc=f"Processing {scene}_{s}_{date}", ): error = future.result() if error: print(error) if __name__ == "__main__": main()