text: # Copyright (c) Microsoft Corporation. All rights reserved.
  # Licensed under the MIT License.
  # commitHash: c76bf579a0d5090ebd32426907d051d499f3e847
  # url: https://github.com/yalaudah/facies_classification_benchmark
  #
  # To Test:
  # python test.py TRAIN.END_EPOCH 1 TRAIN.SNAPSHOTS 1 --cfg "configs/seresnet_unet.yaml" --debug
  #
  # /* spell-checker: disable */
  """
  Modified version of the Alaudah testing script
  Runs only on single GPU
  """
  
  import itertools
  import json
  import logging
  import logging.config
  import os
  from os import path
  
  import fire
  import numpy as np
  import torch
  import torch.nn.functional as F
  from albumentations import Compose, Normalize, PadIfNeeded, Resize
  from toolz import compose, curry, itertoolz, pipe, take
  from torch.utils import data
  
  from cv_lib.segmentation import models
  from cv_lib.segmentation.dutchf3.utils import current_datetime, git_branch, git_hash
  
  from cv_lib.utils import load_log_configuration, mask_to_disk, generate_path, image_to_disk
  from deepseismic_interpretation.dutchf3.data import add_patch_depth_channels, get_test_loader
  from default import _C as config
  from default import update_config
  
  _CLASS_NAMES = [
      "upper_ns",
      "middle_ns",
      "lower_ns",
      "rijnland_chalk",
      "scruff",
      "zechstein",
  ]
  
  
  class runningScore(object):
      def __init__(self, n_classes):
          self.n_classes = n_classes
          self.confusion_matrix = np.zeros((n_classes, n_classes))
  
      def _fast_hist(self, label_true, label_pred, n_class):
          mask = (label_true >= 0) & (label_true < n_class)
          hist = np.bincount(n_class * label_true[mask].astype(int) + label_pred[mask], minlength=n_class ** 2,).reshape(
              n_class, n_class
          )
          return hist
  
      def update(self, label_trues, label_preds):
          for lt, lp in zip(label_trues, label_preds):
              self.confusion_matrix += self._fast_hist(lt.flatten(), lp.flatten(), self.n_classes)
  
      def get_scores(self):
          """Returns accuracy score evaluation result.
              - overall accuracy
              - mean accuracy
              - mean IU
              - fwavacc
          """
          hist = self.confusion_matrix
          acc = np.diag(hist).sum() / hist.sum()
          acc_cls = np.diag(hist) / hist.sum(axis=1)
          mean_acc_cls = np.nanmean(acc_cls)
          iu = np.diag(hist) / (hist.sum(axis=1) + hist.sum(axis=0) - np.diag(hist))
          mean_iu = np.nanmean(iu)
          freq = hist.sum(axis=1) / hist.sum()  # fraction of the pixels that come from each class
          fwavacc = (freq[freq > 0] * iu[freq > 0]).sum()
          cls_iu = dict(zip(range(self.n_classes), iu))
  
          return (
              {
                  "Pixel Acc: ": acc,
                  "Class Accuracy: ": acc_cls,
                  "Mean Class Acc: ": mean_acc_cls,
                  "Freq Weighted IoU: ": fwavacc,
                  "Mean IoU: ": mean_iu,
                  "confusion_matrix": self.confusion_matrix,
              },
              cls_iu,
          )
  
      def reset(self):
          self.confusion_matrix = np.zeros((self.n_classes, self.n_classes))
  
  
  def _transform_CHW_to_HWC(numpy_array):
      return np.moveaxis(numpy_array, 0, -1)
  
  
  def _transform_HWC_to_CHW(numpy_array):
      return np.moveaxis(numpy_array, -1, 0)
  
  
  @curry
  def _apply_augmentation3D(aug, numpy_array):
      assert len(numpy_array.shape) == 3, "This method only accepts 3D arrays"
      patch = _transform_CHW_to_HWC(numpy_array)
      patch = aug(image=patch)["image"]
      return _transform_HWC_to_CHW(patch)
  
  
  @curry
  def _apply_augmentation2D(aug, numpy_array):
      assert len(numpy_array.shape) == 2, "This method only accepts 2D arrays"
      return aug(image=numpy_array)["image"]
  
  
  _AUGMENTATION = {3: _apply_augmentation3D, 2: _apply_augmentation2D}
  
  
  @curry
  def _apply_augmentation(aug, image):
      if isinstance(image, torch.Tensor):
          image = image.numpy()
  
      if aug is not None:
          return _AUGMENTATION[len(image.shape)](aug, image)
      else:
          return image
  
  
  def _add_depth(image):
      if isinstance(image, torch.Tensor):
          image = image.numpy()
      return add_patch_depth_channels(image)
  
  
  def _to_torch(image):
      if isinstance(image, torch.Tensor):
          return image
      else:
          return torch.from_numpy(image).to(torch.float32)
  
  
  def _expand_dims_if_necessary(torch_tensor):
      if len(torch_tensor.shape) == 2:
          return torch_tensor.unsqueeze(dim=0)
      else:
          return torch_tensor
  
  
  @curry
  def _extract_patch(hdx, wdx, ps, patch_size, img_p):
      if len(img_p.shape) == 2:  # 2D
          return img_p[hdx + ps : hdx + ps + patch_size, wdx + ps : wdx + ps + patch_size]
      else:  # 3D
          return img_p[
              :, hdx + ps : hdx + ps + patch_size, wdx + ps : wdx + ps + patch_size,
          ]
  
  
  def _compose_processing_pipeline(depth, aug=None):
      steps = []
      if aug is not None:
          steps.append(_apply_augmentation(aug))
  
      if depth == "patch":
          steps.append(_add_depth)
  
      steps.append(_to_torch)
      steps.append(_expand_dims_if_necessary)
      steps.reverse()
      return compose(*steps)
  
  
  def _generate_batches(h, w, ps, patch_size, stride, batch_size=64):
      hdc_wdx_generator = itertools.product(range(0, h - patch_size + ps, stride), range(0, w - patch_size + ps, stride),)
      for batch_indexes in itertoolz.partition_all(batch_size, hdc_wdx_generator):
          yield batch_indexes
  
  
  @curry
  def _output_processing_pipeline(config, output):
      output = output.unsqueeze(0)
      _, _, h, w = output.shape
      if config.TEST.POST_PROCESSING.SIZE != h or config.TEST.POST_PROCESSING.SIZE != w:
          output = F.interpolate(
              output, size=(config.TEST.POST_PROCESSING.SIZE, config.TEST.POST_PROCESSING.SIZE,), mode="bilinear",
          )
  
      if config.TEST.POST_PROCESSING.CROP_PIXELS > 0:
          _, _, h, w = output.shape
          output = output[
              :,
              :,
              config.TEST.POST_PROCESSING.CROP_PIXELS : h - config.TEST.POST_PROCESSING.CROP_PIXELS,
              config.TEST.POST_PROCESSING.CROP_PIXELS : w - config.TEST.POST_PROCESSING.CROP_PIXELS,
          ]
      return output.squeeze()
  
  
  def _patch_label_2d(
      model,
      img,
      pre_processing,
      output_processing,
      patch_size,
      stride,
      batch_size,
      device,
      num_classes,
      split,
      debug,
      MIN,
      MAX,
  ):
      """Processes a whole section
      """
  
      img = torch.squeeze(img)
      h, w = img.shape[-2], img.shape[-1]  # height and width
  
      # Pad image with patch_size/2:
      ps = int(np.floor(patch_size / 2))  # pad size
      img_p = F.pad(img, pad=(ps, ps, ps, ps), mode="constant", value=0)
      output_p = torch.zeros([1, num_classes, h + 2 * ps, w + 2 * ps])
  
      # generate output:
      for batch_indexes in _generate_batches(h, w, ps, patch_size, stride, batch_size=batch_size):
          batch = torch.stack(
              [pipe(img_p, _extract_patch(hdx, wdx, ps, patch_size), pre_processing,) for hdx, wdx in batch_indexes],
              dim=0,
          )
  
          model_output = model(batch.to(device))
  
          for (hdx, wdx), output in zip(batch_indexes, model_output.detach().cpu()):
              output = output_processing(output)
              output_p[:, :, hdx + ps : hdx + ps + patch_size, wdx + ps : wdx + ps + patch_size,] += output
  
          # dump the data right before it's being put into the model and after scoring
          if debug:
              outdir = f"debug/test/batch_{split}"
              generate_path(outdir)
              for i in range(batch.shape[0]):
                  path_prefix = f"{outdir}/{batch_indexes[i][0]}_{batch_indexes[i][1]}"
                  model_output = model_output.detach().cpu()
                  # save image:
                  image_to_disk(np.array(batch[i, 0, :, :]), path_prefix + "_img.png", MIN, MAX)
                  # dump model prediction:
                  mask_to_disk(model_output[i, :, :, :].argmax(dim=0).numpy(), path_prefix + "_pred.png", num_classes)
                  # dump model confidence values
                  for nclass in range(num_classes):
                      image_to_disk(
                          model_output[i, nclass, :, :].numpy(), path_prefix + f"_class_{nclass}_conf.png", MIN, MAX
                      )
  
      # crop the output_p in the middle
      output = output_p[:, :, ps:-ps, ps:-ps]
      return output
  
  
  def _evaluate_split(
      split,
      section_aug,
      model,
      pre_processing,
      output_processing,
      device,
      running_metrics_overall,
      config,
      data_flow,
      debug=False,
  ):
      logger = logging.getLogger(__name__)
  
      TestSectionLoader = get_test_loader(config)
  
      test_set = TestSectionLoader(config, split=split, is_transform=True, augmentations=section_aug, debug=debug,)
  
      n_classes = test_set.n_classes
  
      if debug:
          data_flow[split] = dict()
          data_flow[split]["test_section_loader_length"] = len(test_set)
          data_flow[split]["test_input_shape"] = test_set.seismic.shape
          data_flow[split]["test_label_shape"] = test_set.labels.shape
          data_flow[split]["n_classes"] = n_classes
  
      test_loader = data.DataLoader(test_set, batch_size=1, num_workers=config.WORKERS, shuffle=False)
  
      if debug:
          data_flow[split]["test_loader_length"] = len(test_loader)
          logger.info("Running in Debug/Test mode")
          take_n = 2
          test_loader = take(take_n, test_loader)
          data_flow[split]["take_n_sections"] = take_n
          pred_list, gt_list, img_list = [], [], []
  
      try:
          output_dir = generate_path(
              f"{config.OUTPUT_DIR}/test/{split}", git_branch(), git_hash(), config.MODEL.NAME, current_datetime(),
          )
      except:
          output_dir = generate_path(f"{config.OUTPUT_DIR}/test/{split}", config.MODEL.NAME, current_datetime(),)
  
      running_metrics_split = runningScore(n_classes)
  
      # evaluation mode:
      with torch.no_grad():  # operations inside don't track history
          model.eval()
          for i, (images, labels) in enumerate(test_loader):
              logger.info(f"split: {split}, section: {i}")
              outputs = _patch_label_2d(
                  model,
                  images,
                  pre_processing,
                  output_processing,
                  config.TRAIN.PATCH_SIZE,
                  config.TEST.TEST_STRIDE,
                  config.VALIDATION.BATCH_SIZE_PER_GPU,
                  device,
                  n_classes,
                  split,
                  debug,
                  config.DATASET.MIN,
                  config.DATASET.MAX,
              )
  
              pred = outputs.detach().max(1)[1].numpy()
              gt = labels.numpy()
              if debug:
                  pred_list.append((pred.shape, len(np.unique(pred))))
                  gt_list.append((gt.shape, len(np.unique(gt))))
                  img_list.append(images.numpy().shape)
  
              running_metrics_split.update(gt, pred)
              running_metrics_overall.update(gt, pred)
  
              #  dump images to disk for review
              mask_to_disk(pred.squeeze(), os.path.join(output_dir, f"{i}_pred.png"), n_classes)
              mask_to_disk(gt.squeeze(), os.path.join(output_dir, f"{i}_gt.png"), n_classes)
  
      if debug:
          data_flow[split]["pred_shape"] = pred_list
          data_flow[split]["gt_shape"] = gt_list
          data_flow[split]["img_shape"] = img_list
  
      # get scores
      score, class_iou = running_metrics_split.get_scores()
  
      # Log split results
      logger.info(f'Pixel Acc: {score["Pixel Acc: "]:.3f}')
      if debug:
          for cdx in range(n_classes):
              logger.info(f'  Class_{cdx}_accuracy {score["Class Accuracy: "][cdx]:.3f}')
      else:
          for cdx, class_name in enumerate(_CLASS_NAMES):
              logger.info(f'  {class_name}_accuracy {score["Class Accuracy: "][cdx]:.3f}')
  
      logger.info(f'Mean Class Acc: {score["Mean Class Acc: "]:.3f}')
      logger.info(f'Freq Weighted IoU: {score["Freq Weighted IoU: "]:.3f}')
      logger.info(f'Mean IoU: {score["Mean IoU: "]:0.3f}')
      running_metrics_split.reset()
  
  
  def _write_section_file(labels, section_file):
      # define indices of the array
      irange, xrange, depth = labels.shape
  
      if config.TEST.INLINE:
          i_list = list(range(irange))
          i_list = ["i_" + str(inline) for inline in i_list]
      else:
          i_list = []
  
      if config.TEST.CROSSLINE:
          x_list = list(range(xrange))
          x_list = ["x_" + str(crossline) for crossline in x_list]
      else:
          x_list = []
  
      list_test = i_list + x_list
  
      file_object = open(section_file, "w")
      file_object.write("\n".join(list_test))
      file_object.close()
  
  
  def test(*options, cfg=None, debug=False):
  
      update_config(config, options=options, config_file=cfg)
      n_classes = config.DATASET.NUM_CLASSES
  
      # Start logging
      load_log_configuration(config.LOG_CONFIG)
      logger = logging.getLogger(__name__)
      device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
      log_dir, _ = os.path.split(config.TEST.MODEL_PATH)
  
      # load model:
      model = getattr(models, config.MODEL.NAME).get_seg_model(config)
      trained_model = torch.load(config.TEST.MODEL_PATH)
      trained_model = {k.replace("module.", ""): v for (k, v) in trained_model.items()}
      model.load_state_dict(trained_model, strict=True)
      model = model.to(device)
  
      running_metrics_overall = runningScore(n_classes)
  
      # Augmentation
      section_aug = Compose([Normalize(mean=(config.TRAIN.MEAN,), std=(config.TRAIN.STD,), max_pixel_value=1,)])
  
      # TODO: make sure that this is consistent with how normalization and agumentation for train.py
      # issue: https://github.com/microsoft/seismic-deeplearning/issues/270
      patch_aug = Compose(
          [
              Resize(
                  config.TRAIN.AUGMENTATIONS.RESIZE.HEIGHT, config.TRAIN.AUGMENTATIONS.RESIZE.WIDTH, always_apply=True,
              ),
              PadIfNeeded(
                  min_height=config.TRAIN.AUGMENTATIONS.PAD.HEIGHT,
                  min_width=config.TRAIN.AUGMENTATIONS.PAD.WIDTH,
                  border_mode=config.OPENCV_BORDER_CONSTANT,
                  always_apply=True,
                  mask_value=255,
              ),
          ]
      )
  
      pre_processing = _compose_processing_pipeline(config.TRAIN.DEPTH, aug=patch_aug)
      output_processing = _output_processing_pipeline(config)
  
      splits = ["test1", "test2"] if "Both" in config.TEST.SPLIT else [config.TEST.SPLIT]
      data_flow = dict()
      for sdx, split in enumerate(splits):
          labels = np.load(path.join(config.DATASET.ROOT, "test_once", split + "_labels.npy"))
          section_file = path.join(config.DATASET.ROOT, "splits", "section_" + split + ".txt")
          _write_section_file(labels, section_file)
          _evaluate_split(
              split,
              section_aug,
              model,
              pre_processing,
              output_processing,
              device,
              running_metrics_overall,
              config,
              data_flow,
              debug=debug,
          )
  
      if debug:
          config_file_name = "default_config" if not cfg else cfg.split("/")[-1].split(".")[0]
  
          fname = f"data_flow_test_{config_file_name}_{config.TRAIN.MODEL_DIR}.json"
          with open(fname, "w") as f:
              json.dump(data_flow, f, indent=1)
  
      # FINAL TEST RESULTS:
      score, class_iou = running_metrics_overall.get_scores()
  
      logger.info("--------------- FINAL RESULTS -----------------")
      logger.info(f'Pixel Acc: {score["Pixel Acc: "]:.4f}')
  
      if debug:
          for cdx in range(n_classes):
              logger.info(f'  Class_{cdx}_accuracy {score["Class Accuracy: "][cdx]:.3f}')
      else:
          for cdx, class_name in enumerate(_CLASS_NAMES):
              logger.info(f'     {class_name}_accuracy {score["Class Accuracy: "][cdx]:.4f}')
  
      logger.info(f'Mean Class Acc: {score["Mean Class Acc: "]:.4f}')
      logger.info(f'Freq Weighted IoU: {score["Freq Weighted IoU: "]:.4f}')
      logger.info(f'Mean IoU: {score["Mean IoU: "]:0.4f}')
  
      # Save confusion matrix:
      confusion = score["confusion_matrix"]
      np.savetxt(path.join(log_dir, "confusion.csv"), confusion, delimiter=" ")
  
      if debug:
          fname = f"metrics_test_{config_file_name}_{config.TRAIN.MODEL_DIR}.json"
          with open(fname, "w") as fid:
              json.dump(
                  {
                      metric: score[metric]
                      for metric in ["Pixel Acc: ", "Mean Class Acc: ", "Freq Weighted IoU: ", "Mean IoU: "]
                  },
                  fid,
              )
  
  
  if __name__ == "__main__":
      fire.Fire(test)
  