text: import numpy as np
  import torch
  import imageio
  import os
  import glob
  from torch.utils import data
  from torch.utils.data import Dataset
  from PIL import Image
  
  
  class  ImageFolder(Dataset):
      """docstring for ArtDataset"""
      def __init__(self, root, transform=None):
          super( ImageFolder, self).__init__()
          self.root = root
  
          self.frame = self._parse_frame()
          self.transform = transform
  
      def _parse_frame(self):
          frame = []
          folders = glob.glob(os.path.join(self.root, '*'))
          for folder in folders:
              for f in glob.glob(os.path.join(folder, '*.png')):
                  frame.append(f)
          return frame
  
      def __len__(self):
          return len(self.frame)
  
      def __getitem__(self, idx):
          file = self.frame[idx]
          img = Image.open(file).convert('RGB')
  
          if self.transform:
              img = self.transform(img)
  
          return img
  
  def InfiniteSampler(n):
      """Data sampler"""
      i = n - 1
      order = np.random.permutation(n)
      while True:
          yield order[i]
          i += 1
          if i >= n:
              np.random.seed()
              order = np.random.permutation(n)
              i = 0
  
  class InfiniteSamplerWrapper(data.sampler.Sampler):
      """Data sampler wrapper"""
      def __init__(self, data_source):
          self.num_samples = len(data_source)
  
      def __iter__(self):
          return iter(InfiniteSampler(self.num_samples))
  
      def __len__(self):
          return 2 ** 15
  
  def get_nsamples(data_loader, N):
    x = []
    n = 0
    while n < N:
      x_next = next(data_loader)
      x_next = x_next.cuda(non_blocking=True)
      x.append(x_next)
      n += x_next.size(0)
    x = torch.cat(x, dim=0)[:N]
    return x
  
  
  def count_trainable_parameters(model):
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    return sum([np.prod(p.size()) for p in model_parameters])
  
  
  def save_video(imgs, fname, as_gif=False, fps=24, quality=8):
      # convert to np.uint8
      imgs = (255 * np.clip(imgs.permute(0, 2, 3, 1).detach().cpu().numpy() / 2 + 0.5, 0, 1)).astype(np.uint8)
      imageio.mimwrite(fname, imgs, fps=fps, quality=quality)
      
      if as_gif:  # save as gif, too
          os.system(f'ffmpeg -i {fname} -r 15 '
                    f'-vf "scale=512:-1,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse" {os.path.splitext(fname)[0] + ".gif"}')
  
  
  def color_depth_map(depths, scale=None):
      """
      Color an input depth map.
  
      Arguments:
          depths -- HxW numpy array of depths
          [scale=None] -- scaling the values (defaults to the maximum depth)
  
      Returns:
          colored_depths -- HxWx3 numpy array visualizing the depths
      """
  
      _color_map_depths = np.array([
        [0, 0, 0],  # 0.000
        [0, 0, 255],  # 0.114
        [255, 0, 0],  # 0.299
        [255, 0, 255],  # 0.413
        [0, 255, 0],  # 0.587
        [0, 255, 255],  # 0.701
        [255, 255, 0],  # 0.886
        [255, 255, 255],  # 1.000
        [255, 255, 255],  # 1.000
      ]).astype(float)
      _color_map_bincenters = np.array([
        0.0,
        0.114,
        0.299,
        0.413,
        0.587,
        0.701,
        0.886,
        1.000,
        2.000,  # doesn't make a difference, just strictly higher than 1
      ])
    
      if scale is None:
        scale = depths.max()
    
      values = np.clip(depths.flatten() / scale, 0, 1)
      # for each value, figure out where they fit in in the bincenters: what is the last bincenter smaller than this value?
      lower_bin = ((values.reshape(-1, 1) >= _color_map_bincenters.reshape(1, -1)) * np.arange(0, 9)).max(axis=1)
      lower_bin_value = _color_map_bincenters[lower_bin]
      higher_bin_value = _color_map_bincenters[lower_bin + 1]
      alphas = (values - lower_bin_value) / (higher_bin_value - lower_bin_value)
      colors = _color_map_depths[lower_bin] * (1 - alphas).reshape(-1, 1) + _color_map_depths[
        lower_bin + 1] * alphas.reshape(-1, 1)
      return colors.reshape(depths.shape[0], depths.shape[1], 3).astype(np.uint8)
  
  
  # Virtual camera utils
  
  
  def to_sphere(u, v):
      theta = 2 * np.pi * u
      phi = np.arccos(1 - 2 * v)
      cx = np.sin(phi) * np.cos(theta)
      cy = np.sin(phi) * np.sin(theta)
      cz = np.cos(phi)
      s = np.stack([cx, cy, cz])
      return s
  
  
  def polar_to_cartesian(r, theta, phi, deg=True):
      if deg:
          phi = phi * np.pi / 180
          theta = theta * np.pi / 180
      cx = np.sin(phi) * np.cos(theta)
      cy = np.sin(phi) * np.sin(theta)
      cz = np.cos(phi)
      return r * np.stack([cx, cy, cz])
  
  
  def to_uv(loc):
      # normalize to unit sphere
      loc = loc / loc.norm(dim=1, keepdim=True)
  
      cx, cy, cz = loc.t()
      v = (1 - cz) / 2
  
      phi = torch.acos(cz)
      sin_phi = torch.sin(phi)
  
      # ensure we do not divide by zero
      eps = 1e-8
      sin_phi[sin_phi.abs() < eps] = eps
  
      theta = torch.acos(cx / sin_phi)
  
      # check for sign of phi
      cx_rec = sin_phi * torch.cos(theta)
      if not np.isclose(cx.numpy(), cx_rec.numpy(), atol=1e-5).all():
          sin_phi = -sin_phi
  
      # check for sign of theta
      cy_rec = sin_phi * torch.sin(theta)
      if not np.isclose(cy.numpy(), cy_rec.numpy(), atol=1e-5).all():
          theta = -theta
  
      u = theta / (2 * np.pi)
      assert np.isclose(to_sphere(u, v).detach().cpu().numpy(), loc.t().detach().cpu().numpy(), atol=1e-5).all()
  
      return u, v
  
  
  def to_phi(u):
      return 360 * u  # 2*pi*u*180/pi
  
  
  def to_theta(v):
      return np.arccos(1 - 2 * v) * 180. / np.pi
  
  
  def sample_on_sphere(range_u=(0, 1), range_v=(0, 1)):
      u = np.random.uniform(*range_u)
      v = np.random.uniform(*range_v)
      return to_sphere(u, v)
  
  
  def look_at(eye, at=np.array([0, 0, 0]), up=np.array([0, 0, 1]), eps=1e-5):
      at = at.astype(float).reshape(1, 3)
      up = up.astype(float).reshape(1, 3)
  
      eye = eye.reshape(-1, 3)
      up = up.repeat(eye.shape[0] // up.shape[0], axis=0)
      eps = np.array([eps]).reshape(1, 1).repeat(up.shape[0], axis=0)
  
      z_axis = eye - at
      z_axis /= np.max(np.stack([np.linalg.norm(z_axis, axis=1, keepdims=True), eps]))
  
      x_axis = np.cross(up, z_axis)
      x_axis /= np.max(np.stack([np.linalg.norm(x_axis, axis=1, keepdims=True), eps]))
  
      y_axis = np.cross(z_axis, x_axis)
      y_axis /= np.max(np.stack([np.linalg.norm(y_axis, axis=1, keepdims=True), eps]))
  
      r_mat = np.concatenate((x_axis.reshape(-1, 3, 1), y_axis.reshape(-1, 3, 1), z_axis.reshape(-1, 3, 1)), axis=2)
  
      return r_mat
  