text: #!/usr/bin/env python3
  # Copyright (c) Facebook, Inc. and its affiliates.
  # All rights reserved.
  #
  # This source code is licensed under the BSD-style license found in the
  # LICENSE file in the root directory of this source tree.
  
  import collections
  import os
  import pickle
  import warnings
  import sys
  sys.path.insert(0, '/home/roironen/pytorch3d/projects/')
  import hydra
  import numpy as np
  import torch
  from CT.CT.visualization import  visualize_ct_outputs
  from nerf.nerf.dataset import trivial_collate
  from CT.CT.dataset import get_cloud_datasets
  from CT.CT.CTnet import CTnet
  from nerf.nerf.stats import Stats
  from omegaconf import DictConfig
  from visdom import Visdom
  
  
  CONFIG_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs")
  
  
  @hydra.main(config_path=CONFIG_DIR, config_name="basic_training")
  def main(cfg: DictConfig):
  
      # Set the relevant seeds for reproducibility.
      np.random.seed(cfg.seed)
      torch.manual_seed(cfg.seed)
  
      # Device on which to run.
      if torch.cuda.is_available():
          device = "cuda"
      else:
          warnings.warn(
              "Please note that although executing on CPU is supported,"
              + "the training is unlikely to finish in reasonable time."
          )
          device = "cpu"
  
      # Load the training/validation data.
      current_dir = os.path.dirname(os.path.realpath(__file__))
      DATA_DIR = os.path.join(current_dir, "data")
      train_dataset, val_dataset, _, n_cam = get_cloud_datasets(
          dataset_name=cfg.data.dataset_name,
          image_size=cfg.data.image_size,
      )
  
      # Initialize the Radiance Field model.
      model = CTnet(cfg=cfg, n_cam=n_cam)
  
      # Move the model to the relevant device.
      model.to(device)
  
      # Init stats to None before loading.
      stats = None
      optimizer_state_dict = None
      start_epoch = 0
  
      checkpoint_path = os.path.join(hydra.utils.get_original_cwd(), cfg.checkpoint_path)
      if len(cfg.checkpoint_path) > 0:
          # Make the root of the experiment directory.
          checkpoint_dir = os.path.split(checkpoint_path)[0]
          os.makedirs(checkpoint_dir, exist_ok=True)
  
          # Resume training if requested.
          if cfg.resume and os.path.isfile(checkpoint_path):
              print(f"Resuming from checkpoint {checkpoint_path}.")
              loaded_data = torch.load(checkpoint_path)
              model.load_state_dict(loaded_data["model"])
              stats = pickle.loads(loaded_data["stats"])
              print(f"   => resuming from epoch {stats.epoch}.")
              optimizer_state_dict = loaded_data["optimizer"]
              start_epoch = stats.epoch
  
      # Initialize the optimizer.
      optimizer = torch.optim.Adam(
          model.parameters(),
          lr=cfg.optimizer.lr,
      )
  
      # Load the optimizer state dict in case we are resuming.
      if optimizer_state_dict is not None:
          optimizer.load_state_dict(optimizer_state_dict)
          optimizer.last_epoch = start_epoch
  
      # Init the stats object.
      if stats is None:
          stats = Stats(
              ["loss", "relative_error", "lr", "max_memory", "sec/it"],
          )
  
      # Learning rate scheduler setup.
  
      # Following the original code, we use exponential decay of the
      # learning rate: current_lr = base_lr * gamma ** (epoch / step_size)
      def lr_lambda(epoch):
          return cfg.optimizer.lr_scheduler_gamma ** (
              epoch #/ cfg.optimizer.lr_scheduler_step_size
          )
  
      # The learning rate scheduling is implemented with LambdaLR PyTorch scheduler.
      lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
          optimizer, lr_lambda, last_epoch=start_epoch - 1, verbose=False
      )
  
      # Initialize the cache for storing variables needed for visualization.
      visuals_cache = collections.deque(maxlen=cfg.visualization.history_size)
  
      # Init the visualization visdom env.
      if cfg.visualization.visdom:
          viz = Visdom(
              server=cfg.visualization.visdom_server,
              port=cfg.visualization.visdom_port,
              use_incoming_socket=False,
          )
      else:
          viz = None
  
  
  
      # if cfg.data.precache_rays:
      #     # Precache the projection rays.
      #     model.eval()
      #     with torch.no_grad():
      #         for dataset in (train_dataset, val_dataset):
      #             for e in dataset:
      #                 cache_cameras = [cam.to(device) for cam in e["camera"]]
      #                 cache_camera_hashes = e["camera_idx"]#[e["camera_idx"] for e in dataset]
      #                 model.precache_rays(cache_cameras, cache_camera_hashes)
  
      train_dataloader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=1,
          shuffle=True,
          num_workers=0,
          collate_fn=trivial_collate,
      )
  
      # The validation dataloader is just an endless stream of random samples.
      val_dataloader = torch.utils.data.DataLoader(
          val_dataset,
          batch_size=1,
          num_workers=0,
          collate_fn=trivial_collate,
          sampler=torch.utils.data.RandomSampler(
              val_dataset,
              replacement=True,
              num_samples=cfg.optimizer.max_epochs,
          ),
      )
      mse = torch.nn.MSELoss()
      # Set the model to the training mode.
      model.train()
      # Run the main training loop.
      for epoch in range(start_epoch, cfg.optimizer.max_epochs):
          for iteration, batch in enumerate(train_dataloader):
              if iteration % (cfg.stats_print_interval*20) == 0:
                  stats.new_epoch()  # Init a new epoch.
              if iteration in cfg.optimizer.iter_steps:
                  # Adjust the learning rate.
                  lr_scheduler.step()
              image, camera, volume, camera_idx = batch[0].values()
              image = image.to(device)
              volume = volume.to(device)
              # R = torch.FloatTensor().to(device)
              # T = torch.FloatTensor().to(device)
              # for cam in camera:
              #     R = torch.cat((R, cam.R), dim=0)
              #     T = torch.cat((T, cam.T), dim=0)
              # camera = PerspectiveCameras(device=device, R=R, T=T)
              for cam in camera:
                  cam.to(device)
  
              optimizer.zero_grad()
  
              # Run the forward pass of the model.
              out = model(
                  camera,
                  image,
                  volume
              )
  
              # The loss is a sum of coarse and fine MSEs
              loss = mse(out["output"], out["volume"])
  
              # Take the training step.
              loss.backward()
              optimizer.step()
              with torch.no_grad():
                  relative_err = torch.norm(out["output"] - out["volume"],p=1) / (torch.norm(out["volume"],p=1) + 1e-6)
              # Update stats with the current metrics.
              stats.update(
                  {"loss": float(loss), "relative_error": float(relative_err), "lr": lr_scheduler.get_last_lr()[0],
                   "max_memory": float(round(torch.cuda.max_memory_allocated()/1e6))},
                  stat_set="train",
              )
  
              if iteration % cfg.stats_print_interval == 0:
                  stats.print(stat_set="train")
                  # if viz is not None:
                  #     # Plot that loss curves into visdom.
                  #     stats.plot_stats(
                  #         viz=viz,
                  #         visdom_env=cfg.visualization.visdom_env,
                  #         plot_file=None,
                  #     )
  
              # Update the visualization cache.
              if viz is not None:
                  visuals_cache.append(
                      {
                          "camera": camera,
                          "camera_idx": camera_idx,
                          "image": image.cpu().detach(),
                      }
                  )
  
  
          # Validation
          if epoch % cfg.validation_epoch_interval == 0 and epoch > 0:
  
              # Sample a validation data.
  
              val_batch = next(val_dataloader.__iter__())
              val_image, val_camera, val_volume, camera_idx = val_batch[0].values()
              val_image = val_image.to(device)
              val_volume = val_volume.to(device)
              for cam in camera:
                  cam.to(device)
  
              # Activate eval mode of the model (lets us do a full rendering pass).
              model.eval()
              with torch.no_grad():
                  val_out = model(
                      val_camera,
                      val_image,
                      val_volume
                  )
                  loss = mse(val_out["output"], val_out["volume"])
                  relative_err = torch.norm(val_out["output"] - val_out["volume"], p=1) / (torch.norm(val_out["volume"], p=1) + 1e-6)
  
              # Update stats with the validation metrics.
              stats.update({"loss": float(loss), "relative_error": float(relative_err)}, stat_set="val")
              stats.print(stat_set="val")
  
              if viz is not None:
                  # Plot that loss curves into visdom.
                  stats.plot_stats(
                      viz=viz,
                      visdom_env=cfg.visualization.visdom_env,
                      plot_file=None,
                  )
                  # Visualize the intermediate results.
                  visualize_ct_outputs(
                      out, visuals_cache, viz, cfg.visualization.visdom_env
                  )
  
              # Set the model back to train mode.
              model.train()
  
          # Checkpoint.
          if (
              epoch % cfg.checkpoint_epoch_interval == 0
              and len(cfg.checkpoint_path) > 0
              and epoch > 0
          ):
              print(f"Storing checkpoint {checkpoint_path}.")
              data_to_store = {
                  "model": model.state_dict(),
                  "optimizer": optimizer.state_dict(),
                  "stats": pickle.dumps(stats),
              }
              torch.save(data_to_store, checkpoint_path)
  
  
  if __name__ == "__main__":
      main()
  