Update exporter.py to export sh_degree 0 case #3371

[jb-ye]

Recover PR #3371

[brentyi]

Seems overall reasonable to me! Is there anywhere to reference which PLY keys are supported by which third-party tools, though? The f_dc_{0,1,2} behavior introduced in this PR seems standard but I'm not actually sure why/when the colors key from the original behavior is useful.

[jb-ye]

@brentyi

The f_dc_{} and f_rest_{} keys come from the original authors' GS code.

The colors key is not standard either and came from a previous PR (#3111). The common properties of PLY for storing colors are called "red", "green", "blue" in uchars, which was the very originally behavior. But I am kinda of okay about using "colors" instead of storing RGB in separate properties.

[brentyi]

Got it! Maybe I'm still wondering: is there anybody relying on the colors key? Are there actually third-party programs that support it?

Since it's relatively rare to train models without spherical harmonics I'm personally okay with changing the default behavior here.

@akristoffersen Justin commented that you might know?

[jb-ye]

Changing back to "red", "green", "blue" keys instead of "colors" is also an option, since they people can directly load the PLY as point cloud in meshlab or similar tools.

[MLLeKander]

The exporter is actually currently broken when sh_degree == 0. See: #3377

Maybe this is beyond the initial scope of this PR, but it might make sense to add logic to this PR to break apart colors into red/green/blue for the non-use_sh0_renderer case:

                    colors = torch.clamp(model.colors.clone(), 0.0, 1.0).data.cpu().numpy()
                    colors = (colors * 255).astype(np.uint8)

                    map_to_tensors["red"] = colors[:, 0]
                    map_to_tensors["green"] = colors[:, 1]
                    map_to_tensors["blue"] = colors[:, 2]

It would probably also be good to change the first check of write_ply to use tensor.size instead of len(tensor):

        if not all(tensor.size == count for tensor in map_to_tensors.values()):

[jb-ye]

@MLLeKander Thanks very much, let me update my PR with the proposed fix.

[brentyi]

One suggestion: what if we have a field ply_color_mode: Literal["sh_coeffs", "rgb"] where...

• The default is "sh_coeffs"
• If ply_color_mode == "sh_coeffs", then:
  • Write f_dc_{0,1,2} regardless of SH degree
  • Only write f_rest_{0,1,2} when the model's SH degree > 0
• If ply_color_mode == "rgb", then:
  • Write red, green, blue, regardless of SH degree

?

I'm not sure about the naming, but something like this seems slightly more consistent to me. Since then by default we can support all of the viewers that support f_dc_*/f_rest_*, and then red/green/blue outputs can be opted into. And both can be independent of the SH config of the original splatfacto model.

[MLLeKander]

I don't see a reason to not respect the ply_color_mode when sh_degree > 0. In fact, I think it would be better to have SplatfactoModel handle the sh_degree special cases, and simplify the exporter logic. It seems like SplatfactoModel.features_dc has a different interpretation in the sh_degree == 0 case, but this is not reflected in SplatfactoModel.shs_0.

In splatfacto.py:

    @property
    def colors(self):
        if self.config.sh_degree > 0:
            return torch.clamp(SH2RGB(self.shs_0), 0.0, 1.0)
        else:
            return torch.clamp(torch.sigmoid(self.features_dc), 0.0, 1.0)

    @property
    def shs_0(self):
        if self.config.sh_degree > 0:
            return self.features_dc
        else:
            return RGB2SH(self.colors)

It seems like these properties are only used in the exporter.

And in exporter.py:

            if self.ply_color_mode == "sh_coeffs":
                shs_0 = model.shs_0.data.cpu().numpy()
                for i in range(shs_0.shape[1]):
                    map_to_tensors[f"f_dc_{i}"] = shs_0[:, i]

            elif self.ply_color_mode == "rgb":
                colors = model.colors.data.cpu().numpy()
                colors = (colors * 255).astype(np.uint8)

                map_to_tensors["red"] = colors[:, 0]
                map_to_tensors["green"] = colors[:, 1]
                map_to_tensors["blue"] = colors[:, 2]

            shs_rest = model.shs_rest.transpose(1, 2).contiguous().cpu().numpy()
            shs_rest = shs_rest.reshape((n, -1))
            for i in range(shs_rest.shape[-1]):
                map_to_tensors[f"f_rest_{i}"] = shs_rest[:, i, None]

Note that shs_rest will be of shape (n, 0), so we won't enter the loop when sh_degree == 0.

I haven't gotten a chance to test/verify this code, so take it with a grain of salt.

[jb-ye]

@MLLeKander Could you take a look at my latest revision (maybe test your color-only model)?

[MLLeKander]

The rgb color mode looks great. However, it seems that the result when ply_color_mode == "sh_coeffs" and sh_degree == 0 is incorrect. When sh_degree > 0, SplatfactoModel.features_dc correspond to the 0-degree SH coefficients, but when sh_degree == 0, it is the logit of the RGB color. However, the SplatfactoModel.sh_0 property just returns SplatfactoModel.features_dc.

In the sh_degree == 0 case, we will need to convert from features_dc to SH coefficients, presumably with RGB2SH. I think it makes most sense to put this logic in models/splatfacto.py (mirroring the conditional in .colors), but we could also add a special case to exporter.py.

[jb-ye]

@MLLeKander @brentyi I think it is now ready to merge. Let me know if you have other suggestions.

[MLLeKander]

Just verified, both ply_color_modes looks great on my model!

[brentyi]

lgtm too, thanks @jb-ye @MLLeKander!

[brentyi]

also cc @mayaz8h who has been fixing up #2951