join_scene fix for TexturesUV

Summary: Fix issue #826. This is a correction to the joining of TexturesUV into a single scene.

Reviewed By: nikhilaravi

Differential Revision: D30767092

fbshipit-source-id: 03ba6a1d2f22e569d1b3641cd13ddbb8dcb87ec7

pytorch3d/renderer/mesh/textures.py

@@ -1157,22 +1157,46 @@ def join_scene(self) -> "TexturesUV":
                     new_uvs = torch.Tensor(new_uvs)
 
             # If align_corners is True, then an index of x (where x is in
-            # the range 0 .. map_.shape[]-1) in one of the input maps
-            # was hit by a u of x/(map_.shape[]-1).
-            # That x is located at the index loc[] + x in the single_map, and
-            # to hit that we need u to equal (loc[] + x) / (total_size[]-1)
+            # the range 0 .. map_.shape[1]-1) in one of the input maps
+            # was hit by a u of x/(map_.shape[1]-1).
+            # That x is located at the index loc[1] + x in the single_map, and
+            # to hit that we need u to equal (loc[1] + x) / (total_size[1]-1)
             # so the old u should be mapped to
-            #   { u*(map_.shape[]-1) + loc[] } / (total_size[]-1)
+            #   { u*(map_.shape[1]-1) + loc[1] } / (total_size[1]-1)
+
+            # Also, an index of y (where y is in
+            # the range 0 .. map_.shape[0]-1) in one of the input maps
+            # was hit by a v of 1 - y/(map_.shape[0]-1).
+            # That y is located at the index loc[0] + y in the single_map, and
+            # to hit that we need v to equal 1 - (loc[0] + y) / (total_size[0]-1)
+            # so the old v should be mapped to
+            #   1 - { (1-v)*(map_.shape[0]-1) + loc[0] } / (total_size[0]-1)
+            # =
+            # { v*(map_.shape[0]-1) + total_size[0] - map.shape[0] - loc[0] }
+            #        / (total_size[0]-1)
 
             # If align_corners is False, then an index of x (where x is in
-            # the range 1 .. map_.shape[]-2) in one of the input maps
-            # was hit by a u of (x+0.5)/(map_.shape[]).
-            # That x is located at the index loc[] + 1 + x in the single_map,
+            # the range 1 .. map_.shape[1]-2) in one of the input maps
+            # was hit by a u of (x+0.5)/(map_.shape[1]).
+            # That x is located at the index loc[1] + 1 + x in the single_map,
             # (where the 1 is for the border)
-            # and to hit that we need u to equal (loc[] + 1 + x + 0.5) / (total_size[])
+            # and to hit that we need u to equal (loc[1] + 1 + x + 0.5) / (total_size[1])
             # so the old u should be mapped to
-            #   { loc[] + 1 + u*map_.shape[]-0.5 + 0.5 } / (total_size[])
-            #  = { loc[] + 1 + u*map_.shape[] } / (total_size[])
+            #   { loc[1] + 1 + u*map_.shape[1]-0.5 + 0.5 } / (total_size[1])
+            #  = { loc[1] + 1 + u*map_.shape[1] } / (total_size[1])
+
+            # Also, an index of y (where y is in
+            # the range 1 .. map_.shape[0]-2) in one of the input maps
+            # was hit by a v of 1 - (y+0.5)/(map_.shape[0]).
+            # That y is located at the index loc[0] + 1 + y in the single_map,
+            # (where the 1 is for the border)
+            # and to hit that we need v to equal 1 - (loc[0] + 1 + y + 0.5) / (total_size[0])
+            # so the old v should be mapped to
+            #   1 - { loc[0] + 1 + (1-v)*map_.shape[0]-0.5 + 0.5 } / (total_size[0])
+            #  = { total_size[0] - loc[0] -1 - (1-v)*map_.shape[0]  }
+            #         / (total_size[0])
+            #  = { total_size[0] - loc[0] - map.shape[0] - 1 + v*map_.shape[0] }
+            #         / (total_size[0])
 
             # We change the y's in new_uvs for the scaling of height,
             # and the x's for the scaling of width.
