EditMP/datasets.py at main · wangwhTj/EditMP · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36


import torch
# import torchvision
# import torchvision.transforms as transforms
from PIL import Image
from torch.utils.data import Dataset
# import string
import numpy as np


class Demostration(Dataset):
    def __init__(self, path):
        Expert_Trajs = np.load(path)
        print(Expert_Trajs.shape)
        #self.Expert_Traj = Expert_Trajs[:512]
        self.Expert_Traj = Expert_Trajs
        # for i in range(512):
        #     self.Expert_Traj[i] = Expert_Trajs[i*6]


    def __getitem__(self, index):
        Expert_Trajs = self.Expert_Traj[index]


        return torch.FloatTensor(Expert_Trajs)

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


if __name__ == "__main__":
    d = Demostration('/home/xuxh/Level_recognition/data/val_list.txt')

    print("000")