class_docstr: Prediction pipeline:
  - receives a patched patient image (n_patches, c, y, x, (z)) from patient data loader.
  - forwards patches through model in chunks of batch_size. (method: batch_tiling_forward)
  - unmolds predictions (boxes and segmentations) to original patient coordinates. (method: spatial_tiling_forward)
  
  Ensembling (mode == 'test'):
  - for inference, forwards 4 mirrored versions of image to through model and unmolds predictions afterwards
    accordingly (method: data_aug_forward)
  - for inference, loads multiple parameter-sets of the trained model corresponding to different epochs. for each
    parameter-set loops over entire test set, runs prediction pipeline for each patient. (method: predict_test_set)
  
  Consolidation of predictions:
  - consolidates a patient's predictions (boxes, segmentations) collected over patches, data_aug- and temporal ensembling,
    performs clustering and weighted averaging (external function: apply_wbc_to_patient) to obtain consistent outptus.
  - for 2D networks, consolidates box predictions to 3D cubes via clustering (adaption of non-maximum surpression).
    (external function: merge_2D_to_3D_preds_per_patient)
  
  Ground truth handling:
  - dissmisses any ground truth boxes returned by the model (happens in validation mode, patch-based groundtruth)
  - if provided by data loader, adds 3D ground truth to the final predictions to be passed to the evaluator.
