luni/reconstruction
1
0
Fork 0
Code Issues 6 Pull Requests Packages Projects 1 Releases Wiki Activity

get+normals work on batches

Browse Source
This commit is contained in:
luni
2023-01-14 17:10:58 +00:00
parent c7ba2f1b54
commit e1d7cafeab
2 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/dataset.py
+1 -1
View File
@@ -85,7 +85,7 @@ class Dataset(Dataset):
start = time.perf_counter() start = time.perf_counter()
# TODO: is it sure I need this K, and not some transformed one? # 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}') print(f'normals time: {time.perf_counter() - start}')
depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask
utils/data_utils.py
+3 -2
View File
@@ -46,10 +46,11 @@ def read_frames_file(file):
def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic): 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 = 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 = 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]) 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 = kornia.geometry.depth.depth_to_normals(depth_dilated, camera_intrinsic, normalize_points=True)
normals = torch.permute(normals, [0,2,3,1]) * depth_mask[:,:,:,None] normals = torch.permute(normals, [0,2,3,1]) * depth_mask[:,:,:,None]