get+normals work on batches

This commit is contained in:
2023-01-14 17:10:58 +00:00
parent c7ba2f1b54
commit e1d7cafeab
2 changed files with 4 additions and 3 deletions
+3 -2
View File
@@ -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]