# Copyright 2022 The Plenoptix Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Semantic NeRF implementation. """ from typing import Dict, Optional, Tuple import torch from torch import nn from torchtyping import TensorType from nerfactory.cameras.rays import RayBundle, RaySamples from nerfactory.data.structs import Semantics from nerfactory.fields.base import Field from nerfactory.fields.modules.encoding import Encoding, Identity, NeRFEncoding from nerfactory.fields.modules.field_heads import ( DensityFieldHead, FieldHeadNames, RGBFieldHead, SemanticFieldHead, ) from nerfactory.fields.modules.mlp import MLP from nerfactory.fields.nerf_field import NeRFField from nerfactory.graphs.vanilla_nerf import NeRFGraph from nerfactory.renderers.renderers import SemanticRenderer from nerfactory.utils import writer class SemanticNerfField(Field): """Semantic-NeRF field""" def __init__( self, num_semantic_classes: int, position_encoding: Encoding = Identity(in_dim=3), direction_encoding: Encoding = Identity(in_dim=3), base_mlp_num_layers: int = 8, base_mlp_layer_width: int = 256, head_mlp_num_layers: int = 2, head_mlp_layer_width: int = 128, skip_connections: Tuple[int] = (4,), ) -> None: super().__init__() self.num_semantic_classes = num_semantic_classes self.position_encoding = position_encoding self.direction_encoding = direction_encoding self.mlp_base = MLP( in_dim=self.position_encoding.get_out_dim(), num_layers=base_mlp_num_layers, layer_width=base_mlp_layer_width, skip_connections=skip_connections, ) self.mlp_head = MLP( in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim(), num_layers=head_mlp_num_layers, layer_width=head_mlp_layer_width, ) self.mlp_semantic = MLP( in_dim=self.mlp_head.get_out_dim(), layer_width=self.mlp_head.layer_width // 2, num_layers=1, activation=nn.ReLU(), ) self.field_head_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) self.field_head_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim()) self.field_head_semantic = SemanticFieldHead( in_dim=self.mlp_semantic.get_out_dim(), num_classes=self.num_semantic_classes ) def get_density(self, ray_samples: RaySamples): encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions()) base_mlp_out = self.mlp_base(encoded_xyz) density = self.field_head_density(base_mlp_out) return density, base_mlp_out def get_outputs( self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: encoded_dir = self.direction_encoding(ray_samples.frustums.directions) mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding], dim=-1)) outputs = {} # rgb outputs[self.field_head_rgb.field_head_name] = self.field_head_rgb(mlp_out) # semantic mlp_out_sem = self.mlp_semantic(mlp_out) outputs[self.field_head_semantic.field_head_name] = self.field_head_semantic(mlp_out_sem) return outputs class SemanticNerfGraph(NeRFGraph): """Semantic-NeRF graph""" def __init__(self, intrinsics=None, camera_to_world=None, semantics: Semantics = None, **kwargs) -> None: self.stuff_classes = semantics.stuff_classes self.stuff_colors = semantics.stuff_colors super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) def populate_fields(self): """Set the fields.""" position_encoding = NeRFEncoding( in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True ) direction_encoding = NeRFEncoding( in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True ) self.field_coarse = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding) num_semantic_classes = len(self.stuff_classes) self.field_fine = SemanticNerfField( num_semantic_classes, position_encoding=position_encoding, direction_encoding=direction_encoding ) def populate_misc_modules(self): super().populate_misc_modules() self.renderer_semantic = SemanticRenderer() self.cross_entropy_loss = nn.CrossEntropyLoss(reduction="mean") def get_outputs(self, ray_bundle: RayBundle): # uniform sampling ray_samples_uniform = self.sampler_uniform(ray_bundle) # coarse field field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform) weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( rgb=field_outputs_coarse[FieldHeadNames.RGB], weights=weights_coarse, ) accumulation_coarse = self.renderer_accumulation(weights_coarse) depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform) # pdf sampling ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) # fine field field_outputs_fine = self.field_fine.forward(ray_samples_pdf) weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) rgb_fine = self.renderer_rgb( rgb=field_outputs_fine[FieldHeadNames.RGB], weights=weights_fine, ) accumulation_fine = self.renderer_accumulation(weights_fine) depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf) semantic_fine = self.renderer_semantic(field_outputs_fine[FieldHeadNames.SEMANTICS], weights=weights_fine) outputs = { "rgb_coarse": rgb_coarse, "rgb_fine": rgb_fine, "accumulation_coarse": accumulation_coarse, "accumulation_fine": accumulation_fine, "depth_coarse": depth_coarse, "depth_fine": depth_fine, "semantic_fine": semantic_fine, } return outputs def get_loss_dict(self, outputs, batch): image = batch["image"] rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) semantic_logits = outputs["semantic_fine"] semantic_classes = batch["semantics"][..., 0].long() semantic_loss_fine = self.cross_entropy_loss(semantic_logits, semantic_classes) loss_dict = { "rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine, "semantic_loss_fine": semantic_loss_fine, } return loss_dict def log_test_image_outputs(self, image_idx, step, batch, outputs): super().log_test_image_outputs(image_idx, step, batch, outputs) semantic_logits = outputs["semantic_fine"] semantic_labels = torch.argmax(torch.nn.functional.softmax(semantic_logits, dim=-1), dim=-1) semantic_labels_image = self.stuff_colors[semantic_labels] writer.put_image(name=f"semantics/image_idx_{image_idx}", image=semantic_labels_image, step=step) mask = batch["mask"].repeat(1, 1, 3) writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step)