Merge pull request #71 from castacks/pr-branch

Pr branch

Author: Amigoshan

docs/examples.rst

@@ -73,6 +73,23 @@ Multi-Thread Download via Python API
                              unzip = True,
                              num_workers = 8)
 
+Download TartanGround
+............................
+
+.. code-block:: python
+
+    import tartanair as ta
+
+    # Initialize TartanAir.
+    tartanground_data_root = '/my/path/to/root/folder/for/tartanair-v2'
+    ta.init(tartanground_data_root)
+
+    # Download a trajectory.
+    ta.download_ground(env = ['OldTownSummer', 'DesertGasStation'], 
+                version = ['v1', 'v2', 'v3_anymal'], # available versions are: 'v1', 'v2', 'v3_anymal'
+                modality = ['image', 'depth', 'seg', 'lidar', 'imu'], # available modalities are: image', 'depth', 'seg', 'imu', 'lidar'
+                camera_name = ['lcam_front', 'lcam_bottom'], # available cameras are: 'lcam_front', 'lcam_left', 'lcam_right', 'lcam_back', 'lcam_top', 'lcam_bottom'
+                unzip = True) # unzip files autonomously after download
 
 Customization Example
 -------------------------------------
@@ -295,11 +312,36 @@ The tool supports sampling dense correspondence between any combination of pinho
 
 Given a pair of RGB and depth image cubes and two direction, the tool will compute dense correspondense represented as optical flow and a occlusion mask signaling is the pixel directly observable.
 
-Please refer to flow_sampling_example.py for a complete example.
+.. code-block:: python
+
+    import tartanair as ta
+
+    # For help with rotations.
+    from scipy.spatial.transform import Rotation
+
+    # Initialize TartanAir.
+    tartanair_data_root = 'your/path/to_tav2'
+    ta.init(tartanair_data_root)
+
+    # Customize the dataset.
+    ta.customize_flow(env = 'ArchVizTinyHouseDay',
+                difficulty = 'easy',
+                trajectory_id = ['P000'],
+                cam_sides=["left", "right"],
+                num_workers = 4,
+                frame_sep = 1,
+                device = "cuda") # or cpu
+
+Please refer to flow_resampling_example.py for a complete example.
 
 .. image:: images/flow_sampling/flow_sampling_example.png
    :alt: Algorithm
    :align: center
    :width: 80%
 
-To learn more about how the resampling happens, see `Flow Sampling <flow_sampling.html>`_.
+To learn more about how the resampling happens, see flow_sampling. 
+
+.. toctree::
+   :maxdepth: 1
+
+   flow_sampling

docs/modalities.rst

@@ -69,6 +69,18 @@ The fisheye and Panorama data are sampled from the raw pinhole data, thus contai
 Optical flow
 --------------
 
+Due to the high serving demand of the dataset, we are moving away from pre-computing the flow labels for download, as they can simply be generated from pose, intrinsics, and depth locally. Please refer to examples/flow_resampling_example.py for the code to generate optical flow from the raw data.
+
+Also, we are changing the storage format due to increased precision of flow models. For a pair of images, the generated flow information is stored as a npz file containing the following fields:
+
+    - flow_fwd/bwd: x,y coordinate of the flow in float32 format
+    - fov_mask_fwd/bwd: boolean mask indicating if the pixel have a valid projection in the forward/backward image, i.e., is it out of the image space
+    - covisible_mask_fwd/bwd: boolean mask indicating if the pixel is covisible in the forward/backward image, i.e., is it occluded by other objects in the scene
+
+
+Old Optical flow format
+-----------------------
+
 Same as TartanAir V1, the optical flow is calculated for the static environments by image warping, using the camera pose and depth images. The biggest upgrades are that we accelerate the code by a Cuda implementation and provide tools for generating optical flow across any type of camera model (e.g. between pinhole and fisheye). For the convenience of the downloading, we compress the optical flow to use 8-bit representation. The decoding code is as follows.
 
 .. code-block:: python

examples/flow_resampling_example.py

@@ -0,0 +1,20 @@
+'''
+Author: Yuchen Zhang
+Date: 2024-11-19
+
+Example script for sample new correspondence from the TartanAir dataset.
+'''
+import tartanair as ta
+
+# Initialize TartanAir.
+tartanair_data_root = 'your/path/to_tav2'
+ta.init(tartanair_data_root)
+
+# Customize the dataset.
+ta.customize_flow(env = 'ArchVizTinyHouseDay',
+             difficulty = 'easy',
+             trajectory_id = ['P000'],
+             camera_name=["lcam_left", "lcam_back"],
+             num_workers = 4,
+             frame_sep = 1,
+             device = "cuda") # or cpu