67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
import os
|
|
import sys
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
|
|
|
|
import math
|
|
import time
|
|
import torch
|
|
import torchvision as tv
|
|
import numpy as np
|
|
import scipy.io as sio
|
|
import plotly.express as ex
|
|
|
|
from torch.utils.data import Dataset
|
|
|
|
from render import Render
|
|
|
|
|
|
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'))
|
|
|
|
self.render = Render(max_depth=10)
|
|
self.body_faces = np.load('custom_data/smpl_faces.npy')
|
|
|
|
|
|
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']
|
|
|
|
id = self.frames[idx].split('_')[0]
|
|
faces = np.load(os.path.join(self.folder, 'faces', id + '.npy'))
|
|
mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy'))
|
|
body = np.load(os.path.join(self.folder, 'mesh_body_posed', self.frames[idx][:-1] + 'y'))
|
|
|
|
info_folder = '/home/workspace/data/dataset/' + ('train' if self.folder.split('/')[-1] == 'train' else 'validation') + '/' + id + '/info.mat'
|
|
info = sio.loadmat(info_folder)
|
|
combined = np.concatenate((body, mesh_posed[mask, :]))
|
|
combined_faces = np.concatenate((self.body_faces, faces + 14475))
|
|
|
|
start = time.perf_counter()
|
|
depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
|
|
t = time.perf_counter() - start
|
|
|
|
# image = self.transforms(image)
|
|
return { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mask': mask }
|
|
|
|
|
|
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, shuffle=True)
|
|
|
|
for batch in train_dataloader:
|
|
batch
|
|
pass
|
|
|
|
pass |