Skip to content

Replaying

FolderReplayer

Source code in mbodied/data/replaying.py 
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
class FolderReplayer:
    def __init__(self, path: str) -> None:
        """Initialize the FolderReplayer class with the given path.

        Args:
            path (str): Path to the folder containing HDF5 files.
        """
        self.path = path

    def __iter__(self):
        """Iterate through the HDF5 files in the folder."""
        for f in os.listdir(self.path):
            if f.endswith(".h5"):
                r = Replayer(f"{self.path}/{f}")
                for _i, sample in enumerate(r):
                    observation = sample[0]
                    action = sample[1]
                    state = sample[2] if len(sample) > 2 else None
                    image = np.asarray(observation["image"])
                    instruction = observation["instruction"]
                    yield {
                        "observation": {"image": image, "instruction": instruction},
                        "action": action,
                        "state": state,
                    }

__init__(path)

Initialize the FolderReplayer class with the given path.

Parameters:

Name Type Description Default
path str

Path to the folder containing HDF5 files.

 required
Source code in mbodied/data/replaying.py 
293
294
295
296
297
298
299
def __init__(self, path: str) -> None:
    """Initialize the FolderReplayer class with the given path.

    Args:
        path (str): Path to the folder containing HDF5 files.
    """
    self.path = path

__iter__()

Iterate through the HDF5 files in the folder.

Source code in mbodied/data/replaying.py 
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
def __iter__(self):
    """Iterate through the HDF5 files in the folder."""
    for f in os.listdir(self.path):
        if f.endswith(".h5"):
            r = Replayer(f"{self.path}/{f}")
            for _i, sample in enumerate(r):
                observation = sample[0]
                action = sample[1]
                state = sample[2] if len(sample) > 2 else None
                image = np.asarray(observation["image"])
                instruction = observation["instruction"]
                yield {
                    "observation": {"image": image, "instruction": instruction},
                    "action": action,
                    "state": state,
                }

Replayer

Replays datasets recorded by Recorder.

This class provides methods to read, process, and analyze HDF5 files recorded by the Recorder class.

Example

replayer = Replayer("data.h5") for sample in replayer: observation, action, state = sample ...

