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
+1 -1
View File
@@ -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
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]