diff --git a/.gitignore b/.gitignore index f8a9584..58aac05 100644 --- a/.gitignore +++ b/.gitignore @@ -148,3 +148,4 @@ docs/make.bat docs/Makefile examples/ta_data_spec.txt *.txt +examples/*.png \ No newline at end of file diff --git a/docs/environments.rst b/docs/environments.rst new file mode 100644 index 0000000..1241eeb --- /dev/null +++ b/docs/environments.rst @@ -0,0 +1,7 @@ +Environments +===================================== + +We definitely have bugs. Please report them in `GitHub issues`_ or submit a pull request! + +.. _`GitHub issues`: https://github.com/castacks/tartanairpy/issues + diff --git a/docs/examples.rst b/docs/examples.rst index c34d78e..a525693 100644 --- a/docs/examples.rst +++ b/docs/examples.rst @@ -20,11 +20,11 @@ Download via Python API ta.init(tartanair_data_root) # Download a trajectory. - ta.download(env = "DesertGasStation", - difficulty = ['easy', 'hard'], - modality = ['image', 'depth'], - camera_name = ['lcam_front'], - unzip = True) + ta.download(env = "ArchVizTinyHouseDay", + difficulty = ['easy'], # this can be 'easy', and/or 'hard' + modality = ['image', 'depth', 'seg', 'imu'], # available modalities are: image', 'depth', 'seg', 'imu', 'lidar', 'flow', 'pose' + camera_name = ['lcam_front', 'lcam_left', 'lcam_right', 'lcam_back', 'lcam_top', 'lcam_bottom'], + unzip = True) # unzip files autonomously after download Download via a yaml config file @@ -37,10 +37,10 @@ The config file if of the following format: .. code-block:: yaml - env: ['DesertGasStation'] + env: ['ArchVizTinyHouseDay'] difficulty: ['easy'] - modality: ['image'] - camera_name: ['lcam_front'] + modality: ['image', 'depth'] + camera_name: ['lcam_front', 'lcam_left', 'lcam_right', 'lcam_back', 'lcam_top', 'lcam_bottom'] unzip: True Customization Example @@ -54,7 +54,6 @@ TartanAir V2 allows you to synthesize your own dataset by modifying the raw data # For help with rotations. from scipy.spatial.transform import Rotation - import numpy as np # Initialize TartanAir. tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' @@ -77,8 +76,8 @@ TartanAir V2 allows you to synthesize your own dataset by modifying the raw data R_raw_new1 = Rotation.from_euler('xyz', [45, 0, 0], degrees=True).as_matrix().tolist() cam_model_1 = {'name': 'doublesphere', - 'raw_side': 'left', - 'params': + 'raw_side': 'left', + 'params': {'fx': 300, 'fy': 300, 'cx': 500, @@ -88,15 +87,16 @@ TartanAir V2 allows you to synthesize your own dataset by modifying the raw data 'alpha': 0.6, 'xi': -0.2, 'fov_degree': 195}, - 'R_raw_new': R_raw_new1} + 'R_raw_new': R_raw_new1} # Customize the dataset. - ta.customize(env = 'SupermarketExposure', + ta.customize(env = 'ArchVizTinyHouseDay', difficulty = 'easy', trajectory_id = ['P000'], - modality = ['image'], - new_camera_models_params=[cam_model_0, cam_model_1], - num_workers = 2) + modality = ['image', 'depth'], + new_camera_models_params=[cam_model_1, cam_model_0], + num_workers = 4, + device = "cuda") # or cpu Dataloader Example ------------------------------------- @@ -113,33 +113,33 @@ TartanAir-V2 includes a powerful parallelized dataloader. It can be used to load ta.init(tartanair_data_root) # Specify the environments, difficulties, and trajectory ids to load. - envs = ['ArchVizTinyHouseDayExposure'] - difficulties = ['hard'] - trajectory_ids = [] #['P000', 'P001'] + envs = ['ArchVizTinyHouseDay'] + difficulties = ['easy'] + trajectory_ids = ['P000', 'P001'] # Specify the modalities to load. - modalities = ['image', 'depth', 'pose'] - camnames = ['lcam_front', 'lcam_left'] + modalities = ['image', 'pose', 'imu'] + camnames = ['lcam_front', 'lcam_left', 'lcam_right', 'lcam_back', 'lcam_top', 'lcam_bottom'] # Specify the dataloader parameters. new_image_shape_hw = [640, 640] # If None, no resizing is performed. If a value is passed, then the image is resized to this shape. - subset_framenum = 364 # This is the number of frames in a subset. Notice that this is an upper bound on the batch size. Ideally, make this number large to utilize your RAM efficiently. Information about the allocated memory will be provided in the console. - seq_length = {'image': 2, 'depth': 1} # This is the length of the data-sequences. For example, if the sequence length is 2, then the dataloader will load pairs of images. + subset_framenum = 200 # This is the number of frames in a subset. Notice that this is an upper bound on the batch size. Ideally, make this number large to utilize your RAM efficiently. Information about the allocated memory will be provided in the console. + seq_length = {'image': 2, 'pose': 2, 'imu': 10} # This is the length of the data-sequences. For example, if the sequence length is 2, then the dataloader will load pairs of images. seq_stride = 1 # This is the stride between the data-sequences. For example, if the sequence length is 2 and the stride is 1, then the dataloader will load pairs of images [0,1], [1,2], [2,3], etc. If the stride is 2, then the dataloader will load pairs of images [0,1], [2,3], [4,5], etc. frame_skip = 0 # This is the number of frames to skip between each frame. For example, if the frame skip is 2 and the sequence length is 3, then the dataloader will load frames [0, 3, 6], [1, 4, 7], [2, 5, 8], etc. batch_size = 8 # This is the number of data-sequences in a mini-batch. num_workers = 4 # This is the number of workers to use for loading the data. - shuffle = False # Whether to shuffle the data. Let's set this to False for now, so that we can see the data loading in a nice video. Yes it is nice don't argue with me please. Just look at it! So nice. :) + shuffle = True # Whether to shuffle the data. Let's set this to False for now, so that we can see the data loading in a nice video. Yes it is nice don't argue with me please. Just look at it! So nice. :) # Create a dataloader object. dataloader = ta.dataloader(env = envs, difficulty = difficulties, trajectory_id = trajectory_ids, modality = modalities, - camear_name = camnames, + camera_name = camnames, new_image_shape_hw = new_image_shape_hw, - subset_framenum = subset_framenum, seq_length = seq_length, + subset_framenum = subset_framenum, seq_stride = seq_stride, frame_skip = frame_skip, batch_size = batch_size, @@ -151,21 +151,37 @@ TartanAir-V2 includes a powerful parallelized dataloader. It can be used to load for i in range(100): # Get the next batch. batch = dataloader.load_sample() - # Check if the batch is None. - if batch is None: - break - print("Batch number: {}".format(i), "Loaded {} samples so far.".format(i * batch_size)) + # Visualize some images. + # The shape of an image batch is (B, S, H, W, C), where B is the batch size, S is the sequence length, H is the height, W is the width, and C is the number of channels. + + print("Batch number: {}".format(i+1), "Loaded {} samples so far.".format((i+1) * batch_size)) for b in range(batch_size): - # Visualize some images. - # The shape of an image batch is (B, S, H, W, C), where B is the batch size, S is the sequence length, H is the height, W is the width, and C is the number of channels. - img0 = batch['rgb_lcam_front'][b][0] - img1 = batch['rgb_lcam_front'][b][1] - - # Visualize the images. - outimg = np.concatenate((img0, img1), axis = 1) - cv2.imshow('outimg', outimg) - cv2.waitKey(1) + + # Create image cross. + left = batch['image_lcam_left'][b][0].numpy().transpose(1,2,0) + front = batch['image_lcam_front'][b][0].numpy().transpose(1,2,0) + right = batch['image_lcam_right'][b][0].numpy().transpose(1,2,0) + back = batch['image_lcam_back'][b][0].numpy().transpose(1,2,0) + top = batch['image_lcam_top'][b][0].numpy().transpose(1,2,0) + bottom = batch['image_lcam_bottom'][b][0].numpy().transpose(1,2,0) + cross_mid = np.concatenate([left, front, right, back], axis=1) + cross_top = np.concatenate([np.zeros_like(top), top, np.zeros_like(top), np.zeros_like(top)], axis=1) + cross_bottom = np.concatenate([np.zeros_like(bottom), bottom, np.zeros_like(bottom), np.zeros_like(bottom)], axis=1) + cross = np.concatenate([cross_top, cross_mid, cross_bottom], axis=0) + + pose = batch['pose_lcam_front'].numpy() + imu = batch['imu'].numpy() + + # Resize. + cross = cv2.resize(cross, (cross.shape[1]//4, cross.shape[0]//4)) + + # Show the image cross. + cv2.imshow('cross', cross) + cv2.waitKey(100) + + print(" Pose: ", pose[0][0]) + print(" IMU: ", imu[0][0]) dataloader.stop_cachers() @@ -184,11 +200,11 @@ Create a data iterator to get samples from the TartanAir V2 dataset. The samples ta.init(tartanair_data_root) # Create iterator. - ta_iterator = ta.iterator(env = 'ConstructionSite', - difficulty = 'easy', - trajectory_id = 'P000', - modality = 'image', - camera_name = 'lcam_front') + ta_iterator = ta.iterator(env = ['ArchVizTinyHouseDay'], + difficulty = 'easy', + trajectory_id = [], + modality = 'image', + camera_name = ['lcam_left']) for i in range(100): sample = next(ta_iterator) @@ -208,7 +224,7 @@ TartanAir also provides tools for evaluating estimated trajectories against the ta.init(tartanair_data_root) # Create an example trajectory. This is a noisy version of the ground truth trajectory. - env = 'AbandonedCableExposure' + env = 'ArchVizTinyHouseDay' difficulty = 'easy' trajectory_id = 'P002' camera_name = 'lcam_front' @@ -216,11 +232,26 @@ TartanAir also provides tools for evaluating estimated trajectories against the est_traj = gt_traj + np.random.normal(0, 0.1, gt_traj.shape) # Pass the ground truth trajectory directly to the evaluation function. - results = ta.evaluate_traj(est_traj, gt_traj = gt_traj, enforce_length = True, plot = True, plot_out_path = plot_out_path, do_scale = True, do_align = True) + results = ta.evaluate_traj(est_traj, + gt_traj = gt_traj, + enforce_length = True, + plot = True, + plot_out_path = plot_out_path, + do_scale = True, + do_align = True) # Or pass the environment, difficulty, and trajectory id to the evaluation function. plot_out_path = "evaluator_example_plot.png" - results = ta.evaluate_traj(est_traj, env, difficulty, trajectory_id, camera_name, enforce_length = True, plot = True, plot_out_path = plot_out_path, do_scale = True, do_align = True) + results = ta.evaluate_traj(est_traj, + env = env, + difficulty = difficulty, + trajectory_id = trajectory_id, + camera_name = camera_name, + enforce_length = True, + plot = True, + plot_out_path = plot_out_path, + do_scale = True, + do_align = True) Flow Sampling Example ------------------------------------- diff --git a/docs/index.rst b/docs/index.rst index 579185e..8ffc7e6 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -18,7 +18,7 @@ Welcome to TartanAir V2! Let's go on an adventure to beautiful mountains, to dark caves, to stylish homes, to the Moon 🚀, and to other exciting places. And there is more! You, your models, and your robots can experience these worlds via a variety of sensors: LiDAR, IMU, optical cameras with any lense configuration you want (we provide customizable fisheye, pinhole, and equirectangular camera models), depth cameras, segmentation "cameras", and event cameras. .. image:: - web_cover_figure.png + _images/tartanair_ending.gif All the environments have recorded trajectories that were designed to be challenging and realistic. Can we improve the state of the art in SLAM, navigation, and robotics? @@ -30,6 +30,8 @@ All the environments have recorded trajectories that were designed to be challen installation examples usage + modalities + environments troubleshooting License diff --git a/docs/modalities.rst b/docs/modalities.rst new file mode 100644 index 0000000..501a512 --- /dev/null +++ b/docs/modalities.rst @@ -0,0 +1,7 @@ +Modalities +===================================== + +Images +Depth +Flow +LiDAR diff --git a/examples/customization_example.py b/examples/customization_example.py index cb8d5bf..ba736e9 100644 --- a/examples/customization_example.py +++ b/examples/customization_example.py @@ -7,7 +7,6 @@ # General imports. import sys -import numpy as np from scipy.spatial.transform import Rotation # Local imports. @@ -15,14 +14,14 @@ import tartanair as ta # Create a TartanAir object. -tartanair_data_root = '/media/yoraish/overflow/data/tartanair-v2' +tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' ta.init(tartanair_data_root) # Create the requested camera models and their parameters. R_raw_new0 = Rotation.from_euler('y', 90, degrees=True).as_matrix().tolist() cam_model_0 = {'name': 'pinhole', - 'raw_side': 'left', # TartanAir has two cameras, one on the left and one on the right. This parameter specifies which camera to use. + 'raw_side': 'left', # TartanAir has two cameras, one on the left and one on the right. This parameter specifies which camera to use. 'params': {'fx': 320, 'fy': 320, 'cx': 320, 'cy': 320, 'width': 640, 'height': 640}, 'R_raw_new': R_raw_new0} @@ -43,4 +42,10 @@ 'fov_degree': 195}, 'R_raw_new': R_raw_new1} -ta.customize(env = 'ModularNeighborhoodIntExt', difficulty = 'easy', trajectory_id = ['P000'], modality = ['image'], new_camera_models_params=[cam_model_1, cam_model_0], num_workers = 2, device='cuda') # Or cpu. +ta.customize(env = 'ArchVizTinyHouseDay', + difficulty = 'easy', + trajectory_id = ['P000'], + modality = ['image', 'depth'], + new_camera_models_params=[cam_model_1, cam_model_0], + num_workers = 4, + device = "cuda") # or cpu diff --git a/examples/dataloader_example.py b/examples/dataloader_example.py index 710f714..4f37d99 100644 --- a/examples/dataloader_example.py +++ b/examples/dataloader_example.py @@ -15,20 +15,19 @@ import tartanair as ta # Create a TartanAir object. -# tartanair_data_root = './data/tartanair-v2' -tartanair_data_root = '/media/yoraish/overflow/data/tartanair-v2' +tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' ta.init(tartanair_data_root) # Specify the environments, difficulties, and trajectory ids to load. envs = [ - "ArchVizTinyHouseDayExposure", + "ArchVizTinyHouseDay", ] -difficulties = [] -trajectory_ids = ['P000'] +difficulties = ['easy'] +trajectory_ids = ['P000','P001'] # Specify the modalities to load. -modalities = ['image', 'pose', 'imu_gyro', 'imu_time'] +modalities = ['image', 'pose', 'imu'] camnames = ['lcam_front', 'lcam_left', 'lcam_right', 'lcam_back', 'lcam_top', 'lcam_bottom'] # Specify dataloader configuration. The dataloader operates in the following way: @@ -38,8 +37,8 @@ new_image_shape_hw = [640, 640] # If None, no resizing is performed. If a value is passed, then the image is resized to this shape. subset_framenum = 200 # This is the number of frames in a subset. Notice that this is an upper bound on the batch size. Ideally, make this number large to utilize your RAM efficiently. Information about the allocated memory will be provided in the console. +seq_length = {'image': 2, 'pose': 2, 'imu': 10} # This is the length of the data-sequences. For example, if the sequence length is 2, then the dataloader will load pairs of images. frame_skip = 0 # This is the number of frames to skip between each frame. For example, if the frame skip is 2 and the sequence length is 3, then the dataloader will load frames [0, 3, 6], [1, 4, 7], [2, 5, 8], etc. -seq_length = {'image': 2, 'pose': 2, 'imu_gyro': 10, 'imu_time': 10} # This is the length of the data-sequences. For example, if the sequence length is 2, then the dataloader will load pairs of images. seq_stride = 1 # This is the stride between the data-sequences. For example, if the sequence length is 2 and the stride is 1, then the dataloader will load pairs of images [0,1], [1,2], [2,3], etc. If the stride is 2, then the dataloader will load pairs of images [0,1], [2,3], [4,5], etc. batch_size = 4 # This is the number of data-sequences in a mini-batch. num_workers = 8 # This is the number of workers to use for loading the data. @@ -52,8 +51,8 @@ modality = modalities, camera_name = camnames, new_image_shape_hw = new_image_shape_hw, - subset_framenum = subset_framenum, seq_length = seq_length, + subset_framenum = subset_framenum, seq_stride = seq_stride, frame_skip = frame_skip, batch_size = batch_size, @@ -65,32 +64,39 @@ # Iterate over the dataloader and visualize the output. # Iterate over the batches. -for i in range(1000): +for i in range(100): # Get the next batch. batch = dataloader.load_sample() - # Check if the batch is None. - if batch is None: - break # Visualize some images. # The shape of an image batch is (B, S, H, W, C), where B is the batch size, S is the sequence length, H is the height, W is the width, and C is the number of channels. - # Create image cross. - left = batch['rgb_lcam_left'][0][0] - front = batch['rgb_lcam_front'][0][0] - right = batch['rgb_lcam_right'][0][0] - back = batch['rgb_lcam_back'][0][0] - top = batch['rgb_lcam_top'][0][0] - bottom = batch['rgb_lcam_bottom'][0][0] - cross_mid = np.concatenate([left, front, right, back], axis=1) - cross_top = np.concatenate([np.zeros_like(top), top, np.zeros_like(top), np.zeros_like(top)], axis=1) - cross_bottom = np.concatenate([np.zeros_like(bottom), bottom, np.zeros_like(bottom), np.zeros_like(bottom)], axis=1) - cross = np.concatenate([cross_top, cross_mid, cross_bottom], axis=0) - - # Resize. - cross = cv2.resize(cross, (cross.shape[1]//4, cross.shape[0]//4)) - - # Show the image cross. - cv2.imshow('cross', cross) - cv2.waitKey(1) + print("Batch number: {}".format(i+1), "Loaded {} samples so far.".format((i+1) * batch_size)) + + for b in range(batch_size): + + # Create image cross. + left = batch['image_lcam_left'][b][0].numpy().transpose(1,2,0) + front = batch['image_lcam_front'][b][0].numpy().transpose(1,2,0) + right = batch['image_lcam_right'][b][0].numpy().transpose(1,2,0) + back = batch['image_lcam_back'][b][0].numpy().transpose(1,2,0) + top = batch['image_lcam_top'][b][0].numpy().transpose(1,2,0) + bottom = batch['image_lcam_bottom'][b][0].numpy().transpose(1,2,0) + cross_mid = np.concatenate([left, front, right, back], axis=1) + cross_top = np.concatenate([np.zeros_like(top), top, np.zeros_like(top), np.zeros_like(top)], axis=1) + cross_bottom = np.concatenate([np.zeros_like(bottom), bottom, np.zeros_like(bottom), np.zeros_like(bottom)], axis=1) + cross = np.concatenate([cross_top, cross_mid, cross_bottom], axis=0) + + pose = batch['pose_lcam_front'].numpy() + imu = batch['imu'].numpy() + + # Resize. + cross = cv2.resize(cross, (cross.shape[1]//4, cross.shape[0]//4)) + + # Show the image cross. + cv2.imshow('cross', cross) + cv2.waitKey(100) + + print(" Pose: ", pose[0][0]) + print(" IMU: ", imu[0][0]) dataloader.stop_cachers() \ No newline at end of file diff --git a/examples/dataset_example.py b/examples/dataset_example.py index c603380..30a2083 100644 --- a/examples/dataset_example.py +++ b/examples/dataset_example.py @@ -13,7 +13,7 @@ import tartanair as ta # Create a TartanAir object. -tartanair_data_root = '/media/yoraish/overflow/data/tartanair-v2_training_data' +tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' ta.init(tartanair_data_root) ##################### @@ -21,7 +21,7 @@ ##################### # Set up the dataset. -dataset = ta.create_image_dataset(env = ['ConstructionSite', 'SupermarketExposure'], difficulty = [], trajectory_id = [], modality = ['image', 'depth'], camera_name = ['lcam_front', 'lcam_back', 'lcam_right', 'lcam_left', 'lcam_top', 'lcam_bottom'], transform = None, num_workers=10) +dataset = ta.create_image_dataset(env = ['ArchVizTinyHouseDay'], difficulty = [], trajectory_id = [], modality = ['image', 'depth'], camera_name = ['lcam_front', 'lcam_back', 'lcam_right', 'lcam_left', 'lcam_top', 'lcam_bottom'], transform = None, num_workers=10) # Print the dataset. print(dataset) @@ -74,7 +74,7 @@ ]) # Set up a dataset. -dataset = ta.create_image_dataset(env = [], difficulty = 'easy', trajectory_id = [], modality = ['image', 'depth'], camera_name = ['lcam_fish', 'lcam_front'], transform = transform) +dataset = ta.create_image_dataset(env = ['ArchVizTinyHouseDay'], difficulty = 'easy', trajectory_id = [], modality = ['image', 'depth'], camera_name = ['lcam_fish', 'lcam_front'], transform = transform) dataloader = DataLoader(dataset, batch_size = 3, shuffle = True, num_workers = 0) @@ -104,37 +104,3 @@ if i_batch == 5: break - -#################### -# Ask for all the data from an environment. -#################### -dataset = ta.create_image_dataset(env = 'ConstructionSite') - -dataloader = DataLoader(dataset, batch_size = 3, shuffle = True, num_workers = 0) - -# Show a few images. - -for i_batch, sample_batched in enumerate(dataloader): - print(i_batch, sample_batched['lcam_fish']['image_0'].size()) - - # Show the batch side by side. - import cv2 - import numpy as np - imgs1 = sample_batched['lcam_fish']['image_0'].numpy() - imgs2 = sample_batched['lcam_fish']['image_1'].numpy() - - # Move the color channel to the end. - imgs1 = np.transpose(imgs1, (0, 2, 3, 1)) - imgs2 = np.transpose(imgs2, (0, 2, 3, 1)) - - - img = np.concatenate((imgs1[0], imgs1[1], imgs1[2]), axis = 1) - img = np.concatenate((img, np.concatenate((imgs2[0], imgs2[1], imgs2[2]), axis = 1)), axis = 0) - img = cv2.resize(img, (0,0), fx = 0.5, fy = 0.5) - img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) - - cv2.imshow('image transform', img) - cv2.waitKey(0) - - if i_batch == 5: - break \ No newline at end of file diff --git a/examples/download_config.yaml b/examples/download_config.yaml index 88007f5..7f16e9e 100644 --- a/examples/download_config.yaml +++ b/examples/download_config.yaml @@ -1,7 +1,7 @@ # Config specifying which trajectories to download from the remote. -env: ['AmericanDiner'] -difficulty: ['easy'] -modality: ['image', 'depth'] -camera_name: ['lcam_fish', 'lcam_front'] -unzip: True + env: ['ArchVizTinyHouseDay'] + difficulty: ['easy'] + modality: ['image', 'depth', 'seg', 'imu'] + camera_name: ['lcam_fish', 'lcam_front', 'lcam_left', 'lcam_right', 'lcam_back', 'lcam_top', 'lcam_bottom'] + unzip: True diff --git a/examples/download_example.py b/examples/download_example.py index b9f35c8..e640996 100644 --- a/examples/download_example.py +++ b/examples/download_example.py @@ -13,25 +13,21 @@ import tartanair as ta # Create a TartanAir object. -tartanair_data_root = '/media/yoraish/overflow/data/tartanair-v2' -azure_token = "?sv=2021-10-04&st=2023-03-01T16%3A34%3A50Z&se=2023-03-30T15%3A34%3A00Z&sr=c&sp=rl&sig=LojCTa60TcA9ApMiMofliedxujeX0AOZdoC3O5u5cxg%3D" - -ta.init(tartanair_data_root, azure_token) - -# Download a trajectory. -# env = [ -# "HQWesternSaloonExposure", -# "ModularNeighborhoodIntExt", -# "PolarSciFiExposure", -# "PrisonExposure", -# "RuinsExposure", -# "TerrainBlendingExposure", -# "UrbanConstructionExposure", -# "VictorianStreetExposure", -# "WaterMillDayExposure", -# "WaterMillNightExposure" -# ] -# ta.download(env = "OldScandinaviaExposure", difficulty = ['easy', 'hard'], trajectory_id = ["P004", "P005", "P003", "P008"], modality = ['image', 'depth'], camera_name = ['lcam_front', 'lcam_right', 'lcam_back', 'lcam_left', 'lcam_top', 'lcam_bottom']) +tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' + +ta.init(tartanair_data_root) + +# Download data from following environments. +env = [ "Prison", + "Ruins", + "UrbanConstruction", +] + +ta.download(env = env, + difficulty = ['easy', 'hard'], + modality = ['image', 'depth'], + camera_name = ['lcam_front', 'lcam_right', 'lcam_back', 'lcam_left', 'lcam_top', 'lcam_bottom'], + unzip = True) # Can also download via a yaml config file. ta.download(config = 'download_config.yaml') diff --git a/examples/evaluator_example.py b/examples/evaluator_example.py index fc994a7..8690151 100644 --- a/examples/evaluator_example.py +++ b/examples/evaluator_example.py @@ -14,13 +14,13 @@ import tartanair as ta # Create a TartanAir object. -tartanair_data_root = '/media/yoraish/overflow/data/tartanair-v2' - +tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' + # Initialize the toolbox. ta.init(tartanair_data_root) # Create an example trajectory. This is a noisy version of the ground truth trajectory. -env = 'AbandonedCableExposure' +env = 'ArchVizTinyHouseDay' difficulty = 'easy' trajectory_id = 'P002' camera_name = 'lcam_front' @@ -31,7 +31,16 @@ # Get the evaluation results. plot_out_path = "evaluator_example.png" -results = ta.evaluate_traj(est_traj, env, difficulty, trajectory_id, camera_name, enforce_length = True, plot = True, plot_out_path = plot_out_path, do_scale = True, do_align = True) +results = ta.evaluate_traj(est_traj, + env = env, + difficulty = difficulty, + trajectory_id = trajectory_id, + camera_name = camera_name, + enforce_length = True, + plot = True, + plot_out_path = plot_out_path, + do_scale = True, + do_align = True) # Optionally pass the ground truth trajectory directly to the evaluation function. results = ta.evaluate_traj(est_traj, gt_traj = gt_traj, enforce_length = True, plot = True, plot_out_path = plot_out_path, do_scale = True, do_align = True) \ No newline at end of file diff --git a/examples/get_traj_example.py b/examples/get_traj_example.py index 84f94e3..046edd8 100644 --- a/examples/get_traj_example.py +++ b/examples/get_traj_example.py @@ -14,12 +14,15 @@ import tartanair as ta # Create a TartanAir object. -tartanair_data_root = '/media/yoraish/overflow/data/tartanair-v2' +tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' ta.init(tartanair_data_root) # List available trajectories. -traj = ta.get_traj_np(env = 'ConstructionSite', difficulty = 'easy', trajectory_id = "P000", camera_name = 'lcam_front') +traj = ta.get_traj_np(env = 'ArchVizTinyHouseDay', + difficulty = 'easy', + trajectory_id = "P000", + camera_name = 'lcam_front') print(traj.shape) np.set_printoptions(precision=3, suppress=True) print(traj[0:10, :]) \ No newline at end of file diff --git a/examples/iterator_example.py b/examples/iterator_example.py index 5208f68..3940594 100644 --- a/examples/iterator_example.py +++ b/examples/iterator_example.py @@ -13,12 +13,16 @@ import tartanair as ta # Create a TartanAir object. -tartanair_data_root = '/media/yoraish/overflow/data/tartanair-v2' - +tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' + ta.init(tartanair_data_root) # List available trajectories. -ta_iterator = ta.iterator(env = ['PolarSciFiExposure'], difficulty = 'easy', trajectory_id = [], modality = 'image', camera_name = ['lcam_left']) +ta_iterator = ta.iterator(env = ['ArchVizTinyHouseDay'], + difficulty = 'easy', + trajectory_id = [], + modality = 'image', + camera_name = ['lcam_left']) for i in range(100): print(next(ta_iterator)) \ No newline at end of file diff --git a/examples/list_example.py b/examples/list_example.py index d41be16..09629b1 100644 --- a/examples/list_example.py +++ b/examples/list_example.py @@ -13,8 +13,8 @@ import tartanair as ta # Create a TartanAir object. -tartanair_data_root = '/media/yoraish/overflow/data/tartanair-v2' - +tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' + ta.init(tartanair_data_root) # List available trajectories. diff --git a/examples/visualization_example.py b/examples/visualization_example.py index 5d39ab2..4e233e4 100644 --- a/examples/visualization_example.py +++ b/examples/visualization_example.py @@ -13,9 +13,13 @@ import tartanair as ta # Create a TartanAir object. -tartanair_data_root = '/media/yoraish/overflow/data/tartanair-v2_training_data' - +tartanair_data_root = '/my/path/to/root/folder/for/tartanair-v2' + ta.init(tartanair_data_root) # List available trajectories. -ta.visualize('ConstructionSite', difficulty='easy', trajectory_id = 'P000', modality = ['image', 'depth', 'seg'], camera_name = ['lcam_front', 'lcam_right', 'lcam_back', 'lcam_left']) \ No newline at end of file +ta.visualize('ArchVizTinyHouseDay', + difficulty='easy', + trajectory_id = 'P000', + modality = ['image', 'depth', 'seg'], + camera_name = ['lcam_front', 'lcam_right', 'lcam_back', 'lcam_left']) \ No newline at end of file diff --git a/tartanair/customizer.py b/tartanair/customizer.py index b10371c..591d17e 100644 --- a/tartanair/customizer.py +++ b/tartanair/customizer.py @@ -21,13 +21,15 @@ # Local imports. from .tartanair_module import TartanAirModule +from .reader import TartanAirImageReader # Image resampling. -from .image_resampling.image_sampler import SixPlanarNumba +from .image_resampling.image_sampler import SixPlanarNumba, SixPlanarTorch from .image_resampling.mvs_utils.camera_models import Pinhole, DoubleSphere, LinearSphere, Equirectangular #, PinholeRadTanFast, EUCM from .image_resampling.mvs_utils.shape_struct import ShapeStruct -from .image_resampling.image_sampler.blend_function import BlendBy2ndOrderGradTorch +from .image_resampling.mvs_utils.ftensor import FTensor +from .image_resampling.image_sampler.blend_function import BlendBy2ndOrderGradTorch, BlendBy2ndOrderGradOcv class TartanAirCustomizer(TartanAirModule): def __init__(self, tartanair_data_root): @@ -136,9 +138,10 @@ def customize(self, env, difficulty = [], trajectory_id = [], modality = [], new # Some mappings between attributes and parameters. - modality_to_reader = {"image": self.read_rgb, "depth": self.read_dist, "seg": self.read_seg} + self.reader = TartanAirImageReader() + modality_to_reader = {"image": self.reader.read_bgr, "depth": self.reader.read_dist, "seg": self.reader.read_seg} modality_to_interpolation = {"image": "linear", "seg": "nearest", "depth": "blend"} - modality_to_writer = {"image": self.write_as_is, "seg": self.write_as_is, "depth": self.write_float_depth} + modality_to_writer = {"image": self.reader.write_as_is, "seg": self.reader.write_as_is, "depth": self.reader.write_float_depth} ############################### # Enumerate the trajectories. @@ -293,10 +296,14 @@ def customize(self, env, difficulty = [], trajectory_id = [], modality = [], new def sample_image_worker(self, argslist): frame_ix, new_cam_model_object, R_raw_new, modality, new_cam_model_name, cam_name, side_to_frame_gfps, new_data_dir_path, modality_to_reader, modality_to_interpolation, modality_to_writer = argslist - sampler = SixPlanarNumba(new_cam_model_object.fov_degree, new_cam_model_object, R_raw_new) + # sampler = SixPlanarNumba(new_cam_model_object.fov_degree, new_cam_model_object, R_raw_new) + R_raw_new = FTensor(R_raw_new, f0='raw', f1='fisheye', rotation=True) # TODO: double check + sampler = SixPlanarTorch(new_cam_model_object, R_raw_new) + # sampler = SixPlanarNumba(new_cam_model_object, R_raw_new) + sampler.device = self.device create_figures = False - + raw_images = {} for side in ['front', 'back', 'left', 'right', 'top', 'bottom']: # Revert to below: @@ -318,13 +325,13 @@ def sample_image_worker(self, argslist): cv2.putText(img, side, (280, 320), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 2) raw_images[side] = img - + # Resample the six raw images to the new camera model. if modality_to_interpolation[modality] == "blend": blend_func = BlendBy2ndOrderGradTorch(0.01) # hard code new_image, new_image_valid_mask = sampler.blend_interpolation(raw_images, blend_func, invalid_pixel_value=0) else: - new_image, new_image_valid_mask = sampler(raw_images, interpolation= modality_to_interpolation[modality]) + new_image, new_image_valid_mask = sampler(raw_images, interpolation= modality_to_interpolation[modality], invalid_pixel_value=0) # Write the new image. out_fn = str(frame_ix).zfill(6) + "_{}_{}_{}.png".format(cam_name, modality, new_cam_model_name) @@ -390,8 +397,8 @@ def check_six_images_exist(self, env = [], difficulty = [], trajectory_id = [], # Iterate trajectory. if not trajectory_id: diff_path = os.path.join(self.data_root, env_folder, difficulty_folder) - env_trajectory_id = os.listdir(diff_path) - for traj_id_folder in env_trajectory_id: + trajectory_id = os.listdir(diff_path) + for traj_id_folder in trajectory_id: # Iterate modality. if not modality: @@ -415,101 +422,3 @@ def check_six_images_exist(self, env = [], difficulty = [], trajectory_id = [], print("Success: {env_folder}/{difficulty_folder}/{traj_id_folder} All files are available for {modality} on {raw_side} side.".format(env_folder = env_folder, difficulty_folder = difficulty_folder, traj_id_folder = traj_id_folder, modality = modality, raw_side = raw_side)) return True - - ############################### - # Data enumeration. - ############################### - def enumerate_trajs(self, data_folders = ['Data_easy','Data_hard']): - ''' - Return a dict: - res['env0']: ['Data_easy/P000', 'Data_easy/P001', ...], - res['env1']: ['Data_easy/P000', 'Data_easy/P001', ...], - ''' - env_folders = os.listdir(self.data_root) - env_folders = [ee for ee in env_folders if os.path.isdir(os.path.join(self.data_root, ee))] - env_folders.sort() - trajlist = {} - for env_folder in env_folders: - env_dir = os.path.join(self.data_root, env_folder) - trajlist[env_folder] = [] - for data_folder in data_folders: - datapath = os.path.join(env_dir, data_folder) - if not os.path.isdir(datapath): - continue - - trajfolders = os.listdir(datapath) - trajfolders = [ os.path.join(data_folder, tf) for tf in trajfolders if tf[0]=='P' ] - trajfolders.sort() - trajlist[env_folder].extend(trajfolders) - return trajlist - - ############################### - # Data reading, writing, and processing. - ############################### - def depth_to_dist(self, depth): - ''' - assume: fov = 90 on both x and y axes, and optical center is at image center. - ''' - # import ipdb;ipdb.set_trace() - if self.depth_shape is None or \ - depth.shape != self.depth_shape or \ - self.conv_matrix is None: # only calculate once if the depth shape has not changed - hh, ww = depth.shape - f = ww/2 - wIdx = np.linspace(0, ww - 1, ww, dtype=np.float32) + 0.5 - ww/2 # put the optical center at the middle of the image - hIdx = np.linspace(0, hh - 1, hh, dtype=np.float32) + 0.5 - hh/2 # put the optical center at the middle of the image - u, v = np.meshgrid(wIdx, hIdx) - dd = np.sqrt(u * u + v * v + f * f)/f - self.conv_matrix = dd - self.depth_shape = depth.shape - disp = self.conv_matrix * depth - return disp - - def ocv_read(self, fn ): - image = cv2.imread(fn, cv2.IMREAD_UNCHANGED) - assert image is not None, \ - f'{fn} read error. ' - return image - - def read_rgb(self, fn ): - return self.ocv_read(fn) - - def read_dist(self, fn ): # read a depth image and convert it to distance - depth = self.read_dep(fn) - return self.depth_to_dist(depth) - - def read_dep(self, fn ): - image = self.ocv_read(fn) - depth = np.squeeze( image.view(' None: @@ -77,6 +67,7 @@ def download(self, filelist, destination_path): - str: A message indicating success or failure. """ + from botocore.exceptions import NoCredentialsError target_filelist = [] for source_file_name in filelist: target_file_name = join(destination_path, source_file_name.replace('/', '_')) @@ -115,64 +106,8 @@ class TartanAirDownloader(TartanAirModule): def __init__(self, tartanair_data_root): super().__init__(tartanair_data_root) - # The modalities that have a camera associated with them and that we'll download the pose file along with. - self.cam_modalities = ['image', 'depth', 'seg'] # the modalities that support all camera names - self.flow_camlist = ['lcam_front'] # valid camera name for the flow modality - self.downloader = CloudFlareDownloader() - def check_env_valid(self, envlist): - invalid_list = [] - for env in envlist: - if not env in self.env_names: - invalid_list.append(env) - - if len(invalid_list) == 0: - return True - - print_error(f"The following envs are invalid: {invalid_list}") - print_warn(f"The available envs are: {self.env_names}") - return False - - def check_modality_valid(self, modlist): - invalid_list = [] - for mod in modlist: - if not mod in self.modality_names: - invalid_list.append(mod) - - if len(invalid_list) == 0: - return True - - print_error(f"The following modalities are invalid: {invalid_list}") - print_warn(f"The available modalities are: {self.modality_names}") - return False - - def check_camera_valid(self, camlist): - invalid_list = [] - for cam in camlist: - if not cam in self.camera_names: - invalid_list.append(cam) - - if len(invalid_list) == 0: - return True - - print_error(f"The following camera names are invalid: {invalid_list}") - print_warn(f"The available camera names are: {self.camera_names}") - return False - - def check_difficulty_valid(self, difflist): - invalid_list = [] - for diff in difflist: - if not diff in self.difficulty_names: - invalid_list.append(diff) - - if len(invalid_list) == 0: - return True - - print_error(f"The following difficulties are invalid: {invalid_list}") - print_warn(f"The available difficulties are: {self.difficulty_names}") - return False - def generate_filelist(self, envs, difficulties, modalities, camera_names): ''' Return a list of zipfiles to be downloaded @@ -189,21 +124,10 @@ def generate_filelist(self, envs, difficulties, modalities, camera_names): envstr = env + '/' for difficulty in difficulties: diffstr = envstr + 'Data_' + difficulty + '/' - for mod in modalities: - if mod in self.cam_modalities: - for camname in camera_names: - zipfile = diffstr + mod + '_' + camname + '.zip' - zipfilelist.append(zipfile) - elif mod == 'flow': - for camname in camera_names: - if camname in self.flow_camlist: - zipfile = diffstr + mod + '_' + camname + '.zip' - zipfilelist.append(zipfile) - else: - print_warn("Warn: flow modality doesn't have {}! We only have flow for {}".format(camname, self.flow_camlist)) - else: # for lidar and imu - zipfile = diffstr + mod + '.zip' - zipfilelist.append(zipfile) + folderlist = self.compile_modality_and_cameraname(modalities, camera_names) + zipfiles = [diffstr + fl + '.zip' for fl in folderlist] + zipfilelist.extend(zipfiles) + return zipfilelist def doublecheck_filelist(self, filelist, gtfile=''): diff --git a/tartanair/evaluator.py b/tartanair/evaluator.py index 0bb17a9..331335e 100644 --- a/tartanair/evaluator.py +++ b/tartanair/evaluator.py @@ -44,7 +44,6 @@ def evaluate_traj(self, if gt_traj is None: raise ValueError("Please pass a ground truth trajectory or a trajectory specification (env, difficulty, trajectory_id, camera_name) to the evaluation method.") - # Make sure that the input is a numpy array with a correct shape. if isinstance(est_traj, list): print(Fore.CYAN + "Warning: Converting the estimated trajectory (list) to a numpy array." + Style.RESET_ALL) diff --git a/tartanair/image_resampling/image_sampler b/tartanair/image_resampling/image_sampler index aca2b45..620f75b 160000 --- a/tartanair/image_resampling/image_sampler +++ b/tartanair/image_resampling/image_sampler @@ -1 +1 @@ -Subproject commit aca2b45cc6b9d04a6c9c45bbc9af9311d5f164e3 +Subproject commit 620f75b06b83a43a426e01fa7922fb14a4ccd584 diff --git a/tartanair/lister.py b/tartanair/lister.py index 19442a9..6e48e5f 100644 --- a/tartanair/lister.py +++ b/tartanair/lister.py @@ -11,6 +11,7 @@ # Local imports. from .tartanair_module import TartanAirModule +from os.path import isdir, join class TartanAirLister(TartanAirModule): def __init__(self, tartanair_data_root): @@ -22,6 +23,7 @@ def list_envs(self, verbose = True): ''' # Get the local environments. local_envs = os.listdir(self.tartanair_data_root) + local_envs = [pp for pp in local_envs if isdir(join(self.tartanair_data_root, pp))] # Get the remote environments. remote_envs = self.env_names diff --git a/tartanair/reader.py b/tartanair/reader.py index 9eec446..6cb2a52 100644 --- a/tartanair/reader.py +++ b/tartanair/reader.py @@ -88,7 +88,7 @@ def read_flow(self, flowpath, scale = 1): return flownp def read_dist(self, fn ): # read a depth image and convert it to distance - depth = self.read_dep(fn) + depth = self.read_depth(fn) return self.depth_to_dist(depth) def depth_to_dist(self, depth): @@ -109,6 +109,78 @@ def depth_to_dist(self, depth): disp = self.conv_matrix * depth return disp + + ############################### + # Data reading, writing, and processing. + ############################### + def depth_to_dist(self, depth): + ''' + assume: fov = 90 on both x and y axes, and optical center is at image center. + ''' + # import ipdb;ipdb.set_trace() + if self.depth_shape is None or \ + depth.shape != self.depth_shape or \ + self.conv_matrix is None: # only calculate once if the depth shape has not changed + hh, ww = depth.shape + f = ww/2 + wIdx = np.linspace(0, ww - 1, ww, dtype=np.float32) + 0.5 - ww/2 # put the optical center at the middle of the image + hIdx = np.linspace(0, hh - 1, hh, dtype=np.float32) + 0.5 - hh/2 # put the optical center at the middle of the image + u, v = np.meshgrid(wIdx, hIdx) + dd = np.sqrt(u * u + v * v + f * f)/f + self.conv_matrix = dd + self.depth_shape = depth.shape + disp = self.conv_matrix * depth + return disp + + # def ocv_read(self, fn ): + # image = cv2.imread(fn, cv2.IMREAD_UNCHANGED) + # assert image is not None, \ + # f'{fn} read error. ' + # return image + + # def read_rgb(self, fn ): + # return self.ocv_read(fn) + + # def read_dist(self, fn ): # read a depth image and convert it to distance + # depth = self.read_dep(fn) + # return self.depth_to_dist(depth) + + # def read_dep(self, fn ): + # image = self.ocv_read(fn) + # depth = np.squeeze( image.view('