Source code in mbodied/data/replaying.py 
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
class Replayer:
    """Replays datasets recorded by Recorder.

    This class provides methods to read, process, and analyze HDF5 files recorded by the Recorder class.

    Example:
        replayer = Replayer("data.h5")
        for sample in replayer:
            observation, action, state = sample
            ...
    """

    def __init__(
        self,
        path: str,
        file_keys: List[str] = None,
        image_keys_to_save: List[str] = None,
    ) -> None:
        """Initialize the Replayer class with the given parameters.

        Args:
            path (str): Path to the HDF5 file.
            file_keys (List[str], optional): List of keys in the file. Defaults to None.
            image_keys_to_save (List[str], optional): List of image keys to save. Defaults to None.
        """
        self.path = path
        self.frames_path = f"{Path(self.path).parent}/{Path(self.path).stem}_frames"
        self.file = h5py.File(path, "r")

        self.file_keys = file_keys or list(self.file.keys())
        self._reorder_keys()
        self.image_keys_to_save = image_keys_to_save or []
        if self.image_keys_to_save:
            Path(self.frames_path).mkdir(parents=True, exist_ok=True)
        self.size = self._get_size()

    def _reorder_keys(self) -> None:
        """Reorder keys to have 'observation' first, 'action' last, and 'supervision' last but one."""
        if "observation" in self.file_keys:
            self.file_keys.remove("observation")
            self.file_keys = ["observation"] + self.file_keys
        if "action" in self.file_keys:
            self.file_keys.remove("action")
            self.file_keys = self.file_keys + ["action"]
        if "state" in self.file_keys:
            self.file_keys.remove("state")
            self.file_keys = self.file_keys + ["state"]
        if "supervision" in self.file_keys:
            self.file_keys.remove("supervision")
            self.file_keys = self.file_keys + ["supervision"]

    def _get_size(self) -> int:
        """Get the size attribute from the HDF5 file or infer it from the first key."""
        size = self.file.attrs.get("size")
        if size is None:
            logging.warning(
                "No 'size' attribute found in the HDF5 file. The number of samples will be inferred from the first key.",
            )
            try:
                key = self.file_keys[0]
                while isinstance(self.file[key], h5py.Group):
                    key = list(self.file[key].keys())[0]
                size = len(self.file[key])
            except Exception:
                logging.warning("No size attribute and empty dataset found. Setting size to 0.")
                size = 0
        return size

    def get_frames_path(self) -> str | None:
        """Get the path to the frames directory."""
        return self.frames_path

    def recursive_do(self, do: Callable, key="", prefix="", **kwargs) -> Any:
        """Recursively perform a function on each key in the HDF5 file.

        Args:
            do (Callable): Function to perform.
            key (str, optional): Key in the HDF5 file. Defaults to ''.
            prefix (str, optional): Prefix for the key. Defaults to ''.
            **kwargs: Additional arguments to pass to the function.

        Returns:
            Any: Result of the function.
        """
        full_key = f"{prefix}/{key}".strip("/")
        if full_key and isinstance(self.file[full_key], h5py.Dataset):
            return do(full_key, **kwargs)

        keys = self.file_keys if not full_key else self.file[full_key].keys()

        result = {}
        for sub_key in keys:
            result[sub_key] = self.recursive_do(do, key=sub_key, prefix=full_key, **kwargs)
        return result

    def get_unique_items(self, key: str) -> List[str]:
        """Get unique items for a given key.

        Args:
            key (str): Key in the HDF5 file.

        Returns:
            List[str]: List of unique items.
        """
        if self.file[key].dtype == string_dtype():
            return list({item.decode() for item in self.file[key][: self.size]})
        return list(set(self.file[key][: self.size]))

    def read_sample(self, index: int) -> Tuple[dict, ...]:
        """Read a sample from the HDF5 file at a given index.

        Args:
            index (int): Index of the sample.

        Returns:
            Tuple[dict, ...]: Tuple of dictionaries containing the sample data.
        """

        def read_do(full_key: str, index: int) -> Any:
            if index >= self.file[full_key].shape[0]:
                logging.warning(f"Index {index} is out of bounds for key {full_key}. Skipping...")
                return None
            value = self.file[full_key][index]
            if full_key in self.image_keys_to_save and len(value.shape) == 3:
                image_path = f"{self.frames_path}/{Path(full_key).name}_{index}.png"
                PILImage.fromarray(value).save(image_path)
            if self.file[full_key].dtype == string_dtype():
                if self.file[full_key].attrs.get("tuple_length") is not None:
                    return Sample.model_validate_json(value.decode()).unpack(to_dicts=True)
                try:
                    value = value.decode()
                except Exception as e:
                    logging.error(f"Error decoding string for key {full_key}: {e}")
                    value = value[0].decode()
            return value

        return tuple(self.recursive_do(read_do, key=key, index=index) for key in self.file_keys)

    def get_structure(self, key="", prefix="") -> dict:
        """Get the structure of the HDF5 file.

        Args:
            key (str, optional): Key in the HDF5 file. Defaults to ''.
            prefix (str, optional): Prefix for the key. Defaults to ''.

        Returns:
            dict: Structure of the HDF5 file.
        """

        def structure_do(full_key: str) -> tuple:
            return self.file[full_key][: self.size].shape, self.file[full_key][: self.size].dtype

        return self.recursive_do(structure_do, key, prefix)

    def pack_one(self, index: int) -> Sample:
        """Pack a single sample into a Sample object.

        Args:
            index (int): Index of the sample.

        Returns:
            Sample: Sample object.
        """
        return Sample(**{key: self.read_sample(index)[i] for i, key in enumerate(self.file_keys)})

    def pack(self) -> Sample:
        """Pack all samples into a Sample object with attributes being lists of samples.

        Returns:
            Sample: Sample object containing all samples.
        """
        return Sample.pack_from([self.pack_one(i) for i in range(self.size)])

    def sample(self, index: int | slice | None = None, n: int = 1) -> Sample:
        """Get a sample from the HDF5 file.

        Args:
            index (Optional[Union[int, slice]], optional): Index or slice of the sample. Defaults to None.
            n (int, optional): Number of samples to get. Defaults to 1.

        Returns:
            Sample: Sample object.
        """
        if index is None:
            index = np.random.randint(0, self.size, n)
        if isinstance(index, int):
            return self.pack_one(index)
        if isinstance(index, slice):
            return [self.read_sample(i) for i in range(index.start, index.stop, index.step)]
        return None

    def get_stats(self, key="", prefix="") -> dict:
        """Get statistics for a given key in the HDF5 file.

        Args:
            key (str, optional): Key in the HDF5 file. Defaults to ''.
            prefix (str, optional): Prefix for the key. Defaults to ''.

        Returns:
            dict: Statistics for the given key.
        """

        def stats_do(full_key: str) -> dict:
            if self.file[full_key].dtype == h5py.string_dtype():
                return {
                    "unique": len({item.decode() for item in self.file[full_key][: self.size]}),
                    "shape": self.file[full_key][: self.size].shape,
                    "dtype": self.file[full_key][: self.size].dtype,
                }
            data = self.file[full_key][: self.size]
            return {
                "min": np.min(data, axis=0),
                "max": np.max(data, axis=0),
                "mean": np.mean(data, axis=0),
                "std": np.std(data, axis=0),
                "median": np.median(data, axis=0),
                "shape": self.file[full_key][: self.size].shape,
                "dtype": self.file[full_key][: self.size].dtype,
            }

        return {key: self.recursive_do(stats_do, key, prefix)}

    def __iter__(self):
        """Iterate over the HDF5 file."""
        self.index = -1
        return self

    def __next__(self) -> Tuple[dict, ...]:
        """Get the next sample from the HDF5 file.

        Returns:
            Tuple[dict, ...]: Tuple of dictionaries containing the sample data.
        """
        if self.index >= self.size - 1:
            raise StopIteration
        self.index += 1
        return self.read_sample(self.index)

    def close(self) -> None:
        """Close the HDF5 file."""
        self.file.close()

