Class MotionDatasetRollout

diffmimic/utils/data.py:112–152

class MotionDatasetRollout(data.Dataset):
    def __init__(self, motion_dir, seg_info=None):
        if 'kit' in motion_dir:
            list_path = 'data/kitml_datalist.txt'
        elif 'human' in motion_dir:
            list_path = 'data/humanml3d_datalist.txt'
        else:
            raise NotImplementedError
        with open(list_path) as f:
            self.motion_list = f.read().splitlines()
            if 'kit' in motion_dir:
                # self.motion_list = ['00586.npy', '00603.npy', '00606.npy', '02916.npy','02351.npy', '02593.npy']
                # self.motion_list = ['02916.npy','02351.npy', '02593.npy'] # 360 jump
                # self.motion_list = ['00607.npy', '00608.npy', '00609.npy', '00610.npy','00611.npy', '00612.npy', '01236.npy'] # kick
                self.motion_list = ['03329.npy', '02797.npy', '00653.npy', '01317.npy', '02320.npy', '02772.npy', '02355.npy', '02625.npy', '02916.npy', '03299.npy'] # acrobatic
            #    self.motion_list = ['02285.npy', '02291.npy', '02292.npy', '02626.npy', '03449.npy', '03311.npy', '03304.npy', '02823.npy', '02628.npy', '02283.npy', '01234.npy'] # dance
            else:
                self.motion_list = ['000000.npy', '000002.npy', '000264.npy', '000499.npy']
            # self.motion_list = [fname for fname in self.motion_list if not os.path.isfile(os.path.join('/mnt/lustre/share/jwren/cvae_kitml_traj_valid_v3/', fname))]
            # self.motion_list = [fname for fname in self.motion_list if not os.path.isfile(os.path.join('/mnt/lustre/share/jwren/cvae_kitml_traj_invalid_v3/', fname))]


            if seg_info is not None:
                num, idx = seg_info
                total_len = len(self.motion_list)
                seg_len = (total_len // num) + 1
                idx_start = seg_len * idx
                idx_end = seg_len * (idx + 1)
                self.motion_list = self.motion_list[idx_start:idx_end]
            motion_path_list = [os.path.join(motion_dir, fname) for fname in self.motion_list]
            self.motion_path_list = [x for x in motion_path_list if os.path.isfile(x)]
        print(len(self.motion_path_list))

    def __len__(self):
        return len(self.motion_path_list)

    def __getitem__(self, idx):
        fpath, fname = self.motion_path_list[idx], self.motion_list[idx]
        demo_traj = np.array(np.load(fpath))

        return demo_traj, fname


class InfiniteSampler(data.sampler.Sampler):