text: """Training/evaluation utilities module.
  
  This module contains utilities and tools used to instantiate training sessions. It also contains
  the prediction consumer interface used by metrics and loggers to receive iteration data during
  training. See :mod:`thelper.optim.metrics` for more information on metrics.
  """
  
  import json
  import logging
  import os
  from typing import Any, AnyStr, Dict, List, Optional, Union  # noqa: F401
  
  import cv2 as cv
  import numpy as np
  import sklearn.metrics
  import torch
  
  import thelper.concepts
  import thelper.ifaces
  import thelper.typedefs  # noqa: F401
  import thelper.utils
  from thelper.ifaces import ClassNamesHandler, FormatHandler, PredictionConsumer
  from thelper.optim.eval import compute_iou
  from thelper.tasks.detect import BoundingBox
  
  logger = logging.getLogger(__name__)
  
  
  class PredictionCallback(PredictionConsumer):
      """Callback function wrapper compatible with the consumer interface.
  
      This interface is used to hide user-defined callbacks into the list of prediction consumers given
      to trainer implementations. The callbacks must always be compatible with the list of arguments
      defined by ``thelper.typedefs.IterCallbackParams``, but may also receive extra arguments defined in
      advance and passed to the constructor of this class.
  
      Attributes:
          callback_func: user-defined function to call on every update from the trainer.
          callback_kwargs: user-defined extra arguments to provide to the callback function.
      """
  
      def __init__(self, callback_func, callback_kwargs=None):
          # type: (thelper.typedefs.IterCallbackType, thelper.typedefs.IterCallbackParams) -> None
          assert callback_func is not None and \
              (isinstance(callback_func, str) or callable(callback_func)), \
              "invalid callback function, must be importable string or callable object"
          if isinstance(callback_func, str):
              callback_func = thelper.utils.import_function(callback_func)
          thelper.utils.check_func_signature(callback_func, thelper.typedefs.IterCallbackParams)
          assert callback_kwargs is None or \
              (isinstance(callback_kwargs, dict) and
               not any([p in callback_kwargs for p in thelper.typedefs.IterCallbackParams])), \
              "invalid callback kwargs (must be dict, and not contain overlap with default args)"
          self.callback_func = callback_func
          self.callback_kwargs = callback_kwargs
  
      def __repr__(self):
          """Returns a generic print-friendly string containing info about this consumer."""
          return self.__class__.__module__ + "." + self.__class__.__qualname__ + \
              f"(callback_func={repr(self.callback_func)}, callback_kwargs={repr(self.callback_kwargs)})"
  
      def update(self, *args, **kwargs):
          """Forwards the latest prediction data from the training session to the user callback."""
          return self.callback_func(*args, **kwargs, **self.callback_kwargs)
  
  
  @thelper.concepts.classification
  class ClassifLogger(PredictionConsumer, ClassNamesHandler, FormatHandler):
      """Classification output logger.
  
      This class provides a simple logging interface for accumulating and saving the predictions of a classifier.
      By default, all predictions will be logged. However, a confidence threshold can be set to focus on "hard"
      samples if necessary. It also optionally offers tensorboardX-compatible output images that can be saved
      locally or posted to tensorboard for browser-based visualization.
  
      Usage examples inside a session configuration file::
  
          # ...
          # lists all metrics to instantiate as a dictionary
          "metrics": {
              # ...
              # this is the name of the example consumer; it is used for lookup/printing only
              "logger": {
                  # this type is used to instantiate the confusion matrix report object
                  "type": "thelper.train.utils.ClassifLogger",
                  "params": {
                      # log the three 'best' predictions for each sample
                      "top_k": 3,
                      # keep updating a set of 10 samples for visualization via tensorboardX
                      "viz_count": 10
                  }
              },
              # ...
          }
          # ...
  
      Attributes:
          top_k: number of 'best' predictions to keep for each sample (along with the gt label).
          conf_threshold: threshold used to eliminate uncertain predictions.
          class_names: holds the list of class label names provided by the dataset parser. If it is not
              provided when the constructor is called, it will be set by the trainer at runtime.
          target_name: name of the targeted label (may be 'None' if all classes are used).
          target_idx: index of the targeted label (may be 'None' if all classes are used).
          viz_count: number of tensorboardX images to generate and update at each epoch.
          report_count: number of samples to print in reports (use 'None' if all samples must be printed).
          log_keys: list of metadata field keys to copy from samples into the log for each prediction.
          force_softmax: specifies whether a softmax operation should be applied to the prediction scores
              obtained from the trainer.
          score: array used to store prediction scores for logging.
          true: array used to store groundtruth labels for logging.
          meta: array used to store metadata pulled from samples for logging.
          format: output format of the produced log (supports: text, CSV)
      """
  
      def __init__(self,
                   top_k=1,               # type: int
                   conf_threshold=None,   # type: Optional[thelper.typedefs.Number]
                   class_names=None,      # type: Optional[List[AnyStr]]
                   target_name=None,      # type: Optional[AnyStr]
                   viz_count=0,           # type: int
                   report_count=None,     # type: Optional[int]
                   log_keys=None,         # type: Optional[List[AnyStr]]
                   force_softmax=True,    # type: bool
                   format=None,           # type: Optional[AnyStr]
                   ):                     # type: (...) -> None
          """Receives the logging parameters & the optional class label names used to decorate the log."""
          assert isinstance(top_k, int) and top_k > 0, "invalid top-k value"
          assert conf_threshold is None or (isinstance(conf_threshold, (float, int)) and 0 < conf_threshold <= 1), \
              "classification confidence threshold should be 'None' or float in ]0, 1]"
          assert isinstance(viz_count, int) and viz_count >= 0, "invalid image count to visualize"
          assert report_count is None or (isinstance(report_count, int) and report_count >= 0), \
              "invalid report sample count"
          assert log_keys is None or isinstance(log_keys, list), "invalid list of sample keys to log"
          self.top_k = top_k
          self.target_name = target_name
          self.target_idx = None
          self.conf_threshold = conf_threshold
          self.viz_count = viz_count
          self.report_count = report_count
          self.log_keys = log_keys if log_keys is not None else []
          self.force_softmax = force_softmax
          self.score = None
          self.true = None
          self.meta = None
          ClassNamesHandler.__init__(self, class_names)
          FormatHandler.__init__(self, format)
  
      def __repr__(self):
          """Returns a generic print-friendly string containing info about this consumer."""
          return self.__class__.__module__ + "." + self.__class__.__qualname__ + \
              f"(top_k={repr(self.top_k)}, conf_threshold={repr(self.conf_threshold)}, " + \
              f"class_names={repr(self.class_names)}, target_name={repr(self.target_name)}, " + \
              f"viz_count={repr(self.viz_count)}, report_count={repr(self.report_count)}, " + \
              f"log_keys={repr(self.log_keys)}, force_softmax={repr(self.force_softmax)})"
  
      @ClassNamesHandler.class_names.setter
      def class_names(self, class_names):
          """Sets the class label names that must be predicted by the model."""
          ClassNamesHandler.class_names.fset(self, class_names)
          if self.target_name is not None and self.class_names is not None:
              assert self.target_name in self.class_indices, \
                  f"could not find target name {repr(self.target_name)} in class names list"
              self.target_idx = self.class_indices[self.target_name]
          else:
              self.target_idx = None
  
      def update(self,        # see `thelper.typedefs.IterCallbackParams` for more info
                 task,        # type: thelper.tasks.utils.Task
                 input,       # type: thelper.typedefs.InputType
                 pred,        # type: thelper.typedefs.ClassificationPredictionType
                 target,      # type: thelper.typedefs.ClassificationTargetType
                 sample,      # type: thelper.typedefs.SampleType
                 loss,        # type: Optional[float]
                 iter_idx,    # type: int
                 max_iters,   # type: int
                 epoch_idx,   # type: int
                 max_epochs,  # type: int
                 **kwargs,    # type: Any
                 ):           # type: (...) -> None
          """Receives the latest predictions and target values from the training session.
  
          The exact signature of this function should match the one of the callbacks defined in
          :class:`thelper.train.base.Trainer` and specified by ``thelper.typedefs.IterCallbackParams``.
          """
          assert len(kwargs) == 0, "unexpected extra arguments present in update call"
          assert isinstance(task, thelper.tasks.Classification), "classif report only impl for classif tasks"
          assert iter_idx is not None and max_iters is not None and iter_idx < max_iters, \
              "bad iteration indices given to update function"
          if self.score is None or self.score.size != max_iters:
              self.score = np.asarray([None] * max_iters)
              self.true = np.asarray([None] * max_iters)
              self.meta = {key: np.asarray([None] * max_iters) for key in self.log_keys}
          if task.class_names != self.class_names:
              self.class_names = task.class_names
          assert pred.dim() == 2, "current classif logger impl only supports 2D outputs (BxC)"
          assert pred.shape[1] == len(self.class_names), "unexpected prediction class dimension size"
          if target is None or target.numel() == 0:
              self.true[iter_idx] = None
          else:
              assert target.dim() == 1, "gt should be batched (1D) tensor"
              assert pred.shape[0] == target.shape[0], "prediction/gt tensors batch size mismatch"
              self.true[iter_idx] = target.numpy()
          if self.force_softmax:
              with torch.no_grad():
                  pred = torch.nn.functional.softmax(pred, dim=1)
          self.score[iter_idx] = pred.numpy()
          for meta_key in self.log_keys:
              assert meta_key in sample, f"could not extract sample field with key {repr(meta_key)}"
              val = sample[meta_key]
              assert isinstance(val, (list, np.ndarray, torch.Tensor)), f"field {repr(meta_key)} should be batched"
              self.meta[meta_key][iter_idx] = val if isinstance(val, list) else val.tolist()
  
      def render(self):
          """Returns an image of predicted outputs as a numpy-compatible RGBA image drawn by pyplot."""
          if self.viz_count == 0:
              return None
          if self.score is None or self.true is None:
              return None
          raise NotImplementedError  # TODO
  
      def report_text(self):
          # type: () -> Optional[AnyStr]
          return self.report_csv()
  
      def report_csv(self):
          # type: () -> Optional[AnyStr]
          """Returns the logged metadata of predicted samples.
  
          The returned object is a print-friendly CSV string that can be consumed directly by tensorboardX. Note
          that this string might be very long if the dataset is large (i.e. it will contain one line per sample).
          """
          if isinstance(self.report_count, int) and self.report_count <= 0:
              return None
          if self.score is None or self.true is None:
              return None
          pack = list(zip(*[(*pack, )
                            for packs in zip(self.score, self.true, *self.meta.values())
                            for pack in zip(*packs)]))
          logdata = {key: np.stack(val, axis=0) for key, val in zip(["pred", "target", *self.meta.keys()], pack)}
          assert all([len(val) == len(logdata["target"]) for val in logdata.values()]), "messed up unpacking"
          header = "target_name,target_score"
          for k in range(self.top_k):
              header += f",pred_{k + 1}_name,pred_{k + 1}_score"
          for meta_key in self.log_keys:
              header += f",{str(meta_key)}"
          lines = []
          for sample_idx in range(len(logdata["target"])):
              gt_label_idx = int(logdata["target"][sample_idx])
              pred_scores = logdata["pred"][sample_idx]
              sorted_score_idxs = np.argsort(pred_scores)[::-1]
              sorted_scores = pred_scores[sorted_score_idxs]
              if self.conf_threshold is None or gt_label_idx is None or \
                      pred_scores[gt_label_idx] >= self.conf_threshold:
                  if gt_label_idx is not None:
                      entry = f"{self.class_names[gt_label_idx]},{pred_scores[gt_label_idx]:2.4f}"
                  else:
                      entry = f"<unknown>,{0.0:2.4f}"
                  for k in range(self.top_k):
                      entry += f",{self.class_names[sorted_score_idxs[k]]},{sorted_scores[k]:2.4f}"
                  for meta_key in self.log_keys:
                      entry += f",{str(logdata[meta_key][sample_idx])}"
                  lines.append(entry)
                  if isinstance(self.report_count, int) and len(lines) >= self.report_count:
                      break
          return "\n".join([header, *lines])
  
      def reset(self):
          """Toggles a reset of the internal state, emptying storage arrays."""
          self.score = None
          self.true = None
          self.meta = None
  
  
  @thelper.concepts.classification
  class ClassifReport(PredictionConsumer, ClassNamesHandler, FormatHandler):
      """Classification report interface.
  
      This class provides a simple interface to ``sklearn.metrics.classification_report`` so that all
      count-based metrics can be reported at once under a string-based representation.
  
      Usage example inside a session configuration file::
  
          # ...
          # lists all metrics to instantiate as a dictionary
          "metrics": {
              # ...
              # this is the name of the example consumer; it is used for lookup/printing only
              "report": {
                  # this type is used to instantiate the classification report object
                  "type": "thelper.train.utils.ClassifReport",
                  # we do not need to provide any parameters to the constructor, defaults are fine
                  "params": {
                      # optional parameter that will indicate output as JSON is desired, plain 'text' otherwise
                      "format": "json"
                  }
              },
              # ...
          }
          # ...
  
      Attributes:
          class_names: holds the list of class label names provided by the dataset parser. If it is not
              provided when the constructor is called, it will be set by the trainer at runtime.
          pred: queue used to store the top-1 (best) predicted class indices at each iteration.
          format: output format of the produced log (supports: text, JSON)
      """
  
      def __init__(self, class_names=None, sample_weight=None, digits=4, format=None):
          """Receives the optional class names and arguments passed to the report generator function.
  
          Args:
              class_names: holds the list of class label names provided by the dataset parser. If it is not
                  provided when the constructor is called, it will be set by the trainer at runtime.
              sample_weight: sample weights, forwarded to ``sklearn.metrics.classification_report``.
              digits: metrics output digit count, forwarded to ``sklearn.metrics.classification_report``.
              format: output format of the produced log.
          """
          self.class_names = None
          self.sample_weight = sample_weight
          self.digits = digits
          self.pred = None
          self.target = None
          ClassNamesHandler.__init__(self, class_names)
          FormatHandler.__init__(self, format)
  
      def __repr__(self):
          """Returns a generic print-friendly string containing info about this consumer."""
          return f"{self.__class__.__module__}.{self.__class__.__qualname__}" + \
                 f"(class_names={repr(self.class_names)}, sample_weight={repr(self.sample_weight)}, " + \
                 f"digits={repr(self.digits)})"
  
      def update(self,        # see `thelper.typedefs.IterCallbackParams` for more info
                 task,        # type: thelper.tasks.utils.Task
                 input,       # type: thelper.typedefs.InputType
                 pred,        # type: thelper.typedefs.ClassificationPredictionType
                 target,      # type: thelper.typedefs.ClassificationTargetType
                 sample,      # type: thelper.typedefs.SampleType
                 loss,        # type: Optional[float]
                 iter_idx,    # type: int
                 max_iters,   # type: int
                 epoch_idx,   # type: int
                 max_epochs,  # type: int
                 **kwargs,    # type: Any
                 ):           # type: (...) -> None
          """Receives the latest predictions and target values from the training session.
  
          The exact signature of this function should match the one of the callbacks defined in
          :class:`thelper.train.base.Trainer` and specified by ``thelper.typedefs.IterCallbackParams``.
          """
          assert len(kwargs) == 0, "unexpected extra arguments present in update call"
          assert isinstance(task, thelper.tasks.Classification), "classif report only impl for classif tasks"
          assert iter_idx is not None and max_iters is not None and iter_idx < max_iters, \
              "bad iteration indices given to update function"
          if self.pred is None or self.pred.size != max_iters:
              self.pred = np.asarray([None] * max_iters)
              self.target = np.asarray([None] * max_iters)
          if task.class_names != self.class_names:
              self.class_names = task.class_names
          if target is None or target.numel() == 0:
              # only accumulate results when groundtruth is available
              self.pred[iter_idx] = None
              self.target[iter_idx] = None
              return
          assert pred.dim() == 2 or target.dim() == 1, "current classif report impl only supports batched 1D outputs"
          assert pred.shape[0] == target.shape[0], "prediction/gt tensors batch size mismatch"
          assert pred.shape[1] == len(self.class_names), "unexpected prediction class dimension size"
          self.pred[iter_idx] = pred.topk(1, dim=1)[1].view(pred.shape[0]).tolist()
          self.target[iter_idx] = target.view(target.shape[0]).tolist()
  
      def gen_report(self, as_dict=False):
          # type: (bool) -> Union[AnyStr, thelper.typedefs.JSON]
          if self.pred is None or self.target is None:
              return None
          pred, target = zip(*[(pred, target) for preds, targets in zip(self.pred, self.target)
                               if targets is not None for pred, target in zip(preds, targets)])
          y_true = np.asarray(target)
          y_pred = np.asarray(pred)
          _y_true = [self.class_names[classid] for classid in y_true]
          _y_pred = [self.class_names[classid] if (0 <= classid < len(self.class_names)) else "<unset>"
                     for classid in y_pred]
          return sklearn.metrics.classification_report(_y_true, _y_pred, sample_weight=self.sample_weight,
                                                       digits=self.digits, output_dict=as_dict)
  
      def report_text(self):
          # type: () -> Optional[AnyStr]
          """Returns the classification report as a multi-line print-friendly string."""
          report = self.gen_report(as_dict=False)
          return f"\n{report}" if isinstance(report, str) else None
  
      def report_json(self):
          # type: () -> Optional[AnyStr]
          """Returns the classification report as a JSON formatted string."""
          report = self.gen_report(as_dict=True)
          return json.dumps(report, indent=4) if isinstance(report, str) else None
  
      def reset(self):
          """Toggles a reset of the metric's internal state, emptying queues."""
          self.pred = None
          self.target = None
  
  
  @thelper.concepts.detection
  class DetectLogger(PredictionConsumer, ClassNamesHandler, FormatHandler):
      """Detection output logger.
  
      This class provides a simple logging interface for accumulating and saving the bounding boxes of an
      object detector. By default, all detections will be logged. However, a confidence threshold can be set
      to focus on strong predictions if necessary.
  
      .. todo::
          It also optionally offers tensorboardX-compatible output images that can be saved
          locally or posted to tensorboard for browser-based visualization.
  
      Usage examples inside a session configuration file::
  
          # ...
          # lists all metrics to instantiate as a dictionary
          "metrics": {
              # ...
              # this is the name of the example consumer; it is used for lookup/printing only
              "logger": {
                  # this type is used to instantiate the confusion matrix report object
                  "type": "thelper.train.utils.DetectLogger",
                  "params": {
                      # (optional) log the three 'best' detections for each target
                      "top_k": 3
                  }
              },
              # ...
          }
          # ...
  
      Attributes:
          top_k: number of 'best' detections to keep for each target bbox (along with the target label).
              If omitted, lists all bounding box predictions by the model after applying IoU and confidence thresholds.
          conf_threshold: threshold used to eliminate uncertain predictions (if they support confidence).
              If confidence is not supported by the model bbox predictions, this parameter is ignored.
          iou_threshold: threshold used to eliminate predictions too far from target (regardless of confidence).
              If omitted, will ignore only completely non-overlapping predicted bounding boxes (:math:`IoU=0`).
              If no target bounding boxes are provided (prediction-only), this parameter is ignored.
          class_names: holds the list of class label names provided by the dataset parser. If it is not
              provided when the constructor is called, it will be set by the trainer at runtime.
          target_name: name of the targeted label (may be 'None' if all classes are used).
          target_idx: index of the targeted label (may be 'None' if all classes are used).
          viz_count: number of tensorboardX images to generate and update at each epoch.
          report_count: number of samples to print in reports (use 'None' if all samples must be printed).
          log_keys: list of metadata field keys to copy from samples into the log for each prediction.
          bbox: array used to store prediction bounding boxes for logging.
          true: array used to store groundtruth labels for logging.
          meta: array used to store metadata pulled from samples for logging.
          format: output format of the produced log (supports: text, CSV, JSON)
      """
  
      def __init__(self,
                   top_k=None,            # type: Optional[int]
                   conf_threshold=None,   # type: Optional[thelper.typedefs.Number]
                   iou_threshold=None,    # type: Optional[thelper.typedefs.Number]
                   class_names=None,      # type: Optional[List[AnyStr]]
                   target_name=None,      # type: Optional[AnyStr]
                   viz_count=0,           # type: int
                   report_count=None,     # type: Optional[int]
                   log_keys=None,         # type: Optional[List[AnyStr]]
                   format=None,           # type: Optional[AnyStr]
                   ):                     # type: (...) -> None
          """Receives the logging parameters & the optional class label names used to decorate the log."""
          assert top_k is None or isinstance(top_k, int) and top_k > 0, "invalid top-k value"
          assert conf_threshold is None or (isinstance(conf_threshold, (float, int)) and 0 <= conf_threshold <= 1), \
              "detection confidence threshold should be 'None' or number in [0, 1]"
          assert iou_threshold is None or (isinstance(iou_threshold, (int, float)) and 0 <= iou_threshold <= 1), \
              "detection IoU threshold should be 'None' or number in [0, 1]"
          assert isinstance(viz_count, int) and viz_count >= 0, "invalid image count to visualize"
          assert report_count is None or (
              isinstance(report_count, int) and report_count >= 0), "invalid report sample count"
          assert log_keys is None or isinstance(log_keys, list), "invalid list of sample keys to log"
          self.top_k = top_k
          self.target_name = target_name
          self.target_idx = None
          self.conf_threshold = conf_threshold
          self.iou_threshold = iou_threshold
          self.viz_count = viz_count
          self.report_count = report_count
          self.log_keys = log_keys if log_keys is not None else []
          self.bbox = None    # type: Optional[thelper.typedefs.DetectionPredictionType]
          self.true = None    # type: Optional[thelper.typedefs.DetectionTargetType]
          self.meta = None    # type: Optional[Dict[AnyStr, List[Any]]]
          ClassNamesHandler.__init__(self, class_names)
          FormatHandler.__init__(self, format)
  
      def __repr__(self):
          """Returns a generic print-friendly string containing info about this consumer."""
          return f"{self.__class__.__module__}.{self.__class__.__qualname__}" + \
                 f"(top_k={repr(self.top_k)}, conf_threshold={repr(self.conf_threshold)}, " + \
                 f"class_names={repr(self.class_names)}, target_name={repr(self.target_name)}, " + \
                 f"viz_count={repr(self.viz_count)}, report_count={repr(self.report_count)}, " + \
                 f"log_keys={repr(self.log_keys)})"
  
      @ClassNamesHandler.class_names.setter
      def class_names(self, class_names):
          """Sets the class label names that must be predicted by the model."""
          ClassNamesHandler.class_names.fset(self, class_names)
          if self.target_name is not None and self.class_names is not None:
              assert self.target_name in self.class_indices, \
                  f"could not find target name {repr(self.target_name)} in class names list"
              self.target_idx = self.class_indices[self.target_name]
          else:
              self.target_idx = None
  
      def update(self,        # see `thelper.typedefs.IterCallbackParams` for more info
                 task,        # type: thelper.tasks.utils.Task
                 input,       # type: thelper.typedefs.InputType
                 pred,        # type: thelper.typedefs.DetectionPredictionType
                 target,      # type: thelper.typedefs.DetectionTargetType
                 sample,      # type: thelper.typedefs.SampleType
                 loss,        # type: Optional[float]
                 iter_idx,    # type: int
                 max_iters,   # type: int
                 epoch_idx,   # type: int
                 max_epochs,  # type: int
                 **kwargs,    # type: Any
                 ):           # type: (...) -> None
          """Receives the latest predictions and target values from the training session.
  
          The exact signature of this function should match the one of the callbacks defined in
          :class:`thelper.train.base.Trainer` and specified by ``thelper.typedefs.IterCallbackParams``.
          """
          assert len(kwargs) == 0, "unexpected extra arguments present in update call"
          assert isinstance(task, thelper.tasks.Detection), "detect report only impl for detection tasks"
          assert iter_idx is not None and max_iters is not None and iter_idx < max_iters, \
              "bad iteration indices given to update function"
          if self.bbox is None or self.bbox.size != max_iters:
              self.bbox = np.asarray([None] * max_iters)
              self.true = np.asarray([None] * max_iters)
              self.meta = {key: np.asarray([None] * max_iters) for key in self.log_keys}
          if task.class_names != self.class_names:
              self.class_names = task.class_names
          if target is None or len(target) == 0 or all(len(t) == 0 for t in target):
              target = [None] * len(pred)   # simplify unpacking during report generation
          else:
              assert len(pred) == len(target), "prediction/target bounding boxes list batch size mismatch"
              for gt in target:
                  assert all(isinstance(bbox, BoundingBox) for bbox in gt), \
                      "detect logger only supports 2D lists of bounding box targets"
          for det in pred:
              assert all(isinstance(bbox, BoundingBox) for bbox in det), \
                  "detect logger only supports 2D lists of bounding box predictions"
          self.bbox[iter_idx] = pred
          self.true[iter_idx] = target
          for meta_key in self.log_keys:
              assert meta_key in sample, f"could not extract sample field with key {repr(meta_key)}"
              val = sample[meta_key]
              assert isinstance(val, (list, np.ndarray, torch.Tensor)), f"field {repr(meta_key)} should be batched"
              self.meta[meta_key][iter_idx] = val if isinstance(val, list) else val.tolist()
  
      def render(self):
          """Returns an image of predicted outputs as a numpy-compatible RGBA image drawn by pyplot."""
          if self.viz_count == 0:
              return None
          if self.bbox is None or self.true is None:
              return None
          raise NotImplementedError  # TODO
  
      def group_bbox(self,
                     target_bboxes,   # type: List[Optional[BoundingBox]]
                     detect_bboxes,   # type: List[BoundingBox]
                     ):               # type: (...) -> List[Dict[AnyStr, Union[BoundingBox, float, None]]]
          """Groups a sample's detected bounding boxes with target bounding boxes according to configuration parameters.
  
          Returns a list of detections grouped by target(s) with following format::
  
              [
                  {
                      "target": <associated-target-bbox>,
                      "detect": [
                          {
                              "bbox": <detection-bbox>,
                              "iou": <IoU(detect-bbox, target-bbox)>
                          },
                          ...
                      ]
                  },
                  ...
              ]
  
          The associated target bounding box and :math:`IoU` can be ``None`` if no target was provided
          (ie: during inference). In this case, the returned list will have only a single element with all detections
          associated to it. A single element list can also be returned if only one target was specified for this sample.
  
          When multiple ground truth targets were provided, the returned list will have the same length and ordering as
          these targets. The associated detected bounding boxes will depend on IoU between target/detection combinations.
  
          All filtering thresholds specified as configuration parameter will be applied for the returned list. Detected
          bounding boxes will also be sorted by highest confidence (if available) or by highest IoU as fallback.
          """
          # remove low confidence and sort by highest
          group_bboxes = sorted(
              [bbox for bbox in (detect_bboxes if detect_bboxes else []) if isinstance(bbox, BoundingBox)],
              key=lambda b: b.confidence if b.confidence is not None else 0, reverse=True)
          if self.conf_threshold:
              group_bboxes = [b for b in group_bboxes if b.confidence is not None and b.confidence >= self.conf_threshold]
          sort_by_iou = all(bbox.confidence is None for bbox in group_bboxes)
          # group according to target count
          target_count = len(target_bboxes)
          if target_count == 0:
              group_bboxes = [{"target": None, "detect": [{"bbox": bbox, "iou": None} for bbox in group_bboxes]}]
          elif target_count == 1:
              sorted_detect = [{"bbox": bbox, "iou": compute_iou(bbox, target_bboxes[0])} for bbox in group_bboxes]
              if sort_by_iou:
                  sorted_detect = list(sorted(sorted_detect, key=lambda d: d["iou"], reverse=True))
              group_bboxes = [{"target": target_bboxes[0], "detect": sorted_detect}]
          else:
              # regroup by highest IoU
              target_detects = [[] for _ in range(target_count)]
              for det in group_bboxes:
                  # FIXME:
                  #  should we do something different if all IoU = 0 (ie: false positive detection)
                  #  for now, they will all be stored in the first target, but can be tracked with IoU = 0
                  det_target_iou = [compute_iou(det, t) for t in target_bboxes]
                  best_iou_idx = int(np.argmax(det_target_iou))
                  target_detects[best_iou_idx].append({"bbox": det, "iou": det_target_iou[best_iou_idx]})
              group_bboxes = [{"target": target_bboxes[i],
                               "detect": list(sorted(
                                   target_detects[i], key=lambda d: d["iou"], reverse=True))
                                   if sort_by_iou else target_detects[i]
                              } for i in range(target_count)]     # noqa: PEP8
          # apply filters on grouped results
          if self.iou_threshold:
              for grp in group_bboxes:
                  grp["detect"] = [d for d in grp["detect"] if d["iou"] is None or d["iou"] >= self.iou_threshold]
          if self.top_k:
              for grp in group_bboxes:
                  grp["detect"] = grp["detect"][:self.top_k]
          return list(group_bboxes)
  
      def gen_report(self):
          # type: () -> Optional[List[Dict[AnyStr, Any]]]
          """Returns the logged metadata of predicted bounding boxes per sample target in a JSON-like structure.
  
          For every target bounding box, the corresponding *best*-sorted detections are returned.
          Sample metadata is appended to every corresponding sub-target if any where requested.
  
          If ``report_count`` was specified, the returned report will be limited to that requested amount of targets.
  
          .. seealso::
              | :meth:`DetectLogger.group_bbox` for formatting, sorting and filtering details.
          """
          if isinstance(self.report_count, int) and self.report_count <= 0:
              return None
          if self.bbox is None or self.true is None:
              return None
          # flatten batches/samples
          pack = list(zip(*[(*pack,)
                            for packs in zip(self.bbox, self.true, *self.meta.values())
                            for pack in zip(*packs)]))
          data = {key: val for key, val in zip(["detect", "target", *self.meta.keys()], pack)}
          assert all([len(val) == len(data["target"]) for val in data.values()]), "messed up unpacking"
  
          # flatten targets per batches/samples
          all_targets = []
          sample_count = len(data["target"])
          for sample_idx in range(sample_count):
              if isinstance(self.report_count, int) and self.report_count >= len(all_targets):
                  logger.warning(f"report max count {self.report_count} reached at {len(all_targets)} targets "
                                 f"(sample {sample_idx}/{sample_count} processed)")
                  break
              sample_targets = data["target"][sample_idx]
              sample_detects = data["detect"][sample_idx]
              sample_report = self.group_bbox(sample_targets, sample_detects)
              for target in sample_report:
                  for k in self.meta:
                      target[k] = self.meta[k][sample_idx]
                  target["target"] = {
                      "bbox": target["target"],
                      "class_name": self.class_names[target["target"].class_id]
                  }
              all_targets.extend(sample_report)
          # format everything nicely as json
          for item in all_targets:
              if isinstance(item["target"]["bbox"], BoundingBox):
                  item["target"].update(item["target"].pop("bbox").json())
                  item["target"]["class_name"] = self.class_names[item["target"]["class_id"]]
              for det in item["detect"]:
                  if isinstance(det["bbox"], BoundingBox):
                      det.update(det.pop("bbox").json())
                      det["class_name"] = self.class_names[det["class_id"]]
          return all_targets
  
      def report_json(self):
          # type: () -> Optional[AnyStr]
          """Returns the logged metadata of predicted bounding boxes as a JSON formatted string."""
          report = self.gen_report()
          if not report:
              return None
          return json.dumps(report, indent=4)
  
      def report_text(self):
          # type: () -> Optional[AnyStr]
          return self.report_csv()
  
      def report_csv(self):
          # type: () -> Optional[AnyStr]
          r"""Returns the logged metadata of predicted bounding boxes.
  
          The returned object is a print-friendly CSV string.
  
          Note that this string might be very long if the dataset is large or if the model tends to generate a lot of
          detections. The string will contain at least :math:`N_sample \cdot N_target` lines and each line will have
          up to :math:`N_bbox` detections, unless limited by configuration parameters.
          """
          report = self.gen_report()
          if not report:
              return None
  
          none_str = "unknown"
  
          def patch_none(to_patch, number_format='2.4f'):  # type: (Any, str) -> str
              if to_patch is None:
                  return none_str
              if isinstance(to_patch, float):
                  s = f"{{:{number_format}}}"
                  return s.format(to_patch)
              return str(to_patch)
  
          header = "sample,target_name,target_bbox"
          for meta_key in self.log_keys:
              header += f",{str(meta_key)}"
          if self.top_k:
              for k in range(self.top_k):
                  header += f",detect_{k + 1}_name,detect_{k + 1}_bbox,detect_{k + 1}_conf,detect_{k + 1}_iou"
          else:
              # unknown count total detections (can be variable)
              header += f",detect_name[N],detect_bbox[N],detect_conf[N],detect_iou[N],(...)[N]"
          lines = [""] * len(report)
          for i, result in enumerate(report):
              target = result["target"]
              detect = result["detect"]
              if not target:
                  entry = f"{none_str},{none_str},{none_str}"
              else:
                  entry = f"{target['image_id']},{patch_none(target['class_name'])},{patch_none(target['bbox'])}"
              for meta_key in self.log_keys:
                  entry += f",{str(target[meta_key])}"
              for det in detect:
                  entry += f",{det['class_name']},{det['bbox']},{patch_none(det['confidence'])},{patch_none(det['iou'])}"
              lines[i] = entry
          return "\n".join([header, *lines])
  
      def reset(self):
          """Toggles a reset of the internal state, emptying storage arrays."""
          self.bbox = None
          self.true = None
          self.meta = None
  
  
  @thelper.concepts.classification
  @thelper.concepts.segmentation
  class ConfusionMatrix(PredictionConsumer, ClassNamesHandler):
      """Confusion matrix report interface.
  
      This class provides a simple interface to ``sklearn.metrics.confusion_matrix`` so that a full
      confusion matrix can be easily reported under a string-based representation. It also offers a
      tensorboardX-compatible output image that can be saved locally or posted to tensorboard for
      browser-based visualization.
  
      Usage example inside a session configuration file::
  
          # ...
          # lists all metrics to instantiate as a dictionary
          "metrics": {
              # ...
              # this is the name of the example consumer; it is used for lookup/printing only
              "confmat": {
                  # this type is used to instantiate the confusion matrix report object
                  "type": "thelper.train.utils.ConfusionMatrix",
                  # we do not need to provide any parameters to the constructor, defaults are fine
                  "params": {
                      # optional parameter that will indicate output as JSON is desired, plain 'text' otherwise
                      "format": "json"
                  }
              },
              # ...
          }
          # ...
  
      Attributes:
          matrix: report generator function, called at evaluation time to generate the output string.
          class_names: holds the list of class label names provided by the dataset parser. If it is not
              provided when the constructor is called, it will be set by the trainer at runtime.
          draw_normalized: defines whether rendered confusion matrices should be normalized or not.
          pred: queue used to store the top-1 (best) predicted class indices at each iteration.
      """
  
      def __init__(self, class_names=None, draw_normalized=True):
          """Receives the optional class label names used to decorate the output string.
  
          Args:
              class_names: holds the list of class label names provided by the dataset parser. If it is not
                  provided when the constructor is called, it will be set by the trainer at runtime.
              draw_normalized: defines whether rendered confusion matrices should be normalized or not.
          """
  
          def gen_matrix(y_true, y_pred, _class_names):
              _y_true = [_class_names[classid] for classid in y_true]
              _y_pred = [_class_names[classid] if (0 <= classid < len(_class_names)) else "<unset>" for classid in y_pred]
              return sklearn.metrics.confusion_matrix(_y_true, _y_pred, labels=_class_names)
  
          self.matrix = gen_matrix
          self.draw_normalized = draw_normalized
          self.pred = None
          self.target = None
          ClassNamesHandler.__init__(self, class_names)
  
      def __repr__(self):
          """Returns a generic print-friendly string containing info about this consumer."""
          return self.__class__.__module__ + "." + self.__class__.__qualname__ + \
              f"(class_names={repr(self.class_names)}, draw_normalized={repr(self.draw_normalized)})"
  
      def update(self,        # see `thelper.typedefs.IterCallbackParams` for more info
                 task,        # type: thelper.tasks.utils.Task
                 input,       # type: thelper.typedefs.InputType
                 pred,        # type: thelper.typedefs.ClassificationPredictionType
                 target,      # type: thelper.typedefs.ClassificationTargetType
                 sample,      # type: thelper.typedefs.SampleType
                 loss,        # type: Optional[float]
                 iter_idx,    # type: int
                 max_iters,   # type: int
                 epoch_idx,   # type: int
                 max_epochs,  # type: int
                 **kwargs,    # type: Any
                 ):           # type: (...) -> None
          """Receives the latest predictions and target values from the training session.
  
          The exact signature of this function should match the one of the callbacks defined in
          :class:`thelper.train.base.Trainer` and specified by ``thelper.typedefs.IterCallbackParams``.
          """
          assert len(kwargs) == 0, "unexpected extra arguments present in update call"
          assert isinstance(task, thelper.tasks.Classification), "confmat only impl for classif tasks"
          assert iter_idx is not None and max_iters is not None and iter_idx < max_iters, \
              "bad iteration indices given to update function"
          if self.pred is None or self.pred.size != max_iters:
              self.pred = np.asarray([None] * max_iters)
              self.target = np.asarray([None] * max_iters)
          if task.class_names != self.class_names:
              self.class_names = task.class_names
          if target is None or target.numel() == 0:
              # only accumulate results when groundtruth is available
              self.pred[iter_idx] = None
              self.target[iter_idx] = None
              return
          assert pred.dim() == 2 or target.dim() == 1, "current confmat impl only supports batched 1D outputs"
          assert pred.shape[0] == target.shape[0], "prediction/gt tensors batch size mismatch"
          assert pred.shape[1] == len(self.class_names), "unexpected prediction class dimension size"
          self.pred[iter_idx] = pred.topk(1, dim=1)[1].view(pred.shape[0]).tolist()
          self.target[iter_idx] = target.view(target.shape[0]).tolist()
  
      def report(self):
          """Returns the confusion matrix as a multi-line print-friendly string."""
          if self.pred is None or self.target is None:
              return None
          pred, target = zip(*[(pred, target) for preds, targets in zip(self.pred, self.target)
                               if targets is not None for pred, target in zip(preds, targets)])
          confmat = self.matrix(np.asarray(target), np.asarray(pred), self.class_names)
          return "\n" + thelper.utils.stringify_confmat(confmat, self.class_names)
  
      def render(self):
          """Returns the confusion matrix as a numpy-compatible RGBA image drawn by pyplot."""
          if self.pred is None or self.target is None:
              return None
          pred, target = zip(*[(pred, target) for preds, targets in zip(self.pred, self.target)
                               if targets is not None for pred, target in zip(preds, targets)])
          confmat = self.matrix(np.asarray(target), np.asarray(pred), self.class_names)
          try:
              fig, ax = thelper.draw.draw_confmat(confmat, self.class_names, normalize=self.draw_normalized)
              array = thelper.draw.fig2array(fig)
              return array
          except AttributeError as e:
              logger.warning(f"failed to render confusion matrix; caught exception:\n{str(e)}")
              # return None if rendering fails (probably due to matplotlib on display-less server)
              return None
  
      def reset(self):
          """Toggles a reset of the metric's internal state, emptying queues."""
          self.pred = None
          self.target = None
  
  
  def create_consumers(config):
      """Instantiates and returns the prediction consumers defined in the configuration dictionary.
  
      All arguments are expected to be handed in through the configuration via a dictionary named 'params'.
      """
      assert isinstance(config, dict), "config should be provided as a dictionary"
      consumers = {}
      for name, consumer_config in config.items():
          assert isinstance(consumer_config, dict), "consumer config should be provided as a dictionary"
          assert "type" in consumer_config and consumer_config["type"], "consumer config missing 'type' field"
          consumer_type = thelper.utils.import_class(consumer_config["type"])
          consumer_params = thelper.utils.get_key_def(["params", "parameters"], consumer_config, {})
          try:
              consumer = consumer_type(**consumer_params)
          except Exception:
              logger.error(f"failed to create consumer {consumer_config['type']} with params:\n\t{str(consumer_params)}")
              raise
          assert isinstance(consumer, PredictionConsumer), \
              "invalid consumer type, must derive from PredictionConsumer interface"
          consumers[name] = consumer
      return consumers
  
  
  def create_trainer(session_name,    # type: AnyStr
                     save_dir,        # type: AnyStr
                     config,          # type: thelper.typedefs.ConfigDict
                     model,           # type: thelper.typedefs.ModelType
                     task,            # type: thelper.tasks.Task
                     loaders,         # type: thelper.typedefs.MultiLoaderType
                     ckptdata=None    # type: Optional[thelper.typedefs.CheckpointContentType]
                     ):               # type: (...) -> thelper.train.Trainer
      """Instantiates the trainer object based on the type contained in the config dictionary.
  
      The trainer type is expected to be in the configuration dictionary's `trainer` field, under the `type` key. For more
      information on the configuration, refer to :class:`thelper.train.base.Trainer`. The instantiated type must be
      compatible with the constructor signature of :class:`thelper.train.base.Trainer`. The object's constructor will
      be given the full config dictionary and the checkpoint data for resuming the session (if available).
  
      If the trainer type is missing, it will be automatically deduced based on the task object.
  
      Args:
          session_name: name of the training session used for printing and to create internal tensorboardX directories.
          save_dir: path to the session directory where logs and checkpoints will be saved.
          config: full configuration dictionary that will be parsed for trainer parameters and saved in checkpoints.
          model: model to train/evaluate; should be compatible with :class:`thelper.nn.utils.Module`.
          task: global task interface defining the type of model and training goal for the session.
          loaders: a tuple containing the training/validation/test data loaders (a loader can be ``None`` if empty).
          ckptdata: raw checkpoint to parse data from when resuming a session (if ``None``, will start from scratch).
  
      Returns:
          The fully-constructed trainer object, ready to begin model training/evaluation.
  
      .. seealso::
          | :class:`thelper.train.base.Trainer`
  
      """
      assert "trainer" in config and config["trainer"], "session configuration dictionary missing 'trainer' section"
      trainer_config = config["trainer"]
      if "type" not in trainer_config:
          if isinstance(task, thelper.tasks.Classification):
              trainer_type = thelper.train.ImageClassifTrainer
          elif isinstance(task, thelper.tasks.Detection):
              trainer_type = thelper.train.ObjDetectTrainer
          elif isinstance(task, thelper.tasks.Regression):
              trainer_type = thelper.train.RegressionTrainer
          elif isinstance(task, thelper.tasks.Segmentation):
              trainer_type = thelper.train.ImageSegmTrainer
          else:
              raise AssertionError(f"unknown trainer type required for task '{str(task)}'")
      else:
          trainer_type = thelper.utils.import_class(trainer_config["type"])
      return trainer_type(session_name, save_dir, model, task, loaders, config, ckptdata=ckptdata)
  
  
  # noinspection PyUnusedLocal
  def _draw_wrapper(task,         # type: thelper.tasks.utils.Task
                    input,        # type: thelper.typedefs.InputType
                    pred,         # type: thelper.typedefs.AnyPredictionType
                    target,       # type: thelper.typedefs.AnyTargetType
                    sample,       # type: thelper.typedefs.SampleType
                    loss,         # type: Optional[float]
                    iter_idx,     # type: int
                    max_iters,    # type: int
                    epoch_idx,    # type: int
                    max_epochs,   # type: int
                    # extra params added by callback interface below
                    output_path,  # type: str
                    save,         # type: bool
                    # all extra params will be forwarded to the display call
                    **kwargs,     # type: Any
                    ):            # type: (...) -> None
      """Wrapper to :func:`thelper.draw.draw` used as a callback entrypoint for trainers."""
      res = thelper.draw.draw(task=task, input=input, pred=pred, target=target, **kwargs)
      if save:
          assert isinstance(res, tuple) and len(res) == 2, "unexpected redraw output (should be 2-elem tuple)"
          if isinstance(res[1], np.ndarray):
              assert "path" in sample and isinstance(sample["path"], list) and len(sample["path"]) == 1, \
                  "unexpected draw format (current implementation needs batch size = 1, and path metadata)"
              os.makedirs(output_path, exist_ok=True)
              filepath = os.path.join(output_path, os.path.basename(sample["path"][0]))
              cv.imwrite(filepath, res[1])
          else:
              # we're displaying with matplotlib, and have no clue on how to save the output
              raise NotImplementedError
  