__init__(path, file_keys=None, image_keys_to_save=None)

Initialize the Replayer class with the given parameters.

Parameters:

Name Type Description Default
path str

Path to the HDF5 file.

 required
file_keys List[str]

List of keys in the file. Defaults to None.

 None
image_keys_to_save List[str]

List of image keys to save. Defaults to None.

 None
Source code in mbodied/data/replaying.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(
    self,
    path: str,
    file_keys: List[str] = None,
    image_keys_to_save: List[str] = None,
) -> None:
    """Initialize the Replayer class with the given parameters.

    Args:
        path (str): Path to the HDF5 file.
        file_keys (List[str], optional): List of keys in the file. Defaults to None.
        image_keys_to_save (List[str], optional): List of image keys to save. Defaults to None.
    """
    self.path = path
    self.frames_path = f"{Path(self.path).parent}/{Path(self.path).stem}_frames"
    self.file = h5py.File(path, "r")

    self.file_keys = file_keys or list(self.file.keys())
    self._reorder_keys()
    self.image_keys_to_save = image_keys_to_save or []
    if self.image_keys_to_save:
        Path(self.frames_path).mkdir(parents=True, exist_ok=True)
    self.size = self._get_size()

__iter__()

Iterate over the HDF5 file.

Source code in mbodied/data/replaying.py 
255
256
257
258
def __iter__(self):
    """Iterate over the HDF5 file."""
    self.index = -1
    return self

__next__()

Get the next sample from the HDF5 file.

Returns:

Type Description
Tuple[dict, ...]

Tuple[dict, ...]: Tuple of dictionaries containing the sample data.
Source code in mbodied/data/replaying.py 
260
261
262
263
264
265
266
267
268
269
def __next__(self) -> Tuple[dict, ...]:
    """Get the next sample from the HDF5 file.

    Returns:
        Tuple[dict, ...]: Tuple of dictionaries containing the sample data.
    """
    if self.index >= self.size - 1:
        raise StopIteration
    self.index += 1
    return self.read_sample(self.index)

close()

Close the HDF5 file.

Source code in mbodied/data/replaying.py 
271
272
273
def close(self) -> None:
    """Close the HDF5 file."""
    self.file.close()

get_frames_path()

Get the path to the frames directory.

Source code in mbodied/data/replaying.py 
101
102
103
def get_frames_path(self) -> str | None:
    """Get the path to the frames directory."""
    return self.frames_path

get_stats(key='', prefix='')

Get statistics for a given key in the HDF5 file.

Parameters:

Name Type Description Default
key str

Key in the HDF5 file. Defaults to ''.

 ''
prefix str

Prefix for the key. Defaults to ''.

 ''

Returns:

Name Type Description
dict dict

