37 lines
904 B
Python
37 lines
904 B
Python
import os
|
|
import sys
|
|
|
|
import math
|
|
import torch
|
|
import torchvision as tv
|
|
import numpy as np
|
|
|
|
from torch.utils.data import Dataset
|
|
|
|
class Dataset(Dataset):
|
|
def __init__(self, folder):
|
|
self.folder = folder
|
|
self.transforms = tv.transforms.ToTensor()
|
|
self.frames = os.listdir(os.path.join(folder, 'mesh_posed'))
|
|
|
|
def __len__(self):
|
|
return len(self.frames)
|
|
|
|
def __getitem__(self, idx):
|
|
img_path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
|
|
image = np.load(img_path)['mesh_posed']
|
|
# image = self.transforms(image)
|
|
return image
|
|
|
|
|
|
if __name__ == '__main__':
|
|
from torch.utils.data import DataLoader
|
|
|
|
dataset = Dataset('/home/workspace/local-data/train')
|
|
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, num_workers=2)
|
|
|
|
for batch in train_dataloader:
|
|
batch
|
|
pass
|
|
|
|
pass |