get+normals work on batches
This commit is contained in:
+1
-1
@@ -85,7 +85,7 @@ class Dataset(Dataset):
|
||||
|
||||
start = time.perf_counter()
|
||||
# TODO: is it sure I need this K, and not some transformed one?
|
||||
normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
|
||||
normal_image = get_normals_from_depth(depth_image[None,None,:,:], depth_mask[None,:,:], K)
|
||||
print(f'normals time: {time.perf_counter() - start}')
|
||||
|
||||
depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask
|
||||
|
||||
+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