get+normals work on batches

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