diff --git a/data/dataset.py b/data/dataset.py index 73daa26..652a092 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -13,7 +13,7 @@ import plotly.express as ex from torch.utils.data import Dataset -from render import Render +from render import Render, DepthRenderer class Dataset(Dataset): @@ -22,7 +22,7 @@ class Dataset(Dataset): self.transforms = tv.transforms.ToTensor() self.frames = os.listdir(os.path.join(folder, 'mesh_posed')) - self.render = Render(max_depth=10) + self.render = DepthRenderer() self.body_faces = np.load('custom_data/smpl_faces.npy') @@ -41,12 +41,15 @@ class Dataset(Dataset): 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)) + combined = torch.Tensor(np.concatenate((body, mesh_posed[mask, :]))).to('cuda') + combined_faces = torch.Tensor(np.concatenate((self.body_faces, faces + 14475))).to('cuda') start = time.perf_counter() - depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0] + # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0] + # depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0] + depth_image = self.render(combined, combined_faces, info['camLoc'][0]) t = time.perf_counter() - start + print(t) # image = self.transforms(image) return { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mask': mask } diff --git a/data/render.py b/data/render.py index 9928771..516d7d1 100644 --- a/data/render.py +++ b/data/render.py @@ -81,6 +81,7 @@ def intrinsic(): [0, fy_px, v], [0, 0, 1]], np.float32) + def extrinsic(camLoc): R_w2bc = np.array([[0, 1, 0], [0, 0, 1], @@ -121,4 +122,57 @@ class Render: depth[~mask] = 0.0 mask = mask.astype(np.float32) - return depth \ No newline at end of file + return depth + + +from torch import nn +from pytorch3d.renderer import ( + look_at_view_transform, + FoVPerspectiveCameras, + PerspectiveCameras, + PointLights, + DirectionalLights, + Materials, + RasterizationSettings, + MeshRenderer, + MeshRasterizer, + SoftPhongShader, + TexturesUV, + TexturesVertex +) +from pytorch3d.structures import Meshes + +class DepthRenderer(nn.Module): + def __init__(self) -> None: + super().__init__() + + + + def forward(self, verts, faces, cam): + R, T = look_at_view_transform(2.7, 0, 0) + cameras = FoVPerspectiveCameras(device='cuda', R=R, T=T) + ext = extrinsic(cam) + cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3], focal_length=50) + + raster_settings = RasterizationSettings( + image_size=512, + blur_radius=0.0, + faces_per_pixel=1, + bin_size=0 + ) + + rasterizer = MeshRasterizer( + cameras=cameras, + raster_settings=raster_settings + ) + depth = rasterizer(Meshes(verts[None,:,:], faces[None,:,:])).zbuf[0,:,:,0] + + mask = depth != -1 + # 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