generate normal image

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