Trouble updating all pixels from texture that contribute to render

[andrewharp]

I'm trying to implement the mesh fitting/texture updating example, but I want to use a TexturesUV object rather than TexturesVertex. I've implemented a custom shader so I can do this without lighting.

What I'm finding is that after switching the optimizer to Adam I can get the training renders to exactly match the ground truth targets, however the texture map is not being updated as I hoped. It seems like it is each pixel in the rendered image only gets mapped to a single pixel in the texture map. Since I have a finite set of images in the orbit that I generated, that means that not all pixels in the texture get touched.

This is the result of training with a 128x128 render. The texture is 1024x1024, so it makes sense that much of it is untouched on the obliquely angled parts of the cow if backprop only updates one pixel per pixel.

The problem decreases dramatically if I increase the render size, which increases the sampling resolution, but I'd like to understand what I'm missing to make it update all pixels in a texture that theoretically should contribute to pixels in a render.

I've tried switching from hard_rgb_blend to softmax_rgb_blend, changing faces_per_pixel, changing blur_radius (which seems to make it only update the texture along triangle seams??), etc. Apart from blur_radius (which has a negative effect) nothing really seems to make a difference.

Here's the relevant code

class SimpleShader(nn.Module):
    def __init__(self, device="gpu", blend_params=None, cameras=None):
        super().__init__()
        self.blend_params = blend_params if blend_params is not None else BlendParams()

    def forward(self, fragments, meshes: Meshes, **kwargs) -> torch.Tensor:
        blend_params = kwargs.get("blend_params", self.blend_params)
        texels = meshes.sample_textures(fragments)
        images = softmax_rgb_blend(texels, fragments, blend_params)
        return images  # (N, H, W, 4) RGBA image

...

        raster_settings_soft = RasterizationSettings(
            image_size=image_size,
            blur_radius=0, 
            faces_per_pixel=1,
            perspective_correct=False, 
        )
        rasterizer = MeshRasterizer(
            cameras=cameras,
            raster_settings=raster_settings_soft,

        )
        shader = SimpleShader(
            device=device,
            cameras=cameras,
            blend_params=BlendParams(sigma=1e-4, gamma=1e-4, background_color=torch.Tensor([1.0, 0.0, 0.0]).to(device)),
        )
        renderer_textured = MeshRenderer(
            rasterizer=rasterizer,
            shader=shader,
        )   

        if update_textures:
            assert hasattr(mesh.textures, "maps_padded"), "Mesh does not have a texture, but update_textures is enabled"

            move_vert_uvs = False
            texture_image = mesh.textures.maps_padded().clone().detach()            
            texture_image = torch.ones_like(texture_image) * 0.5

            texture_image.requires_grad = True
            logging.info(f"Texture image shape: {texture_image.shape}")

            verts_uvs = mesh.textures.verts_uvs_padded().clone().detach()
            verts_uvs.requires_grad = move_vert_uvs

            faces_uvs = mesh.textures.faces_uvs_padded().clone().detach()
            faces_uvs.requires_grad = False

            mesh.textures = TexturesUV(
                    maps=texture_image,
                    faces_uvs=faces_uvs,
                    verts_uvs=verts_uvs,
                    align_corners=False,
                    sampling_mode="bilinear",
                )

            logging.info(f"Texture image shape: {texture_image.shape}, verts_uvs shape: {verts_uvs.shape}")
            logging.info(f"verts: {verts_uvs} faces: {faces_uvs}")

...

        if update_textures:
            variables.append(texture_image)

            if move_vert_uvs:
                variables.append(verts_uvs)
            losses["rgb"] = {"weight": 10.0, "values": []}
...

        for i in range(num_iterations):
            # Initialize optimizer
            optimizer.zero_grad()

            new_src_mesh = mesh.offset_verts(deform_verts)

            # Losses to smooth/regularize the mesh shape
            loss = {k: torch.tensor(0.0, device=device) for k in losses}

            if update_positions:
                update_mesh_shape_prior_losses(new_src_mesh, loss)

            for j in np.random.permutation(num_views).tolist()[:num_views_per_iteration]:
                if update_positions:                
                    predicted_silhouette = renderer_silhouette(new_src_mesh, cameras=cameras[j])[..., 3]
                    loss_silhouette = ((predicted_silhouette - target_silhouette[j]) ** 2).mean()
                    loss["silhouette"] += loss_silhouette / num_views_per_iteration

                if update_textures:
                    predicted_rgb = renderer_textured(new_src_mesh, cameras=cameras[j])[..., :3]

                    if i % viz_schedule == 0:
                        predicted_rgbs.append(predicted_rgb.clone().detach())

                    loss_rgb = ((predicted_rgb - target_rgb[j]) ** 2).mean()
                    loss["rgb"] += loss_rgb / num_views_per_iteration

Thanks for any assistance!

[bottler]

This is a model problem rather than a software problem. Gradient descent optimization won't change unused inputs. We can't help.