55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
import os
|
|
import sys
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
|
|
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
|
# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
|
|
|
|
import torch
|
|
from torch.utils.data import Dataset
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
from time import time
|
|
from cloth3d.read import DataReader
|
|
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
|
|
from superSMPL.smpl_torch import SMPL
|
|
from data.pose_renderer import PoseRenderer
|
|
|
|
|
|
class Dataset(Dataset):
|
|
def __init__(self, device='cuda'):
|
|
self.device = device
|
|
self.smpl = SMPL(path='superSMPL', device=device)
|
|
|
|
|
|
def __len__(self):
|
|
return int(1e6)
|
|
|
|
|
|
def __getitem__(self, idx):
|
|
# generate a random tenssor of size 3 with values betweeen -pi and pi
|
|
angles = torch.rand(3) * 2 * np.pi - np.pi
|
|
angles = angles.to(self.device)
|
|
matrix = self.smpl.rodrigues(angles[None, None, :]).to(self.device)[0]
|
|
|
|
data = {
|
|
'angles': angles,
|
|
'matrix': matrix,
|
|
}
|
|
|
|
|
|
return data
|
|
|
|
|
|
if __name__ == '__main__':
|
|
from torch.utils.data import DataLoader
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
dataset = Dataset(device=device)
|
|
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
|
|
|
|
for batch in train_dataloader:
|
|
batch
|
|
pass
|
|
|
|
pass |