import json import numpy as np from tqdm import tqdm from PIL import Image import matplotlib.pyplot as plt from scipy.spatial.transform import Rotation as R import os import concurrent.futures path_to_droid_repo = "/workspace/externals/droid/" #input IoU_threshold = 0.75 reprojection_error_threshold = 5.0 # Load the extrinsics cam2base_extrinsics_path = f"{path_to_droid_repo}/cam2base_extrinsic_superset.json" with open(cam2base_extrinsics_path, "r") as f: cam2base_extrinsics = json.load(f) # # Load the intrinsics intrinsics_path = f"{path_to_droid_repo}/intrinsics.json" with open(intrinsics_path, "r") as f: intrinsics = json.load(f) # Load camera serials camera_serials_path = f"{path_to_droid_repo}/camera_serials.json" with open(camera_serials_path, "r") as f: camera_serials = json.load(f) ## data sample # "AUTOLab+5d05c5aa+2023-07-07-10h-00m-27s": { # "relative_path": "AUTOLab/failure/2023-07-07/Fri_Jul__7_10:00:27_2023", # "24400334": [ # 0.2596757315060087, # -0.36626259649963777, # 0.24849304837972613, # -1.742115402153725, # -0.0012127426938948194, # -0.7149867215760838 # ], # "24400334_metric_type": "IoU", # "24400334_quality_metric": 0.6712943316156411, # "24400334_source": "GT", # "22008760": [ # 0.4039752945788883, # 0.47318839256292644, # 0.27170584157181743, # -1.6827143529296786, # 0.07550227077885108, # -2.668292283962724 # ], # "22008760_metric_type": "IoU", # "22008760_quality_metric": 0.6186910915171785, # "22008760_source": "GT" # }, def convert_extrinsics_to_matrix(extracted_extrinsics): pos = extracted_extrinsics[0:3] # translation rot_mat = R.from_euler("xyz", extracted_extrinsics[3:6]).as_matrix() # rotation # Make homogenous transformation matrix cam_to_base_extrinsics_matrix = np.eye(4) cam_to_base_extrinsics_matrix[:3, :3] = rot_mat cam_to_base_extrinsics_matrix[:3, 3] = pos return cam_to_base_extrinsics_matrix def convert_intrinsics_to_matrix(extracted_intrinsics): fx, cx, fy, cy = extracted_intrinsics if fx == 0 or fy == 0: return None intrinsics_matrix = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) return intrinsics_matrix def save_lowdim(episode_count, intrinsics_matrix, extrinsics_matrix, prompts, date_time, img_h, img_w, relative_path, episode_id, output_lowdim_path): for i in range(episode_count): lowdim = dict() lowdim.update({"intrinsics": intrinsics_matrix}) lowdim.update({"shape": [img_h, img_w]}) if prompts is not None: lowdim.update({"prompt": prompts[i]}) lowdim.update({"time": date_time}) lowdim.update({"pose": extrinsics_matrix}) lowdim.update({"relative_path": relative_path}) lowdim.update({"original_episode_id": episode_id}) lowdim.update({"IOU_threshold_for_calib": IoU_threshold}) np.savez(os.path.join(output_lowdim_path, f"frame_{i:06d}.npz"), **lowdim) labwise_episodes_with_good_extrinsics = {} for episode_id, episode_metadata in cam2base_extrinsics.items(): if 'relative_path' in episode_metadata: # remove failure trajectories as they are not present in the raw data path if "failure" in episode_metadata['relative_path']: # remove failure trajectories continue else: # print ("No relative path for episode, cannot map, skipping ", episode_id) continue episode_IoUs = [] reprojection_errors = [] for key, value in episode_metadata.items(): if "metric_type" in key and episode_metadata[key] == "IoU": episode_IoUs.append(episode_metadata[key.replace("metric_type", "quality_metric")]) if "metric_type" in key and episode_metadata[key] == "Reprojection_error": reprojection_errors.append(episode_metadata[key.replace("metric_type", "quality_metric")]) # IOUs for both cameras are provided and above threshold if (len(episode_IoUs) >= 2 and all(x > IoU_threshold for x in episode_IoUs)): print ("Keeping episode: ", episode_id) lab_name = episode_metadata['relative_path'].split("/")[0] if lab_name not in labwise_episodes_with_good_extrinsics: labwise_episodes_with_good_extrinsics[lab_name] = [] labwise_episodes_with_good_extrinsics[lab_name].append([episode_id, episode_metadata]) # Alternative ways to filter episodes: # Option-1: IOU + reprojection error criteria # if (len(episode_IoUs)+len(reprojection_errors)) >= 2 and all(x > IoU_threshold for x in episode_IoUs) and all(x < reprojection_error_threshold for x in reprojection_errors): # # episodes_with_good_extrinsics[episode_id] = episode_metadata # lab_name = episode_metadata['relative_path'].split("/")[0] # if lab_name not in labwise_episodes_with_good_extrinsics: # labwise_episodes_with_good_extrinsics[lab_name] = [] # labwise_episodes_with_good_extrinsics[lab_name].append([episode_id, episode_metadata]) # Option-2: Separately deal with IOU and reprojection error criteria # if len(episode_IoUs) >= 1 and all(x > IoU_threshold for x in episode_IoUs): # episodes_with_good_extrinsics[episode_id] = episode_metadata # if len(reprojection_errors) >= 1 and all(x < reprojection_error_threshold for x in reprojection_errors): # episodes_with_good_extrinsics[episode_id] = episode_metadata # print("Number of episodes with good extrinsics: ", len(labwise_episodes_with_good_extrinsics)) print ("Lab names: ", labwise_episodes_with_good_extrinsics.keys()) print ("Labwise episode count: ", {k: len(v) for k, v in labwise_episodes_with_good_extrinsics.items()}) print ("Total episodes with good extrinsics: ", sum([len(v) for k, v in labwise_episodes_with_good_extrinsics.items()])) left_raw_path = "/workspace/cv_datasets/raw/DROIDraw/" depth_raw_path = "/workspace/cv_datasets/raw/DROIDraw/" output_path = "/workspace/datasets/DroidProcessedV4/" # Output (ProceessedDataset format): # lab/episode_no/ # - depth/camera_no/frame_no.npz # - rgb/camera_no/frame_no.jpg (using left_rgb from input) # - lowdim/camera_no/frame_no.npz def process_episode(lab_name, i, episode_id, episode_metadata): camera_serials_to_name = {v: k for k, v in camera_serials[episode_id].items()} extracted_extrinsics = {} extracted_intrinsics = {} img_h = {} img_w = {} relative_path = episode_metadata['relative_path'] date_time = episode_id.split("+")[-1] if not os.path.exists(os.path.join(left_raw_path, relative_path)): print ("WARN: Raw input data path does not exist, skipping ", os.path.join(left_raw_path, relative_path)) return for key, value in episode_metadata.items(): if key.isdigit(): camera_serial = key calib_camera_name = camera_serials_to_name[camera_serial] if calib_camera_name == "ext1_cam_serial": camera_name = "camera_1" elif calib_camera_name == "ext2_cam_serial": camera_name = "camera_2" else: print(f"WARN: Unknown camera name: {calib_camera_name}") continue calculated_intrinsics = convert_intrinsics_to_matrix(intrinsics[episode_id][camera_serial]['cameraMatrix']) if calculated_intrinsics is None: print ("WARN: Skipping camera with invalid intrinsics: ", episode_id, calib_camera_name, intrinsics[episode_id][camera_serial]['cameraMatrix']) continue extracted_intrinsics[camera_name] = calculated_intrinsics extracted_extrinsics[camera_name] = convert_extrinsics_to_matrix(value) img_h[camera_name] = intrinsics[episode_id][camera_serial]['height'] img_w[camera_name] = intrinsics[episode_id][camera_serial]['width'] # Less than 2 cameras with valid intrinsics, skip if len(extracted_intrinsics.keys()) < 2: print ("WARN: Less than 2 cameras with valid intrinsics for episode, skipping ", episode_id, "aka", f"episode_{i:06d}") return if len(extracted_extrinsics.keys()) < 2: print ("WARN: Less than 2 cameras with valid extrinsics for episode, skipping ", episode_id, "aka", f"episode_{i:06d}") return # print (f"Processing episode {i} / {len(labwise_episodes_with_good_extrinsics[lab_name])} : {episode_id} with {len(extracted_intrinsics.keys())} cameras, into {os.path.join(output_path, lab_name, f'episode_{i:06d}')}") os.makedirs(os.path.join(output_path, lab_name, f"episode_{i:06d}"), exist_ok=True) for camera_name in extracted_extrinsics.keys(): img_input_path = os.path.join(left_raw_path, relative_path, "images/left_rgb/", camera_name) left_img_files = sorted(os.listdir(img_input_path)) # print (f"Image Input path: {img_input_path}, found {len(left_img_files)} images") depth_input_path = os.path.join(depth_raw_path, relative_path, "images/tri_depth/", camera_name) depth_files = sorted(os.listdir(depth_input_path)) # print ("Depth input path: ", depth_input_path, "found", len(depth_files), "depth files") if len(left_img_files) != len(depth_files): print (f"WARN: Number of images and depth files do not match for episode {episode_id}, camera {camera_name}, skipping this camera") continue if len(left_img_files) == 0 or len(depth_files) == 0: print (f"WARN: Either image or depth not found for episode {episode_id}, camera {camera_name}, skipping this camera") continue os.makedirs(os.path.join(output_path, lab_name, f"episode_{i:06d}", "rgb", camera_name), exist_ok=True) for j, img_file in enumerate(left_img_files): img = Image.open(os.path.join(img_input_path, img_file)) img.save(os.path.join(output_path, lab_name, f"episode_{i:06d}", "rgb", camera_name, f"frame_{j:06d}.jpg")) os.makedirs(os.path.join(output_path, lab_name, f"episode_{i:06d}", "depth", camera_name), exist_ok=True) for j, depth_file in enumerate(depth_files): depth = np.load(os.path.join(depth_input_path, depth_file))/1000.0 # convert mm to meters if depth.shape[0] == 0 or depth.shape[1] == 0: print ("WARN: Empty depth file, skipping ", os.path.join(depth_input_path, depth_file)) continue np.savez_compressed(os.path.join(output_path, lab_name, f"episode_{i:06d}", "depth", camera_name, f"frame_{j:06d}.npz"), data=depth) output_lowdim_path = os.path.join(output_path, lab_name, f"episode_{i:06d}", "lowdim", camera_name) os.makedirs(output_lowdim_path, exist_ok=True) save_lowdim(len(left_img_files), extracted_intrinsics[camera_name], extracted_extrinsics[camera_name], None, date_time, img_h[camera_name], img_w[camera_name], relative_path, episode_id, output_lowdim_path) # if relative_path == "ILIAD/success/2023-05-31/Wed_May_31_18:14:33_2023": # print (extracted_intrinsics["camera_1"], extracted_intrinsics["camera_2"]) # print ("diff between intrinsics: ", extracted_intrinsics["camera_1"] - extracted_intrinsics["camera_2"]) # Episode ID: AUTOLab+0d4edc83+2023-12-02-15h-11m-45s # Relative path: AUTOLab/success/2023-12-02/Sat_Dec__2_15:11:45_2023 # Lab name: AUTOLab # print ("Episode ID: ", episode_id) # print ("Relative path: ", relative_path) # print ("Lab name: ", lab_name) for lab_name, episodes in labwise_episodes_with_good_extrinsics.items(): if lab_name == "RPL": print ("Skipping RPL lab for now, done already") continue print (f"Processing lab: {lab_name} with {len(episodes)} episodes into {os.path.join(output_path, lab_name)}") os.makedirs(os.path.join(output_path, lab_name), exist_ok=True) with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: futures = [] for i, (episode_id, episode_metadata) in enumerate(episodes): future = executor.submit(process_episode, lab_name, i, episode_id, episode_metadata) futures.append(future) # Show progress as tasks complete for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc=f"Processing episodes for {lab_name}"): try: result = future.result() # Get result or raise exception except Exception as e: print(f"Episode failed: {e}\n")