Merge branch 'main' into multi_thread_download

.gitignore

@@ -148,3 +148,4 @@ docs/make.bat
 docs/Makefile
 examples/ta_data_spec.txt
 *.txt
+examples/*.png

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
+

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,19 +224,34 @@ 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'
     gt_traj = ta.get_traj_np(env, difficulty, trajectory_id, camera_name)
     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
 -------------------------------------

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

docs/modalities.rst

@@ -0,0 +1,7 @@
+Modalities
+=====================================
+
+Images
+Depth
+Flow
+LiDAR

examples/customization_example.py

@@ -7,22 +7,21 @@
 
 # General imports.
 import sys
-import numpy as np
 from scipy.spatial.transform import Rotation
 
 # Local imports.
 sys.path.append('..')
 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