generate normal image

data/dataset.py

@@ -6,6 +6,7 @@ os.environ["CUDA_VISIBLE_DEVICES"] = '0'
 import math
 import time
 import torch
+import kornia
 import torchvision as tv
 import numpy as np
 import scipy.io as sio
@@ -17,7 +18,7 @@ from render import Render, DepthRenderer
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from superSMPL.IO import readPC2Frame, writePC2
-from utils.data_utils import Frame, read_frames_file, mask_faces
+from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth
 from cloth3d.read import DataReader
 from cloth3d.util import quads2tris
 
@@ -71,10 +72,16 @@ class Dataset(Dataset):
         t = time.perf_counter() - start
         print(t)
         start = time.perf_counter()
-        depth_image, v2i = self.render(combined, combined_faces, info['camLoc'])
+        depth_image, depth_mask, v2i, K = self.render(combined, combined_faces, info['camLoc'])
         t = time.perf_counter() - start
         print(t)
 
+        pc = kornia.geometry.depth.depth_to_3d(depth_image[None,None,:,:], K, normalize_points=True)
+        pc = pc[0,:,depth_mask].cpu().numpy()
+
+        normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
+
+
         depth_image = depth_image.cpu().numpy()
         dist_img = depth_image[v2i[idx, 1], v2i[idx, 0]]
         depth_image[v2i[idx, 1], v2i[idx, 0]] = 100
@@ -84,12 +91,16 @@ class Dataset(Dataset):
 
         import plotly.express as ex
         import plotly.graph_objects as go
-        ex.imshow(deodr_depth[:,:,0]).show()
+        # ex.imshow(deodr_depth[:,:,0]).show()
-        ex.imshow(depth_image).show()
+        # ex.imshow(depth_image).show()
+        ex.imshow(normal_image[0].cpu().numpy()).show()
 
-        go.Figure(data = [go.Scatter3d(x=combined[:, 0], y=combined[:, 1], z=combined[:, 2], marker=dict(size=3, line=dict(width=0)), mode='markers'),
+        # go.Figure(data = [go.Scatter3d(x=combined[:, 0], y=combined[:, 1], z=combined[:, 2], marker=dict(size=3, line=dict(width=0)), mode='markers'),
-                          go.Scatter3d(x=[combined[idx, 0]], y=[combined[idx, 1]], z=[combined[idx, 2]], marker=dict(size=10, line=dict(width=0)), mode='markers'),
+        #                   go.Scatter3d(x=[combined[idx, 0]], y=[combined[idx, 1]], z=[combined[idx, 2]], marker=dict(size=10, line=dict(width=0)), mode='markers'),
-                         ]).show()
+        #                  ]).show()
+
+        # go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),
+        #                  ]).show()
 
         # v2i_coverage = np.zeros((480, 640), dtype=np.float32)
         # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1

data/render.py

@@ -163,13 +163,14 @@ class DepthRenderer(nn.Module):
 
         ext = torch.tensor(extrinsic(cam).astype(np.float32), device='cuda')
         # ext = proj(cam)
-        fp = np.array([240 * (480 / 36), 320 * (640 / 36)], dtype=np.float32)
+        # fp = np.array([240 * (480 / 36), 320 * (640 / 36)], dtype=np.float32)
         img_size = np.array([480, 640], dtype=np.float32)
         # K = np.eye(4, dtype=np.float32)
         # K[:3,:3] = intrinsic()
+        K = torch.tensor(intrinsic()[None, :, :], device='cuda')
         # cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3], focal_length=1,)# principal_point=fp[None, :], image_size=img_size[None, :])
 
-        cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=torch.tensor(intrinsic()[None, :, :], device='cuda'), image_size=torch.tensor(img_size[None, :], device='cuda'))
+        cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(img_size[None, :], device='cuda'))
 
         raster_settings = RasterizationSettings(
                             image_size=(480, 640), 
@@ -197,4 +198,4 @@ class DepthRenderer(nn.Module):
         # import plotly.express as ex
         # ex.imshow(depth.cpu().numpy()).show()
 
-        return depth, v2i
+        return depth, mask, v2i, K

utils/data_utils.py

@@ -1,9 +1,12 @@
 import os
 import sys
+import kornia
+import torch
 import numpy as np
 import scipy.io as sio
 from math import cos, sin
 
+
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from superSMPL.IO import readPC2Frame, readOBJ
 
@@ -34,3 +37,14 @@ def read_frames_file(file):
             id, frame_nr, garment = line.split()
             frames.append(Frame(id, int(frame_nr), garment))
     return frames
+
+
+def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic):
+    depth_dilated = kornia.morphology.dilation(depth_image[:,None,:,:], torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
+    depth_dilated = kornia.morphology.dilation(depth_dilated, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
+    depth_dilated = torch.permute(depth_dilated, [0,2,3,1])
+    depth_dilated[depth_mask,:] = depth_image[:,:,:,None][depth_mask,:]
+    depth_dilated = torch.permute(depth_dilated, [0,3,1,2])
+    normals = kornia.geometry.depth.depth_to_normals(depth_dilated, camera_intrinsic, normalize_points=True)
+    normals = torch.permute(normals, [0,2,3,1]) * depth_mask[:,:,:,None]
+    return normals