|
|
if __name__ == "__main__": |
|
|
import sys |
|
|
import os |
|
|
import pathlib |
|
|
|
|
|
ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent) |
|
|
sys.path.append(ROOT_DIR) |
|
|
|
|
|
import multiprocessing |
|
|
import os |
|
|
import shutil |
|
|
import click |
|
|
import pathlib |
|
|
import h5py |
|
|
from tqdm import tqdm |
|
|
import numpy as np |
|
|
import collections |
|
|
import pickle |
|
|
from equi_diffpo.common.robomimic_util import RobomimicObsConverter |
|
|
|
|
|
multiprocessing.set_start_method('spawn', force=True) |
|
|
|
|
|
def worker(x): |
|
|
path, idx = x |
|
|
converter = RobomimicObsConverter(path) |
|
|
obss = converter.convert_idx(idx) |
|
|
return obss |
|
|
|
|
|
@click.command() |
|
|
@click.option('-i', '--input', required=True, help='input hdf5 path') |
|
|
@click.option('-o', '--output', required=True, help='output hdf5 path. Parent directory must exist') |
|
|
@click.option('-n', '--num_workers', default=None, type=int) |
|
|
def main(input, output, num_workers): |
|
|
|
|
|
input = pathlib.Path(input).expanduser() |
|
|
assert input.is_file() |
|
|
output = pathlib.Path(output).expanduser() |
|
|
assert output.parent.is_dir() |
|
|
assert not output.is_dir() |
|
|
|
|
|
converter = RobomimicObsConverter(input) |
|
|
|
|
|
|
|
|
print('Copying hdf5') |
|
|
shutil.copy(str(input), str(output)) |
|
|
|
|
|
|
|
|
idx = 0 |
|
|
while idx < len(converter): |
|
|
with multiprocessing.Pool(num_workers) as pool: |
|
|
end = min(idx + num_workers, len(converter)) |
|
|
results = pool.map(worker, [(input, i) for i in range(idx, end)]) |
|
|
|
|
|
|
|
|
print('Writing {} to {}'.format(idx, end)) |
|
|
with h5py.File(output, 'r+') as out_file: |
|
|
for i in tqdm(range(idx, end), desc="Writing to output"): |
|
|
obss = results[i - idx] |
|
|
demo = out_file[f'data/demo_{i}'] |
|
|
del demo['obs'] |
|
|
for k in obss: |
|
|
demo.create_dataset("obs/{}".format(k), data=np.array(obss[k]), compression="gzip") |
|
|
|
|
|
idx = end |
|
|
del results |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
|
