text: import warnings
  from typing import Callable
  
  import cv2
  import matplotlib.pyplot as plt
  import numpy as np
  import torch
  from catalyst.dl import Callback, RunnerState, CallbackOrder
  from catalyst.dl.callbacks import TensorboardLogger
  from tensorboardX import SummaryWriter
  
  from pytorch_toolbelt.utils.torch_utils import rgb_image_from_tensor, to_numpy
  from pytorch_toolbelt.utils.torch_utils import tensor_from_rgb_image
  
  __all__ = [
      "get_tensorboard_logger",
      "ShowPolarBatchesCallback",
      "draw_binary_segmentation_predictions",
      "draw_semantic_segmentation_predictions",
  ]
  
  
  def get_tensorboard_logger(state: RunnerState) -> SummaryWriter:
      for logger_name, logger in state.loggers.items():
          if isinstance(logger, TensorboardLogger):
              return logger.loggers[state.loader_name]
      raise RuntimeError(f"Cannot find Tensorboard logger for loader {state.loader_name}")
  
  
  class ShowPolarBatchesCallback(Callback):
      """
      Visualize best and worst batch based in metric in Tensorboard
      """
  
      def __init__(
          self,
          visualize_batch: Callable,
          metric: str = "loss",
          minimize: bool = True,
          min_delta: float = 1e-6,
          targets="tensorboard",
      ):
          """
  
          :param visualize_batch: Visualization function that must return list of images.
                 It's takes two arguments: (batch input, predicted output).
          :param metric:
          :param minimize:
          :param min_delta:
          :param targets: Str 'tensorboard' or 'matplotlib, or ['tensorboard', 'matplotlib']
          """
          super().__init__(CallbackOrder.Other)
          assert isinstance(targets, (list, str))
  
          self.best_score = None
          self.best_input = None
          self.best_output = None
  
          self.worst_score = None
          self.worst_input = None
          self.worst_output = None
  
          self.target_metric = metric
          self.num_bad_epochs = 0
          self.is_better = None
          self.visualize_batch = visualize_batch
          self.targets = [targets] if isinstance(targets, str) else targets
  
          if minimize:
              self.is_better = lambda score, best: score <= (best - min_delta)
              self.is_worse = lambda score, worst: score >= (worst + min_delta)
          else:
              self.is_better = lambda score, best: score >= (best + min_delta)
              self.is_worse = lambda score, worst: score <= (worst - min_delta)
  
      def to_cpu(self, data):
          if isinstance(data, dict):
              return dict((key, self.to_cpu(value)) for (key, value) in data.items())
          if isinstance(data, torch.Tensor):
              return data.detach().cpu()
          if isinstance(data, (list, tuple)):
              return [self.to_cpu(value) for value in data]
          if isinstance(data, str):
              return data
          raise ValueError("Unsupported type", type(data))
  
      def on_loader_start(self, state):
          self.best_score = None
          self.best_input = None
          self.best_output = None
  
          self.worst_score = None
          self.worst_input = None
          self.worst_output = None
  
      def on_batch_end(self, state: RunnerState):
          value = state.metrics.batch_values.get(self.target_metric, None)
          if value is None:
              warnings.warn(
                  f"Metric value for {self.target_metric} is not available in state.metrics.batch_values"
              )
              return
  
          if self.best_score is None or self.is_better(value, self.best_score):
              self.best_score = value
              self.best_input = self.to_cpu(state.input)
              self.best_output = self.to_cpu(state.output)
  
          if self.worst_score is None or self.is_worse(value, self.worst_score):
              self.worst_score = value
              self.worst_input = self.to_cpu(state.input)
              self.worst_output = self.to_cpu(state.output)
  
      def on_loader_end(self, state: RunnerState) -> None:
          logger = get_tensorboard_logger(state)
  
          if self.best_score is not None:
              best_samples = self.visualize_batch(self.best_input, self.best_output)
              self._log_samples(best_samples, "best", logger, state.step)
  
          if self.worst_score is not None:
              worst_samples = self.visualize_batch(self.worst_input, self.worst_output)
              self._log_samples(worst_samples, "worst", logger, state.step)
  
      def _log_samples(self, samples, name, logger, step):
          if "tensorboard" in self.targets:
              for i, image in enumerate(samples):
                  logger.add_image(
                      f"{self.target_metric}/{name}/{i}",
                      tensor_from_rgb_image(image),
                      step,
                  )
  
          if "matplotlib" in self.targets:
              for i, image in enumerate(samples):
                  plt.figure()
                  plt.imshow(image)
                  plt.tight_layout()
                  plt.axis("off")
                  plt.show()
  
  
  def draw_binary_segmentation_predictions(
      input: dict,
      output: dict,
      image_key="features",
      image_id_key="image_id",
      targets_key="targets",
      outputs_key="logits",
      mean=(0.485, 0.456, 0.406),
      std=(0.229, 0.224, 0.225),
  ):
      images = []
      image_id_input = (
          input[image_id_key]
          if image_id_key is not None
          else [None] * len(input[image_key])
      )
  
      for image, target, image_id, logits in zip(
          input[image_key], input[targets_key], image_id_input, output[outputs_key]
      ):
          image = rgb_image_from_tensor(image, mean, std)
          target = to_numpy(target).squeeze(0)
          logits = to_numpy(logits).squeeze(0)
  
          overlay = image.copy()
          true_mask = target > 0
          pred_mask = logits > 0
  
          overlay[true_mask & pred_mask] = np.array(
              [0, 250, 0], dtype=overlay.dtype
          )  # Correct predictions (Hits) painted with green
          overlay[true_mask & ~pred_mask] = np.array(
              [250, 0, 0], dtype=overlay.dtype
          )  # Misses painted with red
          overlay[~true_mask & pred_mask] = np.array(
              [250, 250, 0], dtype=overlay.dtype
          )  # False alarm painted with yellow
          overlay = cv2.addWeighted(image, 0.5, overlay, 0.5, 0, dtype=cv2.CV_8U)
  
          if image_id is not None:
              cv2.putText(
                  overlay,
                  str(image_id),
                  (10, 15),
                  cv2.FONT_HERSHEY_PLAIN,
                  1,
                  (250, 250, 250),
              )
  
          images.append(overlay)
      return images
  
  
  def draw_semantic_segmentation_predictions(
      input: dict,
      output: dict,
      class_colors,
      mode="overlay",
      image_key="features",
      image_id_key="image_id",
      targets_key="targets",
      outputs_key="logits",
      mean=(0.485, 0.456, 0.406),
      std=(0.229, 0.224, 0.225),
  ):
      assert mode in {"overlay", "side-by-side"}
  
      images = []
      image_id_input = (
          input[image_id_key]
          if image_id_key is not None
          else [None] * len(input[image_key])
      )
  
      for image, target, image_id, logits in zip(
          input[image_key], input[targets_key], image_id_input, output[outputs_key]
      ):
          image = rgb_image_from_tensor(image, mean, std)
          logits = to_numpy(logits).argmax(axis=0)
          target = to_numpy(target)
  
          if mode == "overlay":
              overlay = image.copy()
              for class_index, class_color in enumerate(class_colors):
                  overlay[logits == class_index, :] = class_color
  
              overlay = cv2.addWeighted(image, 0.5, overlay, 0.5, 0, dtype=cv2.CV_8U)
  
              if image_id is not None:
                  cv2.putText(
                      overlay,
                      str(image_id),
                      (10, 15),
                      cv2.FONT_HERSHEY_PLAIN,
                      1,
                      (250, 250, 250),
                  )
          elif mode == "side-by-side":
  
              true_mask = np.zeros_like(image)
              for class_index, class_color in enumerate(class_colors):
                  true_mask[target == class_index, :] = class_color
  
              pred_mask = np.zeros_like(image)
              for class_index, class_color in enumerate(class_colors):
                  pred_mask[logits == class_index, :] = class_color
  
              overlay = np.hstack((image, true_mask, pred_mask))
          else:
              raise ValueError(mode)
  
          images.append(overlay)
  
      return images
  