content: from torch.utils.data import Dataset
  import os
  import json
  import numpy as np
  import cv2
  from PIL import Image
  import torch
  from torchvision import transforms as T
  from data.scene import get_boundingbox
  
  from models.rays import gen_rays_from_single_image, gen_random_rays_from_single_image
  from kornia import create_meshgrid
  
  def get_ray_directions(H, W, focal, center=None):
      """
      Get ray directions for all pixels in camera coordinate.
      Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
                 ray-tracing-generating-camera-rays/standard-coordinate-systems
      Inputs:
          H, W, focal: image height, width and focal length
      Outputs:
          directions: (H, W, 3), the direction of the rays in camera coordinate
      """
      grid = create_meshgrid(H, W, normalized_coordinates=False)[0] + 0.5 # 1xHxWx2
  
      i, j = grid.unbind(-1)
      # the direction here is without +0.5 pixel centering as calibration is not so accurate
      # see https://github.com/bmild/nerf/issues/24
      cent = center if center is not None else [W / 2, H / 2]
      directions = torch.stack([(i - cent[0]) / focal[0], (j - cent[1]) / focal[1], torch.ones_like(i)], -1)  # (H, W, 3)
  
      return directions
  
  def load_K_Rt_from_P(filename, P=None):
      if P is None:
          lines = open(filename).read().splitlines()
          if len(lines) == 4:
              lines = lines[1:]
          lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
          P = np.asarray(lines).astype(np.float32).squeeze()
  
      out = cv2.decomposeProjectionMatrix(P)
      K = out[0]
      R = out[1]
      t = out[2]
  
      K = K / K[2, 2]
      intrinsics = np.eye(4)
      intrinsics[:3, :3] = K
  
      pose = np.eye(4, dtype=np.float32)
      pose[:3, :3] = R.transpose() 
      pose[:3, 3] = (t[:3] / t[3])[:, 0]
  
      return intrinsics, pose  # ! return cam2world matrix here
  
  
  # ! load one ref-image with multiple src-images in camera coordinate system
  class BlenderPerView(Dataset):
      def __init__(self, root_dir, split, img_wh=(256, 256), downSample=1.0,
                   N_rays=512,
                   vol_dims=[128, 128, 128], batch_size=1,
                   clean_image=False, importance_sample=False,
                   specific_dataset_name = 'GSO'
                   ):
  
  
          self.root_dir = root_dir
          self.split = split
  
          self.specific_dataset_name = specific_dataset_name
          self.N_rays = N_rays
          self.batch_size = batch_size  # - used for construct new metas for gru fusion training
          
          self.clean_image = clean_image
          self.importance_sample = importance_sample
          self.scale_factor = 1.0
          self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
          assert self.split == 'val' or 'export_mesh', 'only support val or export_mesh'
          # find all subfolders
          main_folder = os.path.join(root_dir, self.specific_dataset_name)
          self.shape_list = [""] # os.listdir(main_folder) # MODIFIED
          self.shape_list.sort()
          
          self.lvis_paths = []
          for shape_name in self.shape_list:
              self.lvis_paths.append(os.path.join(main_folder, shape_name))
  
          if img_wh is not None:
              assert img_wh[0] % 32 == 0 and img_wh[1] % 32 == 0, \
                  'img_wh must both be multiples of 32!'
  
          # * bounding box for rendering
          self.bbox_min = np.array([-1.0, -1.0, -1.0])
          self.bbox_max = np.array([1.0, 1.0, 1.0])
  
          # - used for cost volume regularization
          self.voxel_dims = torch.tensor(vol_dims, dtype=torch.float32)
          self.partial_vol_origin = torch.tensor([-1., -1., -1.], dtype=torch.float32)
          
  
      def define_transforms(self):
          self.transform = T.Compose([T.ToTensor()])
  
  
      def load_cam_info(self):
          for vid, img_id in enumerate(self.img_ids):
              intrinsic, extrinsic, near_far = self.intrinsic, np.linalg.inv(self.c2ws[vid]), self.near_far
              self.all_intrinsics.append(intrinsic)
              self.all_extrinsics.append(extrinsic)
              self.all_near_fars.append(near_far)
  
      def read_mask(self, filename):
          mask_h = cv2.imread(filename, 0)
          mask_h = cv2.resize(mask_h, None, fx=self.downSample, fy=self.downSample,
                              interpolation=cv2.INTER_NEAREST)
          mask = cv2.resize(mask_h, None, fx=0.25, fy=0.25,
                            interpolation=cv2.INTER_NEAREST)
  
          mask[mask > 0] = 1  # the masks stored in png are not binary
          mask_h[mask_h > 0] = 1
  
          return mask, mask_h
  
      def cal_scale_mat(self, img_hw, intrinsics, extrinsics, near_fars, factor=1.):
  
          center, radius, bounds = get_boundingbox(img_hw, intrinsics, extrinsics, near_fars)
  
          radius = radius * factor
          scale_mat = np.diag([radius, radius, radius, 1.0])
          scale_mat[:3, 3] = center.cpu().numpy()
          scale_mat = scale_mat.astype(np.float32)
  
          return scale_mat, 1. / radius.cpu().numpy()
  
      def __len__(self):
          return len(self.lvis_paths)
  
      def __getitem__(self, idx):
          sample = {}
          origin_idx = idx
          imgs, depths_h, masks_h = [], [], []  # full size (256, 256)
          intrinsics, w2cs, c2ws, near_fars = [], [], [], []  # record proj-mats between views
  
          folder_path = self.lvis_paths[idx]
          target_idx = 0
          # last subdir name
          shape_name = os.path.split(folder_path)[-1]
  
          pose_json_path = os.path.join(folder_path, "pose.json")
          with open(pose_json_path, 'r') as f:
              meta = json.load(f)
          
          self.img_ids = list(meta["c2ws"].keys()) # e.g. "view_0", "view_7", "view_0_2_10"
          self.img_wh = (256, 256)
          self.input_poses = np.array(list(meta["c2ws"].values()))
          intrinsic = np.eye(4)
          intrinsic[:3, :3] = np.array(meta["intrinsics"])
          self.intrinsic = intrinsic
          self.near_far = np.array(meta["near_far"])
          self.define_transforms()
          self.blender2opencv = np.array(
              [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
          )
          
          self.c2ws = []
          self.w2cs = []
          self.all_intrinsics = []  # the cam info of the whole scene
          self.all_extrinsics = []
          self.all_near_fars = []
  
          for idx, img_id in enumerate(self.img_ids):
              pose = self.input_poses[idx]
              c2w = pose @ self.blender2opencv
              self.c2ws.append(c2w)
              self.all_intrinsics.append(self.intrinsic)
              self.all_near_fars.append(self.near_far)
              self.all_extrinsics.append(np.linalg.inv(c2w))
              self.w2cs.append(np.linalg.inv(c2w))
          self.c2ws = np.stack(self.c2ws, axis=0)
          self.w2cs = np.stack(self.w2cs, axis=0)
  
  
          # target view
          c2w = self.c2ws[target_idx]
          w2c = np.linalg.inv(c2w)
          w2c_ref = w2c
          w2c_ref_inv = np.linalg.inv(w2c_ref)
  
          w2cs.append(w2c @ w2c_ref_inv)
          c2ws.append(np.linalg.inv(w2c @ w2c_ref_inv))
  
          img_filename = os.path.join(folder_path, 'stage1_8', f'{self.img_ids[target_idx]}')
  
          img = Image.open(img_filename)
          img = self.transform(img)  # (4, h, w)
          
  
          if img.shape[0] == 4:
              img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
          imgs += [img]
  
  
          depth_h = torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
          depth_h = depth_h.fill_(-1.0)
          mask_h = torch.ones((img.shape[1], img.shape[2]), dtype=torch.int32)
  
  
          depths_h.append(depth_h)
          masks_h.append(mask_h)
          
          intrinsic = self.intrinsic
          intrinsics.append(intrinsic)
      
          near_fars.append(self.all_near_fars[target_idx])
          image_perm = 0  # only supervised on reference view
  
          mask_dilated = None
  
          src_views = range(8, 8 + 8 * 4)
  
          for vid in src_views:
  
              img_filename = os.path.join(folder_path, 'stage2_8', f'{self.img_ids[vid]}')
              img = Image.open(img_filename)
              img_wh = self.img_wh
  
              img = self.transform(img)
              if img.shape[0] == 4:
                  img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
  
              imgs += [img]
              depth_h = np.ones(img.shape[1:], dtype=np.float32)
              depths_h.append(depth_h)
              masks_h.append(np.ones(img.shape[1:], dtype=np.int32))
  
              near_fars.append(self.all_near_fars[vid])
              intrinsics.append(self.all_intrinsics[vid])
  
              w2cs.append(self.all_extrinsics[vid] @ w2c_ref_inv)
  
          
          # ! estimate scale_mat
          scale_mat, scale_factor = self.cal_scale_mat(
                                                       img_hw=[img_wh[1], img_wh[0]],
                                                       intrinsics=intrinsics, extrinsics=w2cs,
                                                       near_fars=near_fars, factor=1.1
                                                       )
  
  
          new_near_fars = []
          new_w2cs = []
          new_c2ws = []
          new_affine_mats = []
          new_depths_h = []
          for intrinsic, extrinsic, near_far, depth in zip(intrinsics, w2cs, near_fars, depths_h):
  
              P = intrinsic @ extrinsic @ scale_mat
              P = P[:3, :4]
              # - should use load_K_Rt_from_P() to obtain c2w
              c2w = load_K_Rt_from_P(None, P)[1]
              w2c = np.linalg.inv(c2w)
              new_w2cs.append(w2c)
              new_c2ws.append(c2w)
              affine_mat = np.eye(4)
              affine_mat[:3, :4] = intrinsic[:3, :3] @ w2c[:3, :4]
              new_affine_mats.append(affine_mat)
  
              camera_o = c2w[:3, 3]
              dist = np.sqrt(np.sum(camera_o ** 2))
              near = dist - 1
              far = dist + 1
  
              new_near_fars.append([0.95 * near, 1.05 * far])
              new_depths_h.append(depth * scale_factor)
  
          imgs = torch.stack(imgs).float()
          depths_h = np.stack(new_depths_h)
          masks_h = np.stack(masks_h)
  
          affine_mats = np.stack(new_affine_mats)
          intrinsics, w2cs, c2ws, near_fars = np.stack(intrinsics), np.stack(new_w2cs), np.stack(new_c2ws), np.stack(
              new_near_fars)
          
          if self.split == 'train':
              start_idx = 0
          else:
              start_idx = 1
  
  
          target_w2cs = []
          target_intrinsics = []
          new_target_w2cs = []
          for i_idx in range(8):
              target_w2cs.append(self.all_extrinsics[i_idx] @ w2c_ref_inv)
              target_intrinsics.append(self.all_intrinsics[i_idx])
  
          for intrinsic, extrinsic in zip(target_intrinsics, target_w2cs):
  
              P = intrinsic @ extrinsic @ scale_mat
              P = P[:3, :4]
              # - should use load_K_Rt_from_P() to obtain c2w
              c2w = load_K_Rt_from_P(None, P)[1]
              w2c = np.linalg.inv(c2w)
              new_target_w2cs.append(w2c)
          target_w2cs = np.stack(new_target_w2cs)
  
  
  
          view_ids = [idx] + list(src_views)
          sample['origin_idx'] = origin_idx
          sample['images'] = imgs  # (V, 3, H, W)
          sample['depths_h'] = torch.from_numpy(depths_h.astype(np.float32))  # (V, H, W)
          sample['masks_h'] = torch.from_numpy(masks_h.astype(np.float32))  # (V, H, W)
          sample['w2cs'] = torch.from_numpy(w2cs.astype(np.float32))  # (V, 4, 4)
          sample['c2ws'] = torch.from_numpy(c2ws.astype(np.float32))  # (V, 4, 4)
          sample['target_candidate_w2cs'] = torch.from_numpy(target_w2cs.astype(np.float32))  # (8, 4, 4)
          sample['near_fars'] = torch.from_numpy(near_fars.astype(np.float32))  # (V, 2)
          sample['intrinsics'] = torch.from_numpy(intrinsics.astype(np.float32))[:, :3, :3]  # (V, 3, 3)
          sample['view_ids'] = torch.from_numpy(np.array(view_ids))
          sample['affine_mats'] = torch.from_numpy(affine_mats.astype(np.float32))  # ! in world space
  
          sample['scan'] = shape_name
  
          sample['scale_factor'] = torch.tensor(scale_factor)
          sample['img_wh'] = torch.from_numpy(np.array(img_wh))
          sample['render_img_idx'] = torch.tensor(image_perm)
          sample['partial_vol_origin'] = self.partial_vol_origin
          sample['meta'] = str(self.specific_dataset_name) + '_' + str(shape_name) + "_refview" + str(view_ids[0])
          # print("meta: ", sample['meta'])
  
          # - image to render
          sample['query_image'] = sample['images'][0]
          sample['query_c2w'] = sample['c2ws'][0]
          sample['query_w2c'] = sample['w2cs'][0]
          sample['query_intrinsic'] = sample['intrinsics'][0]
          sample['query_depth'] = sample['depths_h'][0]
          sample['query_mask'] = sample['masks_h'][0]
          sample['query_near_far'] = sample['near_fars'][0]
  
          sample['images'] = sample['images'][start_idx:]  # (V, 3, H, W)
          sample['depths_h'] = sample['depths_h'][start_idx:]  # (V, H, W)
          sample['masks_h'] = sample['masks_h'][start_idx:]  # (V, H, W)
          sample['w2cs'] = sample['w2cs'][start_idx:]  # (V, 4, 4)
          sample['c2ws'] = sample['c2ws'][start_idx:]  # (V, 4, 4)
          sample['intrinsics'] = sample['intrinsics'][start_idx:]  # (V, 3, 3)
          sample['view_ids'] = sample['view_ids'][start_idx:]
          sample['affine_mats'] = sample['affine_mats'][start_idx:]  # ! in world space
  
          sample['scale_mat'] = torch.from_numpy(scale_mat)
          sample['trans_mat'] = torch.from_numpy(w2c_ref_inv)
  
          # - generate rays
          if ('val' in self.split) or ('test' in self.split):
              sample_rays = gen_rays_from_single_image(
                  img_wh[1], img_wh[0],
                  sample['query_image'],
                  sample['query_intrinsic'],
                  sample['query_c2w'],
                  depth=sample['query_depth'],
                  mask=sample['query_mask'] if self.clean_image else None)
          else:
              sample_rays = gen_random_rays_from_single_image(
                  img_wh[1], img_wh[0],
                  self.N_rays,
                  sample['query_image'],
                  sample['query_intrinsic'],
                  sample['query_c2w'],
                  depth=sample['query_depth'],
                  mask=sample['query_mask'] if self.clean_image else None,
                  dilated_mask=mask_dilated,
                  importance_sample=self.importance_sample)
  
  
          sample['rays'] = sample_rays
  
          return sample
  