From cf9c557b7bf69ee9870d21ea96b58dee39c8437a Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 6 May 2022 18:07:12 -0700 Subject: [PATCH] Add colormaps to logs --- .vscode/settings.json | 3 +++ environment/requirements.txt | 1 + mattport/nerf/graph/vanilla_nerf.py | 9 +++++++-- mattport/nerf/trainer.py | 4 ++-- mattport/utils/visualization.py | 24 ++++++++++++++++++++++++ mattport/utils/writer.py | 4 ++-- 6 files changed, 39 insertions(+), 6 deletions(-) create mode 100644 mattport/utils/visualization.py diff --git a/.vscode/settings.json b/.vscode/settings.json index 6f7f45846..2196d04e5 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -10,6 +10,9 @@ }, "editor.formatOnSave": true, "python.formatting.provider": "black", + "python.formatting.blackArgs": [ + "--line-length=120" + ], "python.linting.pylintEnabled": true, "python.linting.flake8Enabled": false, "python.linting.enabled": true, diff --git a/environment/requirements.txt b/environment/requirements.txt index d26e36758..56624bc05 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -5,6 +5,7 @@ imageio==2.16.1 ipywidgets>=7.6 jupyterlab kaleido +matplotlib mediapy plotly==5.7.0 pylint==2.13.4 diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 16014b239..7c0bf5600 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -20,6 +20,7 @@ from mattport.nerf.renderers import AccumulationRenderer, DisparityRenderer, RGB from mattport.nerf.sampler import PDFSampler, UniformSampler from mattport.structures import colors from mattport.structures.rays import RaySamples +from mattport.utils import visualization class NeRFField(nn.Module): @@ -162,12 +163,16 @@ class NeRFGraph(Graph): fine_renderer_accumulation = self.renderer_accumulation(fine_weights) # RendererOutputs fine_renderer_disparity = self.renderer_disparity(fine_weights, pdf_ray_samples.ts) + # TODO refactor this into "vis" section. Doesn't need to be run during training. + coarse_renderer_accumulation = visualization.apply_colormap(coarse_renderer_accumulation.accumulation) + fine_renderer_accumulation = visualization.apply_colormap(fine_renderer_accumulation.accumulation) + # outputs: outputs = { "rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb, - "accumulation_coarse": coarse_renderer_accumulation.accumulation, - "accumulation_fine": fine_renderer_accumulation.accumulation, + "accumulation_coarse": coarse_renderer_accumulation, + "accumulation_fine": fine_renderer_accumulation, "disparity_coarse": coarse_renderer_disparity.disparity, "disparity_fine": fine_renderer_disparity.disparity, } diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 5ef5c149c..de289a29f 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -260,8 +260,8 @@ class Trainer: disparity_fine.append(graph_outputs["disparity_fine"]) rgb_coarse = torch.cat(rgb_coarse).view(image_height, image_width, 3).detach().cpu() rgb_fine = torch.cat(rgb_fine).view(image_height, image_width, 3).detach().cpu() - accumulation_coarse = torch.cat(accumulation_coarse).view(image_height, image_width, 1).detach().cpu() - accumulation_fine = torch.cat(accumulation_fine).view(image_height, image_width, 1).detach().cpu() + accumulation_coarse = torch.cat(accumulation_coarse).view(image_height, image_width, 3).detach().cpu() + accumulation_fine = torch.cat(accumulation_fine).view(image_height, image_width, 3).detach().cpu() disparity_coarse = torch.cat(disparity_coarse).view(image_height, image_width, 1).detach().cpu() disparity_fine = torch.cat(disparity_fine).view(image_height, image_width, 1).detach().cpu() diff --git a/mattport/utils/visualization.py b/mattport/utils/visualization.py new file mode 100644 index 000000000..df9f8de05 --- /dev/null +++ b/mattport/utils/visualization.py @@ -0,0 +1,24 @@ +""" Helper functions for visualizing outputs """ + +import torch +from matplotlib import cm +from torchtyping import TensorType + + +def apply_colormap(image: TensorType[..., 1], cmap="viridis") -> TensorType[..., 3]: + """Convert single channel to a color image. + + Args: + image (TensorType[..., 1]): Single channel image. + cmap (str, optional): Colormap for image. Defaults to 'turbo'. + + Returns: + TensorType[..., 3]: Colored image + """ + + colormap = cm.get_cmap(cmap) + colormap = torch.tensor(colormap.colors).to(image.device) + + image = (image * 255).long() + + return colormap[image] diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index 02eb1eddb..b73068b6a 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -89,13 +89,13 @@ class TensorboardWriter(Writer): self.tb_writer = SummaryWriter(log_dir=self.save_dir) def write_image( - self, name: str, x: TensorType["H", "W", 3], step: int, group: str = None, prefix: str = None + self, name: str, x: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None ) -> None: """_summary_ Args: name (str): data identifier - x (TensorType["H", "W", 3]): rendered image to write + x (TensorType["H", "W", "C"]): rendered image to write """ x = to8b(x) tensorboard_name = get_tensorboard_name(name, group=group, prefix=prefix)