Statistics for the given key.
Source code in mbodied/data/replaying.py 
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
def get_stats(self, key="", prefix="") -> dict:
    """Get statistics for a given key in the HDF5 file.

    Args:
        key (str, optional): Key in the HDF5 file. Defaults to ''.
        prefix (str, optional): Prefix for the key. Defaults to ''.

    Returns:
        dict: Statistics for the given key.
    """

    def stats_do(full_key: str) -> dict:
        if self.file[full_key].dtype == h5py.string_dtype():
            return {
                "unique": len({item.decode() for item in self.file[full_key][: self.size]}),
                "shape": self.file[full_key][: self.size].shape,
                "dtype": self.file[full_key][: self.size].dtype,
            }
        data = self.file[full_key][: self.size]
        return {
            "min": np.min(data, axis=0),
            "max": np.max(data, axis=0),
            "mean": np.mean(data, axis=0),
            "std": np.std(data, axis=0),
            "median": np.median(data, axis=0),
            "shape": self.file[full_key][: self.size].shape,
            "dtype": self.file[full_key][: self.size].dtype,
        }

    return {key: self.recursive_do(stats_do, key, prefix)}

get_structure(key='', prefix='')

Get the structure of the HDF5 file.

Parameters:

Name Type Description Default
key str

Key in the HDF5 file. Defaults to ''.

 ''
prefix str

Prefix for the key. Defaults to ''.

 ''

Returns:

Name Type Description
dict dict

Structure of the HDF5 file.
Source code in mbodied/data/replaying.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def get_structure(self, key="", prefix="") -> dict:
    """Get the structure of the HDF5 file.

    Args:
        key (str, optional): Key in the HDF5 file. Defaults to ''.
        prefix (str, optional): Prefix for the key. Defaults to ''.

    Returns:
        dict: Structure of the HDF5 file.
    """

    def structure_do(full_key: str) -> tuple:
        return self.file[full_key][: self.size].shape, self.file[full_key][: self.size].dtype

    return self.recursive_do(structure_do, key, prefix)

get_unique_items(key)

Get unique items for a given key.

Parameters:

Name Type Description Default
key str

Key in the HDF5 file.

 required

Returns:

Type Description
List[str]

List[str]: List of unique items.
Source code in mbodied/data/replaying.py 
128
129
130
131
132
133
134
135
136
137
138
139
def get_unique_items(self, key: str) -> List[str]:
    """Get unique items for a given key.

    Args:
        key (str): Key in the HDF5 file.

    Returns:
        List[str]: List of unique items.
    """
    if self.file[key].dtype == string_dtype():
        return list({item.decode() for item in self.file[key][: self.size]})
    return list(set(self.file[key][: self.size]))

pack()

Pack all samples into a Sample object with attributes being lists of samples.

Returns:

Name Type Description
Sample Sample

Sample object containing all samples.
Source code in mbodied/data/replaying.py 
198
199
200
201
202
203
204
def pack(self) -> Sample:
    """Pack all samples into a Sample object with attributes being lists of samples.

    Returns:
        Sample: Sample object containing all samples.
    """
    return Sample.pack_from([self.pack_one(i) for i in range(self.size)])

pack_one(index)

Pack a single sample into a Sample object.

Parameters:

Name Type Description Default
index int

Index of the sample.

 required

Returns:

Name Type Description
Sample Sample

Sample object.
Source code in mbodied/data/replaying.py 
187
188
189
190
191
192
193
194
195
196
def pack_one(self, index: int) -> Sample:
    """Pack a single sample into a Sample object.

    Args:
        index (int): Index of the sample.

    Returns:
        Sample: Sample object.
    """
    return Sample(**{key: self.read_sample(index)[i] for i, key in enumerate(self.file_keys)})

read_sample(index)

Read a sample from the HDF5 file at a given index.

Parameters:

Name Type Description Default
index int

Index of the sample.

 required

Returns:

Type Description
Tuple[dict, ...]

Tuple[dict, ...]: Tuple of dictionaries containing the sample data.
Source code in mbodied/data/replaying.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def read_sample(self, index: int) -> Tuple[dict, ...]:
    """Read a sample from the HDF5 file at a given index.

    Args:
        index (int): Index of the sample.

    Returns:
        Tuple[dict, ...]: Tuple of dictionaries containing the sample data.
    """

    def read_do(full_key: str, index: int) -> Any:
        if index >= self.file[full_key].shape[0]:
            logging.warning(f"Index {index} is out of bounds for key {full_key}. Skipping...")
            return None
        value = self.file[full_key][index]
        if full_key in self.image_keys_to_save and len(value.shape) == 3:
            image_path = f"{self.frames_path}/{Path(full_key).name}_{index}.png"
            PILImage.fromarray(value).save(image_path)
        if self.file[full_key].dtype == string_dtype():
            if self.file[full_key].attrs.get("tuple_length") is not None:
                return Sample.model_validate_json(value.decode()).unpack(to_dicts=True)
            try:
                value = value.decode()
            except Exception as e:
                logging.error(f"Error decoding string for key {full_key}: {e}")
                value = value[0].decode()
        return value

    return tuple(self.recursive_do(read_do, key=key, index=index) for key in self.file_keys)

