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

data/preprocess.py

@@ -0,0 +1,25 @@
+import os
+import sys
+
+import numpy as np
+from tqdm import tqdm
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from utils.utils import mask_faces
+
+def generate_faces(folder):
+    frames = os.listdir(os.path.join(folder, 'mesh_mask'))
+    F = np.load('custom_data/smpl_faces.npy')
+    os.makedirs(os.path.join(folder, 'faces'), exist_ok=True)
+
+    for frame in tqdm(frames):
+        mask = np.load(os.path.join(folder, 'mesh_mask', frame))
+        faces = mask_faces(F, mask)
+        np.save(os.path.join(folder, 'faces', frame), faces)
+
+
+if __name__ == '__main__':
+    # folder = '/home/workspace/local-data/train'
+    folder = '/home/workspace/local-data/valid'
+
+    generate_faces(folder)

data/render.py

@@ -0,0 +1,124 @@
+import numpy as np
+
+
+from deodr import Camera, Scene3D
+from deodr.triangulated_mesh import TriMesh
+
+
+class RenderCore:
+    def __init__(self, max_depth=0, depth_scale=1):
+        self.mesh = None
+        self.depthScale = depth_scale
+        self.scene = Scene3D(0)
+        self.set_max_depth(max_depth)
+		
+    def set_mesh(self, vertices, faces):
+        self.mesh = TriMesh(
+            faces, vertices
+        )  # we do a copy to avoid negative stride not support by Tensorflow
+		
+        object_center = vertices.mean(axis=0)
+        object_radius = np.max(np.std(vertices, axis=0))
+        self.camera_center = object_center + np.array([-0.5, 0, 5]) * object_radius
+		
+        self.scene.set_mesh(self.mesh)
+
+    def set_max_depth(self, max_depth):
+        self.scene.max_depth = max_depth
+        self.scene.set_background_color([max_depth])
+
+    def set_depth_scale(self, depth_scale):
+        self.depthScale = depth_scale
+
+    def set_image(self, width, height, intrinsic, extrinsic):
+        self.width = width
+        self.height = height
+		
+        # focal = None
+        # if focal is None:
+            # focal = 2 * self.width
+
+        # rot = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
+        # trans = -rot.T.dot(self.camera_center)
+        # intrinsic = np.array(
+            # [[focal, 0, self.width / 2], [0, focal, self.height / 2], [0, 0, 1]]
+        # )
+        # extrinsic = np.column_stack((rot, trans))
+		
+        self.camera = Camera(
+            extrinsic=extrinsic,
+            intrinsic=intrinsic,
+            width=self.width,
+            height=self.height,
+            distortion=None,
+        )
+        
+    def render(self):
+        depth_scale = 1 * self.depthScale
+        depth = self.scene.render_depth(
+            self.camera,
+            depth_scale=depth_scale,
+        )
+        #depth = tf.clip_by_value(depth, 0, self.scene.max_depth)
+        return depth
+
+
+""" CAMERA """
+def intrinsic():
+	RES_X = 640
+	RES_Y = 480
+	f_mm             = 50 # blender default
+	sensor_w_mm      = 36 # blender default
+	sensor_h_mm = sensor_w_mm * RES_Y / RES_X
+
+	fx_px = f_mm * RES_X / sensor_w_mm
+	fy_px = f_mm * RES_Y / sensor_h_mm
+
+	u = RES_X / 2
+	v = RES_Y / 2
+
+	return np.array([[fx_px, 0,     u],
+					 [0,     fy_px, v],
+					 [0,     0,     1]], np.float32)
+				
+def extrinsic(camLoc):
+	R_w2bc = np.array([[0, 1, 0],
+					   [0, 0, 1],
+					   [1, 0, 0]], np.float32)
+	
+	T_w2bc = -1 * R_w2bc.dot(camLoc)
+	
+	R_bc2cv = np.array([[1,  0,  0],
+						[0, -1,  0],
+						[0,  0, -1]], np.float32)
+	
+	R_w2cv = R_bc2cv.dot(R_w2bc)
+	T_w2cv = R_bc2cv.dot(T_w2bc)
+
+	return np.concatenate((R_w2cv, T_w2cv[:,None]), axis=1)
+	
+def proj(camLoc):
+	return intrinsic().dot(extrinsic(camLoc))
+
+
+class Render:
+    def __init__(self, max_depth=0, depth_scale=1, size=(640, 480)):
+        self.render_core = RenderCore(max_depth=max_depth)
+        self.size = size
+        self.max_depth = max_depth
+
+    def render(self, mesh, faces, cam):
+        self.render_core.set_mesh(mesh, faces)
+        self.render_core.set_image(self.size[0], self.size[1], intrinsic(), extrinsic(cam))
+        depth = self.render_core.render()
+
+        thresh = self.max_depth - 1
+        mask = depth<thresh
+        # dmin = depth.min()
+        dmean = depth[mask].mean()
+        # dmax = depth[mask].max()
+        depth[mask] = depth[mask] - dmean
+        depth[~mask] = 0.0
+        mask = mask.astype(np.float32)
+
+        return depth