content: import json
  import math
  import os
  import re
  import shutil
  from typing import List, Optional, Union
  
  import cv2
  import imageio
  import matplotlib.pyplot as plt
  import numpy as np
  import torch
  
  # import wandb
  from matplotlib import cm
  from matplotlib.colors import LinearSegmentedColormap
  from PIL import Image, ImageDraw
  
  from .typing import *
  
  
  def image_to_tensor(
      images: List[Image.Image], device: str = "cuda"
  ) -> torch.FloatTensor:
      tensors = []
      for image in images:
          tensor = torch.from_numpy(np.array(image)).float() / 255.0
          tensor = tensor.to(device)
          tensors.append(tensor)
      return torch.stack(tensors)
  
  
  def tensor_to_image(
      data: Union[Image.Image, torch.Tensor, np.ndarray],
      batched: bool = False,
      format: str = "HWC",
  ) -> Union[Image.Image, List[Image.Image]]:
      if isinstance(data, Image.Image):
          return data
      if isinstance(data, torch.Tensor):
          data = data.detach().cpu().numpy()
      if data.dtype == np.float32 or data.dtype == np.float16:
          data = (data * 255).astype(np.uint8)
      elif data.dtype == np.bool_:
          data = data.astype(np.uint8) * 255
      assert data.dtype == np.uint8
      if format == "CHW":
          if batched and data.ndim == 4:
              data = data.transpose((0, 2, 3, 1))
          elif not batched and data.ndim == 3:
              data = data.transpose((1, 2, 0))
  
      if batched:
          return [Image.fromarray(d) for d in data]
      return Image.fromarray(data)
  
  
  def largest_factor_near_sqrt(n: int) -> int:
      """
      Finds the largest factor of n that is closest to the square root of n.
  
      Args:
          n (int): The integer for which to find the largest factor near its square root.
  
      Returns:
          int: The largest factor of n that is closest to the square root of n.
      """
      sqrt_n = int(math.sqrt(n))  # Get the integer part of the square root
  
      # First, check if the square root itself is a factor
      if sqrt_n * sqrt_n == n:
          return sqrt_n
  
      # Otherwise, find the largest factor by iterating from sqrt_n downwards
      for i in range(sqrt_n, 0, -1):
          if n % i == 0:
              return i
  
      # If n is 1, return 1
      return 1
  
  
  def make_image_grid(
      images: List[Image.Image],
      rows: Optional[int] = None,
      cols: Optional[int] = None,
      resize: Optional[int] = None,
  ) -> Image.Image:
      """
      Prepares a single grid of images. Useful for visualization purposes.
      """
      if rows is None and cols is not None:
          assert len(images) % cols == 0
          rows = len(images) // cols
      elif cols is None and rows is not None:
          assert len(images) % rows == 0
          cols = len(images) // rows
      elif rows is None and cols is None:
          rows = largest_factor_near_sqrt(len(images))
          cols = len(images) // rows
  
      assert len(images) == rows * cols
  
      if resize is not None:
          images = [img.resize((resize, resize)) for img in images]
  
      w, h = images[0].size
      grid = Image.new("RGB", size=(cols * w, rows * h))
  
      for i, img in enumerate(images):
          grid.paste(img, box=(i % cols * w, i // cols * h))
      return grid
  
  
  class SaverMixin:
      _save_dir: Optional[str] = None
      _wandb_logger: Optional[Any] = None
  
      def set_save_dir(self, save_dir: str):
          self._save_dir = save_dir
  
      def get_save_dir(self):
          if self._save_dir is None:
              raise ValueError("Save dir is not set")
          return self._save_dir
  
      def convert_data(self, data):
          if data is None:
              return None
          elif isinstance(data, np.ndarray):
              return data
          elif isinstance(data, torch.Tensor):
              if data.dtype in [torch.float16, torch.bfloat16]:
                  data = data.float()
              return data.detach().cpu().numpy()
          elif isinstance(data, list):
              return [self.convert_data(d) for d in data]
          elif isinstance(data, dict):
              return {k: self.convert_data(v) for k, v in data.items()}
          else:
              raise TypeError(
                  "Data must be in type numpy.ndarray, torch.Tensor, list or dict, getting",
                  type(data),
              )
  
      def get_save_path(self, filename):
          save_path = os.path.join(self.get_save_dir(), filename)
          os.makedirs(os.path.dirname(save_path), exist_ok=True)
          return save_path
  
      DEFAULT_RGB_KWARGS = {"data_format": "HWC", "data_range": (0, 1)}
      DEFAULT_UV_KWARGS = {
          "data_format": "HWC",
          "data_range": (0, 1),
          "cmap": "checkerboard",
      }
      DEFAULT_GRAYSCALE_KWARGS = {"data_range": None, "cmap": "jet"}
      DEFAULT_GRID_KWARGS = {"align": "max"}
  
      def get_rgb_image_(self, img, data_format, data_range, rgba=False):
          img = self.convert_data(img)
          assert data_format in ["CHW", "HWC"]
          if data_format == "CHW":
              img = img.transpose(1, 2, 0)
          if img.dtype != np.uint8:
              img = img.clip(min=data_range[0], max=data_range[1])
              img = (
                  (img - data_range[0]) / (data_range[1] - data_range[0]) * 255.0
              ).astype(np.uint8)
          nc = 4 if rgba else 3
          imgs = [img[..., start : start + nc] for start in range(0, img.shape[-1], nc)]
          imgs = [
              (
                  img_
                  if img_.shape[-1] == nc
                  else np.concatenate(
                      [
                          img_,
                          np.zeros(
                              (img_.shape[0], img_.shape[1], nc - img_.shape[2]),
                              dtype=img_.dtype,
                          ),
                      ],
                      axis=-1,
                  )
              )
              for img_ in imgs
          ]
          img = np.concatenate(imgs, axis=1)
          if rgba:
              img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
          else:
              img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
          return img
  
      def _save_rgb_image(
          self,
          filename,
          img,
          data_format,
          data_range,
          name: Optional[str] = None,
          step: Optional[int] = None,
      ):
          img = self.get_rgb_image_(img, data_format, data_range)
          cv2.imwrite(filename, img)
          if name and self._wandb_logger:
              self._wandb_logger.log_image(
                  key=name, images=[self.get_save_path(filename)], step=step
              )
  
      def save_rgb_image(
          self,
          filename,
          img,
          data_format=DEFAULT_RGB_KWARGS["data_format"],
          data_range=DEFAULT_RGB_KWARGS["data_range"],
          name: Optional[str] = None,
          step: Optional[int] = None,
      ) -> str:
          save_path = self.get_save_path(filename)
          self._save_rgb_image(save_path, img, data_format, data_range, name, step)
          return save_path
  
      def get_uv_image_(self, img, data_format, data_range, cmap):
          img = self.convert_data(img)
          assert data_format in ["CHW", "HWC"]
          if data_format == "CHW":
              img = img.transpose(1, 2, 0)
          img = img.clip(min=data_range[0], max=data_range[1])
          img = (img - data_range[0]) / (data_range[1] - data_range[0])
          assert cmap in ["checkerboard", "color"]
          if cmap == "checkerboard":
              n_grid = 64
              mask = (img * n_grid).astype(int)
              mask = (mask[..., 0] + mask[..., 1]) % 2 == 0
              img = np.ones((img.shape[0], img.shape[1], 3), dtype=np.uint8) * 255
              img[mask] = np.array([255, 0, 255], dtype=np.uint8)
              img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
          elif cmap == "color":
              img_ = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
              img_[..., 0] = (img[..., 0] * 255).astype(np.uint8)
              img_[..., 1] = (img[..., 1] * 255).astype(np.uint8)
              img_ = cv2.cvtColor(img_, cv2.COLOR_RGB2BGR)
              img = img_
          return img
  
      def save_uv_image(
          self,
          filename,
          img,
          data_format=DEFAULT_UV_KWARGS["data_format"],
          data_range=DEFAULT_UV_KWARGS["data_range"],
          cmap=DEFAULT_UV_KWARGS["cmap"],
      ) -> str:
          save_path = self.get_save_path(filename)
          img = self.get_uv_image_(img, data_format, data_range, cmap)
          cv2.imwrite(save_path, img)
          return save_path
  
      def get_grayscale_image_(self, img, data_range, cmap):
          img = self.convert_data(img)
          img = np.nan_to_num(img)
          if data_range is None:
              img = (img - img.min()) / (img.max() - img.min())
          else:
              img = img.clip(data_range[0], data_range[1])
              img = (img - data_range[0]) / (data_range[1] - data_range[0])
          assert cmap in [None, "jet", "magma", "spectral"]
          if cmap == None:
              img = (img * 255.0).astype(np.uint8)
              img = np.repeat(img[..., None], 3, axis=2)
          elif cmap == "jet":
              img = (img * 255.0).astype(np.uint8)
              img = cv2.applyColorMap(img, cv2.COLORMAP_JET)
          elif cmap == "magma":
              img = 1.0 - img
              base = cm.get_cmap("magma")
              num_bins = 256
              colormap = LinearSegmentedColormap.from_list(
                  f"{base.name}{num_bins}", base(np.linspace(0, 1, num_bins)), num_bins
              )(np.linspace(0, 1, num_bins))[:, :3]
              a = np.floor(img * 255.0)
              b = (a + 1).clip(max=255.0)
              f = img * 255.0 - a
              a = a.astype(np.uint16).clip(0, 255)
              b = b.astype(np.uint16).clip(0, 255)
              img = colormap[a] + (colormap[b] - colormap[a]) * f[..., None]
              img = (img * 255.0).astype(np.uint8)
          elif cmap == "spectral":
              colormap = plt.get_cmap("Spectral")
  
              def blend_rgba(image):
                  image = image[..., :3] * image[..., -1:] + (
                      1.0 - image[..., -1:]
                  )  # blend A to RGB
                  return image
  
              img = colormap(img)
              img = blend_rgba(img)
              img = (img * 255).astype(np.uint8)
              img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
          return img
  
      def _save_grayscale_image(
          self,
          filename,
          img,
          data_range,
          cmap,
          name: Optional[str] = None,
          step: Optional[int] = None,
      ):
          img = self.get_grayscale_image_(img, data_range, cmap)
          cv2.imwrite(filename, img)
          if name and self._wandb_logger:
              self._wandb_logger.log_image(
                  key=name, images=[self.get_save_path(filename)], step=step
              )
  
      def save_grayscale_image(
          self,
          filename,
          img,
          data_range=DEFAULT_GRAYSCALE_KWARGS["data_range"],
          cmap=DEFAULT_GRAYSCALE_KWARGS["cmap"],
          name: Optional[str] = None,
          step: Optional[int] = None,
      ) -> str:
          save_path = self.get_save_path(filename)
          self._save_grayscale_image(save_path, img, data_range, cmap, name, step)
          return save_path
  
      def get_image_grid_(self, imgs, align):
          if isinstance(imgs[0], list):
              return np.concatenate(
                  [self.get_image_grid_(row, align) for row in imgs], axis=0
              )
          cols = []
          for col in imgs:
              assert col["type"] in ["rgb", "uv", "grayscale"]
              if col["type"] == "rgb":
                  rgb_kwargs = self.DEFAULT_RGB_KWARGS.copy()
                  rgb_kwargs.update(col["kwargs"])
                  cols.append(self.get_rgb_image_(col["img"], **rgb_kwargs))
              elif col["type"] == "uv":
                  uv_kwargs = self.DEFAULT_UV_KWARGS.copy()
                  uv_kwargs.update(col["kwargs"])
                  cols.append(self.get_uv_image_(col["img"], **uv_kwargs))
              elif col["type"] == "grayscale":
                  grayscale_kwargs = self.DEFAULT_GRAYSCALE_KWARGS.copy()
                  grayscale_kwargs.update(col["kwargs"])
                  cols.append(self.get_grayscale_image_(col["img"], **grayscale_kwargs))
  
          if align == "max":
              h = max([col.shape[0] for col in cols])
          elif align == "min":
              h = min([col.shape[0] for col in cols])
          elif isinstance(align, int):
              h = align
          else:
              raise ValueError(
                  f"Unsupported image grid align: {align}, should be min, max, or int"
              )
  
          for i in range(len(cols)):
              if cols[i].shape[0] != h:
                  w = int(cols[i].shape[1] * h / cols[i].shape[0])
                  cols[i] = cv2.resize(cols[i], (w, h), interpolation=cv2.INTER_CUBIC)
          return np.concatenate(cols, axis=1)
  
      def save_image_grid(
          self,
          filename,
          imgs,
          align=DEFAULT_GRID_KWARGS["align"],
          name: Optional[str] = None,
          step: Optional[int] = None,
          texts: Optional[List[float]] = None,
      ):
          save_path = self.get_save_path(filename)
          img = self.get_image_grid_(imgs, align=align)
  
          if texts is not None:
              img = Image.fromarray(img)
              draw = ImageDraw.Draw(img)
              black, white = (0, 0, 0), (255, 255, 255)
              for i, text in enumerate(texts):
                  draw.text((2, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
                  draw.text((0, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
                  draw.text((2, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
                  draw.text((0, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
                  draw.text((1, (img.size[1] // len(texts)) * i), f"{text}", black)
              img = np.asarray(img)
  
          cv2.imwrite(save_path, img)
          if name and self._wandb_logger:
              self._wandb_logger.log_image(key=name, images=[save_path], step=step)
          return save_path
  
      def save_image(self, filename, img) -> str:
          save_path = self.get_save_path(filename)
          img = self.convert_data(img)
          assert img.dtype == np.uint8 or img.dtype == np.uint16
          if img.ndim == 3 and img.shape[-1] == 3:
              img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
          elif img.ndim == 3 and img.shape[-1] == 4:
              img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
          cv2.imwrite(save_path, img)
          return save_path
  
      def save_cubemap(self, filename, img, data_range=(0, 1), rgba=False) -> str:
          save_path = self.get_save_path(filename)
          img = self.convert_data(img)
          assert img.ndim == 4 and img.shape[0] == 6 and img.shape[1] == img.shape[2]
  
          imgs_full = []
          for start in range(0, img.shape[-1], 3):
              img_ = img[..., start : start + 3]
              img_ = np.stack(
                  [
                      self.get_rgb_image_(img_[i], "HWC", data_range, rgba=rgba)
                      for i in range(img_.shape[0])
                  ],
                  axis=0,
              )
              size = img_.shape[1]
              placeholder = np.zeros((size, size, 3), dtype=np.float32)
              img_full = np.concatenate(
                  [
                      np.concatenate(
                          [placeholder, img_[2], placeholder, placeholder], axis=1
                      ),
                      np.concatenate([img_[1], img_[4], img_[0], img_[5]], axis=1),
                      np.concatenate(
                          [placeholder, img_[3], placeholder, placeholder], axis=1
                      ),
                  ],
                  axis=0,
              )
              imgs_full.append(img_full)
  
          imgs_full = np.concatenate(imgs_full, axis=1)
          cv2.imwrite(save_path, imgs_full)
          return save_path
  
      def save_data(self, filename, data) -> str:
          data = self.convert_data(data)
          if isinstance(data, dict):
              if not filename.endswith(".npz"):
                  filename += ".npz"
              save_path = self.get_save_path(filename)
              np.savez(save_path, **data)
          else:
              if not filename.endswith(".npy"):
                  filename += ".npy"
              save_path = self.get_save_path(filename)
              np.save(save_path, data)
          return save_path
  
      def save_state_dict(self, filename, data) -> str:
          save_path = self.get_save_path(filename)
          torch.save(data, save_path)
          return save_path
  
      def save_img_sequence(
          self,
          filename,
          img_dir,
          matcher,
          save_format="mp4",
          fps=30,
          name: Optional[str] = None,
          step: Optional[int] = None,
      ) -> str:
          assert save_format in ["gif", "mp4"]
          if not filename.endswith(save_format):
              filename += f".{save_format}"
          save_path = self.get_save_path(filename)
          matcher = re.compile(matcher)
          img_dir = os.path.join(self.get_save_dir(), img_dir)
          imgs = []
          for f in os.listdir(img_dir):
              if matcher.search(f):
                  imgs.append(f)
          imgs = sorted(imgs, key=lambda f: int(matcher.search(f).groups()[0]))
          imgs = [cv2.imread(os.path.join(img_dir, f)) for f in imgs]
  
          if save_format == "gif":
              imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
              imageio.mimsave(save_path, imgs, fps=fps, palettesize=256)
          elif save_format == "mp4":
              imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
              imageio.mimsave(save_path, imgs, fps=fps)
          if name and self._wandb_logger:
              from .core import warn
  
              warn("Wandb logger does not support video logging yet!")
          return save_path
  
      def save_img_sequences(
          self,
          seq_dir,
          matcher,
          save_format="mp4",
          fps=30,
          delete=True,
          name: Optional[str] = None,
          step: Optional[int] = None,
      ):
          seq_dir_ = os.path.join(self.get_save_dir(), seq_dir)
          for f in os.listdir(seq_dir_):
              img_dir_ = os.path.join(seq_dir_, f)
              if not os.path.isdir(img_dir_):
                  continue
              try:
                  self.save_img_sequence(
                      os.path.join(seq_dir, f),
                      os.path.join(seq_dir, f),
                      matcher,
                      save_format=save_format,
                      fps=fps,
                      name=f"{name}_{f}",
                      step=step,
                  )
                  if delete:
                      shutil.rmtree(img_dir_)
              except:
                  from .core import warn
  
                  warn(f"Video saving for directory {seq_dir_} failed!")
  
      def save_file(self, filename, src_path, delete=False) -> str:
          save_path = self.get_save_path(filename)
          shutil.copyfile(src_path, save_path)
          if delete:
              os.remove(src_path)
          return save_path
  
      def save_json(self, filename, payload) -> str:
          save_path = self.get_save_path(filename)
          with open(save_path, "w") as f:
              f.write(json.dumps(payload))
          return save_path
  