class_docstr_tok: ["Prediction","pipeline",":","-","receives","a","patched","patient","image","(","n_patches",",","c",",","y",",","x",",","(","z",")",")","from","patient","data","loader",".","-","forwards","patches","through","model","in","chunks","of","batch_size",".","(","method",":","batch_tiling_forward",")","-","unmolds","predictions","(","boxes","and","segmentations",")","to","original","patient","coordinates",".","(","method",":","spatial_tiling_forward",")","Ensembling","(","mode","==","'test","'",")",":","-","for","inference",",","forwards","4","mirrored","versions","of","image","to","through","model","and","unmolds","predictions","afterwards","accordingly","(","method",":","data_aug_forward",")","-","for","inference",",","loads","multiple","parameter-sets","of","the","trained","model","corresponding","to","different","epochs",".","for","each","parameter-set","loops","over","entire","test","set",",","runs","prediction","pipeline","for","each","patient",".","(","method",":","predict_test_set",")","Consolidation","of","predictions",":","-","consolidates","a","patient","'s","predictions","(","boxes",",","segmentations",")","collected","over","patches",",","data_aug-","and","temporal","ensembling",",","performs","clustering","and","weighted","averaging","(","external","function",":","apply_wbc_to_patient",")","to","obtain","consistent","outptus",".","-","for","2D","networks",",","consolidates","box","predictions","to","3D","cubes","via","clustering","(","adaption","of","non-maximum","surpression",")",".","(","external","function",":","merge_2D_to_3D_preds_per_patient",")","Ground","truth","handling",":","-","dissmisses","any","ground","truth","boxes","returned","by","the","model","(","happens","in","validation","mode",",","patch-based","groundtruth",")","-","if","provided","by","data","loader",",","adds","3D","ground","truth","to","the","final","predictions","to","be","passed","to","the","evaluator","."]
docstr: wrapper around test method, which loads multiple (or one) epoch parameters (temporal ensembling), loops through
  the test set and collects predictions per patient. Also flattens the results per patient and epoch
  and adds optional ground truth boxes for evaluation. Saves out the raw result list for later analysis and
  optionally consolidates and returns predictions immediately.
  :return: (optionally) list_of_results_per_patient: list over patient results. each entry is a dict with keys:
           - 'boxes': list over batch elements. each element is a list over boxes, where each box is
                      one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions
                      (if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.
           - 'seg_preds': not implemented yet. todo for evaluation of instance/semantic segmentation.
docstr_tok: ["wrapper","around","test","method",",","which","loads","multiple","(","or","one",")","epoch","parameters","(","temporal","ensembling",")",",","loops","through","the","test","set","and","collects","predictions","per","patient",".","Also","flattens","the","results","per","patient","and","epoch","and","adds","optional","ground","truth","boxes","for","evaluation",".","Saves","out","the","raw","result","list","for","later","analysis","and","optionally","consolidates","and","returns","predictions","immediately",".",":","return",":","(","optionally",")","list_of_results_per_patient",":","list","over","patient","results",".","each","entry","is","a","dict","with","keys",":","-","'boxes","'",":","list","over","batch","elements",".","each","element","is","a","list","over","boxes",",","where","each","box","is","one","dictionary",":","[","[","box_0",",","...","]",",","[","box_n",",","...","]","]",".","batch","elements","are","slices","for","2D","predictions","(","if","not","merged","to","3D",")",",","and","a","dummy","batch","dimension","of","1","for","3D","predictions",".","-","'seg_preds","'",":","not","implemented","yet",".","todo","for","evaluation","of","instance\/semantic","segmentation","."]
code: def predict_test_set(self, batch_gen, return_results=True):
      """
      wrapper around test method, which loads multiple (or one) epoch parameters (temporal ensembling), loops through
      the test set and collects predictions per patient. Also flattens the results per patient and epoch
      and adds optional ground truth boxes for evaluation. Saves out the raw result list for later analysis and
      optionally consolidates and returns predictions immediately.
      :return: (optionally) list_of_results_per_patient: list over patient results. each entry is a dict with keys:
               - 'boxes': list over batch elements. each element is a list over boxes, where each box is
                          one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions
                          (if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.
               - 'seg_preds': not implemented yet. todo for evaluation of instance/semantic segmentation.
      """
      dict_of_patient_results = OrderedDict()
  
      # get paths of all parameter sets to be loaded for temporal ensembling. (or just one for no temp. ensembling).
      weight_paths = [
          os.path.join(
              self.cf.fold_dir,
              "{}_best_checkpoint".format(epoch),
              "params.pth",
          )
          for epoch in self.epoch_ranking
      ]
  
      for rank_ix, weight_path in enumerate(weight_paths):
          self.logger.info(
              (
                  "tmp ensembling over rank_ix:{} epoch:{}".format(
                      rank_ix, weight_path
                  )
              )
          )
          self.net.load_state_dict(torch.load(weight_path))
          self.net.eval()
          self.rank_ix = str(
              rank_ix
          )  # get string of current rank for unique patch ids.
  
          with torch.no_grad():
              for _ in range(batch_gen["n_test"]):
                  batch = next(batch_gen["test"])
  
                  # store batch info in patient entry of results dict.
                  if rank_ix == 0:
                      dict_of_patient_results[batch["pid"]] = {}
                      dict_of_patient_results[batch["pid"]][
                          "results_list"
                      ] = []
                      dict_of_patient_results[batch["pid"]][
                          "patient_bb_target"
                      ] = batch["patient_bb_target"]
                      dict_of_patient_results[batch["pid"]][
                          "patient_roi_labels"
                      ] = batch["patient_roi_labels"]
  
                  # call prediction pipeline and store results in dict.
                  results_dict = self.predict_patient(batch)
                  dict_of_patient_results[batch["pid"]][
                      "results_list"
                  ].append(results_dict["boxes"])
  
      self.logger.info(
          "finished predicting test set. starting post-processing of predictions."
      )
      list_of_results_per_patient = []
  
      # loop over patients again to flatten results across epoch predictions.
      # if provided, add ground truth boxes for evaluation.
      for pid, p_dict in dict_of_patient_results.items():
          tmp_ens_list = p_dict["results_list"]
          results_dict = {}
          # collect all boxes/seg_preds of same batch_instance over temporal instances.
          results_dict["boxes"] = [
              [item for d in tmp_ens_list for item in d[batch_instance]]
              for batch_instance in range(len(tmp_ens_list[0]))
          ]
  
          # TODO return for instance segmentation:
          # results_dict['seg_preds'] = np.mean(results_dict['seg_preds'], 1)[:, None]
          # results_dict['seg_preds'] = np.array([[item for d in tmp_ens_list for item in d['seg_preds'][batch_instance]]
          #                                       for batch_instance in range(len(tmp_ens_list[0]['boxes']))])
  
          # add 3D ground truth boxes for evaluation.
          for b in range(p_dict["patient_bb_target"].shape[0]):
              for t in range(len(p_dict["patient_bb_target"][b])):
                  results_dict["boxes"][b].append(
                      {
                          "box_coords": p_dict["patient_bb_target"][b][
                              t
                          ],
                          "box_label": p_dict["patient_roi_labels"][b][
                              t
                          ],
                          "box_type": "gt",
                      }
                  )
  
          list_of_results_per_patient.append(
              [results_dict["boxes"], pid]
          )
  
      # save out raw predictions.
      out_string = (
          "raw_pred_boxes_hold_out_list"
          if self.cf.hold_out_test_set
          else "raw_pred_boxes_list"
      )
      with open(
          os.path.join(
              self.cf.fold_dir, "{}.pickle".format(out_string)
          ),
          "wb",
      ) as handle:
          pickle.dump(list_of_results_per_patient, handle)
  
      if return_results:
          # consolidate predictions.
          self.logger.info(
              "applying wcs to test set predictions with iou = {} and n_ens = {}.".format(
                  self.cf.wcs_iou, self.n_ens
              )
          )
          pool = Pool(processes=6)
          mp_inputs = [
              [
                  ii[0],
                  ii[1],
                  self.cf.class_dict,
                  self.cf.wcs_iou,
                  self.n_ens,
              ]
              for ii in list_of_results_per_patient
          ]
          list_of_results_per_patient = pool.map(
              apply_wbc_to_patient, mp_inputs, chunksize=1
          )
          pool.close()
          pool.join()
  
          # merge 2D boxes to 3D cubes. (if model predicts 2D but evaluation is run in 3D)
          if self.cf.merge_2D_to_3D_preds:
              self.logger.info(
                  "applying 2Dto3D merging to test set predictions with iou = {}.".format(
                      self.cf.merge_3D_iou
                  )
              )
              pool = Pool(processes=6)
              mp_inputs = [
                  [
                      ii[0],
                      ii[1],
                      self.cf.class_dict,
                      self.cf.merge_3D_iou,
                  ]
                  for ii in list_of_results_per_patient
              ]
              list_of_results_per_patient = pool.map(
                  merge_2D_to_3D_preds_per_patient,
                  mp_inputs,
                  chunksize=1,
              )
              pool.close()
              pool.join()
  
          return list_of_results_per_patient
  
imports: import os
  import numpy
  import torch
  from scipy.stats import norm
  from collections import OrderedDict
  from multiprocessing import Pool
  import pickle
  import pandas
label_desc: Use image node_id 3 for calling the Predictor obj's underlying member method code with example usage: obj.predict_test_set(batch_gen, return_results) and returns: list_of_results_per_patient