recursive_do(do, key='', prefix='', **kwargs)

Recursively perform a function on each key in the HDF5 file.

Parameters:

Name Type Description Default
do Callable

Function to perform.

 required
key str

Key in the HDF5 file. Defaults to ''.

 ''
prefix str

Prefix for the key. Defaults to ''.

 ''
**kwargs

Additional arguments to pass to the function.

 {}

Returns:

Name Type Description
Any Any

Result of the function.
Source code in mbodied/data/replaying.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def recursive_do(self, do: Callable, key="", prefix="", **kwargs) -> Any:
    """Recursively perform a function on each key in the HDF5 file.

    Args:
        do (Callable): Function to perform.
        key (str, optional): Key in the HDF5 file. Defaults to ''.
        prefix (str, optional): Prefix for the key. Defaults to ''.
        **kwargs: Additional arguments to pass to the function.

    Returns:
        Any: Result of the function.
    """
    full_key = f"{prefix}/{key}".strip("/")
    if full_key and isinstance(self.file[full_key], h5py.Dataset):
        return do(full_key, **kwargs)

    keys = self.file_keys if not full_key else self.file[full_key].keys()

    result = {}
    for sub_key in keys:
        result[sub_key] = self.recursive_do(do, key=sub_key, prefix=full_key, **kwargs)
    return result

sample(index=None, n=1)

Get a sample from the HDF5 file.

Parameters:

Name Type Description Default
index Optional[Union[int, slice]]

Index or slice of the sample. Defaults to None.

 None
n int

Number of samples to get. Defaults to 1.

 1

Returns:

Name Type Description
Sample Sample

Sample object.
Source code in mbodied/data/replaying.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
def sample(self, index: int | slice | None = None, n: int = 1) -> Sample:
    """Get a sample from the HDF5 file.

    Args:
        index (Optional[Union[int, slice]], optional): Index or slice of the sample. Defaults to None.
        n (int, optional): Number of samples to get. Defaults to 1.

    Returns:
        Sample: Sample object.
    """
    if index is None:
        index = np.random.randint(0, self.size, n)
    if isinstance(index, int):
        return self.pack_one(index)
    if isinstance(index, slice):
        return [self.read_sample(i) for i in range(index.start, index.stop, index.step)]
    return None

clean_folder(folder, image_keys_to_save)

Clean the folder by iterating through the files and asking for deletion.

Parameters:

Name Type Description Default
folder str

Path to the folder.

 required
image_keys_to_save List[str]

List of image keys to save.

 required
Source code in mbodied/data/replaying.py 
276
277
278
279
280
281
282
283
284
285
286
287
288
289
def clean_folder(folder: str, image_keys_to_save: List[str]) -> None:
    """Clean the folder by iterating through the files and asking for deletion.

    Args:
        folder (str): Path to the folder.
        image_keys_to_save (List[str]): List of image keys to save.
    """
    for f in os.listdir(folder):
        r = Replayer(f"{folder}/{f}", image_keys_to_save=image_keys_to_save)
        for _i, sample in enumerate(r):
            obs, act, state = sample
        should_delete = input(f"Delete {f}? (y/n): ")
        if should_delete.lower() == "y":
            Path(f"{folder}/{f}").rmdir()

parse_slice(s)

Parse a string to an integer or slice.

Parameters:

Name Type Description Default
s str

String to parse.

 required

Returns:

Type Description
int | slice

Union[int, slice]: Integer or slice.

Examples:

>>> lst = [0, 1, 2, 3, 4, 5]
>>> lst[parse_slice("1")]
1
>>> lst[parse_slice("1:5:2")]
[1, 3]
Source code in mbodied/data/replaying.py 
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
def parse_slice(s: str) -> int | slice:
    """Parse a string to an integer or slice.

    Args:
        s (str): String to parse.

    Returns:
        Union[int, slice]: Integer or slice.

    Examples:
        >>> lst = [0, 1, 2, 3, 4, 5]
        >>> lst[parse_slice("1")]
        1
        >>> lst[parse_slice("1:5:2")]
        [1, 3]
    """
    if ":" in s or s == "all":
        parts = s.strip().split(":")
        start = int(parts[0]) if parts[0] else None
        stop = int(parts[1]) if len(parts) > 1 and parts[1] else None
        step = int(parts[2]) if len(parts) > 2 and parts[2] else None
        return slice(start, stop, step)
    return int(s)

