import torch, torchvision from safetensors.torch import load_file from utils import log_utils import networks class RCNetModel(object): ''' Image radar fusion to determine correspondence of radar to image Arg(s): input_channels_image : int number of channels in the image input_channels_depth : int number of channels in depth map input_patch_size_image : int patch of image to consider for radar point encoder_type : str encoder type n_filters_encoder_image : list[int] list of filters for each layer in image encoder n_neurons_encoder_image : list[int] list of neurons for each layer in depth encoder decoder_type : str decoder type n_filters_decoder : list[int] list of filters for each layer in decoder weight_initializer : str kaiming_normal, kaiming_uniform, xavier_normal, xavier_uniform activation_func : str activation function for network device : torch.device device for running model ''' def __init__(self, input_channels_image, input_channels_depth, input_patch_size_image, encoder_type, n_filters_encoder_image, n_neurons_encoder_depth, decoder_type, n_filters_decoder, weight_initializer='kaiming_uniform', activation_func='leaky_relu', device=torch.device('cuda')): self.input_patch_size_image = input_patch_size_image self.device = device # height, width = input_patch_size_image # latent_height = np.ceil(height / 32.0).astype(int) # latent_width = np.ceil(width / 32.0).astype(int) height, width = input_patch_size_image latent_height = int((height // 32.0)) latent_width = int((width // 32.0)) latent_size_depth = latent_height * latent_width * n_neurons_encoder_depth[-1] # Build encoder if 'rcnet' in encoder_type: self.encoder = networks.RCNetEncoder( input_channels_image=input_channels_image, input_channels_depth=input_channels_depth, input_patch_size_image=input_patch_size_image, n_filters_encoder_image=n_filters_encoder_image, n_neurons_encoder_depth=n_neurons_encoder_depth, latent_size_depth=latent_size_depth, weight_initializer=weight_initializer, activation_func=activation_func, use_batch_norm='batch_norm' in encoder_type) else: raise ValueError('Encoder type {} not supported.'.format(encoder_type)) # Calculate number of channels for latent and skip connections combining image + depth n_skips = n_filters_encoder_image[:-1] n_skips = n_skips[::-1] + [0] latent_channels = n_filters_encoder_image[-1] + n_neurons_encoder_depth[-1] # Build decoder if 'multiscale' in decoder_type: self.decoder = networks.MultiScaleDecoder( input_channels=latent_channels, output_channels=1, n_resolution=1, n_filters=n_filters_decoder, n_skips=n_skips, weight_initializer=weight_initializer, activation_func=activation_func, output_func='linear', use_batch_norm='batch_norm' in decoder_type, deconv_type='up') else: raise ValueError('Decoder type {} not supported.'.format(decoder_type)) # Move to device self.to(self.device) def forward(self, image, point, bounding_boxes, return_logits=True): ''' Forwards the inputs through the network Arg(s): image : torch.Tensor[float32] N x 3 x H x W image point : torch.Tensor[float32] N x 3 input point return_logits : bool if set, then return logits otherwise sigmoid Returns: torch.Tensor[float32] : N x 1 x H x W logits (correspondence map) ''' latent, skips = self.encoder(image, point, bounding_boxes) logits = self.decoder(x=latent, skips=skips, shape=self.input_patch_size_image)[-1] if return_logits: return logits else: return torch.sigmoid(logits) def parameters(self): ''' Returns the list of parameters in the model Returns: list[torch.Tensor[float32]] : list of parameters ''' parameters = \ list(self.encoder.parameters()) + \ list(self.decoder.parameters()) return parameters def eval(self): ''' Sets model to evaluation mode ''' self.encoder.eval() self.decoder.eval() def to(self, device): ''' Moves model to specified device Arg(s): device : torch.device device for running model ''' # Move to device self.encoder.to(device) self.decoder.to(device) def restore_model(self, checkpoint_path): ''' Restore weights of the model Arg(s): checkpoint_path : str path to checkpoint Returns: int : current step in optimization None : retained for compatibility with the baseline inference call ''' checkpoint = load_file(checkpoint_path, device='cpu') encoder_prefix = 'radarnet_encoder.' decoder_prefix = 'radarnet_decoder.' encoder_state = { key[len(encoder_prefix):]: value for key, value in checkpoint.items() if key.startswith(encoder_prefix) } decoder_state = { key[len(decoder_prefix):]: value for key, value in checkpoint.items() if key.startswith(decoder_prefix) } self.encoder.load_state_dict(encoder_state) self.decoder.load_state_dict(decoder_state) return 0, None def data_parallel(self): ''' Allows multi-gpu split along batch ''' self.encoder = torch.nn.DataParallel(self.encoder) self.decoder = torch.nn.DataParallel(self.decoder) def log_summary(self, summary_writer, tag, step, image=None, output_response=None, output_label=None, output_depth=None, validity_map=None, ground_truth_label=None, ground_truth_depth=None, scalars={}, n_display=4): ''' Logs summary to Tensorboard Arg(s): summary_writer : SummaryWriter Tensorboard summary writer tag : str tag that prefixes names to log step : int current step in training image : torch.Tensor[float32] N x 3 x H x W image output_response : torch.Tensor[float32] N x 1 x H x W soft correspondence map output_label : torch.Tensor[float32] N x 1 x H x W binary correspondence map output_depth : torch.Tensor[float32] N x 1 x H x W depth map validity_map : torch.Tensor[float32] N x 1 x H x W validity map ground_truth_label : torch.Tensor[float32] N x 1 x H x W ground truth label ground_truth_depth : torch.Tensor[float32] N x 1 x H x W ground truth depth map scalars : dict[str, float] dictionary of scalars to log n_display : int number of images to display ''' with torch.no_grad(): display_summary_image = [] display_summary_image_text = tag if image is not None: image_summary = image[0:n_display, ...] display_summary_image_text += '-image' # Add to list of images to log display_summary_image.append(image_summary.cpu()) if output_response is not None: output_response_summary = output_response[0:n_display, ...] display_summary_image_text += '-output_response' # Add to list of images to log display_summary_image.append( log_utils.colorize( output_response_summary.cpu(), colormap='inferno')) # Log distribution of output response summary_writer.add_histogram( tag + '-output_response_distro', output_response_summary, global_step=step) if output_label is not None: output_label_summary = output_label[0:n_display, ...] display_summary_image_text += '-output_label' # Add to list of images to log display_summary_image.append( log_utils.colorize( output_label_summary.cpu(), colormap='inferno')) # Log distribution of output and label summary_writer.add_histogram( tag + '-output_label_distro', output_label_summary, global_step=step) if ground_truth_label is not None: ground_truth_label_summary = ground_truth_label[0:n_display, ...] validity_map_label_summary = torch.where( ground_truth_label_summary > 0, torch.ones_like(ground_truth_label_summary), torch.zeros_like(ground_truth_label_summary)) display_summary_image_text += '-ground_truth_label' if output_label is not None: display_summary_image_text += '-error' # Compute output error w.r.t. ground truth ground_truth_label_error_summary = \ torch.abs(output_label_summary - ground_truth_label_summary) ground_truth_label_error_summary = torch.where( validity_map_label_summary == 1.0, (ground_truth_label_error_summary + 1e-8) / (ground_truth_label_summary + 1e-8), validity_map_label_summary) display_summary_image.append( log_utils.colorize( ground_truth_label_error_summary.cpu(), colormap='inferno')) # Add to list of images to log display_summary_image.append( log_utils.colorize( ground_truth_label_summary.cpu(), colormap='inferno')) # Log distribution of ground truth summary_writer.add_histogram( tag + '_ground_truth_label_distro', ground_truth_label, global_step=step) if validity_map is not None: validity_map_summary = validity_map[0:n_display, ...] display_summary_image_text += '-validity_map' # Add to list of images to log display_summary_image.append( log_utils.colorize( validity_map_summary.cpu(), colormap='inferno')) if output_depth is not None: output_depth_summary = output_depth[0:n_display, ...] display_summary_image_text += '-output_depth' # Add to list of images to log display_summary_image.append( log_utils.colorize( (output_depth_summary / 100.0).cpu(), colormap='viridis')) # Log distribution of output depth summary_writer.add_histogram( tag + '-output_depth_distro', output_depth, global_step=step) if ground_truth_depth is not None: ground_truth_depth = torch.unsqueeze(ground_truth_depth[:, 0, :, :], dim=1) ground_truth_depth_summary = ground_truth_depth[0:n_display, ...] validity_map_summary = torch.where( ground_truth_depth_summary > 0, torch.ones_like(ground_truth_depth_summary), torch.zeros_like(ground_truth_depth_summary)) display_summary_image_text += '-ground_truth_label' if output_depth is not None: display_summary_image_text += '-error' # Compute output error w.r.t. ground truth ground_truth_depth_error_summary = \ torch.abs(output_depth_summary - ground_truth_depth_summary) ground_truth_depth_error_summary = torch.where( validity_map_summary == 1.0, (ground_truth_depth_error_summary + 1e-8) / (ground_truth_depth_summary + 1e-8), validity_map_summary) display_summary_image.append( log_utils.colorize( (ground_truth_depth_error_summary / 0.05).cpu(), colormap='inferno')) # Add to list of images to log display_summary_image.append( log_utils.colorize( (ground_truth_depth_summary / 100.0).cpu(), colormap='viridis')) # Log distribution of ground truth summary_writer.add_histogram( tag + '-ground_truth_distro', ground_truth_depth, global_step=step) # Log scalars to tensorboard for (name, value) in scalars.items(): summary_writer.add_scalar(tag + '-' + name, value, global_step=step) # Log image summaries to tensorboard if len(display_summary_image) > 1: display_summary_image = torch.cat(display_summary_image, dim=2) summary_writer.add_image( display_summary_image_text, torchvision.utils.make_grid(display_summary_image, nrow=n_display), global_step=step)