diff --git a/fog_x/dataset.py b/fog_x/dataset.py index f20d343..016cae7 100644 --- a/fog_x/dataset.py +++ b/fog_x/dataset.py @@ -6,6 +6,7 @@ import numpy as np import polars import pandas +from tqdm import tqdm from fog_x.database import ( DatabaseConnector, @@ -432,12 +433,13 @@ def _build_rtx_episodes_from_tfds_builder( data_type = builder.info.features["steps"] - for tf_episode in ds: - logger.info(tf_episode) + pbar_len = len(ds) + for i, tf_episode in tqdm(enumerate(ds), position=0, desc="Episodes", total=pbar_len): + logger.debug(tf_episode) fog_episode = self.new_episode( metadata=additional_metadata, ) - for step in tf_episode["steps"]: + for step in tqdm(tf_episode["steps"], position=1, desc="Steps"): ret = self._load_rtx_step_data_from_tf_step( step, data_type, ) diff --git a/pyproject.toml b/pyproject.toml index 6d41820..14799dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ "opencv-python", "sqlalchemy==1.4.51", "smart_open", + "tqdm", ] description = "An Efficient and Scalable Data Collection and Management Framework For Robotics Learning" readme = {file = "README.md", content-type = "text/markdown"}