import os
import sys

os.environ["CUDA_VISIBLE_DEVICES"] = '1'

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):
        path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
        mesh_posed = np.load(path)['mesh_posed']
        # image = self.transforms(image)
        return { 'mesh_posed': mesh_posed }


if __name__ == '__main__':
    from torch.utils.data import DataLoader
    
    device = "cuda" if torch.cuda.is_available() else "cpu"

    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