Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Utility changes + IMU gyro reading support #29

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,14 @@
# TartanAir

Hello and welcome to the official TartanAir repository. This repository includes a set of tools that complement the [TartanAir Dataset](https://www.tartanair.org/)).
Hello and welcome to the official TartanAir repository. This repository includes a set of tools that complement the [TartanAir Dataset](https://www.tartanair.org/)).

## Installation Instructions
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for starting this PR, Akash!

Could you please change this to distinguish between pip installation and local installation?

Suggested change
## Installation Instructions
## Installation Instructions
`pip install tartanair`
## Local Installation Instructions


```
git clone --recursive https://github.com/castacks/tartanairpy.git
pip install . --verbose
```
if you forget to clone using `--recursive` then you can update all submodules as follows
```
git submodule update --init --recursive
```
45 changes: 27 additions & 18 deletions tartanair/dataloader.py
Original file line number Diff line number Diff line change
@@ -8,7 +8,7 @@

'''

# TODO(yoraish): there is a notation discrepancy between 'rgb' and 'image'. Should fix this and probably stick with 'image' as this is the name used in the dataset.
# TODO(yoraish): there is a notation discrepancy between 'rgb' and 'image'. Should fix this and probably stick with 'image' as this is the name used in the dataset.
# TODO(yoraish): support naming between 'imu_acc' and 'imu_gyro' etc. Currently only 'imu' is supported and is mapped to 'imu_acc'.

# General imports.
@@ -22,7 +22,7 @@

class TartanAirDataLoader(TartanAirModule):
'''
The TartanAirDataset class contains the _information_ about the TartanAir dataset, and implements no functionality. All functionalities are implemented in inherited classes like the TartanAirDownloader, and the interface is via the TartanAir class.
The TartanAirDataset class contains the _information_ about the TartanAir dataset, and implements no functionality. All functionalities are implemented in inherited classes like the TartanAirDownloader, and the interface is via the TartanAir class.
'''
def __init__(self, tartanair_data_root):
# Call the parent class constructor.
@@ -32,7 +32,7 @@ def __init__(self, tartanair_data_root):
self.modality_name_remaps = {
'rgb': 'rgb',
'image': 'rgb',
'imu': 'imu_acc',
'imu_gyro': 'imu_gyro',
'imu_acc': 'imu_acc',
'depth': 'depth',
'lidar': 'lidar',
@@ -48,6 +48,7 @@ def __init__(self, tartanair_data_root):
'seg': 1,
'flow': 1,
'pose': 1,
'imu_gyro': 10,
'imu_acc': 10,
'lidar': 1,
}
@@ -58,24 +59,25 @@ def __init__(self, tartanair_data_root):
'seg': [640, 640],
'flow': [640, 640],
'pose': [7],
'imu_gyro': [3],
'imu_acc': [3],
'lidar': [3],
}

def get_data_cacher(self,
env,
difficulty = None,
trajectory_id = None,
modality = None,
camera_name = None,
def get_data_cacher(self,
env,
difficulty = None,
trajectory_id = None,
modality = None,
camera_name = None,
new_image_shape_hw = [640, 640],
subset_framenum = 56, # <--- Note in the docs that this is an upper bound on the batch size.
seq_length = 1, # This can also be a dictionary, mapping each modality name to a sequence length.
seq_stride = 1,
frame_skip = 0,
batch_size=1,
num_workers=1,
shuffle=False,
shuffle=False,
verbose=False,):

'''
@@ -140,6 +142,13 @@ def get_data_cacher(self,
print("Building data cacher for env {}...".format(env_name))
# If no difficulty was given, then use all difficulties.
if not difficulty:
difficulty = []
for diff in ['easy', 'hard']:
tartanair_diff_path = os.path.join(self.tartanair_data_root, env_name, 'Data_' + diff)
if os.path.isdir(tartanair_diff_path):
difficulty.append(tartanair_diff_path)
else:
raise ValueError("No data found for env {} and difficulty {}. Please check the path to the TartanAir dataset.".format(env_name, diff))
available_diffs = [diff for diff in ['easy', 'hard'] if os.path.isdir(os.path.join(self.tartanair_data_root, env_name, 'Data_' + diff))]
print([os.path.join(self.tartanair_data_root, env_name, 'Data_' + diff) for diff in ['easy', 'hard']])
print(Fore.GREEN + "WARNING: No difficulty was specified for env {}. Defaulting to all available difficulties: {}".format(env_name, available_diffs), Style.RESET_ALL)
@@ -158,9 +167,9 @@ def get_data_cacher(self,
if traj_id not in available_traj_ids:
print(Fore.RED + "WARNING: Trajectory id {} was specified for env {} and difficulty {}, but it is not available. It is skipped.".format(traj_id, env_name, diff), Style.RESET_ALL)
available_traj_ids = [traj_id for traj_id in available_traj_ids if traj_id in trajectory_id]

for traj_id in available_traj_ids:
# Build the data entry.
# Build the data entry.

# Read the data specification file and concatenate it to the data spec file.
with open(data_spec_fpath, 'a') as f:
@@ -201,9 +210,9 @@ def get_data_cacher(self,

print(config)
# Create the data loader from the config.
trainDataloader = MultiDatasets(config,
'local',
batch= batch_size,
trainDataloader = MultiDatasets(config,
'local',
batch= batch_size,
workernum= num_workers,
shuffle= shuffle,
verbose= verbose)
@@ -233,7 +242,7 @@ def traj_spec_file_to_pseudo_traj_spec_list(self, traj_spec_fpath, pseudo_traj_l

else:
pseudo_traj_spec_list.append(numless_header + str(len(lines) - i) + '\n')

# Append the line.
pseudo_traj_spec_list.append(lines[i])

@@ -269,13 +278,13 @@ def add_modality_entries(self, data_entry, modality, camera_name, new_image_shap
modality_entry['cacher_size'] = new_image_shape_hw
else: # 'lidar', 'imu', 'pose', etc.
modality_entry['cacher_size'] = self.modality_default_cacher_size[mod]

modality_entry['length'] = seq_length[mod]
modality_entry['subset_framenum'] = subset_framenum

# Add the modality entry to the data entry.
data_entry['modality'][type] = modality_entry
data_entry['modality'][type]['type'] = type
data_entry['modality'][type]['type'] = type


def get_available_trajectory_ids(self, env_name, diff):