replay(path, image_keys_to_save, stats, unique, keys, show, clean, upload_name)

Replay command for processing HDF5 files.

This command provides various options to manipulate and analyze HDF5 files including cleaning, uploading, showing samples, and generating statistics or extracting unique items based on provided keys.

Source code in mbodied/data/replaying.py 
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
@click.command("replay")
@click.argument("path", type=click.Path(exists=True))
@click.option("-ik", "--image_keys_to_save", multiple=True, help="Specify image keys to save during processing.")
@click.option("-s", "--stats", is_flag=True, help="Enable statistics generation for specified keys.")
@click.option("-u", "--unique", is_flag=True, help="Enable extraction of unique items from specified keys.")
@click.option("-k", "--keys", multiple=True, help="Specify keys to process for statistics or unique item extraction.")
@click.option(
    "--show-slice",
    "-v",
    default="all",
    type=str,
    help="Display and print samples at the specified index or range. Use 'all' to show all samples.",
)
@click.option("--clean", "-c", is_flag=True, help="Clean the specified folder by removing specified image keys.")
@click.option("--upload-name", "-up", default=None, help="Upload the dataset after processing.")
def replay(
    path: str,
    image_keys_to_save: List[str],
    stats: bool,
    unique: bool,
    keys: List[str],
    show: str,
    clean: bool,
    upload_name: str,
) -> None:
    """Replay command for processing HDF5 files.

    This command provides various options to manipulate and analyze HDF5 files including cleaning, uploading,
    showing samples, and generating statistics or extracting unique items based on provided keys.
    """
    if clean:
        clean_folder(path, image_keys_to_save)
        return

    replayer = Replayer(path, image_keys_to_save=image_keys_to_save)
    if stats:
        if keys:
            for _key in keys:
                pass
        else:
            logging.info(replayer.get_stats())

    if unique:
        if keys:
            for _key in keys:
                pass
        else:
            raise ValueError("Keys must be provided to get unique items.")

    if show:
        slice_ = parse_slice(show)
        for _sample in replayer[slice_]:
            showing = {}
            for key in replayer.file_keys:
                if key in replayer.image_keys_to_save:
                    img = MbImage(_sample[key])
                    img.save(f"{replayer.get_frames_path()}/{key}.png")
                    showing[key] = img
                else:
                    showing[key] = _sample[key]
            logging.info(showing)
    replayer.close()

    if upload_name:
        to_dataset(path, name=upload_name)
        return

to_dataset(folder, name, description=None, **kwargs)

Convert the folder of HDF5 files to a Hugging Face dataset.

Parameters:

Name Type Description Default
folder str

Path to the folder containing HDF5 files.

 required
name str

Name of the dataset.

 required
description str

Description of the dataset. Defaults to None.

 None
**kwargs

Additional arguments to pass to the Dataset.push_to_hub method.

 {}
Source code in mbodied/data/replaying.py 
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
def to_dataset(folder: str, name: str, description: str = None, **kwargs) -> None:
    """Convert the folder of HDF5 files to a Hugging Face dataset.

    Args:
        folder (str): Path to the folder containing HDF5 files.
        name (str): Name of the dataset.
        description (str, optional): Description of the dataset. Defaults to None.
        **kwargs: Additional arguments to pass to the Dataset.push_to_hub method.
    """
    r = FolderReplayer(folder)
    data = list(r.__iter__())
    features = Features(
        {
            "observation": {"image": Image(), "instruction": Value("string")},
            "action": infer_features(data[0]["action"]),
            "state": infer_features(data[0]["state"]),
            "episode": Value("int32"),
        },
    )

    # Convert list of dicts to dict of lists, preserving order
    data_dict = {key: [item[key] for item in data] for key in data[0]}
    info = DatasetInfo(
        description=description,
        license="Apache-2.0",
        citation="None",
        size_in_bytes=8000000,
        features=features,
    )

    ds = Dataset.from_dict(data_dict, info=info)
    ds = ds.with_format("pandas")
    ds.push_to_hub(name, private=True, **kwargs)