@@ -1184,7 +1208,9 @@ def join_scene(self) -> "TexturesUV":
             denom_y = merging_plan.total_size[1] - one_if_align
             scale_y = y_shape - one_if_align
             new_uvs[:, 1] *= scale_x / denom_x
-            new_uvs[:, 1] += (loc.x + one_if_not_align) / denom_x
+            new_uvs[:, 1] += (
+                merging_plan.total_size[0] - x_shape - loc.x - one_if_not_align
+            ) / denom_x
             new_uvs[:, 0] *= scale_y / denom_y
             new_uvs[:, 0] += (loc.y + one_if_not_align) / denom_y
 

pytorch3d/vis/texture_vis.py

@@ -36,7 +36,8 @@ def texturesuv_image_matplotlib(
         color: any matplotlib-understood color for the circles.
         subsample: if not None, number of points to plot.
                 Otherwise all points are plotted.
-        origin: "upper" or "lower" like matplotlib.imshow
+        origin: "upper" or "lower" like matplotlib.imshow .
+            upper (the default) matches texturesuv_image_PIL.
     """
 
     import matplotlib.pyplot as plt

tests/test_render_meshes.py

@@ -753,7 +753,7 @@ def test_join_uvs(self):
                 Image.fromarray((output.numpy() * 255).astype(np.uint8)).save(
                     DATA_DIR / f"test_joinuvs{i}_final_.png"
                 )
-                Image.fromarray((output.numpy() * 255).astype(np.uint8)).save(
+                Image.fromarray((merged.numpy() * 255).astype(np.uint8)).save(
                     DATA_DIR / f"test_joinuvs{i}_merged.png"
                 )
 
@@ -782,10 +782,41 @@ def test_join_uvs(self):
                     )
                 ).save(DATA_DIR / f"test_joinuvs{i}_map3.png")
 
-            self.assertClose(output, merged, atol=0.015)
-            self.assertClose(output, image_ref, atol=0.05)
+            self.assertClose(output, merged)
+            self.assertClose(output, image_ref, atol=0.005)
             self.assertClose(mesh.textures.maps_padded()[0].cpu(), map_ref, atol=0.05)
 
+    def test_join_uvs_simple(self):
+        # Example from issue #826
+        a = TexturesUV(
+            maps=torch.full((1, 4000, 4000, 3), 0.8),
+            faces_uvs=torch.arange(300).reshape(1, 100, 3),
+            verts_uvs=torch.rand(1, 300, 2) * 0.4 + 0.1,
+        )
+        b = TexturesUV(
+            maps=torch.full((1, 2000, 2000, 3), 0.7),
+            faces_uvs=torch.arange(150).reshape(1, 50, 3),
+            verts_uvs=torch.rand(1, 150, 2) * 0.2 + 0.3,
+        )
+        c = a.join_batch([b]).join_scene()
+
+        color = c.faces_verts_textures_packed()
+        color1 = color[:100, :, 0].flatten()
+        color2 = color[100:, :, 0].flatten()
+        expect1 = color1.new_tensor(0.8)
+        expect2 = color2.new_tensor(0.7)
+        self.assertClose(color1.min(), expect1)
+        self.assertClose(color1.max(), expect1)
+        self.assertClose(color2.min(), expect2)
+        self.assertClose(color2.max(), expect2)
+
+        if DEBUG:
+            from pytorch3d.vis.texture_vis import texturesuv_image_PIL as PI
+
+            PI(a, radius=5).save(DATA_DIR / "test_join_uvs_simple_a.png")
+            PI(b, radius=5).save(DATA_DIR / "test_join_uvs_simple_b.png")
+            PI(c, radius=5).save(DATA_DIR / "test_join_uvs_simple_c.png")
+
     def test_join_verts(self):
         """Meshes with TexturesVertex joined into a scene"""
         # Test the result of rendering two tori with separate textures.