diff --git a/custom_data/smpl_faces.npy b/custom_data/smpl_faces.npy new file mode 100644 index 0000000..d1d74c9 Binary files /dev/null and b/custom_data/smpl_faces.npy differ diff --git a/data/dataset.py b/data/dataset.py index a8f2eb7..73daa26 100644 --- a/data/dataset.py +++ b/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 diff --git a/data/preprocess.py b/data/preprocess.py new file mode 100644 index 0000000..2b6c934 --- /dev/null +++ b/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) diff --git a/data/render.py b/data/render.py new file mode 100644 index 0000000..9928771 --- /dev/null +++ b/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