PyTorch3D - Avoid flip in TexturesAtlas

Summary: Performance improvement: Use torch.lerp to map uv coordinates to the range needed for grid_sample (i.e. map [0, 1] to [-1, 1] and invert the y-axis)

Reviewed By: bottler

Differential Revision: D51961728

fbshipit-source-id: db19a5e3f482e9af7b96b20f88a1e5d0076dac43

Author: eyoung2015

pytorch3d/renderer/mesh/textures.py

@@ -995,9 +995,13 @@ def sample_textures(self, fragments, **kwargs) -> torch.Tensor:
         #   is the left-top pixel of input, and values x = 1, y = 1 is the
         #   right-bottom pixel of input.
 
-        pixel_uvs = pixel_uvs * 2.0 - 1.0
+        # map to a range of [-1, 1] and flip the y axis
+        pixel_uvs = torch.lerp(
+            pixel_uvs.new_tensor([-1.0, 1.0]),
+            pixel_uvs.new_tensor([1.0, -1.0]),
+            pixel_uvs,
+        )
 
-        texture_maps = torch.flip(texture_maps, [2])  # flip y axis of the texture map
         if texture_maps.device != pixel_uvs.device:
             texture_maps = texture_maps.to(pixel_uvs.device)
         texels = F.grid_sample(
@@ -1035,8 +1039,12 @@ def faces_verts_textures_packed(self) -> torch.Tensor:
         texture_maps = self.maps_padded()  # NxHxWxC
         texture_maps = texture_maps.permute(0, 3, 1, 2)  # NxCxHxW
 
-        faces_verts_uvs = faces_verts_uvs * 2.0 - 1.0
-        texture_maps = torch.flip(texture_maps, [2])  # flip y axis of the texture map
+        # map to a range of [-1, 1] and flip the y axis
+        faces_verts_uvs = torch.lerp(
+            faces_verts_uvs.new_tensor([-1.0, 1.0]),
+            faces_verts_uvs.new_tensor([1.0, -1.0]),
+            faces_verts_uvs,
+        )
 
         textures = F.grid_sample(
             texture_maps,

tests/test_render_meshes.py

@@ -932,7 +932,7 @@ def _join_uvs(self, rasterizer_type):
                     )
                 ).save(DATA_DIR / f"DEBUG_test_joinuvs{i}_map3.png")
 
-            self.assertClose(output, merged)
+            self.assertClose(output, merged, atol=0.005)
             self.assertClose(output, image_ref, atol=0.005)
             self.assertClose(mesh.textures.maps_padded()[0].cpu(), map_ref, atol=0.05)