renderer beta

data/dataset.py

@@ -4,26 +4,52 @@ 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 }
+        return { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mask': mask }
 
 
 if __name__ == '__main__':
@@ -32,7 +58,7 @@ if __name__ == '__main__':
     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)
+    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True)
 
     for batch in train_dataloader:
         batch