content: import numpy as np
  import torch
  from torch.utils.data import Dataset
  from PIL import Image
  from typing import Literal, Tuple, Optional, Any
  import cv2
  
  import os
  import math
  
  
  import cv2
  import numpy as np
  
  def add_margin(pil_img, color=0, size=256):
      width, height = pil_img.size
      result = Image.new(pil_img.mode, (size, size), color)
      result.paste(pil_img, ((size - width) // 2, (size - height) // 2))
      return result
  
  def scale_and_place_object(image, scale_factor):
      assert np.shape(image)[-1]==4  # RGBA
  
      # Extract the alpha channel (transparency) and the object (RGB channels)
      alpha_channel = image[:, :, 3]
  
      # Find the bounding box coordinates of the object
      coords = cv2.findNonZero(alpha_channel)
      x, y, width, height = cv2.boundingRect(coords)
  
      # Calculate the scale factor for resizing
      original_height, original_width = image.shape[:2]
  
      if width > height:
          size = width
          original_size = original_width
      else:
          size = height
          original_size = original_height
  
      scale_factor = min(scale_factor, size / (original_size+0.0))
  
      new_size = scale_factor * original_size
      scale_factor = new_size / size
  
      # Calculate the new size based on the scale factor
      new_width = int(width * scale_factor)
      new_height = int(height * scale_factor)
  
      center_x = original_width // 2
      center_y = original_height // 2
  
      paste_x = center_x - (new_width // 2)
      paste_y = center_y - (new_height // 2)
  
      # Resize the object (RGB channels) to the new size
      rescaled_object = cv2.resize(image[y:y+height, x:x+width], (new_width, new_height))
  
      # Create a new RGBA image with the resized image
      new_image = np.zeros((original_height, original_width, 4), dtype=np.uint8)
  
      new_image[paste_y:paste_y + new_height, paste_x:paste_x + new_width] = rescaled_object
  
      return new_image
  
  class SingleImageDataset(Dataset):
      def __init__(self,
          fix_cam_pose_dir: str,
          num_views: int,
          img_wh: Tuple[int, int],
          bg_color: str,
          crop_size: int = 224,
          single_image: Optional[Image.Image] = None,
          num_validation_samples: Optional[int] = None,
          root_dir: Optional[str] = '',
          filepaths: Optional[list] = None,
          cond_type: Optional[str] = None
          ) -> None:
          """Create a dataset from a folder of images.
          If you pass in a root directory it will be searched for images
          ending in ext (ext can be a list)
          """
          self.root_dir = root_dir
          self.num_views = num_views
          self.img_wh = img_wh
          self.crop_size = crop_size
          self.bg_color = bg_color
          self.cond_type = cond_type
  
          if self.num_views == 4:
              self.view_types  = ['front', 'right', 'back', 'left']
          elif self.num_views == 5:
              self.view_types  = ['front', 'front_right', 'right', 'back', 'left']
          elif self.num_views == 6:
              self.view_types  = ['front', 'front_right', 'right', 'back', 'left', 'front_left']
          
          self.fix_cam_pose_dir = fix_cam_pose_dir
          
          self.fix_cam_poses = self.load_fixed_poses()  # world2cam matrix
  
          if single_image is None:
              if filepaths is None:
                  # Get a list of all files in the directory
                  file_list = os.listdir(self.root_dir)
              else:
                  file_list = filepaths
  
              # Filter the files that end with .png or .jpg
              self.file_list = [file for file in file_list if file.endswith(('.png', '.jpg'))]
          else:
              self.file_list = None
  
          # load all images
          self.all_images = []
          self.all_alphas = []
          bg_color = self.get_bg_color()
  
          if single_image is not None:
              image, alpha = self.load_image(None, bg_color, return_type='pt', Imagefile=single_image)
              self.all_images.append(image)
              self.all_alphas.append(alpha)
          else:
              for file in self.file_list:
                  print(os.path.join(self.root_dir, file))
                  image, alpha = self.load_image(os.path.join(self.root_dir, file), bg_color, return_type='pt')
                  self.all_images.append(image)
                  self.all_alphas.append(alpha)
  
          self.all_images = self.all_images[:num_validation_samples]
          self.all_alphas = self.all_alphas[:num_validation_samples]
  
  
      def __len__(self):
          return len(self.all_images)
  
      def load_fixed_poses(self):
          poses = {}
          for face in self.view_types:
              RT = np.loadtxt(os.path.join(self.fix_cam_pose_dir,'%03d_%s_RT.txt'%(0, face)))
              poses[face] = RT
  
          return poses
          
      def cartesian_to_spherical(self, xyz):
          ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
          xy = xyz[:,0]**2 + xyz[:,1]**2
          z = np.sqrt(xy + xyz[:,2]**2)
          theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
          #ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
          azimuth = np.arctan2(xyz[:,1], xyz[:,0])
          return np.array([theta, azimuth, z])
  
      def get_T(self, target_RT, cond_RT):
          R, T = target_RT[:3, :3], target_RT[:, -1]
          T_target = -R.T @ T # change to cam2world
  
          R, T = cond_RT[:3, :3], cond_RT[:, -1]
          T_cond = -R.T @ T
  
          theta_cond, azimuth_cond, z_cond = self.cartesian_to_spherical(T_cond[None, :])
          theta_target, azimuth_target, z_target = self.cartesian_to_spherical(T_target[None, :])
          
          d_theta = theta_target - theta_cond
          d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
          d_z = z_target - z_cond
          
          # d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
          return d_theta, d_azimuth
  
      def get_bg_color(self):
          if self.bg_color == 'white':
              bg_color = np.array([1., 1., 1.], dtype=np.float32)
          elif self.bg_color == 'black':
              bg_color = np.array([0., 0., 0.], dtype=np.float32)
          elif self.bg_color == 'gray':
              bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
          elif self.bg_color == 'random':
              bg_color = np.random.rand(3)
          elif isinstance(self.bg_color, float):
              bg_color = np.array([self.bg_color] * 3, dtype=np.float32)
          else:
              raise NotImplementedError
          return bg_color
      
      
      def load_image(self, img_path, bg_color, return_type='np', Imagefile=None):
          # pil always returns uint8
          if Imagefile is None:
              image_input = Image.open(img_path)
          else:
              image_input = Imagefile
          image_size = self.img_wh[0]
  
          if self.crop_size!=-1:
              alpha_np = np.asarray(image_input)[:, :, 3]
              coords = np.stack(np.nonzero(alpha_np), 1)[:, (1, 0)]
              min_x, min_y = np.min(coords, 0)
              max_x, max_y = np.max(coords, 0)
              ref_img_ = image_input.crop((min_x, min_y, max_x, max_y))
              h, w = ref_img_.height, ref_img_.width
              scale = self.crop_size / max(h, w)
              h_, w_ = int(scale * h), int(scale * w)
              ref_img_ = ref_img_.resize((w_, h_))
              image_input = add_margin(ref_img_, size=image_size)
          else:
              image_input = add_margin(image_input, size=max(image_input.height, image_input.width))
              image_input = image_input.resize((image_size, image_size))
  
          # img = scale_and_place_object(img, self.scale_ratio)
          img = np.array(image_input)
          img = img.astype(np.float32) / 255. # [0, 1]
          assert img.shape[-1] == 4 # RGBA
  
          alpha = img[...,3:4]
          img = img[...,:3] * alpha + bg_color * (1 - alpha)
  
          if return_type == "np":
              pass
          elif return_type == "pt":
              img = torch.from_numpy(img)
              alpha = torch.from_numpy(alpha)
          else:
              raise NotImplementedError
          
          return img, alpha
      
  
      def __len__(self):
          return len(self.all_images)
  
      def __getitem__(self, index):
  
          image = self.all_images[index%len(self.all_images)]
          alpha = self.all_alphas[index%len(self.all_images)]
          if self.file_list is not None:
              filename = self.file_list[index%len(self.all_images)].replace(".png", "")
          else:
              filename = 'null'
  
          cond_w2c = self.fix_cam_poses['front']
  
          tgt_w2cs = [self.fix_cam_poses[view] for view in self.view_types]
  
          elevations = []
          azimuths = []
  
          img_tensors_in = [
              image.permute(2, 0, 1)
          ] * self.num_views
  
          alpha_tensors_in = [
              alpha.permute(2, 0, 1)
          ] * self.num_views
  
          for view, tgt_w2c in zip(self.view_types, tgt_w2cs):
              # evelations, azimuths
              elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
              elevations.append(elevation)
              azimuths.append(azimuth)
  
          img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
          alpha_tensors_in = torch.stack(alpha_tensors_in, dim=0).float() # (Nv, 3, H, W)
  
          elevations = torch.as_tensor(elevations).float().squeeze(1)
          azimuths = torch.as_tensor(azimuths).float().squeeze(1)
          elevations_cond = torch.as_tensor([0] * self.num_views).float()
  
          normal_class = torch.tensor([1, 0]).float()
          normal_task_embeddings = torch.stack([normal_class]*self.num_views, dim=0)  # (Nv, 2)
          color_class = torch.tensor([0, 1]).float()
          color_task_embeddings = torch.stack([color_class]*self.num_views, dim=0)  # (Nv, 2)
  
          camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1) # (Nv, 3)
  
          out =  {
              'elevations_cond': elevations_cond,
              'elevations_cond_deg': torch.rad2deg(elevations_cond),
              'elevations': elevations,
              'azimuths': azimuths,
              'elevations_deg': torch.rad2deg(elevations),
              'azimuths_deg': torch.rad2deg(azimuths),
              'imgs_in': img_tensors_in,
              'alphas': alpha_tensors_in,
              'camera_embeddings': camera_embeddings,
              'normal_task_embeddings': normal_task_embeddings,
              'color_task_embeddings': color_task_embeddings,
              'filename': filename,
          }
  
          return out
  
          
  