From b0d05bb0485101dd05370c50886ccb8a7e6fe9e7 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 25 Apr 2023 02:54:12 +0000 Subject: [PATCH] compositing for all relevancies --- lerf/lerf.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/lerf/lerf.py b/lerf/lerf.py index 60d8420..b2912e9 100644 --- a/lerf/lerf.py +++ b/lerf/lerf.py @@ -214,10 +214,11 @@ def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle) -> Dic # TODO: handle lists of tensors as well continue outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1) # type: ignore - p_i = torch.clip(outputs["relevancy_0"] - 0.5, 0, 1) - outputs["composited"] = apply_colormap(p_i / (p_i.max() + 1e-6), "turbo") - mask = (outputs["relevancy_0"] < 0.5).squeeze() - outputs["composited"][mask, :] = outputs["rgb"][mask, :] + for i in range(len(self.image_encoder.positives)): + p_i = torch.clip(outputs[f"relevancy_{i}"] - 0.5, 0, 1) + outputs[f"composited_{i}"] = apply_colormap(p_i / (p_i.max() + 1e-6), "turbo") + mask = (outputs["relevancy_0"] < 0.5).squeeze() + outputs[f"composited_{i}"][mask, :] = outputs["rgb"][mask, :] return outputs def _get_outputs_nerfacto(self, ray_samples: RaySamples):