get+normals work on batches
This commit is contained in:
+3
-2
@@ -46,10 +46,11 @@ def read_frames_file(file):
|
||||
|
||||
|
||||
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_image, 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_image = torch.permute(depth_image, [0,2,3,1])
|
||||
depth_dilated[depth_mask,:] = depth_image[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]
|
||||
|
||||
Reference in New Issue
Block a user