| 38 | |
| 39 | |
| 40 | class MotionDataset(data.Dataset): |
| 41 | def __init__(self, motion_dir, seq_len=60, subset=None, all=False, resample=False): |
| 42 | # motion_path_list = glob.glob(os.path.join(motion_dir, "*.npy")) |
| 43 | if 'kit' in motion_dir: |
| 44 | list_path = 'data/kitml_datalist.txt' |
| 45 | text_dir = 'data/kit_texts' |
| 46 | elif 'human' in motion_dir: |
| 47 | # list_path = '/scratch/users/ntu/cheeguan/jwren/data/humanml3d.txt' |
| 48 | list_path = 'data/humanml3d_datalist.txt' |
| 49 | text_dir = 'data/humanml3d_text' |
| 50 | else: |
| 51 | raise NotImplementedError |
| 52 | with open(list_path) as f: |
| 53 | motion_list = f.read().splitlines() |
| 54 | motion_path_list = [os.path.join(motion_dir, fname) for fname in motion_list] |
| 55 | motion_path_list = [x for x in motion_path_list if os.path.isfile(x)] |
| 56 | if subset is not None: |
| 57 | motion_path_list = motion_path_list[:subset] |
| 58 | self.frame_list = [] |
| 59 | if all: |
| 60 | for fpath in motion_path_list: |
| 61 | valid_start_frames = [(fpath, 0)] |
| 62 | self.frame_list += valid_start_frames |
| 63 | else: |
| 64 | if resample: |
| 65 | for fpath in motion_path_list: |
| 66 | motion = np.array(np.load(fpath)) |
| 67 | text_path = os.path.join(text_dir, os.path.basename(fpath).replace('.npy', '.txt')) |
| 68 | valid_start_frames = [(fpath, x) for x in range(max(0, motion.shape[0]-seq_len)+1)] |
| 69 | with open(text_path) as f: |
| 70 | text_description = f.read().splitlines() |
| 71 | text_description = [x.split('#')[0] for x in text_description] |
| 72 | text_description = ' '.join(text_description) |
| 73 | if 'walk' not in text_description: |
| 74 | self.frame_list += valid_start_frames * 10 |
| 75 | else: |
| 76 | self.frame_list += valid_start_frames |
| 77 | else: |
| 78 | for fpath in motion_path_list: |
| 79 | motion = np.array(np.load(fpath)) |
| 80 | valid_start_frames = [(fpath, x) for x in range(max(0, motion.shape[0]-seq_len)+1)] |
| 81 | self.frame_list += valid_start_frames |
| 82 | |
| 83 | |
| 84 | # self.motion_path_list = glob.glob(os.path.join(motion_dir, "*.npy")) |
| 85 | # self.motion_path_list = [fn for fn in self.motion_path_list if np.array(np.load(fn)).shape[0] >= seq_len] |
| 86 | self.seq_len = seq_len |
| 87 | print(len(self.frame_list)) |
| 88 | |
| 89 | def __len__(self): |
| 90 | return len(self.frame_list) |
| 91 | |
| 92 | def __getitem__(self, idx): |
| 93 | fpath, start_idx = self.frame_list[idx] |
| 94 | demo_traj = np.array(np.load(fpath)) |
| 95 | demo_traj = demo_traj[start_idx:start_idx+self.seq_len] |
| 96 | mask = np.ones(self.seq_len) |
| 97 | demo_size = demo_traj.shape[0] |