text: import torch
  import numpy as np
  import matplotlib.pyplot as plt
  from typing import Dict, List
  from train_eval.initialization import initialize_prediction_model, initialize_dataset, get_specific_args
  from nuscenes.eval.prediction.splits import get_prediction_challenge_split
  from nuscenes.prediction.input_representation.static_layers import StaticLayerRasterizer
  from nuscenes.prediction.input_representation.agents import AgentBoxesWithFadedHistory
  from nuscenes.prediction.input_representation.interface import InputRepresentation
  from nuscenes.prediction.input_representation.combinators import Rasterizer
  import train_eval.utils as u
  import imageio
  import os
  
  
  # Initialize device:
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  
  
  class Visualizer:
      """
      Class for visualizing predictions generated by trained model
      """
      def __init__(self, cfg: Dict, data_root: str, data_dir: str, checkpoint_path: str):
          """
          Initialize evaluator object
          :param cfg: Configuration parameters
          :param data_root: Root directory with data
          :param data_dir: Directory with extracted, pre-processed data
          :param checkpoint_path: Path to checkpoint with trained weights
          """
  
          # Initialize dataset
          ds_type = cfg['dataset'] + '_' + cfg['agent_setting'] + '_' + cfg['input_representation']
          spec_args = get_specific_args(cfg['dataset'], data_root, cfg['version'] if 'version' in cfg.keys() else None)
          test_set = initialize_dataset(ds_type, ['load_data', data_dir, cfg['test_set_args']] + spec_args)
          self.ds = test_set
  
          # Initialize model
          self.model = initialize_prediction_model(cfg['encoder_type'], cfg['aggregator_type'], cfg['decoder_type'],
                                                   cfg['encoder_args'], cfg['aggregator_args'], cfg['decoder_args'])
          self.model = self.model.float().to(device)
          self.model.eval()
  
          # Load checkpoint
          checkpoint = torch.load(checkpoint_path)
          self.model.load_state_dict(checkpoint['model_state_dict'])
  
      def visualize(self, output_dir: str, dataset_type: str):
          """
          Generate visualizations for predictions
          :param output_dir: results directory to dump visualizations in
          :param dataset_type: e.g. 'nuScenes'. Visualizations will vary based on dataset.
          :return:
          """
          if dataset_type == 'nuScenes':
              self.visualize_nuscenes(output_dir)
  
      def visualize_nuscenes(self, output_dir):
          index_list = self.get_vis_idcs_nuscenes()
          if not os.path.isdir(os.path.join(output_dir, 'results', 'gifs')):
              os.mkdir(os.path.join(output_dir, 'results', 'gifs'))
          for n, indices in enumerate(index_list):
              imgs = self.generate_nuscenes_gif(indices)
              filename = os.path.join(output_dir, 'results', 'gifs', 'example' + str(n) + '.gif')
              imageio.mimsave(filename, imgs, format='GIF', fps=2)
  
      def get_vis_idcs_nuscenes(self):
          """
          Returns list of list of indices for generating gifs for the nuScenes val set.
          Instance tokens are hardcoded right now. I'll fix this later (TODO)
          """
          token_list = get_prediction_challenge_split('val', dataroot=self.ds.helper.data.dataroot)
          instance_tokens = [token_list[idx].split("_")[0] for idx in range(len(token_list))]
          unique_instance_tokens = []
          for i_t in instance_tokens:
              if i_t not in unique_instance_tokens:
                  unique_instance_tokens.append(i_t)
  
          instance_tokens_to_visualize = [54, 98, 91, 5, 114, 144, 291, 204, 312, 187, 36, 267, 146]
  
          idcs = []
          for i_t_id in instance_tokens_to_visualize:
              idcs_i_t = [i for i in range(len(instance_tokens)) if instance_tokens[i] == unique_instance_tokens[i_t_id]]
              idcs.append(idcs_i_t)
  
          return idcs
  
      def generate_nuscenes_gif(self, idcs: List[int]):
          """
          Generates gif of predictions for the given set of indices.
          :param idcs: val set indices corresponding to a particular instance token.
          """
  
          # Raster maps for visualization.
          map_extent = self.ds.map_extent
          resolution = 0.1
          static_layer_rasterizer = StaticLayerRasterizer(self.ds.helper,
                                                          resolution=resolution,
                                                          meters_ahead=map_extent[3],
                                                          meters_behind=-map_extent[2],
                                                          meters_left=-map_extent[0],
                                                          meters_right=map_extent[1])
  
          agent_rasterizer = AgentBoxesWithFadedHistory(self.ds.helper, seconds_of_history=1,
                                                        resolution=resolution,
                                                        meters_ahead=map_extent[3],
                                                        meters_behind=-map_extent[2],
                                                        meters_left=-map_extent[0],
                                                        meters_right=map_extent[1])
  
          raster_maps = InputRepresentation(static_layer_rasterizer, agent_rasterizer, Rasterizer())
  
          imgs = []
          for idx in idcs:
  
              # Load data
              data = self.ds[idx]
              data = u.send_to_device(u.convert_double_to_float(u.convert2tensors(data)))
              i_t = data['inputs']['instance_token']
              s_t = data['inputs']['sample_token']
  
              # Get raster map
              hd_map = raster_maps.make_input_representation(i_t, s_t)
              r, g, b = hd_map[:, :, 0] / 255, hd_map[:, :, 1] / 255, hd_map[:, :, 2] / 255
              hd_map_gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
  
              # Predict
              predictions = self.model(data['inputs'])
  
              # Plot
              fig, ax = plt.subplots(1, 3, figsize=(15, 5))
              ax[0].imshow(hd_map, extent=self.ds.map_extent)
              ax[1].imshow(hd_map_gray, cmap='gist_gray', extent=self.ds.map_extent)
              ax[2].imshow(hd_map_gray, cmap='gist_gray', extent=self.ds.map_extent)
  
              for n, traj in enumerate(predictions['traj'][0]):
                  ax[1].plot(traj[:, 0].detach().cpu().numpy(), traj[:, 1].detach().cpu().numpy(), lw=4,
                             color='r', alpha=0.8)
                  ax[1].scatter(traj[-1, 0].detach().cpu().numpy(), traj[-1, 1].detach().cpu().numpy(), 60,
                                color='r', alpha=0.8)
  
              traj_gt = data['ground_truth']['traj'][0]
              ax[2].plot(traj_gt[:, 0].detach().cpu().numpy(), traj_gt[:, 1].detach().cpu().numpy(), lw=4, color='g')
              ax[2].scatter(traj_gt[-1, 0].detach().cpu().numpy(), traj_gt[-1, 1].detach().cpu().numpy(), 60, color='g')
  
              ax[0].axis('off')
              ax[1].axis('off')
              ax[2].axis('off')
              fig.tight_layout(pad=0)
              ax[0].margins(0)
              ax[1].margins(0)
              ax[2].margins(0)
  
              fig.canvas.draw()
              image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
              image_from_plot = image_from_plot.reshape(fig.canvas.get_width_height()[::-1] + (3,))
              imgs.append(image_from_plot)
              plt.close(fig)
  
          return imgs
  