sparseneus and elev est

.gitignore

@@ -1 +1,2 @@
-__pycache__/
+__pycache__/
+*.DS_Store

SparseNeuS_demo_v1/confs/blender_general_lod1_val_new.conf

@@ -0,0 +1,137 @@
+# - for the lod1 geometry network, using adaptive cost for sparse cost regularization network
+#- for lod1 rendering network, using depth-adaptive render
+
+general {
+  base_exp_dir = ./exp/val/1_4_only_narrow_lod1
+
+  recording = [
+    ./,
+    ./data
+    ./ops
+    ./models
+    ./loss
+  ]
+}
+
+dataset {
+  # local path
+  trainpath = /objaverse-processed/zero12345_img/eval_selected
+  valpath = /objaverse-processed/zero12345_img/eval_selected
+  testpath = /objaverse-processed/zero12345_img/eval_selected
+  # trainpath = /objaverse-processed/zero12345_img/zero12345_2stage_5pred_sample/
+  # valpath = /objaverse-processed/zero12345_img/zero12345_2stage_5pred_sample/
+  # testpath = /objaverse-processed/zero12345_img/zero12345_2stage_5pred_sample/
+  imgScale_train = 1.0
+  imgScale_test = 1.0
+  nviews = 5
+  clean_image = True
+  importance_sample = True
+  test_ref_views = [23]
+
+  # test dataset
+  test_n_views = 2
+  test_img_wh = [256, 256]
+  test_clip_wh = [0, 0]
+  test_scan_id = scan110
+  train_img_idx = [49, 50, 52, 53, 54, 56, 58] #[21, 22, 23, 24, 25]  #
+  test_img_idx = [51, 55, 57]  #[32, 33, 34]  #
+
+  test_dir_comment = train
+}
+
+train {
+  learning_rate = 2e-4
+  learning_rate_milestone = [100000, 150000, 200000]
+  learning_rate_factor = 0.5
+  end_iter = 200000
+  save_freq = 5000
+  val_freq = 1
+  val_mesh_freq =1
+  report_freq = 100
+
+  N_rays = 512
+
+  validate_resolution_level = 4
+  anneal_start = 0
+  anneal_end = 25000
+  anneal_start_lod1 = 0
+  anneal_end_lod1 = 15000
+
+  use_white_bkgd = True
+
+  # Loss
+  # ! for training the lod1 network, don't use this regularization in first 10k steps; then use the regularization
+  sdf_igr_weight = 0.1
+  sdf_sparse_weight = 0.02  # 0.002 for lod1 network;  0.02 for lod0 network
+  sdf_decay_param = 100   # cannot be too large, which decide the tsdf range
+  fg_bg_weight = 0.01  # first 0.01
+  bg_ratio = 0.3
+
+  if_fix_lod0_networks = True
+}
+
+model {
+  num_lods = 2
+
+  sdf_network_lod0 {
+    lod = 0,
+    ch_in = 56,  # the channel num of fused pyramid features
+    voxel_size = 0.02105263,  # 0.02083333, should be 2/95
+    vol_dims = [96, 96, 96],
+    hidden_dim = 128,
+    cost_type = variance_mean
+    d_pyramid_feature_compress = 16,
+    regnet_d_out = 16,
+    num_sdf_layers = 4,
+    # position embedding
+    multires = 6
+  }
+
+
+  sdf_network_lod1 {
+    lod = 1,
+    ch_in = 56,  # the channel num of fused pyramid features
+    voxel_size = 0.0104712, #0.01041667, should be 2/191
+    vol_dims = [192, 192, 192],
+    hidden_dim = 128,
+    cost_type = variance_mean
+    d_pyramid_feature_compress = 8,
+    regnet_d_out = 8,
+    num_sdf_layers = 4,
+    # position embedding
+    multires = 6
+  }
+
+
+  variance_network {
+    init_val = 0.2
+  }
+
+  variance_network_lod1 {
+    init_val = 0.2
+  }
+
+  rendering_network {
+    in_geometry_feat_ch = 16
+    in_rendering_feat_ch = 56
+    anti_alias_pooling = True
+  }
+
+  rendering_network_lod1 {
+    in_geometry_feat_ch = 8
+    in_rendering_feat_ch = 56
+    anti_alias_pooling = True
+
+  }
+
+
+  trainer {
+    n_samples_lod0 = 64
+    n_importance_lod0 = 64
+    n_samples_lod1 = 64
+    n_importance_lod1 = 64
+    n_outside = 0  # 128 if render_outside_uniform_sampling
+    perturb = 1.0
+    alpha_type = div
+  }
+}

SparseNeuS_demo_v1/confs/one2345_lod0_val_demo.conf

@@ -0,0 +1,137 @@
+# - for the lod1 geometry network, using adaptive cost for sparse cost regularization network
+#- for lod1 rendering network, using depth-adaptive render
+
+general {
+
+  base_exp_dir = exp/lod0 # !!! where you store the results and checkpoints to be used
+  recording = [
+    ./,
+    ./data
+    ./ops
+    ./models
+    ./loss
+  ]
+}
+
+dataset {
+  trainpath = ../
+  valpath = ../ # !!! where you store the validation data
+  testpath = ../
+
+
+
+  imgScale_train = 1.0
+  imgScale_test = 1.0
+  nviews = 5
+  clean_image = True
+  importance_sample = True
+  test_ref_views = [23]
+
+  # test dataset
+  test_n_views = 2
+  test_img_wh = [256, 256]
+  test_clip_wh = [0, 0]
+  test_scan_id = scan110
+  train_img_idx = [49, 50, 52, 53, 54, 56, 58] #[21, 22, 23, 24, 25]  #
+  test_img_idx = [51, 55, 57]  #[32, 33, 34]  #
+
+  test_dir_comment = train
+}
+
+train {
+  learning_rate = 2e-4
+  learning_rate_milestone = [100000, 150000, 200000]
+  learning_rate_factor = 0.5
+  end_iter = 200000
+  save_freq = 5000
+  val_freq = 1
+  val_mesh_freq = 1
+  report_freq = 100
+
+  N_rays = 512
+
+  validate_resolution_level = 4
+  anneal_start = 0
+  anneal_end = 25000
+  anneal_start_lod1 = 0
+  anneal_end_lod1 = 15000
+
+  use_white_bkgd = True
+
+  # Loss
+  # ! for training the lod1 network, don't use this regularization in first 10k steps; then use the regularization
+  sdf_igr_weight = 0.1
+  sdf_sparse_weight = 0.02  # 0.002 for lod1 network;  0.02 for lod0 network
+  sdf_decay_param = 100   # cannot be too large, which decide the tsdf range
+  fg_bg_weight = 0.01  # first 0.01
+  bg_ratio = 0.3
+
+  if_fix_lod0_networks = False
+}
+
+model {
+  num_lods = 1
+
+  sdf_network_lod0 {
+    lod = 0,
+    ch_in = 56,  # the channel num of fused pyramid features
+    voxel_size = 0.02105263,  # 0.02083333, should be 2/95
+    vol_dims = [96, 96, 96],
+    hidden_dim = 128,
+    cost_type = variance_mean
+    d_pyramid_feature_compress = 16,
+    regnet_d_out = 16,
+    num_sdf_layers = 4,
+    # position embedding
+    multires = 6
+  }
+
+
+  sdf_network_lod1 {
+    lod = 1,
+    ch_in = 56,  # the channel num of fused pyramid features
+    voxel_size = 0.0104712, #0.01041667, should be 2/191
+    vol_dims = [192, 192, 192],
+    hidden_dim = 128,
+    cost_type = variance_mean
+    d_pyramid_feature_compress = 8,
+    regnet_d_out = 16,
+    num_sdf_layers = 4,
+
+    # position embedding
+    multires = 6
+  }
+
+
+  variance_network {
+    init_val = 0.2
+  }
+
+  variance_network_lod1 {
+    init_val = 0.2
+  }
+
+  rendering_network {
+    in_geometry_feat_ch = 16
+    in_rendering_feat_ch = 56
+    anti_alias_pooling = True
+  }
+
+  rendering_network_lod1 {
+    in_geometry_feat_ch = 16 # default 8
+    in_rendering_feat_ch = 56
+    anti_alias_pooling = True
+
+  }
+
+
+  trainer {
+    n_samples_lod0 = 64
+    n_importance_lod0 = 64
+    n_samples_lod1 = 64
+    n_importance_lod1 = 64
+    n_outside = 0  # 128 if render_outside_uniform_sampling
+    perturb = 1.0
+    alpha_type = div
+  }
+}

SparseNeuS_demo_v1/data/blender.py

@@ -0,0 +1,340 @@
+import torch
+from torch.utils.data import Dataset
+import json
+import numpy as np
+import os
+from PIL import Image
+from torchvision import transforms as T
+from kornia import create_meshgrid
+from models.rays import gen_rays_from_single_image, gen_random_rays_from_single_image
+import cv2 as cv
+from data.scene import get_boundingbox
+
+
+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]
+    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 get_rays(directions, c2w):
+    """
+    Get ray origin and normalized directions in world coordinate for all pixels in one image.
+    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
+               ray-tracing-generating-camera-rays/standard-coordinate-systems
+    Inputs:
+        directions: (H, W, 3) precomputed ray directions in camera coordinate
+        c2w: (3, 4) transformation matrix from camera coordinate to world coordinate
+    Outputs:
+        rays_o: (H*W, 3), the origin of the rays in world coordinate
+        rays_d: (H*W, 3), the normalized direction of the rays in world coordinate
+    """
+    # Rotate ray directions from camera coordinate to the world coordinate
+    rays_d = directions @ c2w[:3, :3].T  # (H, W, 3)
+    # rays_d = rays_d / torch.norm(rays_d, dim=-1, keepdim=True)
+    # The origin of all rays is the camera origin in world coordinate
+    rays_o = c2w[:3, 3].expand(rays_d.shape)  # (H, W, 3)
+
+    rays_d = rays_d.view(-1, 3)
+    rays_o = rays_o.view(-1, 3)
+
+    return rays_o, rays_d
+
+
+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 = cv.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()  # ? why need transpose here
+    pose[:3, 3] = (t[:3] / t[3])[:, 0]
+
+    return intrinsics, pose  # ! return cam2world matrix here
+
+
+class BlenderDataset(Dataset):
+    def __init__(self, root_dir, split, scan_id, n_views, train_img_idx=[], test_img_idx=[],
+                 img_wh=[800, 800], clip_wh=[0, 0], original_img_wh=[800, 800],
+                 N_rays=512, h_patch_size=5, near=2.0, far=6.0):
+        self.root_dir = root_dir
+        self.split = split
+        self.img_wh = img_wh
+        self.clip_wh = clip_wh
+        self.define_transforms()
+        self.train_img_idx = train_img_idx
+        self.test_img_idx = test_img_idx
+        self.N_rays = N_rays
+        self.h_patch_size = h_patch_size  # used to extract patch for supervision
+        self.n_views = n_views
+        self.near, self.far = near, far
+        self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
+
+        with open(os.path.join(self.root_dir, f"transforms_{self.split}.json"), 'r') as f:
+            self.meta = json.load(f)
+
+
+        self.read_meta(near, far)
+        # import ipdb; ipdb.set_trace()
+        self.raw_near_fars = np.stack([np.array([self.near, self.far]) for i in range(len(self.meta['frames']))])
+
+
+        # ! estimate scale_mat
+        self.scale_mat, self.scale_factor = self.cal_scale_mat(
+            img_hw=[self.img_wh[1], self.img_wh[0]],
+            intrinsics=self.all_intrinsics[self.train_img_idx],
+            extrinsics=self.all_w2cs[self.train_img_idx],
+            near_fars=self.raw_near_fars[self.train_img_idx],
+            factor=1.1)
+        # self.scale_mat = np.eye(4)
+        # self.scale_factor = 1.0
+        # import ipdb; ipdb.set_trace()
+        # * after scaling and translation, unit bounding box
+        self.scaled_intrinsics, self.scaled_w2cs, self.scaled_c2ws, \
+        self.scaled_affine_mats, self.scaled_near_fars = self.scale_cam_info()
+
+        self.bbox_min = np.array([-1.0, -1.0, -1.0])
+        self.bbox_max = np.array([1.0, 1.0, 1.0])
+        self.partial_vol_origin = torch.Tensor([-1., -1., -1.])
+        self.white_back = True
+
+    def read_meta(self, near=2.0, far=6.0):
+
+
+        self.ref_img_idx = self.train_img_idx[0]
+        ref_c2w = np.array(self.meta['frames'][self.ref_img_idx]['transform_matrix']) @ self.blender2opencv
+        # ref_c2w = torch.FloatTensor(ref_c2w)
+        self.ref_c2w = ref_c2w
+        self.ref_w2c = np.linalg.inv(ref_c2w)
+
+
+        w, h = self.img_wh
+        self.focal = 0.5 * 800 / np.tan(0.5 * self.meta['camera_angle_x'])  # original focal length
+        self.focal *= self.img_wh[0] / 800  # modify focal length to match size self.img_wh
+
+        # bounds, common for all scenes
+        self.near = near
+        self.far = far
+        self.bounds = np.array([self.near, self.far])
+
+        # ray directions for all pixels, same for all images (same H, W, focal)
+        self.directions = get_ray_directions(h, w, [self.focal,self.focal])  # (h, w, 3)
+        intrinsics = np.eye(4)
+        intrinsics[:3, :3] = np.array([[self.focal,0,w/2],[0,self.focal,h/2],[0,0,1]]).astype(np.float32)
+        self.intrinsics = intrinsics
+
+        self.image_paths = []
+        self.poses = []
+        self.all_rays = []
+        self.all_images = []
+        self.all_masks = []
+        self.all_w2cs = []
+        self.all_intrinsics = []
+        for frame in self.meta['frames']:
+            pose = np.array(frame['transform_matrix']) @ self.blender2opencv
+            self.poses += [pose]
+            c2w = torch.FloatTensor(pose)
+            w2c = np.linalg.inv(c2w)
+            image_path = os.path.join(self.root_dir, f"{frame['file_path']}.png")
+            self.image_paths += [image_path]
+            img = Image.open(image_path)
+            img = img.resize(self.img_wh, Image.LANCZOS)
+            img = self.transform(img)  # (4, h, w)
+            
+            self.all_masks += [img[-1:,:]>0]
+            # img = img[:3, :] * img[ -1:,:] + (1 - img[-1:, :])  # blend A to RGB
+            img = img[:3, :] * img[ -1:,:] 
+            img = img.numpy()  # (3, h, w)
+            self.all_images += [img]
+
+
+            self.all_masks += []
+            self.all_intrinsics.append(self.intrinsics)
+            # - transform from world system to ref-camera system
+            self.all_w2cs.append(w2c @ np.linalg.inv(self.ref_w2c))
+            
+        self.all_images = torch.from_numpy(np.stack(self.all_images)).to(torch.float32)
+        self.all_intrinsics = torch.from_numpy(np.stack(self.all_intrinsics)).to(torch.float32)
+        self.all_w2cs = torch.from_numpy(np.stack(self.all_w2cs)).to(torch.float32)
+        # self.img_wh = [self.img_wh[0] - self.clip_wh[0] - self.clip_wh[2],
+        #                self.img_wh[1] - self.clip_wh[1] - self.clip_wh[3]]
+
+    def cal_scale_mat(self, img_hw, intrinsics, extrinsics, near_fars, factor=1.):
+        center, radius, _ = 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 scale_cam_info(self):
+        new_intrinsics = []
+        new_near_fars = []
+        new_w2cs = []
+        new_c2ws = []
+        new_affine_mats = []
+        for idx in range(len(self.all_images)):
+
+            intrinsics = self.all_intrinsics[idx]
+            # import ipdb; ipdb.set_trace()
+            P = intrinsics @ self.all_w2cs[idx] @ self.scale_mat
+            P = P.cpu().numpy()[: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)
+            new_intrinsics.append(intrinsics)
+            affine_mat = np.eye(4)
+            affine_mat[:3, :4] = intrinsics[: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_intrinsics, new_w2cs, new_c2ws, new_affine_mats, new_near_fars = \
+            np.stack(new_intrinsics), np.stack(new_w2cs), np.stack(new_c2ws), \
+            np.stack(new_affine_mats), np.stack(new_near_fars)
+
+        new_intrinsics = torch.from_numpy(np.float32(new_intrinsics))
+        new_w2cs = torch.from_numpy(np.float32(new_w2cs))
+        new_c2ws = torch.from_numpy(np.float32(new_c2ws))
+        new_affine_mats = torch.from_numpy(np.float32(new_affine_mats))
+        new_near_fars = torch.from_numpy(np.float32(new_near_fars))
+
+        return new_intrinsics, new_w2cs, new_c2ws, new_affine_mats, new_near_fars
+
+    def load_poses_all(self, file=f"transforms_train.json"):
+        with open(os.path.join(self.root_dir, file), 'r') as f:
+            meta = json.load(f)
+
+        c2ws = []
+        for i,frame in enumerate(meta['frames']):
+            c2ws.append(np.array(frame['transform_matrix']) @ self.blender2opencv)
+        return np.stack(c2ws)
+
+    def define_transforms(self):
+        self.transform = T.ToTensor()
+
+
+
+    def get_conditional_sample(self):
+        sample = {}
+        support_idxs = self.train_img_idx
+
+        sample['images'] = self.all_images[support_idxs]  # (V, 3, H, W)
+        sample['w2cs'] = self.scaled_w2cs[self.train_img_idx]  # (V, 4, 4)
+        sample['c2ws'] = self.scaled_c2ws[self.train_img_idx]  # (V, 4, 4)
+        sample['near_fars'] = self.scaled_near_fars[self.train_img_idx]  # (V, 2)
+        sample['intrinsics'] = self.scaled_intrinsics[self.train_img_idx][:, :3, :3]  # (V, 3, 3)
+        sample['affine_mats'] = self.scaled_affine_mats[self.train_img_idx]  # ! in world space
+
+        # sample['scan'] = self.scan_id
+        sample['scale_factor'] = torch.tensor(self.scale_factor)
+        sample['scale_mat'] = torch.from_numpy(self.scale_mat)
+        sample['trans_mat'] = torch.from_numpy(np.linalg.inv(self.ref_w2c))
+        sample['img_wh'] = torch.from_numpy(np.array(self.img_wh))
+        sample['partial_vol_origin'] = torch.tensor(self.partial_vol_origin, dtype=torch.float32)
+
+        return sample
+
+
+
+    def __len__(self):
+        if self.split == 'train':
+            return self.n_views * 1000
+        else:
+            return len(self.test_img_idx) * 1000
+
+
+    def __getitem__(self, idx):
+        sample = {}
+
+        if self.split == 'train':
+            render_idx = self.train_img_idx[idx % self.n_views]
+            support_idxs = [idx for idx in self.train_img_idx if idx != render_idx]
+        else:
+            # render_idx = idx % self.n_test_images + self.n_train_images
+            render_idx = self.test_img_idx[idx % len(self.test_img_idx)]
+            support_idxs = [render_idx]
+
+        sample['images'] = self.all_images[support_idxs]  # (V, 3, H, W)
+        sample['w2cs'] = self.scaled_w2cs[support_idxs]  # (V, 4, 4)
+        sample['c2ws'] = self.scaled_c2ws[support_idxs]  # (V, 4, 4)
+        sample['intrinsics'] = self.scaled_intrinsics[support_idxs][:, :3, :3]  # (V, 3, 3)
+        sample['affine_mats'] = self.scaled_affine_mats[support_idxs]  # ! in world space
+        # sample['scan'] = self.scan_id
+        sample['scale_factor'] = torch.tensor(self.scale_factor)
+        sample['img_wh'] = torch.from_numpy(np.array(self.img_wh))
+        sample['partial_vol_origin'] = torch.tensor(self.partial_vol_origin, dtype=torch.float32)
+        sample['img_index'] = torch.tensor(render_idx)
+
+        # - query image
+        sample['query_image'] = self.all_images[render_idx]
+        sample['query_c2w'] = self.scaled_c2ws[render_idx]
+        sample['query_w2c'] = self.scaled_w2cs[render_idx]
+        sample['query_intrinsic'] = self.scaled_intrinsics[render_idx]
+        sample['query_near_far'] = self.scaled_near_fars[render_idx]
+        # sample['meta'] = str(self.scan_id) + "_" + os.path.basename(self.images_list[render_idx])
+        sample['scale_mat'] = torch.from_numpy(self.scale_mat)
+        sample['trans_mat'] = torch.from_numpy(np.linalg.inv(self.ref_w2c))
+        sample['rendering_c2ws'] = self.scaled_c2ws[self.test_img_idx]
+        sample['rendering_imgs_idx'] = torch.Tensor(np.array(self.test_img_idx).astype(np.int32))
+
+        # - generate rays
+        if self.split == 'val' or self.split == 'test':
+            sample_rays = gen_rays_from_single_image(
+                self.img_wh[1], self.img_wh[0],
+                sample['query_image'],
+                sample['query_intrinsic'],
+                sample['query_c2w'],
+                depth=None,
+                mask=None)
+        else:
+            sample_rays = gen_random_rays_from_single_image(
+                self.img_wh[1], self.img_wh[0],
+                self.N_rays,
+                sample['query_image'],
+                sample['query_intrinsic'],
+                sample['query_c2w'],
+                depth=None,
+                mask=None,
+                dilated_mask=None,
+                importance_sample=False,
+                h_patch_size=self.h_patch_size
+            )
+
+        sample['rays'] = sample_rays
+
+        return sample

SparseNeuS_demo_v1/data/blender_general.py

@@ -0,0 +1,432 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        depth_h = np.array(read_pfm(filename)[0], dtype=np.float32)  # (1200, 1600)
+        depth_h = cv2.resize(depth_h, None, fx=0.5, fy=0.5,
+                             interpolation=cv2.INTER_NEAREST)  # (600, 800)
+        depth_h = depth_h[44:556, 80:720]  # (512, 640)
+        depth_h = cv2.resize(depth_h, None, fx=self.downSample, fy=self.downSample,
+                             interpolation=cv2.INTER_NEAREST)
+        depth = cv2.resize(depth_h, None, fx=1.0 / 4, fy=1.0 / 4,
+                           interpolation=cv2.INTER_NEAREST)
+
+        return depth, depth_h
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 8*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        depth_h = cv2.imread(filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 65535 * 1.4 + 0.5
+
+        depth_h[depth_h < near_bound+1e-3] = 0.0
+
+        depth = {}
+        for l in range(3):
+            depth[f"level_{l}"] = cv2.resize(
+                depth_h,
+                None,
+                fx=1.0 / (2**l),
+                fy=1.0 / (2**l),
+                interpolation=cv2.INTER_NEAREST,
+            )
+
+        if self.split == "train":
+            cutout = np.ones_like(depth[f"level_2"])
+            h0 = int(np.random.randint(0, high=cutout.shape[0] // 5, size=1))
+            h1 = int(
+                np.random.randint(
+                    4 * cutout.shape[0] // 5, high=cutout.shape[0], size=1
+                )
+            )
+            w0 = int(np.random.randint(0, high=cutout.shape[1] // 5, size=1))
+            w1 = int(
+                np.random.randint(
+                    4 * cutout.shape[1] // 5, high=cutout.shape[1], size=1
+                )
+            )
+            cutout[h0:h1, w0:w1] = 0
+            depth_aug = depth[f"level_2"] * cutout
+        else:
+            depth_aug = depth[f"level_2"].copy()
+
+        return depth, depth_h, depth_aug
+
+
+    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_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        # idx = idx % 8
+        # uid = 'c40d63d5d740405e91c7f5fce855076e'
+        # folder_id = '000-123'
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        src_views = range(8+idx*4, 8+(idx+1)*4)
+
+
+        for vid in src_views:
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_{vid%4}_10.png')
+
+            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
+                                                     )
+        # print(scale_mat)
+        # print(scale_factor)
+        # ! calculate the new w2cs after scaling
+        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)
+
+        # print(new_near_fars)
+        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
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_12_narrow.py

@@ -0,0 +1,427 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+        self.imgs_per_instance = 12
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/narrow_12_split_upd.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow_8 = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow_8, 'r') as f:
+            narrow_8_meta = json.load(f)
+
+        pose_json_path_narrow_4 = "/objaverse-processed/zero12345_img/zero12345_2stage_12_pose.json"
+        with open(pose_json_path_narrow_4, 'r') as f:
+            narrow_4_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_8_meta["c2ws"].keys()) + list(narrow_4_meta["c2ws"].keys()) # (8 + 8*4) + (4 + 4*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_8_meta["c2ws"].values()) + list(narrow_4_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_8_meta["intrinsics"] == narrow_4_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_8_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_8_meta["near_far"] == narrow_4_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_8_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+        
+        idx = idx % self.imgs_per_instance # [0, 11]
+        if idx < 8:
+            # target view
+            c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+            img = Image.open(img_filename)
+
+            img = self.transform(img)  # (4, h, w)
+        else:
+            # target view
+            c2w = self.c2ws[idx-8+40]
+            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("/objaverse-processed/zero12345_img/zero12345_narrow_12/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow_12/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+            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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+        
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        src_views = range(8, 8 + 8 * 4 + 4 + 4*4)
+        src_views_used = []
+        skipped_idx = [40, 41, 42, 43]
+        for vid in src_views:
+            if vid in skipped_idx:
+                continue
+
+            src_views_used.append(vid)
+            cur_view_id = (vid - 8) // 4  # [0, 7]
+
+            # choose narrow
+            if cur_view_id < 8:
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{cur_view_id}_{vid%4}_10.png')
+            else: # choose 2-stage
+                cur_view_id = cur_view_id - 1
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow_12", folder_id, uid, f'view_{cur_view_id}_{vid%4}.png')
+
+            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)
+
+           
+        
+
+        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)
+
+        # print(new_near_fars)
+        # print("img numeber: ", len(imgs))
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        if view_ids[0] < 8:
+            meta_end = "_narrow"+ "_refview" + str(view_ids[0])
+        else:
+            meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_12_narrow_8.py

@@ -0,0 +1,427 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+        self.imgs_per_instance = 8
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/narrow_12_split_upd.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow_8 = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow_8, 'r') as f:
+            narrow_8_meta = json.load(f)
+
+        pose_json_path_narrow_4 = "/objaverse-processed/zero12345_img/zero12345_2stage_12_pose.json"
+        with open(pose_json_path_narrow_4, 'r') as f:
+            narrow_4_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_8_meta["c2ws"].keys()) + list(narrow_4_meta["c2ws"].keys()) # (8 + 8*4) + (4 + 4*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_8_meta["c2ws"].values()) + list(narrow_4_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_8_meta["intrinsics"] == narrow_4_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_8_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_8_meta["near_far"] == narrow_4_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_8_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+        
+        idx = idx % self.imgs_per_instance # [0, 11]
+        if idx < 8:
+            # target view
+            c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+            img = Image.open(img_filename)
+
+            img = self.transform(img)  # (4, h, w)
+        else:
+            # target view
+            c2w = self.c2ws[idx-8+40]
+            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("/objaverse-processed/zero12345_img/zero12345_narrow_12/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow_12/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+            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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+        
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        src_views = range(8, 8 + 8 * 4 + 4 + 4*4)
+        src_views_used = []
+        skipped_idx = [40, 41, 42, 43]
+        for vid in src_views:
+            if vid in skipped_idx:
+                continue
+
+            src_views_used.append(vid)
+            cur_view_id = (vid - 8) // 4  # [0, 7]
+
+            # choose narrow
+            if cur_view_id < 8:
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{cur_view_id}_{vid%4}_10.png')
+            else: # choose 2-stage
+                cur_view_id = cur_view_id - 1
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow_12", folder_id, uid, f'view_{cur_view_id}_{vid%4}.png')
+
+            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)
+
+           
+        
+
+        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)
+
+        # print(new_near_fars)
+        # print("img numeber: ", len(imgs))
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        if view_ids[0] < 8:
+            meta_end = "_narrow"+ "_refview" + str(view_ids[0])
+        else:
+            meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_360.py

@@ -0,0 +1,412 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_wide_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_0", "view_0_5", "view_1_7"
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        depth_h = np.array(read_pfm(filename)[0], dtype=np.float32)  # (1200, 1600)
+        depth_h = cv2.resize(depth_h, None, fx=0.5, fy=0.5,
+                             interpolation=cv2.INTER_NEAREST)  # (600, 800)
+        depth_h = depth_h[44:556, 80:720]  # (512, 640)
+        depth_h = cv2.resize(depth_h, None, fx=self.downSample, fy=self.downSample,
+                             interpolation=cv2.INTER_NEAREST)
+        depth = cv2.resize(depth_h, None, fx=1.0 / 4, fy=1.0 / 4,
+                           interpolation=cv2.INTER_NEAREST)
+
+        return depth, depth_h
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 36*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+ 
+        imgs, depths_h, masks_h = [], [], []  # full size (256, 256)
+        intrinsics, w2cs, c2ws, near_fars = [], [], [], []  # record proj mats between views
+
+
+        folder_uid_dict = self.lvis_paths[idx//36]
+
+
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        idx = idx % 36 # [0, 35]
+        gt_view_idx = idx // 12 # [0, 2]
+        target_view_idx = idx % 12 # [0, 11]
+
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{gt_view_idx}_{target_view_idx}_gt.png')
+
+        depth_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{gt_view_idx}_{target_view_idx}_gt_depth_mm.png')
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(gt_view_idx * 12, (gt_view_idx + 1) * 12)
+
+        idx_of_12 = idx - 12 * gt_view_idx # idx % 12
+
+        src_views = list(i % 12 + 12 * gt_view_idx for i in range(idx_of_12 - 1-1, idx_of_12 + 2+1))
+        
+
+        for vid in src_views:
+            # if vid == idx:
+            #     continue
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{gt_view_idx}_{target_view_idx}.png')
+
+            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
+                                                     )
+        # print(scale_mat)
+        # print(scale_factor)
+        # ! calculate the new w2cs after scaling
+        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)
+
+        # print(new_near_fars)
+        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
+
+        view_ids = [idx] + list(src_views)
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_360_2_stage_1_3.py

@@ -0,0 +1,406 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_2stage_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_0", "view_0_5", "view_1_7"
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        depth_h = np.array(read_pfm(filename)[0], dtype=np.float32)  # (1200, 1600)
+        depth_h = cv2.resize(depth_h, None, fx=0.5, fy=0.5,
+                             interpolation=cv2.INTER_NEAREST)  # (600, 800)
+        depth_h = depth_h[44:556, 80:720]  # (512, 640)
+        depth_h = cv2.resize(depth_h, None, fx=self.downSample, fy=self.downSample,
+                             interpolation=cv2.INTER_NEAREST)
+        depth = cv2.resize(depth_h, None, fx=1.0 / 4, fy=1.0 / 4,
+                           interpolation=cv2.INTER_NEAREST)
+
+        return depth, depth_h
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 6*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+ 
+        imgs, depths_h, masks_h = [], [], []  # full size (256, 256)
+        intrinsics, w2cs, c2ws, near_fars = [], [], [], []  # record proj mats between views
+
+
+        folder_uid_dict = self.lvis_paths[idx//6]
+        idx = idx % 6
+
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        # idx = idx % 24 # [0, 23]
+
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_0_{idx}_gt.png')
+
+        depth_filename = os.path.join(self.root_dir, folder_id, uid, f'view_0_{idx}_gt_depth_mm.png')
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(gt_view_idx * 12, (gt_view_idx + 1) * 12)
+
+        
+        src_views = range(6+idx*4, 6+(idx+1)*4)
+
+        for vid in src_views:
+            # if vid == idx:
+            #     continue
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_0_{idx}_{vid % 4}.png')
+
+            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
+                                                     )
+        # print(scale_mat)
+        # print(scale_factor)
+        # ! calculate the new w2cs after scaling
+        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)
+
+        # print(new_near_fars)
+        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
+
+        view_ids = [idx] + list(src_views)
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_360_2_stage_1_4.py

@@ -0,0 +1,411 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_2stage_5pred_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_0", "view_0_5", "view_1_7"
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        depth_h = np.array(read_pfm(filename)[0], dtype=np.float32)  # (1200, 1600)
+        depth_h = cv2.resize(depth_h, None, fx=0.5, fy=0.5,
+                             interpolation=cv2.INTER_NEAREST)  # (600, 800)
+        depth_h = depth_h[44:556, 80:720]  # (512, 640)
+        depth_h = cv2.resize(depth_h, None, fx=self.downSample, fy=self.downSample,
+                             interpolation=cv2.INTER_NEAREST)
+        depth = cv2.resize(depth_h, None, fx=1.0 / 4, fy=1.0 / 4,
+                           interpolation=cv2.INTER_NEAREST)
+
+        return depth, depth_h
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 6*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+ 
+        imgs, depths_h, masks_h = [], [], []  # full size (256, 256)
+        intrinsics, w2cs, c2ws, near_fars = [], [], [], []  # record proj mats between views
+
+
+        folder_uid_dict = self.lvis_paths[idx//6]
+        idx = idx % 6
+
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        # idx = idx % 24 # [0, 23]
+
+
+
+        # target view
+        c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_2stage", folder_id, uid, f'view_0_{idx}_gt.png')
+
+        depth_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_2stage", folder_id, uid, f'view_0_{idx}_gt_depth_mm.png')
+        
+
+        img = Image.open(img_filename)
+
+        img = self.transform(img)  # (4, h, w)
+        
+        # print("img_pre", img.shape)
+        if img.shape[0] == 4:
+            img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+        # print("img", img.shape)
+        imgs += [img]
+
+
+        depth_h = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+        # print("depth_h", depth_h.shape)
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(gt_view_idx * 12, (gt_view_idx + 1) * 12)
+
+        
+        src_views = range(6+idx*4, 6+(idx+1)*4)
+
+        for vid in src_views:
+            # if vid == idx:
+            #     continue
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_0_{idx}_{vid % 4 + 1}.png')
+
+            img = Image.open(img_filename)
+            img_wh = self.img_wh
+
+            img = self.transform(img)
+            # print("img shape1: ", img.shape)
+            if img.shape[0] == 4:
+                img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+            # print("img shape2: ", img.shape)
+            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
+                                                     )
+        # print(scale_mat)
+        # print(scale_factor)
+        # ! calculate the new w2cs after scaling
+        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)
+
+        # print(new_near_fars)
+        # print("imgs: ", len(imgs))
+        # print("img1 shape:", imgs[0].shape)
+        # print("img2 shape:", imgs[1].shape)
+        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
+
+        view_ids = [idx] + list(src_views)
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_4_narrow_and_4_2_stage_mix.py

@@ -0,0 +1,480 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+        self.imgs_per_instance = 16
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow, 'r') as f:
+            narrow_meta = json.load(f)
+
+        pose_json_path_two_stage = "/objaverse-processed/zero12345_img/zero12345_2stage_8_pose.json"
+        with open(pose_json_path_two_stage, 'r') as f:
+            two_stage_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_meta["c2ws"].keys()) + list(two_stage_meta["c2ws"].keys()) # (8 + 8*4) + (8 + 4*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_meta["c2ws"].values()) + list(two_stage_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_meta["intrinsics"] == two_stage_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_meta["near_far"] == two_stage_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance * len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        if idx % 2 == 0:
+            valid_list = [0, 2, 4, 6]
+        else:
+            valid_list = [1, 3, 5, 7]
+
+        if idx % 16 < 8:
+            idx = idx % 16 # [0, 7]
+            # target view
+            c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+            
+
+            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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+            mask_h = depth_h > 0
+            # print("valid pixels", np.sum(mask_h))
+            directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+            surface_points = directions * depth_h[..., None]  # [H, W, 3]
+            distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+            depth_h = distance
+
+
+            depths_h.append(depth_h)
+            masks_h.append(mask_h)
+            
+            intrinsic = self.intrinsic
+            intrinsics.append(intrinsic)
+        
+
+            near_fars.append(self.near_fars[idx])
+            image_perm = 0  # only supervised on reference view
+
+            mask_dilated = None
+
+            # src_views = range(8+idx*4, 8+(idx+1)*4)
+
+            src_views = range(8, 8 + 8 * 4)
+            src_views_used = []
+            for vid in src_views:
+                view_dix_to_use = (vid - 8) // 4
+                if view_dix_to_use not in valid_list:
+                    continue
+                src_views_used.append(vid)
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+                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)
+
+        else:
+            idx = idx % 16 - 8 # [0, 7]
+       
+            c2w = self.c2ws[idx + 40]
+            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("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{idx}_0.png')
+
+
+
+            img = Image.open(img_filename)
+
+            img = self.transform(img)  # (4, h, w)
+            
+            # print("img_pre", img.shape)
+            if img.shape[0] == 4:
+                img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+            # print("img", img.shape)
+            imgs += [img]
+
+
+            depth_h =torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+            depth_h = depth_h.fill_(-1.0)
+            # depth_h = torch.fill((img.shape[1], img.shape[2]), -1.0)
+            # print("depth_h", depth_h.shape)
+            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.near_fars[idx])
+            image_perm = 0  # only supervised on reference view
+
+            mask_dilated = None
+
+
+            
+            src_views = range(40+8, 40+8+32)
+            src_views_used = []
+            for vid in src_views:
+                view_dix_to_use = (vid - 40 - 8) // 4
+                if view_dix_to_use not in valid_list:
+                    continue
+                src_views_used.append(vid)
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{idx}_{(vid-48) % 4 + 1}.png')
+
+                img = Image.open(img_filename)
+                img_wh = self.img_wh
+
+                img = self.transform(img)
+                # print("img shape1: ", img.shape)
+                if img.shape[0] == 4:
+                    img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+                # print("img shape2: ", img.shape)
+                imgs += [img]
+                depth_h =torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+                depth_h = depth_h.fill_(-1.0)
+                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)            
+        
+
+        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)
+
+        # print(new_near_fars)
+        # print("img numeber: ", len(imgs))
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        if view_ids[0] < 8:
+            meta_end = "_narrow"+ "_refview" + str(view_ids[0])
+        else:
+            meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_4_narrow_and_6_2_stage_mix.py

@@ -0,0 +1,476 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow, 'r') as f:
+            narrow_meta = json.load(f)
+
+        pose_json_path_two_stage = "/objaverse-processed/zero12345_img/zero12345_2stage_5pred_pose.json"
+        with open(pose_json_path_two_stage, 'r') as f:
+            two_stage_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_meta["c2ws"].keys()) + list(two_stage_meta["c2ws"].keys()) # (8 + 8*4) + (6 + 6*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_meta["c2ws"].values()) + list(two_stage_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_meta["intrinsics"] == two_stage_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_meta["near_far"] == two_stage_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 12*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//12]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        if idx % 12 < 8:
+            idx = idx % 12 # [0, 7]
+            # target view
+            c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+            
+
+            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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+            mask_h = depth_h > 0
+            # print("valid pixels", np.sum(mask_h))
+            directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+            surface_points = directions * depth_h[..., None]  # [H, W, 3]
+            distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+            depth_h = distance
+
+
+            depths_h.append(depth_h)
+            masks_h.append(mask_h)
+            
+            intrinsic = self.intrinsic
+            intrinsics.append(intrinsic)
+        
+
+            near_fars.append(self.near_fars[idx])
+            image_perm = 0  # only supervised on reference view
+
+            mask_dilated = None
+
+            # src_views = range(8+idx*4, 8+(idx+1)*4)
+
+            src_views = range(8, 8 + 8 * 4)
+            src_views_used = []
+            for vid in src_views:
+                if (vid // 4) % 2 != idx % 2:
+                    continue
+                src_views_used.append(vid)
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+                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)
+
+        else:
+            idx = idx % 12 - 8 # [0, 5]
+            valid_list = [0, 2, 3, 5]
+            idx = valid_list[idx]    # [0, 3]
+            c2w = self.c2ws[idx + 40]
+            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("/objaverse-processed/zero12345_img/zero12345_2stage_5pred/", folder_id, uid, f'view_0_{idx}_0.png')
+
+
+
+            img = Image.open(img_filename)
+
+            img = self.transform(img)  # (4, h, w)
+            
+            # print("img_pre", img.shape)
+            if img.shape[0] == 4:
+                img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+            # print("img", img.shape)
+            imgs += [img]
+
+
+            depth_h =torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+            depth_h = depth_h.fill_(-1.0)
+            # depth_h = torch.fill((img.shape[1], img.shape[2]), -1.0)
+            # print("depth_h", depth_h.shape)
+            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.near_fars[idx])
+            image_perm = 0  # only supervised on reference view
+
+            mask_dilated = None
+
+            # src_views = range(gt_view_idx * 12, (gt_view_idx + 1) * 12)
+
+            
+            src_views = range(40+6, 40+6+24)
+            src_views_used = []
+            for vid in src_views:
+                view_dix_to_use = (vid - 40 - 6) // 4
+                if view_dix_to_use not in valid_list:
+                    continue
+                src_views_used.append(vid)
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_2stage_5pred/", folder_id, uid, f'view_0_{idx}_{(vid-46) % 4 + 1}.png')
+
+                img = Image.open(img_filename)
+                img_wh = self.img_wh
+
+                img = self.transform(img)
+                # print("img shape1: ", img.shape)
+                if img.shape[0] == 4:
+                    img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+                # print("img shape2: ", img.shape)
+                imgs += [img]
+                depth_h =torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+                depth_h = depth_h.fill_(-1.0)
+                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)            
+        
+
+        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)
+
+        # print(new_near_fars)
+        # print("img numeber: ", len(imgs))
+        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
+
+        view_ids = [idx_original % 12] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        if view_ids[0] < 8:
+            meta_end = "_narrow"+ "_refview" + str(view_ids[0])
+        else:
+            meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_6_narrow_and_6_2_stage_blend_mix.py

@@ -0,0 +1,449 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+        if self.split == 'train':
+            self.imgs_per_instance = 12
+        else:
+            self.imgs_per_instance = 16
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow, 'r') as f:
+            narrow_meta = json.load(f)
+
+        pose_json_path_two_stage = "/objaverse-processed/zero12345_img/zero12345_2stage_8_pose.json"
+        with open(pose_json_path_two_stage, 'r') as f:
+            two_stage_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_meta["c2ws"].keys()) + list(two_stage_meta["c2ws"].keys()) # (8 + 8*4) + (8 + 4*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_meta["c2ws"].values()) + list(two_stage_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_meta["intrinsics"] == two_stage_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_meta["near_far"] == two_stage_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+        
+        if self.split == 'train':
+            if idx == 4:
+                idx = 5
+            elif idx == 5:
+                idx = 7
+            elif idx == 10:
+                idx = 13
+            elif idx == 11:
+                idx = 15
+
+        if idx % 16 < 8: # narrow image as target
+            idx = idx % 16 # [0, 7]
+            # target view
+            c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+            
+
+            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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+            mask_h = depth_h > 0
+            # print("valid pixels", np.sum(mask_h))
+            directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+            surface_points = directions * depth_h[..., None]  # [H, W, 3]
+            distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+            depth_h = distance
+
+        else:
+            idx = idx % 16 - 8 # [0, 5]
+            c2w = self.c2ws[idx + 40]
+            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("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{idx}_0.png')
+
+            img = Image.open(img_filename)
+            img = self.transform(img)  # (4, h, w)
+            
+            # print("img_pre", img.shape)
+            if img.shape[0] == 4:
+                img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+            # print("img", img.shape)
+            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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+        if_use_narrow = []
+        if self.split == 'train':
+            for i in range(8):
+                if np.random.random() > 0.5:
+                    if_use_narrow.append(True) # use narrow
+                else:
+                    if_use_narrow.append(False) # 2-stage prediction
+            if_use_narrow[origin_idx % 8] = True if origin_idx < 8 else False
+        else:
+            for i in range(8):
+                if_use_narrow.append( True if origin_idx < 8 else False)
+        src_views = range(8, 8 + 8 * 4)
+        src_views_used = []
+        for vid in src_views:
+            if ((vid - 8) // 4 == 4) or ((vid - 8) // 4 == 6):
+                continue
+            src_views_used.append(vid)
+            cur_view_id = (vid - 8) // 4
+            # choose narrow
+            if if_use_narrow[cur_view_id]:
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{cur_view_id}_{vid%4}_10.png')
+            else: # choose 2-stage
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{(vid - 8) // 4}_{(vid-8) % 4 + 1}.png')
+
+            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)
+
+
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        if view_ids[0] < 8:
+            meta_end = "_narrow"+ "_refview" + str(view_ids[0])
+        else:
+            meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_8_2_stage.py

@@ -0,0 +1,396 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+
+        self.imgs_per_instance = 8
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow, 'r') as f:
+            narrow_meta = json.load(f)
+
+        pose_json_path_two_stage = "/objaverse-processed/zero12345_img/zero12345_2stage_8_pose.json"
+        with open(pose_json_path_two_stage, 'r') as f:
+            two_stage_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_meta["c2ws"].keys()) + list(two_stage_meta["c2ws"].keys()) # (8 + 8*4) + (8 + 8*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_meta["c2ws"].values()) + list(two_stage_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_meta["intrinsics"] == two_stage_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_meta["near_far"] == two_stage_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance * len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        idx = idx % self.imgs_per_instance # [0, 7]
+        # target view
+        c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        
+        src_views = range(8, 8+32)
+        src_views_used = []
+        for vid in src_views:
+            view_dix_to_use = (vid - 8) // 4
+            src_views_used.append(vid)
+            img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{idx}_{(vid-8) % 4 + 1}.png')
+
+            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 =torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+            depth_h = depth_h.fill_(-1.0)
+            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)            
+        
+
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_8_4_gt.py

@@ -0,0 +1,396 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+
+        self.imgs_per_instance = 8
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow, 'r') as f:
+            narrow_meta = json.load(f)
+
+        pose_json_path_two_stage = "/objaverse-processed/zero12345_img/zero12345_2stage_8_pose.json"
+        with open(pose_json_path_two_stage, 'r') as f:
+            two_stage_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_meta["c2ws"].keys()) + list(two_stage_meta["c2ws"].keys()) # (8 + 8*4) + (8 + 8*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_meta["c2ws"].values()) + list(two_stage_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_meta["intrinsics"] == two_stage_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_meta["near_far"] == two_stage_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance * len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        idx = idx % self.imgs_per_instance # [0, 7]
+        # target view
+        c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        
+        src_views = range(8, 8+32)
+        src_views_used = []
+        for vid in src_views:
+            view_dix_to_use = (vid - 8) // 4
+            src_views_used.append(vid)
+            img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10_gt.png')
+
+            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 =torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+            depth_h = depth_h.fill_(-1.0)
+            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)            
+        
+
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_8_narrow_and_8_2_stage_blend_3_views.py

@@ -0,0 +1,446 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+        self.imgs_per_instance = 16
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow, 'r') as f:
+            narrow_meta = json.load(f)
+
+        pose_json_path_two_stage = "/objaverse-processed/zero12345_img/zero12345_2stage_8_pose.json"
+        with open(pose_json_path_two_stage, 'r') as f:
+            two_stage_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_meta["c2ws"].keys()) + list(two_stage_meta["c2ws"].keys()) # (8 + 8*4) + (8 + 4*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_meta["c2ws"].values()) + list(two_stage_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_meta["intrinsics"] == two_stage_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_meta["near_far"] == two_stage_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+        
+        if idx % 16 < 8: # narrow image as target
+            idx = idx % self.imgs_per_instance # [0, 7]
+            # target view
+            c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+            
+
+            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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+            mask_h = depth_h > 0
+            # print("valid pixels", np.sum(mask_h))
+            directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+            surface_points = directions * depth_h[..., None]  # [H, W, 3]
+            distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+            depth_h = distance
+
+        else:
+            idx = idx % self.imgs_per_instance - 8 # [0, 5]
+            c2w = self.c2ws[idx + 40]
+            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("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{idx}_0.png')
+
+
+            img = Image.open(img_filename)
+            img = self.transform(img)  # (4, h, w)
+            
+            # print("img_pre", img.shape)
+            if img.shape[0] == 4:
+                img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+            # print("img", img.shape)
+            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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+        if_use_narrow = []
+        if self.split == 'train':
+            for i in range(8):
+                if np.random.random() > 0.5:
+                    if_use_narrow.append(True) # use narrow
+                else:
+                    if_use_narrow.append(False) # 2-stage prediction
+            if_use_narrow[origin_idx % 8] = True if origin_idx < 8 else False
+        else:
+            for i in range(8):
+                if_use_narrow.append( True if origin_idx < 8 else False)
+
+        src_views = list()
+        for i in range(8):
+            # randomly choose 3 different number from [0,3]
+            local_idxs = np.random.choice(4, 3, replace=False)
+            local_idxs = [0,1,2]
+            local_idxs = [8+i*4+local_idx for local_idx in local_idxs]
+            src_views += local_idxs
+        src_views_used = []
+        for vid in src_views:
+            src_views_used.append(vid)
+            cur_view_id = (vid - 8) // 4
+            # choose narrow
+            if if_use_narrow[cur_view_id]:
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{cur_view_id}_{vid%4}_10.png')
+            else: # choose 2-stage
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{(vid - 8) // 4}_{(vid-8) % 4 + 1}.png')
+
+            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)
+
+           
+        
+
+        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)
+
+        # print(new_near_fars)
+        # print("img numeber: ", len(imgs))
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        if view_ids[0] < 8:
+            meta_end = "_narrow"+ "_refview" + str(view_ids[0])
+        else:
+            meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_8_narrow_and_8_2_stage_blend_mix.py

@@ -0,0 +1,439 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+        self.imgs_per_instance = 16
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow, 'r') as f:
+            narrow_meta = json.load(f)
+
+        pose_json_path_two_stage = "/objaverse-processed/zero12345_img/zero12345_2stage_8_pose.json"
+        with open(pose_json_path_two_stage, 'r') as f:
+            two_stage_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_meta["c2ws"].keys()) + list(two_stage_meta["c2ws"].keys()) # (8 + 8*4) + (8 + 8*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_meta["c2ws"].values()) + list(two_stage_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_meta["intrinsics"] == two_stage_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_meta["near_far"] == two_stage_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+        
+        if idx % 16 < 8: # gt image as target
+            idx = idx % self.imgs_per_instance # [0, 7]
+            # target view
+            c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+            
+
+            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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+            mask_h = depth_h > 0
+            # print("valid pixels", np.sum(mask_h))
+            directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+            surface_points = directions * depth_h[..., None]  # [H, W, 3]
+            distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+            depth_h = distance
+
+        else:
+            idx = idx % self.imgs_per_instance - 8 # [0, 7]
+            c2w = self.c2ws[idx + 40]
+            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("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{idx}_0.png')
+
+
+            img = Image.open(img_filename)
+            img = self.transform(img)  # (4, h, w)
+            
+            # print("img_pre", img.shape)
+            if img.shape[0] == 4:
+                img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+            # print("img", img.shape)
+            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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+        if_use_narrow = []
+        if self.split == 'train':
+            for i in range(8):
+                if np.random.random() > 0.5:
+                    if_use_narrow.append(True) # use narrow
+                else:
+                    if_use_narrow.append(False) # 2-stage prediction
+            if_use_narrow[origin_idx % 8] = True if (origin_idx % 16) < 8 else False
+        else:
+            for i in range(8):
+                if_use_narrow.append( True if (origin_idx % 16) < 8 else False)
+        src_views = range(8, 8 + 8 * 4)
+        src_views_used = []
+        for vid in src_views:
+            src_views_used.append(vid)
+            cur_view_id = (vid - 8) // 4  # [0, 7]
+            # choose narrow
+            if if_use_narrow[cur_view_id]:
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{cur_view_id}_{vid%4}_10.png')
+            else: # choose 2-stage
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{cur_view_id}_{(vid) % 4 + 1}.png')
+
+            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)
+
+           
+        
+
+        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)
+
+        # print(new_near_fars)
+        # print("img numeber: ", len(imgs))
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        if view_ids[0] < 8:
+            meta_end = "_narrow"+ "_refview" + str(view_ids[0])
+        else:
+            meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_8_narrow_and_8_2_stage_mix.py

@@ -0,0 +1,470 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+        self.imgs_per_instance = 16
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow, 'r') as f:
+            narrow_meta = json.load(f)
+
+        pose_json_path_two_stage = "/objaverse-processed/zero12345_img/zero12345_2stage_8_pose.json"
+        with open(pose_json_path_two_stage, 'r') as f:
+            two_stage_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_meta["c2ws"].keys()) + list(two_stage_meta["c2ws"].keys()) # (8 + 8*4) + (8 + 8*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_meta["c2ws"].values()) + list(two_stage_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_meta["intrinsics"] == two_stage_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_meta["near_far"] == two_stage_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance * len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        if idx % self.imgs_per_instance < 8:
+            idx = idx % self.imgs_per_instance # [0, 7]
+            # target view
+            c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+            depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+            
+
+            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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+            mask_h = depth_h > 0
+            # print("valid pixels", np.sum(mask_h))
+            directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+            surface_points = directions * depth_h[..., None]  # [H, W, 3]
+            distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+            depth_h = distance
+
+
+            depths_h.append(depth_h)
+            masks_h.append(mask_h)
+            
+            intrinsic = self.intrinsic
+            intrinsics.append(intrinsic)
+        
+
+            near_fars.append(self.near_fars[idx])
+            image_perm = 0  # only supervised on reference view
+
+            mask_dilated = None
+
+            # src_views = range(8+idx*4, 8+(idx+1)*4)
+
+            src_views = range(8, 8 + 8 * 4)
+            src_views_used = []
+            for vid in src_views:
+                src_views_used.append(vid)
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+                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)
+
+        else:
+            idx = idx % self.imgs_per_instance - 8 # [0, 5]
+       
+            c2w = self.c2ws[idx + 40]
+            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("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{idx}_0.png')
+
+
+            img = Image.open(img_filename)
+
+            img = self.transform(img)  # (4, h, w)
+            
+            # print("img_pre", img.shape)
+            if img.shape[0] == 4:
+                img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+            # print("img", img.shape)
+            imgs += [img]
+
+
+            depth_h = torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+            depth_h = depth_h.fill_(-1.0)
+            # depth_h = torch.fill((img.shape[1], img.shape[2]), -1.0)
+            # print("depth_h", depth_h.shape)
+            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.near_fars[idx])
+            image_perm = 0  # only supervised on reference view
+
+            mask_dilated = None
+
+
+            
+            src_views = range(40+8, 40+8+32)
+            src_views_used = []
+            for vid in src_views:
+                view_dix_to_use = (vid - 40 - 8) // 4
+
+                src_views_used.append(vid)
+                img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{idx}_{(vid-48) % 4 + 1}.png')
+
+                img = Image.open(img_filename)
+                img_wh = self.img_wh
+
+                img = self.transform(img)
+                # print("img shape1: ", img.shape)
+                if img.shape[0] == 4:
+                    img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+                # print("img shape2: ", img.shape)
+                imgs += [img]
+                depth_h =torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+                depth_h = depth_h.fill_(-1.0)
+                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)            
+        
+
+        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)
+
+        # print(new_near_fars)
+        # print("img numeber: ", len(imgs))
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        if view_ids[0] < 8:
+            meta_end = "_narrow"+ "_refview" + str(view_ids[0])
+        else:
+            meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_8_wide_from_2_stage.py

@@ -0,0 +1,395 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+
+        self.imgs_per_instance = 8
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/random32_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path_narrow = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+        with open(pose_json_path_narrow, 'r') as f:
+            narrow_meta = json.load(f)
+
+        pose_json_path_two_stage = "/objaverse-processed/zero12345_img/zero12345_2stage_8_pose.json"
+        with open(pose_json_path_two_stage, 'r') as f:
+            two_stage_meta = json.load(f)
+
+
+        self.img_ids = list(narrow_meta["c2ws"].keys()) + list(two_stage_meta["c2ws"].keys()) # (8 + 8*4) + (8 + 8*4)
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(narrow_meta["c2ws"].values()) + list(two_stage_meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        assert narrow_meta["intrinsics"] == two_stage_meta["intrinsics"], "intrinsics not equal"
+        intrinsic[:3, :3] = np.array(narrow_meta["intrinsics"])
+        self.intrinsic = intrinsic
+        assert narrow_meta["near_far"] == two_stage_meta["near_far"], "near_far not equal"
+        self.near_far = np.array(narrow_meta["near_far"])
+        self.near_far[1] = 1.8
+        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.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.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+
+
+
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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 self.imgs_per_instance * len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+        idx_original=idx
+
+        folder_uid_dict = self.lvis_paths[idx//self.imgs_per_instance]
+        
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        idx = idx % self.imgs_per_instance # [0, 7]
+        # target view
+        c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        
+        src_views = range(0, 8)
+        src_views_used = []
+        for vid in src_views:
+            src_views_used.append(vid)
+            img_filename = os.path.join("/objaverse-processed/zero12345_img/zero12345_2stage_8/", folder_id, uid, f'view_0_{vid}_0.png')
+
+            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 =torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+            depth_h = depth_h.fill_(-1.0)
+            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)            
+        
+
+        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
+
+        view_ids = [idx_original % self.imgs_per_instance] + src_views_used
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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
+        meta_end = "_two_stage"+ "_refview" + str(view_ids[0] - 8)
+        sample['meta'] = str(folder_id) + "_" + str(uid) + "_refview" + meta_end
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_4_1_eval_new_data.py

@@ -0,0 +1,418 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[],
+                 specific_dataset_name = 'GSO'
+                 ):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+        # self.specific_dataset_name = 'Realfusion'
+        # self.specific_dataset_name = 'GSO'
+        # self.specific_dataset_name = 'Objaverse' 
+        # self.specific_dataset_name = 'Zero123'
+
+        self.specific_dataset_name = specific_dataset_name
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        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)
+        self.shape_list.sort()
+
+        # self.shape_list = ['barrel_render']
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        # print("lvis_paths: ", self.lvis_paths)
+
+        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_depth(self, filename):
+        pass
+
+    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 8*len(self.lvis_paths)
+        # return len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+        # idx = idx * 8 # to be deleted
+        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//8]
+        idx = idx % 8 # [0, 7]
+
+        # 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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[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_debug', f'{self.img_ids[idx]}')
+        img_filename = os.path.join(folder_path, 'stage1_8', f'{self.img_ids[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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        # src_views = range(8, 8 + 8 * 4)
+        src_views = range(8+idx*4, 8+(idx+1)*4)
+        for vid in src_views:
+
+            # img_filename = os.path.join(folder_path, 'stage2_8_debug', f'{self.img_ids[vid]}')
+            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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        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

SparseNeuS_demo_v1/data/blender_general_narrow_6.py

@@ -0,0 +1,399 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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):
+        if self.split == 'train':
+            return 6*len(self.lvis_paths)
+        else:
+            return 8*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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
+
+        if self.split == 'train':
+            folder_uid_dict = self.lvis_paths[idx//6]
+            idx = idx % 6 # [0, 5]
+            if idx == 4:
+                idx = 5
+            elif idx == 5:
+                idx = 7
+        else:
+            folder_uid_dict = self.lvis_paths[idx//8]
+            idx = idx % 8 # [0, 7]
+
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            if ((vid - 8) // 4 == 4) or ((vid - 8) // 4 == 6):
+                continue
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+            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)
+        # print("len(imges)", len(imgs))
+        
+        # ! 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)
+
+        # print(new_near_fars)
+        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
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_8_3_fixed.py

@@ -0,0 +1,393 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 8*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = list()
+        for i in range(8):
+            # randomly choose 3 different number from [0,3]
+            # local_idxs = np.random.choice(4, 3, replace=False)
+            local_idxs = [0, 2, 3]
+            # local_idxs = np.random.choice(4, 3, replace=False)
+
+            local_idxs = [8 + i * 4 + local_idx for local_idx in local_idxs]
+            src_views += local_idxs
+        for vid in src_views:
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+            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)
+
+        # print("len(imgs)", len(imgs))
+        # ! 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)
+
+        # print(new_near_fars)
+        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
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_8_3_random.py

@@ -0,0 +1,395 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 8*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = list()
+        for i in range(8):
+
+            if self.split == 'train':
+                local_idxs = np.random.choice(4, 3, replace=False)
+            else:
+                local_idxs = [0, 2, 3]
+            # local_idxs = np.random.choice(4, 3, replace=False)
+
+            local_idxs = [8 + i * 4 + local_idx for local_idx in local_idxs]
+            src_views += local_idxs
+        for vid in src_views:
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+            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)
+
+        # print("len(imgs)", len(imgs))
+        # ! 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)
+
+        # print(new_near_fars)
+        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
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_8_4_random_shading.py

@@ -0,0 +1,432 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 8*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+            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)
+
+        if self.split == 'train':
+            # randomly select one view from eight views as reference view
+            idx_to_select = np.random.randint(0, 8)
+            
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{idx_to_select}.png')
+            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[0] = img
+
+            w2c_selected = self.all_extrinsics[idx_to_select] @ w2c_ref_inv
+            P = self.all_intrinsics[idx_to_select] @ w2c_selected @ scale_mat
+            P = P[:3, :4]
+
+            c2w = load_K_Rt_from_P(None, P)[1]
+            w2c = np.linalg.inv(c2w)
+            affine_mat = np.eye(4)
+            affine_mat[:3, :4] = self.all_intrinsics[idx_to_select][:3, :3] @ w2c[:3, :4]
+            new_affine_mats[0] = affine_mat
+            camera_o = c2w[:3, 3]
+            dist = np.sqrt(np.sum(camera_o ** 2))
+            near = dist - 1
+            far = dist + 1
+            new_near_fars[0] = [0.95 * near, 1.05 * far]
+
+            new_w2cs[0] = w2c
+            new_c2ws[0] = c2w
+
+            depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx_to_select}_depth_mm.png'))
+            depth_h = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+            mask_h = depth_h > 0
+            directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+            surface_points = directions * depth_h[..., None]  # [H, W, 3]
+            distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+            depth_h = distance * scale_factor
+
+            new_depths_h[0] = depth_h 
+            masks_h[0] = mask_h 
+
+
+
+        # print(new_near_fars)
+        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
+
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_all.py

@@ -0,0 +1,386 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 8*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+            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)
+
+        # print(new_near_fars)
+        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
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_2_stage.py

@@ -0,0 +1,410 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 8*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        img = Image.open(img_filename)
+
+        img = self.transform(img)  # (4, h, w)
+        
+        # print("img_pre", img.shape)
+        if img.shape[0] == 4:
+            img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+        # print("img", img.shape)
+        imgs += [img]
+
+
+        depth_h = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+        # print("depth_h", depth_h.shape)
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_0_{(vid - 8) // 4}_{vid % 4 + 1}.png')
+
+            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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_2_stage_temp.py

@@ -0,0 +1,411 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 10
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        idx = idx * 8
+        sample = {}
+        origin_idx = idx
+        imgs, depths_h, masks_h = [], [], []  # full size (256, 256)
+        intrinsics, w2cs, c2ws, near_fars = [], [], [], []  # record proj mats between views
+
+
+        folder_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join("/objaverse-processed/zero12345_img/zero12345_narrow/", folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        img = Image.open(img_filename)
+
+        img = self.transform(img)  # (4, h, w)
+        
+        # print("img_pre", img.shape)
+        if img.shape[0] == 4:
+            img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+        # print("img", img.shape)
+        imgs += [img]
+
+
+        depth_h = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+        # print("depth_h", depth_h.shape)
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_0_{(vid - 8) // 4}_{vid % 4 + 1}.png')
+
+            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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data.py

@@ -0,0 +1,418 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[],
+                 specific_dataset_name = 'GSO'
+                 ):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+        # self.specific_dataset_name = 'Realfusion'
+        # self.specific_dataset_name = 'GSO'
+        # self.specific_dataset_name = 'Objaverse' 
+        # self.specific_dataset_name = 'Zero123'
+
+        self.specific_dataset_name = specific_dataset_name
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        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.shape_list = ['barrel_render']
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        # print("lvis_paths: ", self.lvis_paths)
+
+        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_depth(self, filename):
+        pass
+
+    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 8*len(self.lvis_paths)
+        return len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+        idx = idx * 8 # to be deleted
+        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//8]
+        idx = idx % 8 # [0, 7]
+
+        # 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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[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_debug', f'{self.img_ids[idx]}')
+        img_filename = os.path.join(folder_path, 'stage1_8', f'{self.img_ids[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.near_fars[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_debug', f'{self.img_ids[vid]}')
+            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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data3_1.py

@@ -0,0 +1,414 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+        # self.specific_dataset_name = 'Realfusion'
+        self.specific_dataset_name = 'Objaverse'
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        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)
+        self.shape_list.sort()
+
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        # print("lvis_paths: ", self.lvis_paths)
+
+        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_depth(self, filename):
+        pass
+
+    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 8*len(self.lvis_paths)
+        return len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+        idx = idx * 8 # to be deleted
+        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//8]
+        idx = idx % 8 # [0, 7]
+
+        # 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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[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[idx]}')
+        # print(self.img_ids)
+        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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            if vid % 4 == 0:
+                vid = (vid - 8) // 4
+                img_filename = os.path.join(folder_path, 'stage1_8', f'{self.img_ids[vid]}')
+            else:
+                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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data_32_wide.py

@@ -0,0 +1,465 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+
+def calc_pose(phis, thetas, size, radius = 1.2):
+    import torch
+    def normalize(vectors):
+        return vectors / (torch.norm(vectors, dim=-1, keepdim=True) + 1e-10)
+    # device = torch.device('cuda')
+    thetas = torch.FloatTensor(thetas)
+    phis = torch.FloatTensor(phis)
+    
+    centers = torch.stack([
+        radius * torch.sin(thetas) * torch.sin(phis),
+        -radius * torch.cos(thetas) * torch.sin(phis),
+        radius * torch.cos(phis),
+    ], dim=-1) # [B, 3]
+
+    # lookat
+    forward_vector = normalize(centers).squeeze(0)
+    up_vector = torch.FloatTensor([0, 0, 1]).unsqueeze(0).repeat(size, 1) 
+    right_vector = normalize(torch.cross(up_vector, forward_vector, dim=-1))      
+    if right_vector.pow(2).sum() < 0.01:
+        right_vector = torch.FloatTensor([0, 1, 0]).unsqueeze(0).repeat(size, 1)  
+    up_vector = normalize(torch.cross(forward_vector, right_vector, dim=-1))     
+
+    poses = torch.eye(4, dtype=torch.float)[:3].unsqueeze(0).repeat(size, 1, 1)
+    poses[:, :3, :3] = torch.stack((right_vector, up_vector, forward_vector), dim=-1)
+    poses[:, :3, 3] = centers 
+    return poses
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[],
+                 specific_dataset_name = 'GSO'
+                 ):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+        # self.specific_dataset_name = 'Realfusion'
+        # self.specific_dataset_name = 'GSO'
+        # self.specific_dataset_name = 'Objaverse' 
+        # self.specific_dataset_name = 'Zero123'
+
+        self.specific_dataset_name = specific_dataset_name
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        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.shape_list = os.listdir(main_folder)
+        self.shape_list.sort()
+
+        # self.shape_list = ['barrel_render']
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        # print("lvis_paths: ", self.lvis_paths)
+
+        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!'
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+
+        with open(pose_json_path, 'r') as f:
+            meta = json.load(f)
+        intrinsic = np.eye(4)
+        intrinsic[:3, :3] = np.array(meta["intrinsics"])
+        self.intrinsic = intrinsic
+        self.near_far = np.array(meta["near_far"])
+        self.near_far[1] = 1.8
+
+        # * 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 in range(self.input_poses.shape[0]):
+            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_depth(self, filename):
+        pass
+
+    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 8*len(self.lvis_paths)
+        return len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+        idx = idx * 8 # to be deleted
+        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//8]
+        idx = idx % 8 # [0, 7]
+
+        # 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.near_far[1] = 1.8
+        # self.define_transforms()
+        # self.blender2opencv = np.array(
+        #     [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
+        # )
+
+        pose_file = os.path.join(folder_path, '32_random', 'views.npz')
+        pose_array = np.load(pose_file)
+        pose = calc_pose(pose_array['elevations'], pose_array['azimuths'], 32) # [32, 3, 4] c2ws
+  
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(pose)
+        self.input_poses = np.concatenate([self.input_poses, np.tile(np.array([0, 0, 0, 1], dtype=np.float32)[None, None, :], [self.input_poses.shape[0], 1, 1])], axis=1)
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix in range(pose.shape[0]):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[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_debug', f'{self.img_ids[idx]}')
+        img_filename = os.path.join(folder_path, '32_random', f'{idx}.png')
+
+        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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        src_views = range(0, 8 * 4)
+
+        for vid in src_views:
+
+            # img_filename = os.path.join(folder_path, 'stage2_8_debug', f'{self.img_ids[vid]}')
+            img_filename = os.path.join(folder_path, '32_random', f'{vid}.png')
+            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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data_4_4.py

@@ -0,0 +1,419 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[],
+                 specific_dataset_name = 'GSO'
+                 ):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+        # self.specific_dataset_name = 'Realfusion'
+        # self.specific_dataset_name = 'GSO'
+        # self.specific_dataset_name = 'Objaverse' 
+        # self.specific_dataset_name = 'Zero123'
+
+        self.specific_dataset_name = specific_dataset_name
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        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)
+        self.shape_list.sort()
+
+        # self.shape_list = ['barrel_render']
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        # print("lvis_paths: ", self.lvis_paths)
+
+        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_depth(self, filename):
+        pass
+
+    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 8*len(self.lvis_paths)
+        return len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+        idx = idx * 8 # to be deleted
+        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//8]
+        idx = idx % 8 # [0, 7]
+
+        # 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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[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_debug', f'{self.img_ids[idx]}')
+        img_filename = os.path.join(folder_path, 'stage1_8', f'{self.img_ids[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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            if (vid // 4) % 2 != 0:
+                continue
+            # img_filename = os.path.join(folder_path, 'stage2_8_debug', f'{self.img_ids[vid]}')
+            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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data_6_4.py

@@ -0,0 +1,420 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[],
+                 specific_dataset_name = 'GSO'
+                 ):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+        # self.specific_dataset_name = 'Realfusion'
+        # self.specific_dataset_name = 'GSO'
+        # self.specific_dataset_name = 'Objaverse' 
+        # self.specific_dataset_name = 'Zero123'
+
+        self.specific_dataset_name = specific_dataset_name
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        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)
+        self.shape_list.sort()
+
+        # self.shape_list = ['barrel_render']
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        # print("lvis_paths: ", self.lvis_paths)
+
+        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_depth(self, filename):
+        pass
+
+    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 8*len(self.lvis_paths)
+        return len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+        idx = idx * 8 # to be deleted
+        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//8]
+        idx = idx % 8 # [0, 7]
+
+        # 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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[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_debug', f'{self.img_ids[idx]}')
+        img_filename = os.path.join(folder_path, 'stage1_8', f'{self.img_ids[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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            if ((vid - 8) // 4 == 4) or ((vid - 8) // 4 == 6):
+                continue
+
+            # img_filename = os.path.join(folder_path, 'stage2_8_debug', f'{self.img_ids[vid]}')
+            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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data_8_3.py

@@ -0,0 +1,428 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[],
+                 specific_dataset_name = 'GSO'
+                 ):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+        # self.specific_dataset_name = 'Realfusion'
+        # self.specific_dataset_name = 'GSO'
+        # self.specific_dataset_name = 'Objaverse' 
+        # self.specific_dataset_name = 'Zero123'
+
+        self.specific_dataset_name = specific_dataset_name
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        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)
+        self.shape_list.sort()
+
+        # self.shape_list = ['barrel_render']
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        # print("lvis_paths: ", self.lvis_paths)
+
+        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_depth(self, filename):
+        pass
+
+    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 8*len(self.lvis_paths)
+        return len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+        idx = idx * 8 # to be deleted
+        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//8]
+        idx = idx % 8 # [0, 7]
+
+        # 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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[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_debug', f'{self.img_ids[idx]}')
+        img_filename = os.path.join(folder_path, 'stage1_8', f'{self.img_ids[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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        # src_views = range(8, 8 + 8 * 4)
+
+        src_views = list()
+        for i in range(8):
+            # randomly choose 3 different number from [0,3]
+            # local_idxs = np.random.choice(4, 3, replace=False)
+            local_idxs = [0, 2, 3]
+            # local_idxs = np.random.choice(4, 3, replace=False)
+
+            local_idxs = [8 + i * 4 + local_idx for local_idx in local_idxs]
+            src_views += local_idxs
+
+        for vid in src_views:
+
+            # img_filename = os.path.join(folder_path, 'stage2_8_debug', f'{self.img_ids[vid]}')
+            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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data_8_wide.py

@@ -0,0 +1,420 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[],
+                 specific_dataset_name = 'GSO'
+                 ):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+        # self.specific_dataset_name = 'Realfusion'
+        # self.specific_dataset_name = 'GSO'
+        # self.specific_dataset_name = 'Objaverse' 
+        # self.specific_dataset_name = 'Zero123'
+
+        self.specific_dataset_name = specific_dataset_name
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        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)
+        self.shape_list.sort()
+
+        # self.shape_list = ['barrel_render']
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        # print("lvis_paths: ", self.lvis_paths)
+
+        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_depth(self, filename):
+        pass
+
+    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 8*len(self.lvis_paths)
+        return len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+        idx = idx * 8 # to be deleted
+        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//8]
+        idx = idx % 8 # [0, 7]
+
+        # 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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[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_debug', f'{self.img_ids[idx]}')
+        img_filename = os.path.join(folder_path, 'stage1_8', f'{self.img_ids[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.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        src_views = range(8)
+
+
+        for vid in src_views:
+
+            # img_filename = os.path.join(folder_path, 'stage2_8_debug', f'{self.img_ids[vid]}')
+            # img_filename = os.path.join(folder_path, 'stage2_8', f'{self.img_ids[vid]}')
+            img_filename = os.path.join(folder_path, 'stage1_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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        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

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data_temp.py

@@ -0,0 +1,417 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[],
+                 specific_dataset_name = 'GSO'
+                 ):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+        # self.specific_dataset_name = 'Realfusion'
+        # self.specific_dataset_name = 'GSO'
+        # self.specific_dataset_name = 'Objaverse' 
+        self.specific_dataset_name = 'Objaverse_archived'
+
+        # self.specific_dataset_name = specific_dataset_name
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        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)
+        self.shape_list.sort()
+
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        # print("lvis_paths: ", self.lvis_paths)
+
+        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_depth(self, filename):
+        pass
+
+    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 8*len(self.lvis_paths)
+        return len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        sample = {}
+        idx = idx * 8 # to be deleted
+        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//8]
+        idx = idx % 8 # [0, 7]
+
+        # last subdir name
+        shape_name = os.path.split(folder_path)[-1]
+
+        pose_json_path = os.path.join('/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[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_debug', f'{self.img_ids[idx]}')
+        img_filename = os.path.join(folder_path, 'stage1_8', f'{idx}.png')
+
+        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.near_fars[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_debug', f'{self.img_ids[vid]}')
+            img_filename = os.path.join(folder_path, 'stage2_8', f'{(vid-8)//4}_{(vid-8)%4}.png')
+            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)
+
+        # print(new_near_fars)
+        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['light_idx'] = torch.tensor(light_idx)
+        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

SparseNeuS_demo_v1/data/blender_general_narrow_all_no_depth.py

@@ -0,0 +1,388 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 8*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        # directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        # surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        # distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        # depth_h = distance
+
+        depth_h = torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+        depth_h = depth_h.fill_(-1.0)
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+            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)
+
+        # print(new_near_fars)
+        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
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_all_only_4.py

@@ -0,0 +1,389 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 4*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        idx = idx * 2
+        sample = {}
+        origin_idx = idx
+        imgs, depths_h, masks_h = [], [], []  # full size (256, 256)
+        intrinsics, w2cs, c2ws, near_fars = [], [], [], []  # record proj mats between views
+
+
+        folder_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+            if (vid // 4) % 2 != 0:
+                continue
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+            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)
+
+        # print("len(imgs)", len(imgs))
+        # ! 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)
+
+        # print(new_near_fars)
+        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
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_general_narrow_all_only_4_and_4.py

@@ -0,0 +1,395 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/lvis_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+
+
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_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.near_far[1] = 1.8
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+        # * 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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 8*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    def __getitem__(self, idx):
+        idx = 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_uid_dict = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+
+        # target view
+        c2w = self.c2ws[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(self.root_dir, folder_id, uid, f'view_{idx}.png')
+
+        depth_filename = os.path.join(os.path.join(self.root_dir, folder_id, uid, f'view_{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+        # print("valid pixels", np.sum(mask_h))
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+
+        src_views = range(8, 8 + 8 * 4)
+        
+        vid_list = []
+        for vid in src_views:
+            if (vid // 4) % 2 != idx % 2:
+                continue
+            vid_list.append(vid)
+            img_filename = os.path.join(self.root_dir, folder_id, uid, f'view_{(vid - 8) // 4}_{vid%4}_10.png')
+
+            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)
+
+        
+        # print("idx:", idx)
+        # print("len(imgs)", len(imgs))
+        # print("vid_list", vid_list)
+        # ! 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)
+
+        # print(new_near_fars)
+        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
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/blender_gt_32.py

@@ -0,0 +1,419 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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
+import json
+from termcolor import colored
+import imageio
+from kornia import create_meshgrid
+import open3d as o3d
+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
+
+import os, json
+import numpy as np
+def calc_pose(phis, thetas, size, radius = 1.2):
+    import torch
+    def normalize(vectors):
+        return vectors / (torch.norm(vectors, dim=-1, keepdim=True) + 1e-10)
+    # device = torch.device('cuda')
+    thetas = torch.FloatTensor(thetas)
+    phis = torch.FloatTensor(phis)
+    
+    centers = torch.stack([
+        radius * torch.sin(thetas) * torch.sin(phis),
+        -radius * torch.cos(thetas) * torch.sin(phis),
+        radius * torch.cos(phis),
+    ], dim=-1) # [B, 3]
+
+    # lookat
+    forward_vector = normalize(centers).squeeze(0)
+    up_vector = torch.FloatTensor([0, 0, 1]).unsqueeze(0).repeat(size, 1) 
+    right_vector = normalize(torch.cross(up_vector, forward_vector, dim=-1))      
+    if right_vector.pow(2).sum() < 0.01:
+        right_vector = torch.FloatTensor([0, 1, 0]).unsqueeze(0).repeat(size, 1)  
+    up_vector = normalize(torch.cross(forward_vector, right_vector, dim=-1))     
+
+    poses = torch.eye(4, dtype=torch.float)[:3].unsqueeze(0).repeat(size, 1, 1)
+    poses[:, :3, :3] = torch.stack((right_vector, up_vector, forward_vector), dim=-1)
+    poses[:, :3, 3] = centers 
+    return poses
+
+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()  # ? why need transpose here
+    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, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        lvis_json_path = '/objaverse-processed/zero12345_img/random32_split.json' # folder_id and uid
+        with open(lvis_json_path, 'r') as f:
+            lvis_paths = json.load(f)
+            if self.split == 'train':
+                self.lvis_paths = lvis_paths['train']
+            else:
+                self.lvis_paths = lvis_paths['val']
+        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!'
+        
+        pose_json_path = "/objaverse-processed/zero12345_img/zero12345_narrow_pose.json"
+
+        with open(pose_json_path, 'r') as f:
+            meta = json.load(f)
+        intrinsic = np.eye(4)
+        intrinsic[:3, :3] = np.array(meta["intrinsics"])
+        self.intrinsic = intrinsic
+        self.near_far = np.array(meta["near_far"])
+        self.near_far[1] = 1.8
+
+        # * 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 in range(self.input_poses.shape[0]):
+            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_depth(self, filename):
+        pass
+
+    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)
+        # print("center", center)
+        # print("radius", radius)
+        # print("bounds", bounds)
+        # import ipdb; ipdb.set_trace()
+        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 32*len(self.lvis_paths)
+
+
+    def read_depth(self, filename, near_bound, noisy_factor=1.0):
+        pass
+
+
+    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_uid_dict = self.lvis_paths[idx//32]
+        idx = idx % 32 # [0, 7]
+        folder_id = folder_uid_dict['folder_id']
+        uid = folder_uid_dict['uid']
+
+        pose_file = os.path.join('/objaverse-processed/zero12345_img/random32/', folder_id, uid, 'views.npz')
+        pose_array = np.load(pose_file)
+        pose = calc_pose(pose_array['elevations'], pose_array['azimuths'], 32) # [32, 3, 4] c2ws
+  
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(pose)
+        self.input_poses = np.concatenate([self.input_poses, np.tile(np.array([0, 0, 0, 1], dtype=np.float32)[None, None, :], [self.input_poses.shape[0], 1, 1])], axis=1)
+        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.near_fars = []
+        # self.root_dir = root_dir
+        for image_dix in range(pose.shape[0]):
+            pose = self.input_poses[image_dix]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+
+        # target view
+        c2w = self.c2ws[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('/objaverse-processed/zero12345_img/random32/', folder_id, uid, f'{idx}.png')
+
+        depth_filename = os.path.join(os.path.join('/objaverse-processed/zero12345_img/random32/', folder_id, uid, f'{idx}_depth_mm.png'))
+        
+
+        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 = cv2.imread(depth_filename, cv2.IMREAD_UNCHANGED).astype(np.uint16) / 1000.0
+        mask_h = depth_h > 0
+
+        directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsic[0, 0], self.intrinsic[1, 1]])  # [H, W, 3]
+        surface_points = directions * depth_h[..., None]  # [H, W, 3]
+        distance = np.linalg.norm(surface_points, axis=-1)  # [H, W]
+        depth_h = distance
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+        # src_views = range(8+idx*4, 8+(idx+1)*4)
+        src_views = range(0, 8 * 4)
+
+        for vid in src_views:
+            img_filename = os.path.join('/objaverse-processed/zero12345_img/random32/', folder_id, uid, f'{vid}.png')
+
+            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)
+
+        # print(new_near_fars)
+        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
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = folder_id
+
+        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(folder_id) + "_" + str(uid) + "_refview" + str(view_ids[0])
+
+
+        # - 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

SparseNeuS_demo_v1/data/dtu/dtu_pairs.txt

@@ -0,0 +1,93 @@
+46
+0
+10 10 2346.410000 1 2036.530000 9 1243.890000 12 1052.870000 11 1000.840000 13 703.583000 2 604.456000 8 439.759000 14 327.419000 27 249.278000 
+1
+10 9 2850.870000 10 2583.940000 2 2105.590000 0 2052.840000 8 1868.240000 13 1184.230000 14 1017.510000 12 961.966000 7 670.208000 15 657.218000 
+2
+10 8 2501.240000 1 2106.880000 7 1856.500000 9 1782.340000 3 1141.770000 15 1061.760000 14 815.457000 16 762.153000 6 709.789000 10 699.921000 
+3
+10 7 1294.390000 6 1159.130000 2 1134.270000 4 905.717000 8 687.320000 5 600.015000 17 496.958000 16 481.969000 1 379.011000 15 307.450000 
+4
+10 5 1333.740000 6 1145.150000 3 895.254000 7 486.504000 18 446.420000 2 418.517000 17 326.528000 8 161.115000 16 149.154000 1 103.626000 
+5
+10 6 1676.060000 18 1555.060000 4 1335.550000 17 868.416000 3 593.755000 7 467.816000 20 440.579000 19 428.255000 16 242.327000 21 210.253000 
+6
+10 17 2332.350000 7 1848.240000 18 1812.740000 5 1696.070000 16 1273.000000 3 1157.990000 4 1155.410000 20 771.624000 21 744.945000 2 700.368000 
+7
+10 16 2709.460000 8 2439.700000 15 2078.210000 6 1864.160000 2 1846.600000 17 1791.710000 3 1296.860000 22 957.793000 9 879.088000 21 782.277000 
+8
+10 15 3124.010000 9 3099.920000 14 2756.290000 2 2501.220000 7 2449.320000 1 1875.940000 16 1726.040000 13 1325.760000 23 1177.090000 24 1108.820000 
+9
+10 13 3355.620000 14 3226.070000 8 3098.800000 10 3097.070000 1 2861.420000 12 1873.630000 2 1785.980000 15 1753.320000 25 1365.450000 0 1261.590000 
+10
+10 12 3750.700000 9 3085.870000 13 3028.390000 1 2590.550000 0 2369.790000 11 2266.670000 14 1524.160000 26 1448.150000 27 1293.600000 8 1041.840000 
+11
+10 12 3543.760000 27 3056.050000 10 2248.070000 26 1524.280000 28 1273.330000 13 1265.900000 29 1129.550000 0 998.164000 9 591.176000 30 572.919000 
+12
+10 27 3889.870000 10 3754.540000 13 3745.210000 11 3584.260000 26 3574.560000 25 1877.110000 9 1866.340000 29 1482.720000 30 1418.510000 14 1341.860000 
+13
+10 12 3773.140000 26 3699.280000 25 3657.170000 14 3652.040000 9 3356.290000 10 3049.270000 24 2098.910000 27 1900.960000 31 1460.960000 30 1349.620000 
+14
+10 13 3663.520000 24 3610.690000 9 3232.550000 25 3216.400000 15 3128.840000 8 2758.040000 23 2219.910000 26 1567.450000 10 1536.600000 32 1419.330000 
+15
+10 23 3194.920000 14 3126.000000 8 3120.430000 16 2897.020000 24 2562.490000 7 2084.050000 22 2041.630000 9 1752.080000 33 1232.290000 13 1137.550000 
+16
+10 15 2884.140000 7 2713.880000 22 2708.570000 17 2448.500000 21 2173.300000 23 1908.030000 8 1718.790000 6 1281.960000 35 1047.380000 34 980.064000 
+17
+10 21 2632.480000 16 2428.000000 6 2343.570000 18 2250.230000 20 2149.750000 7 1779.420000 22 1380.250000 36 957.046000 5 878.398000 15 789.068000 
+18
+9 17 2219.150000 20 2173.020000 6 1802.390000 19 1575.770000 5 1564.810000 21 1160.130000 16 660.317000 7 589.484000 36 559.983000 
+19
+7 20 1828.970000 18 1564.630000 17 685.249000 36 613.420000 21 572.770000 5 427.597000 6 368.651000 
+20
+8 21 2569.790000 36 2258.330000 18 2186.710000 17 2130.670000 19 1865.060000 35 996.122000 16 799.808000 40 778.721000 
+21
+9 36 2704.590000 35 2639.690000 17 2638.190000 20 2605.430000 22 2604.260000 16 2158.250000 34 1239.250000 18 1178.240000 40 1128.570000 
+22
+10 23 3232.680000 34 3175.150000 35 2831.090000 16 2712.510000 21 2632.190000 15 2033.390000 33 1712.670000 17 1393.860000 36 1290.960000 24 1195.330000 
+23
+10 24 3710.900000 33 3603.070000 22 3244.200000 15 3190.620000 34 3086.490000 14 2220.110000 32 2100.000000 16 1917.100000 35 1359.790000 25 1356.710000 
+24
+10 25 3844.600000 32 3750.750000 23 3710.600000 14 3609.090000 33 3091.040000 15 2559.240000 31 2423.710000 13 2109.360000 26 1440.580000 34 1410.030000 
+25
+10 26 3951.740000 31 3888.570000 24 3833.070000 13 3667.350000 14 3208.210000 32 2993.460000 30 2681.520000 12 1900.230000 45 1484.030000 27 1462.880000 
+26
+10 30 4033.350000 27 3970.470000 25 3925.250000 13 3686.340000 12 3595.590000 29 2943.870000 31 2917.000000 14 1556.340000 11 1554.750000 46 1503.840000 
+27
+10 29 4027.840000 26 3929.940000 12 3875.580000 11 3085.030000 28 2908.600000 30 2792.670000 13 1878.420000 25 1438.550000 47 1425.200000 10 1290.250000 
+28
+10 29 3687.020000 48 3209.130000 27 2872.860000 47 2014.530000 30 1361.950000 11 1273.600000 26 1062.850000 12 840.841000 46 672.985000 31 271.952000 
+29
+10 27 4029.430000 30 3909.550000 28 3739.930000 47 3695.230000 48 3135.870000 26 2910.970000 46 2229.550000 12 1479.160000 31 1430.260000 11 1144.560000 
+30
+10 26 4029.860000 29 3953.720000 31 3811.120000 46 3630.460000 47 3105.960000 27 2824.430000 25 2657.890000 45 2347.750000 32 1459.110000 12 1429.620000 
+31
+10 25 3882.210000 30 3841.880000 32 3808.500000 45 3649.820000 46 3000.670000 26 2939.940000 24 2409.930000 44 2381.300000 13 1467.590000 29 1459.560000 
+32
+10 31 3826.500000 24 3744.140000 33 3613.240000 44 3552.040000 25 3004.600000 45 2884.590000 43 2393.340000 23 2095.270000 30 1478.600000 14 1420.780000 
+33
+10 32 3618.110000 23 3598.100000 34 3530.530000 43 3462.370000 24 3091.530000 44 2608.080000 42 2426.000000 22 1717.940000 31 1407.650000 25 1324.780000 
+34
+10 33 3523.370000 42 3356.550000 35 3210.340000 22 3178.850000 23 3079.030000 43 2396.450000 41 2386.860000 24 1408.020000 32 1301.340000 21 1256.450000 
+35
+10 34 3187.880000 41 3106.440000 36 2866.040000 22 2817.740000 21 2654.870000 40 2416.980000 42 2137.810000 23 1346.860000 33 1150.330000 16 1044.660000 
+36
+8 40 2910.700000 35 2832.660000 21 2689.960000 20 2280.460000 41 1787.970000 22 1268.490000 34 981.636000 17 954.229000 
+40
+7 36 2918.140000 41 2852.620000 35 2392.960000 21 1124.300000 42 1056.480000 34 877.946000 20 788.701000 
+41
+9 35 3111.050000 42 3049.710000 40 2885.360000 34 2371.020000 36 1813.690000 43 1164.710000 22 1126.900000 21 906.536000 33 903.238000 
+42
+10 34 3356.980000 43 3183.000000 41 3070.540000 33 2421.770000 35 2155.080000 44 1278.410000 23 1183.520000 22 1147.070000 40 1077.080000 32 899.646000 
+43
+10 33 3461.240000 44 3380.740000 42 3188.700000 34 2400.600000 32 2399.090000 45 1359.370000 23 1314.080000 41 1176.120000 24 1159.620000 31 901.556000 
+44
+10 32 3550.810000 45 3510.160000 43 3373.110000 33 2602.330000 31 2395.930000 24 1410.430000 46 1386.310000 42 1279.000000 25 1095.240000 34 968.440000 
+45
+10 31 3650.090000 46 3555.090000 44 3491.150000 32 2868.390000 30 2373.590000 25 1485.370000 47 1405.280000 43 1349.540000 33 1104.770000 26 1046.810000 
+46
+10 30 3635.640000 47 3562.170000 45 3524.170000 31 2976.820000 29 2264.040000 26 1508.870000 44 1367.410000 48 1352.100000 32 1211.240000 25 1102.170000 
+47
+10 29 3705.310000 46 3519.760000 48 3450.480000 30 3074.770000 28 2054.630000 27 1434.570000 45 1377.340000 31 1268.230000 26 1223.830000 25 471.111000 
+48
+10 47 3401.950000 28 3224.840000 29 3101.160000 46 1317.100000 30 1306.700000 27 1235.070000 26 537.731000 31 291.919000 45 276.869000 11 258.856000 

SparseNeuS_demo_v1/data/dtu/lists/test.txt

@@ -0,0 +1,15 @@
+scan24
+scan37
+scan40
+scan55
+scan63
+scan65
+scan69
+scan83
+scan97
+scan105
+scan106
+scan110
+scan114
+scan118
+scan122

SparseNeuS_demo_v1/data/dtu/lists/train.txt

@@ -0,0 +1,75 @@
+scan1
+scan4
+scan5
+scan6
+scan8
+scan9
+scan10
+scan11
+scan12
+scan13
+scan14
+scan15
+scan16
+scan17
+scan18
+scan19
+scan20
+scan21
+scan22
+scan23
+scan28
+scan29
+scan30
+scan31
+scan32
+scan33
+scan34
+scan35
+scan36
+scan38
+scan39
+scan41
+scan42
+scan43
+scan44
+scan45
+scan46
+scan47
+scan48
+scan49
+scan50
+scan51
+scan52
+scan59
+scan60
+scan61
+scan62
+scan64
+scan74
+scan75
+scan76
+scan77
+scan84
+scan85
+scan86
+scan87
+scan88
+scan89
+scan90
+scan91
+scan92
+scan93
+scan94
+scan95
+scan96
+scan98
+scan99
+scan100
+scan101
+scan102
+scan103
+scan104
+scan126
+scan127
+scan128

SparseNeuS_demo_v1/data/dtu_fit.py

@@ -0,0 +1,278 @@
+import torch
+import torch.nn as nn
+import cv2 as cv
+import numpy as np
+import re
+import os
+import logging
+from glob import glob
+
+from models.rays import gen_rays_from_single_image, gen_random_rays_from_single_image
+
+from data.scene import get_boundingbox
+
+
+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 = cv.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()  # ? why need transpose here
+    pose[:3, 3] = (t[:3] / t[3])[:, 0]
+
+    return intrinsics, pose  # ! return cam2world matrix here
+
+
+class DtuFit:
+    def __init__(self, root_dir, split, scan_id, n_views, train_img_idx=[], test_img_idx=[],
+                 img_wh=[800, 600], clip_wh=[0, 0], original_img_wh=[1600, 1200],
+                 N_rays=512, h_patch_size=5, near=425, far=900):
+        super(DtuFit, self).__init__()
+        logging.info('Load data: Begin')
+
+        self.root_dir = root_dir
+        self.split = split
+        self.scan_id = scan_id
+        self.n_views = n_views
+
+        self.near = near
+        self.far = far
+
+        if self.scan_id is not None:
+            self.data_dir = os.path.join(self.root_dir, self.scan_id)
+        else:
+            self.data_dir = self.root_dir
+
+        self.img_wh = img_wh
+        self.clip_wh = clip_wh
+
+        if len(self.clip_wh) == 2:
+            self.clip_wh = self.clip_wh + self.clip_wh
+
+        self.original_img_wh = original_img_wh
+        self.N_rays = N_rays
+        self.h_patch_size = h_patch_size  # used to extract patch for supervision
+        self.train_img_idx = train_img_idx
+        self.test_img_idx = test_img_idx
+
+        camera_dict = np.load(os.path.join(self.data_dir, 'cameras.npz'), allow_pickle=True)
+        self.images_list = sorted(glob(os.path.join(self.data_dir, "image/*.png")))
+        # world_mat: projection matrix: world to image
+        self.world_mats_np = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in
+                              range(len(self.images_list))]
+
+        self.raw_near_fars = np.stack([np.array([self.near, self.far]) for i in range(len(self.images_list))])
+
+        # - reference image; transform the world system to the ref-camera system
+        self.ref_img_idx = self.train_img_idx[0]
+        ref_world_mat = self.world_mats_np[self.ref_img_idx]
+        self.ref_w2c = np.linalg.inv(load_K_Rt_from_P(None, ref_world_mat[:3, :4])[1])
+
+        self.all_images = []
+        self.all_intrinsics = []
+        self.all_w2cs = []
+
+        self.load_scene()  # load the scene
+
+        # ! estimate scale_mat
+        self.scale_mat, self.scale_factor = self.cal_scale_mat(
+            img_hw=[self.img_wh[1], self.img_wh[0]],
+            intrinsics=self.all_intrinsics[self.train_img_idx],
+            extrinsics=self.all_w2cs[self.train_img_idx],
+            near_fars=self.raw_near_fars[self.train_img_idx],
+            factor=1.1)
+
+        # * after scaling and translation, unit bounding box
+        self.scaled_intrinsics, self.scaled_w2cs, self.scaled_c2ws, \
+        self.scaled_affine_mats, self.scaled_near_fars = self.scale_cam_info()
+        # import ipdb; ipdb.set_trace()
+        self.bbox_min = np.array([-1.0, -1.0, -1.0])
+        self.bbox_max = np.array([1.0, 1.0, 1.0])
+        self.partial_vol_origin = torch.Tensor([-1., -1., -1.])
+
+        logging.info('Load data: End')
+
+    def load_scene(self):
+
+        scale_x = self.img_wh[0] / self.original_img_wh[0]
+        scale_y = self.img_wh[1] / self.original_img_wh[1]
+
+        for idx in range(len(self.images_list)):
+            image = cv.imread(self.images_list[idx])
+            image = cv.resize(image, (self.img_wh[0], self.img_wh[1])) / 255.
+
+            image = image[self.clip_wh[1]:self.img_wh[1] - self.clip_wh[3],
+                    self.clip_wh[0]:self.img_wh[0] - self.clip_wh[2]]
+            self.all_images.append(np.transpose(image[:, :, ::-1], (2, 0, 1))) # append [3,]
+
+            P = self.world_mats_np[idx]
+            P = P[:3, :4]
+            intrinsics, c2w = load_K_Rt_from_P(None, P)
+            w2c = np.linalg.inv(c2w)
+
+            intrinsics[:1] *= scale_x
+            intrinsics[1:2] *= scale_y
+
+            intrinsics[0, 2] -= self.clip_wh[0]
+            intrinsics[1, 2] -= self.clip_wh[1]
+
+            self.all_intrinsics.append(intrinsics)
+            # - transform from world system to ref-camera system
+            self.all_w2cs.append(w2c @ np.linalg.inv(self.ref_w2c))
+
+
+        self.all_images = torch.from_numpy(np.stack(self.all_images)).to(torch.float32)
+        self.all_intrinsics = torch.from_numpy(np.stack(self.all_intrinsics)).to(torch.float32)
+        self.all_w2cs = torch.from_numpy(np.stack(self.all_w2cs)).to(torch.float32)
+        self.img_wh = [self.img_wh[0] - self.clip_wh[0] - self.clip_wh[2],
+                       self.img_wh[1] - self.clip_wh[1] - self.clip_wh[3]]
+
+    def cal_scale_mat(self, img_hw, intrinsics, extrinsics, near_fars, factor=1.):
+        center, radius, _ = 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 scale_cam_info(self):
+        new_intrinsics = []
+        new_near_fars = []
+        new_w2cs = []
+        new_c2ws = []
+        new_affine_mats = []
+        for idx in range(len(self.all_images)):
+            intrinsics = self.all_intrinsics[idx]
+            P = intrinsics @ self.all_w2cs[idx] @ self.scale_mat
+            P = P.cpu().numpy()[: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)
+            new_intrinsics.append(intrinsics)
+            affine_mat = np.eye(4)
+            affine_mat[:3, :4] = intrinsics[: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_intrinsics, new_w2cs, new_c2ws, new_affine_mats, new_near_fars = \
+            np.stack(new_intrinsics), np.stack(new_w2cs), np.stack(new_c2ws), \
+            np.stack(new_affine_mats), np.stack(new_near_fars)
+
+        new_intrinsics = torch.from_numpy(np.float32(new_intrinsics))
+        new_w2cs = torch.from_numpy(np.float32(new_w2cs))
+        new_c2ws = torch.from_numpy(np.float32(new_c2ws))
+        new_affine_mats = torch.from_numpy(np.float32(new_affine_mats))
+        new_near_fars = torch.from_numpy(np.float32(new_near_fars))
+
+        return new_intrinsics, new_w2cs, new_c2ws, new_affine_mats, new_near_fars
+
+
+    def get_conditional_sample(self):
+        sample = {}
+        support_idxs = self.train_img_idx
+
+        sample['images'] = self.all_images[support_idxs]  # (V, 3, H, W)
+        sample['w2cs'] = self.scaled_w2cs[self.train_img_idx]  # (V, 4, 4)
+        sample['c2ws'] = self.scaled_c2ws[self.train_img_idx]  # (V, 4, 4)
+        sample['near_fars'] = self.scaled_near_fars[self.train_img_idx]  # (V, 2)
+        sample['intrinsics'] = self.scaled_intrinsics[self.train_img_idx][:, :3, :3]  # (V, 3, 3)
+        sample['affine_mats'] = self.scaled_affine_mats[self.train_img_idx]  # ! in world space
+
+        sample['scan'] = self.scan_id
+        sample['scale_factor'] = torch.tensor(self.scale_factor)
+        sample['scale_mat'] = torch.from_numpy(self.scale_mat)
+        sample['trans_mat'] = torch.from_numpy(np.linalg.inv(self.ref_w2c))
+        sample['img_wh'] = torch.from_numpy(np.array(self.img_wh))
+        sample['partial_vol_origin'] = torch.tensor(self.partial_vol_origin, dtype=torch.float32)
+
+        return sample
+
+    def __len__(self):
+        if self.split == 'train':
+            return self.n_views * 1000
+        else:
+            return len(self.test_img_idx) * 1000
+
+    def __getitem__(self, idx):
+        sample = {}
+
+        if self.split == 'train':
+            render_idx = self.train_img_idx[idx % self.n_views]
+            support_idxs = [idx for idx in self.train_img_idx if idx != render_idx]
+        else:
+            # render_idx = idx % self.n_test_images + self.n_train_images
+            render_idx = self.test_img_idx[idx % len(self.test_img_idx)]
+            support_idxs = [render_idx]
+
+        sample['images'] = self.all_images[support_idxs]  # (V, 3, H, W)
+        sample['w2cs'] = self.scaled_w2cs[support_idxs]  # (V, 4, 4)
+        sample['c2ws'] = self.scaled_c2ws[support_idxs]  # (V, 4, 4)
+        sample['intrinsics'] = self.scaled_intrinsics[support_idxs][:, :3, :3]  # (V, 3, 3)
+        sample['affine_mats'] = self.scaled_affine_mats[support_idxs]  # ! in world space
+        sample['scan'] = self.scan_id
+        sample['scale_factor'] = torch.tensor(self.scale_factor)
+        sample['img_wh'] = torch.from_numpy(np.array(self.img_wh))
+        sample['partial_vol_origin'] = torch.tensor(self.partial_vol_origin, dtype=torch.float32)
+        sample['img_index'] = torch.tensor(render_idx)
+
+        # - query image
+        sample['query_image'] = self.all_images[render_idx]
+        sample['query_c2w'] = self.scaled_c2ws[render_idx]
+        sample['query_w2c'] = self.scaled_w2cs[render_idx]
+        sample['query_intrinsic'] = self.scaled_intrinsics[render_idx]
+        sample['query_near_far'] = self.scaled_near_fars[render_idx]
+        sample['meta'] = str(self.scan_id) + "_" + os.path.basename(self.images_list[render_idx])
+        sample['scale_mat'] = torch.from_numpy(self.scale_mat)
+        sample['trans_mat'] = torch.from_numpy(np.linalg.inv(self.ref_w2c))
+        sample['rendering_c2ws'] = self.scaled_c2ws[self.test_img_idx]
+        sample['rendering_imgs_idx'] = torch.Tensor(np.array(self.test_img_idx).astype(np.int32))
+
+        # - generate rays
+        if self.split == 'val' or self.split == 'test':
+            sample_rays = gen_rays_from_single_image(
+                self.img_wh[1], self.img_wh[0],
+                sample['query_image'],
+                sample['query_intrinsic'],
+                sample['query_c2w'],
+                depth=None,
+                mask=None)
+        else:
+            sample_rays = gen_random_rays_from_single_image(
+                self.img_wh[1], self.img_wh[0],
+                self.N_rays,
+                sample['query_image'],
+                sample['query_intrinsic'],
+                sample['query_c2w'],
+                depth=None,
+                mask=None,
+                dilated_mask=None,
+                importance_sample=False,
+                h_patch_size=self.h_patch_size
+            )
+
+        sample['rays'] = sample_rays
+
+        return sample

SparseNeuS_demo_v1/data/dtu_general.py

@@ -0,0 +1,376 @@
+from torch.utils.data import Dataset
+from utils.misc_utils import read_pfm
+import os
+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 termcolor import colored
+import pdb
+import random
+
+
+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()  # ? why need transpose here
+    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 MVSDatasetDtuPerView(Dataset):
+    def __init__(self, root_dir, split, n_views=3, img_wh=(640, 512), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[]):
+
+        self.root_dir = root_dir
+        self.split = split
+
+        self.img_wh = img_wh
+        self.downSample = downSample
+        self.num_all_imgs = 49  # this preprocessed DTU dataset has 49 images
+        self.n_views = n_views
+        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.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+
+        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!'
+
+        self.split_filepath = f'data/dtu/lists/{self.split}.txt' if split_filepath is None else split_filepath
+        self.pair_filepath = f'data/dtu/dtu_pairs.txt' if pair_filepath is None else pair_filepath
+
+        print(colored("loading all scenes together", 'red'))
+        with open(self.split_filepath) as f:
+            self.scans = [line.rstrip() for line in f.readlines()]
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+
+        self.metas, self.ref_src_pairs = self.build_metas()  # load ref-srcs view pairs info of the scene
+
+        self.allview_ids = [i for i in range(self.num_all_imgs)]
+
+        self.load_cam_info()  # load camera info of DTU, and estimate scale_mat
+
+        self.build_remap()
+        self.define_transforms()
+        print(f'==> image down scale: {self.downSample}')
+
+        # * 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.])
+
+    def build_remap(self):
+        self.remap = np.zeros(np.max(self.allview_ids) + 1).astype('int')
+        for i, item in enumerate(self.allview_ids):
+            self.remap[item] = i
+
+    def define_transforms(self):
+        self.transform = T.Compose([T.ToTensor()])
+
+    def build_metas(self):
+        metas = []
+        ref_src_pairs = {}
+        # light conditions 0-6 for training
+        # light condition 3 for testing (the brightest?)
+        light_idxs = [3] if 'train' not in self.split else range(7)
+
+        with open(self.pair_filepath) as f:
+            num_viewpoint = int(f.readline())
+            # viewpoints (49)
+            for _ in range(num_viewpoint):
+                ref_view = int(f.readline().rstrip())
+                src_views = [int(x) for x in f.readline().rstrip().split()[1::2]]
+
+                ref_src_pairs[ref_view] = src_views
+
+        for light_idx in light_idxs:
+            for scan in self.scans:
+                with open(self.pair_filepath) as f:
+                    num_viewpoint = int(f.readline())
+                    # viewpoints (49)
+                    for _ in range(num_viewpoint):
+                        ref_view = int(f.readline().rstrip())
+                        src_views = [int(x) for x in f.readline().rstrip().split()[1::2]]
+
+                        # ! only for validation
+                        if len(self.test_ref_views) > 0 and ref_view not in self.test_ref_views:
+                            continue
+
+                        metas += [(scan, light_idx, ref_view, src_views)]
+
+        return metas, ref_src_pairs
+
+    def read_cam_file(self, filename):
+        with open(filename) as f:
+            lines = [line.rstrip() for line in f.readlines()]
+        # extrinsics: line [1,5), 4x4 matrix
+        extrinsics = np.fromstring(' '.join(lines[1:5]), dtype=np.float32, sep=' ')
+        extrinsics = extrinsics.reshape((4, 4))
+        # intrinsics: line [7-10), 3x3 matrix
+        intrinsics = np.fromstring(' '.join(lines[7:10]), dtype=np.float32, sep=' ')
+        intrinsics = intrinsics.reshape((3, 3))
+        # depth_min & depth_interval: line 11
+        depth_min = float(lines[11].split()[0])
+        depth_max = depth_min + float(lines[11].split()[1]) * 192
+        self.depth_interval = float(lines[11].split()[1])
+        intrinsics_ = np.float32(np.diag([1, 1, 1, 1]))
+        intrinsics_[:3, :3] = intrinsics
+        return intrinsics_, extrinsics, [depth_min, depth_max]
+
+    def load_cam_info(self):
+        for vid in range(self.num_all_imgs):
+            proj_mat_filename = os.path.join(self.root_dir,
+                                             f'Cameras/train/{vid:08d}_cam.txt')
+            intrinsic, extrinsic, near_far = self.read_cam_file(proj_mat_filename)
+            intrinsic[:2] *= 4  # * the provided intrinsics is 4x downsampled, now keep the same scale with image
+            self.all_intrinsics.append(intrinsic)
+            self.all_extrinsics.append(extrinsic)
+            self.all_near_fars.append(near_far)
+
+    def read_depth(self, filename):
+        # import ipdb; ipdb.set_trace()
+        depth_h = np.array(read_pfm(filename)[0], dtype=np.float32)  # (1200, 1600)
+        depth_h = np.ones((1200, 1600))
+        # print(depth_h.shape)
+        depth_h = cv2.resize(depth_h, None, fx=0.5, fy=0.5,
+                             interpolation=cv2.INTER_NEAREST)  # (600, 800)
+        depth_h = depth_h[44:556, 80:720]  # (512, 640)
+        # print(depth_h.shape)
+        # import ipdb; ipdb.set_trace()
+        depth_h = cv2.resize(depth_h, None, fx=self.downSample, fy=self.downSample,
+                             interpolation=cv2.INTER_NEAREST)
+        depth = cv2.resize(depth_h, None, fx=1.0 / 4, fy=1.0 / 4,
+                           interpolation=cv2.INTER_NEAREST)
+
+        return depth, depth_h
+
+    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, _ = 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.metas)
+
+    def __getitem__(self, idx):
+        sample = {}
+        scan, light_idx, ref_view, src_views = self.metas[idx % len(self.metas)]
+
+        # generalized, load some images at once
+        view_ids = [ref_view] + src_views[:self.n_views]
+        # * transform from world system to camera system
+        w2c_ref = self.all_extrinsics[self.remap[ref_view]]
+        w2c_ref_inv = np.linalg.inv(w2c_ref)
+
+        image_perm = 0  # only supervised on reference view
+
+        imgs, depths_h, masks_h = [], [], []  # full size (640, 512)
+        intrinsics, w2cs, near_fars = [], [], []  # record proj mats between views
+        mask_dilated = None
+        for i, vid in enumerate(view_ids):
+            # NOTE that the id in image file names is from 1 to 49 (not 0~48)
+            img_filename = os.path.join(self.root_dir,
+                                        f'Rectified/{scan}_train/rect_{vid + 1:03d}_{light_idx}_r5000.png')
+            depth_filename = os.path.join(self.root_dir,
+                                          f'Depths/{scan}_train/depth_map_{vid:04d}.pfm')
+            # print(depth_filename)
+            mask_filename = os.path.join(self.root_dir,
+                                         f'Masks_clean_dilated/{scan}_train/mask_{vid:04d}.png')
+
+            img = Image.open(img_filename)
+            img_wh = np.round(np.array(img.size) * self.downSample).astype('int')
+            img = img.resize(img_wh, Image.BILINEAR)
+
+            if os.path.exists(mask_filename) and self.clean_image:
+                mask_l, mask_h = self.read_mask(mask_filename)
+            else:
+                # print(self.split, "don't find mask file", mask_filename)
+                mask_h = np.ones([img_wh[1], img_wh[0]])
+            masks_h.append(mask_h)
+
+            if i == 0:
+                kernel_size = 101  # default 101
+                kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
+                mask_dilated = np.float32(cv2.dilate(np.uint8(mask_h * 255), kernel, iterations=1) > 128)
+
+            if self.clean_image:
+                img = np.array(img)
+                img[mask_h < 0.5] = 0.0
+
+            img = self.transform(img)
+
+            imgs += [img]
+
+            index_mat = self.remap[vid]
+            near_fars.append(self.all_near_fars[index_mat])
+            intrinsics.append(self.all_intrinsics[index_mat])
+            
+            w2cs.append(self.all_extrinsics[index_mat] @ w2c_ref_inv)
+
+            # print(depth_filename)
+            if os.path.exists(depth_filename):  # and i == 0
+                # print("file exists")
+                depth_l, depth_h = self.read_depth(depth_filename)
+                depths_h.append(depth_h)
+        # ! 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)
+
+        # ! calculate the new w2cs after scaling
+        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()
+        print(new_near_fars)
+        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 'train' in self.split:
+            start_idx = 0
+        else:
+            start_idx = 1
+
+        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['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['light_idx'] = torch.tensor(light_idx)
+        sample['scan'] = scan
+
+        sample['scale_factor'] = torch.tensor(scale_factor)
+        sample['img_wh'] = torch.from_numpy(img_wh)
+        sample['render_img_idx'] = torch.tensor(image_perm)
+        sample['partial_vol_origin'] = torch.tensor(self.partial_vol_origin, dtype=torch.float32)
+        sample['meta'] = str(scan) + "_light" + str(light_idx) + "_refview" + str(ref_view)
+
+        # - 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

SparseNeuS_demo_v1/data/scene.py

@@ -0,0 +1,102 @@
+import numpy as np
+import torch
+import pdb
+
+
+def rigid_transform(xyz, transform):
+    """Applies a rigid transform (c2w) to an (N, 3) pointcloud.
+    """
+    device = xyz.device
+    xyz_h = torch.cat([xyz, torch.ones((len(xyz), 1)).to(device)], dim=1)  # (N, 4)
+    xyz_t_h = (transform @ xyz_h.T).T  # * checked: the same with the below
+
+    return xyz_t_h[:, :3]
+
+
+def get_view_frustum(min_depth, max_depth, size, cam_intr, c2w):
+    """Get corners of 3D camera view frustum of depth image
+    """
+    device = cam_intr.device
+    im_h, im_w = size
+    im_h = int(im_h)
+    im_w = int(im_w)
+    view_frust_pts = torch.stack([
+        (torch.tensor([0, 0, im_w, im_w, 0, 0, im_w, im_w]).to(device) - cam_intr[0, 2]) * torch.tensor(
+            [min_depth, min_depth, min_depth, min_depth, max_depth, max_depth, max_depth, max_depth]).to(device) /
+        cam_intr[0, 0],
+        (torch.tensor([0, im_h, 0, im_h, 0, im_h, 0, im_h]).to(device) - cam_intr[1, 2]) * torch.tensor(
+            [min_depth, min_depth, min_depth, min_depth, max_depth, max_depth, max_depth, max_depth]).to(device) /
+        cam_intr[1, 1],
+        torch.tensor([min_depth, min_depth, min_depth, min_depth, max_depth, max_depth, max_depth, max_depth]).to(
+            device)
+    ])
+    view_frust_pts = view_frust_pts.type(torch.float32)
+    c2w = c2w.type(torch.float32)
+    view_frust_pts = rigid_transform(view_frust_pts.T, c2w).T
+    return view_frust_pts
+
+
+def set_pixel_coords(h, w):
+    i_range = torch.arange(0, h).view(1, h, 1).expand(1, h, w).type(torch.float32)  # [1, H, W]
+    j_range = torch.arange(0, w).view(1, 1, w).expand(1, h, w).type(torch.float32)  # [1, H, W]
+    ones = torch.ones(1, h, w).type(torch.float32)
+
+    pixel_coords = torch.stack((j_range, i_range, ones), dim=1)  # [1, 3, H, W]
+
+    return pixel_coords
+
+
+def get_boundingbox(img_hw, intrinsics, extrinsics, near_fars):
+    """
+    # get the minimum bounding box of all visual hulls
+    :param img_hw:
+    :param intrinsics:
+    :param extrinsics:
+    :param near_fars:
+    :return:
+    """
+
+    bnds = torch.zeros((3, 2))
+    bnds[:, 0] = np.inf
+    bnds[:, 1] = -np.inf
+
+    if isinstance(intrinsics, list):
+        num = len(intrinsics)
+    else:
+        num = intrinsics.shape[0]
+    # print("num: ", num)
+    view_frust_pts_list = []
+    for i in range(num):
+        if not isinstance(intrinsics[i], torch.Tensor):
+            cam_intr = torch.tensor(intrinsics[i])
+            w2c = torch.tensor(extrinsics[i])
+            c2w = torch.inverse(w2c)
+        else:
+            cam_intr = intrinsics[i]
+            w2c = extrinsics[i]
+            c2w = torch.inverse(w2c)
+        min_depth, max_depth = near_fars[i][0], near_fars[i][1]
+        # todo: check the coresponding points are matched
+
+        view_frust_pts = get_view_frustum(min_depth, max_depth, img_hw, cam_intr, c2w)
+        bnds[:, 0] = torch.min(bnds[:, 0], torch.min(view_frust_pts, dim=1)[0])
+        bnds[:, 1] = torch.max(bnds[:, 1], torch.max(view_frust_pts, dim=1)[0])
+        view_frust_pts_list.append(view_frust_pts)
+    all_view_frust_pts = torch.cat(view_frust_pts_list, dim=1)
+
+    # print("all_view_frust_pts: ", all_view_frust_pts.shape)
+    # distance = torch.norm(all_view_frust_pts, dim=0)
+    # print("distance: ", distance)
+
+    # print("all_view_frust_pts_z: ", all_view_frust_pts[2, :])
+
+    center = torch.tensor(((bnds[0, 1] + bnds[0, 0]) / 2, (bnds[1, 1] + bnds[1, 0]) / 2,
+                           (bnds[2, 1] + bnds[2, 0]) / 2))
+
+    lengths = bnds[:, 1] - bnds[:, 0]
+
+    max_length, _ = torch.max(lengths, dim=0)
+    radius = max_length / 2
+
+    # print("radius: ", radius)
+    return center, radius, bnds

SparseNeuS_demo_v1/evaluation/clean_mesh.py

@@ -0,0 +1,283 @@
+import numpy as np
+import cv2 as cv
+import os
+from glob import glob
+from scipy.io import loadmat
+import trimesh
+import open3d as o3d
+import torch
+from tqdm import tqdm
+
+import sys
+
+sys.path.append("../")
+
+
+def gen_rays_from_single_image(H, W, image, intrinsic, c2w, depth=None, mask=None):
+    """
+    generate rays in world space, for image image
+    :param H:
+    :param W:
+    :param intrinsics: [3,3]
+    :param c2ws: [4,4]
+    :return:
+    """
+    device = image.device
+    ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
+                            torch.linspace(0, W - 1, W))  # pytorch's meshgrid has indexing='ij'
+    p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3
+
+    # normalized ndc uv coordinates, (-1, 1)
+    ndc_u = 2 * xs / (W - 1) - 1
+    ndc_v = 2 * ys / (H - 1) - 1
+    rays_ndc_uv = torch.stack([ndc_u, ndc_v], dim=-1).view(-1, 2).float().to(device)
+
+    intrinsic_inv = torch.inverse(intrinsic)
+
+    p = p.view(-1, 3).float().to(device)  # N_rays, 3
+    p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze()  # N_rays, 3
+    rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # N_rays, 3
+    rays_v = torch.matmul(c2w[None, :3, :3], rays_v[:, :, None]).squeeze()  # N_rays, 3
+    rays_o = c2w[None, :3, 3].expand(rays_v.shape)  # N_rays, 3
+
+    image = image.permute(1, 2, 0)
+    color = image.view(-1, 3)
+    depth = depth.view(-1, 1) if depth is not None else None
+    mask = mask.view(-1, 1) if mask is not None else torch.ones([H * W, 1]).to(device)
+    sample = {
+        'rays_o': rays_o,
+        'rays_v': rays_v,
+        'rays_ndc_uv': rays_ndc_uv,
+        'rays_color': color,
+        # 'rays_depth': depth,
+        'rays_mask': mask,
+        'rays_norm_XYZ_cam': p  # - XYZ_cam, before multiply depth
+    }
+    if depth is not None:
+        sample['rays_depth'] = depth
+
+    return sample
+
+
+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 = cv.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()  # ? why need transpose here
+    pose[:3, 3] = (t[:3] / t[3])[:, 0]
+
+    return intrinsics, pose  # ! return cam2world matrix here
+
+
+def clean_points_by_mask(points, scan, imgs_idx=None, minimal_vis=0, mask_dilated_size=11):
+    cameras = np.load('{}/scan{}/cameras.npz'.format(DTU_DIR, scan))
+    mask_lis = sorted(glob('{}/scan{}/mask/*.png'.format(DTU_DIR, scan)))
+    n_images = 49 if scan < 83 else 64
+    inside_mask = np.zeros(len(points))
+
+    if imgs_idx is None:
+        imgs_idx = [i for i in range(n_images)]
+
+    # imgs_idx = [i for i in range(n_images)]
+    for i in imgs_idx:
+        P = cameras['world_mat_{}'.format(i)]
+        pts_image = np.matmul(P[None, :3, :3], points[:, :, None]).squeeze() + P[None, :3, 3]
+        pts_image = pts_image / pts_image[:, 2:]
+        pts_image = np.round(pts_image).astype(np.int32) + 1
+
+        mask_image = cv.imread(mask_lis[i])
+        kernel_size = mask_dilated_size  # default 101
+        kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (kernel_size, kernel_size))
+        mask_image = cv.dilate(mask_image, kernel, iterations=1)
+        mask_image = (mask_image[:, :, 0] > 128)
+
+        mask_image = np.concatenate([np.ones([1, 1600]), mask_image, np.ones([1, 1600])], axis=0)
+        mask_image = np.concatenate([np.ones([1202, 1]), mask_image, np.ones([1202, 1])], axis=1)
+
+        in_mask = (pts_image[:, 0] >= 0) * (pts_image[:, 0] <= 1600) * (pts_image[:, 1] >= 0) * (
+                pts_image[:, 1] <= 1200) > 0
+        curr_mask = mask_image[(pts_image[:, 1].clip(0, 1201), pts_image[:, 0].clip(0, 1601))]
+
+        curr_mask = curr_mask.astype(np.float32) * in_mask
+
+        inside_mask += curr_mask
+
+    return inside_mask > minimal_vis
+
+
+def clean_mesh_faces_by_mask(mesh_file, new_mesh_file, scan, imgs_idx, minimal_vis=0, mask_dilated_size=11):
+    old_mesh = trimesh.load(mesh_file)
+    old_vertices = old_mesh.vertices[:]
+    old_faces = old_mesh.faces[:]
+    mask = clean_points_by_mask(old_vertices, scan, imgs_idx, minimal_vis, mask_dilated_size)
+    indexes = np.ones(len(old_vertices)) * -1
+    indexes = indexes.astype(np.long)
+    indexes[np.where(mask)] = np.arange(len(np.where(mask)[0]))
+
+    faces_mask = mask[old_faces[:, 0]] & mask[old_faces[:, 1]] & mask[old_faces[:, 2]]
+    new_faces = old_faces[np.where(faces_mask)]
+    new_faces[:, 0] = indexes[new_faces[:, 0]]
+    new_faces[:, 1] = indexes[new_faces[:, 1]]
+    new_faces[:, 2] = indexes[new_faces[:, 2]]
+    new_vertices = old_vertices[np.where(mask)]
+
+    new_mesh = trimesh.Trimesh(new_vertices, new_faces)
+
+    new_mesh.export(new_mesh_file)
+
+
+def clean_mesh_by_faces_num(mesh, faces_num=500):
+    old_vertices = mesh.vertices[:]
+    old_faces = mesh.faces[:]
+
+    cc = trimesh.graph.connected_components(mesh.face_adjacency, min_len=faces_num)
+    mask = np.zeros(len(mesh.faces), dtype=np.bool)
+    mask[np.concatenate(cc)] = True
+
+    indexes = np.ones(len(old_vertices)) * -1
+    indexes = indexes.astype(np.long)
+    indexes[np.where(mask)] = np.arange(len(np.where(mask)[0]))
+
+    faces_mask = mask[old_faces[:, 0]] & mask[old_faces[:, 1]] & mask[old_faces[:, 2]]
+    new_faces = old_faces[np.where(faces_mask)]
+    new_faces[:, 0] = indexes[new_faces[:, 0]]
+    new_faces[:, 1] = indexes[new_faces[:, 1]]
+    new_faces[:, 2] = indexes[new_faces[:, 2]]
+    new_vertices = old_vertices[np.where(mask)]
+
+    new_mesh = trimesh.Trimesh(new_vertices, new_faces)
+
+    return new_mesh
+
+
+def clean_mesh_faces_outside_frustum(old_mesh_file, new_mesh_file, imgs_idx, H=1200, W=1600, mask_dilated_size=11,
+                                     isolated_face_num=500, keep_largest=True):
+    '''Remove faces of mesh which cannot be orserved by all cameras
+    '''
+    # if path_mask_npz:
+    #     path_save_clean = IOUtils.add_file_name_suffix(path_save_clean, '_mask')
+
+    cameras = np.load('{}/scan{}/cameras.npz'.format(DTU_DIR, scan))
+    mask_lis = sorted(glob('{}/scan{}/mask/*.png'.format(DTU_DIR, scan)))
+
+    mesh = trimesh.load(old_mesh_file)
+    intersector = trimesh.ray.ray_pyembree.RayMeshIntersector(mesh)
+
+    all_indices = []
+    chunk_size = 5120
+    for i in imgs_idx:
+        mask_image = cv.imread(mask_lis[i])
+        kernel_size = mask_dilated_size  # default 101
+        kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (kernel_size, kernel_size))
+        mask_image = cv.dilate(mask_image, kernel, iterations=1)
+
+        P = cameras['world_mat_{}'.format(i)]
+
+        intrinsic, pose = load_K_Rt_from_P(None, P[:3, :])
+
+        rays = gen_rays_from_single_image(H, W, torch.from_numpy(mask_image).permute(2, 0, 1).float(),
+                                          torch.from_numpy(intrinsic)[:3, :3].float(),
+                                          torch.from_numpy(pose).float())
+        rays_o = rays['rays_o']
+        rays_d = rays['rays_v']
+        rays_mask = rays['rays_color']
+
+        rays_o = rays_o.split(chunk_size)
+        rays_d = rays_d.split(chunk_size)
+        rays_mask = rays_mask.split(chunk_size)
+
+        for rays_o_batch, rays_d_batch, rays_mask_batch in tqdm(zip(rays_o, rays_d, rays_mask)):
+            rays_mask_batch = rays_mask_batch[:, 0] > 128
+            rays_o_batch = rays_o_batch[rays_mask_batch]
+            rays_d_batch = rays_d_batch[rays_mask_batch]
+
+            idx_faces_hits = intersector.intersects_first(rays_o_batch.cpu().numpy(), rays_d_batch.cpu().numpy())
+            all_indices.append(idx_faces_hits)
+
+    values = np.unique(np.concatenate(all_indices, axis=0))
+    mask_faces = np.ones(len(mesh.faces))
+    mask_faces[values[1:]] = 0
+    print(f'Surfaces/Kept: {len(mesh.faces)}/{len(values)}')
+
+    mesh_o3d = o3d.io.read_triangle_mesh(old_mesh_file)
+    print("removing triangles by mask")
+    mesh_o3d.remove_triangles_by_mask(mask_faces)
+
+    o3d.io.write_triangle_mesh(new_mesh_file, mesh_o3d)
+
+    # # clean meshes
+    new_mesh = trimesh.load(new_mesh_file)
+    cc = trimesh.graph.connected_components(new_mesh.face_adjacency, min_len=500)
+    mask = np.zeros(len(new_mesh.faces), dtype=np.bool)
+    mask[np.concatenate(cc)] = True
+    new_mesh.update_faces(mask)
+    new_mesh.remove_unreferenced_vertices()
+    new_mesh.export(new_mesh_file)
+
+    # meshes = new_mesh.split(only_watertight=False)
+    #
+    # if not keep_largest:
+    #     meshes = [mesh for mesh in meshes if len(mesh.faces) > isolated_face_num]
+    #     # new_mesh = meshes[np.argmax([len(mesh.faces) for mesh in meshes])]
+    #     merged_mesh = trimesh.util.concatenate(meshes)
+    #     merged_mesh.export(new_mesh_file)
+    # else:
+    #     new_mesh = meshes[np.argmax([len(mesh.faces) for mesh in meshes])]
+    #     new_mesh.export(new_mesh_file)
+
+    o3d.io.write_triangle_mesh(new_mesh_file.replace(".ply", "_raw.ply"), mesh_o3d)
+    print("finishing removing triangles")
+
+
+def clean_outliers(old_mesh_file, new_mesh_file):
+    new_mesh = trimesh.load(old_mesh_file)
+
+    meshes = new_mesh.split(only_watertight=False)
+    new_mesh = meshes[np.argmax([len(mesh.faces) for mesh in meshes])]
+
+    new_mesh.export(new_mesh_file)
+
+
+if __name__ == "__main__":
+
+    scans = [24, 37, 40, 55, 63, 65, 69, 83, 97, 105, 106, 110, 114, 118, 122]
+
+    mask_kernel_size = 11
+
+    imgs_idx = [0, 1, 2]
+    # imgs_idx = [42, 43, 44]
+    # imgs_idx = [1, 8, 9]
+
+    DTU_DIR = "/home/xiaoxiao/dataset/DTU_IDR/DTU"
+    # DTU_DIR = "/userhome/cs/xxlong/dataset/DTU_IDR/DTU"
+
+    base_path = "/home/xiaoxiao/Workplace/nerf_reconstruction/Volume_NeuS/neus_camsys/exp/dtu/evaluation_23_24_33_new/volsdf"
+
+    for scan in scans:
+        print("processing scan%d" % scan)
+        dir_path = os.path.join(base_path, "scan%d" % scan)
+
+        old_mesh_file = glob(os.path.join(dir_path, "*.ply"))[0]
+
+        clean_mesh_file = os.path.join(dir_path, "clean_%03d.ply" % scan)
+        final_mesh_file = os.path.join(dir_path, "final_%03d.ply" % scan)
+
+        clean_mesh_faces_by_mask(old_mesh_file, clean_mesh_file, scan, imgs_idx, minimal_vis=1,
+                                 mask_dilated_size=mask_kernel_size)
+        clean_mesh_faces_outside_frustum(clean_mesh_file, final_mesh_file, imgs_idx, mask_dilated_size=mask_kernel_size)
+
+        print("finish processing scan%d" % scan)

SparseNeuS_demo_v1/evaluation/eval_dtu_python.py

@@ -0,0 +1,369 @@
+import numpy as np
+import open3d as o3d
+import sklearn.neighbors as skln
+from tqdm import tqdm
+from scipy.io import loadmat
+import multiprocessing as mp
+import argparse, os, sys
+import cv2 as cv
+
+from pathlib import Path
+
+
+def get_path_components(path):
+    path = Path(path)
+    ppath = str(path.parent)
+    stem = str(path.stem)
+    ext = str(path.suffix)
+    return ppath, stem, ext
+
+
+def sample_single_tri(input_):
+    n1, n2, v1, v2, tri_vert = input_
+    c = np.mgrid[:n1 + 1, :n2 + 1]
+    c += 0.5
+    c[0] /= max(n1, 1e-7)
+    c[1] /= max(n2, 1e-7)
+    c = np.transpose(c, (1, 2, 0))
+    k = c[c.sum(axis=-1) < 1]  # m2
+    q = v1 * k[:, :1] + v2 * k[:, 1:] + tri_vert
+    return q
+
+
+def write_vis_pcd(file, points, colors):
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points)
+    pcd.colors = o3d.utility.Vector3dVector(colors)
+    o3d.io.write_point_cloud(file, pcd)
+
+
+def eval_cloud(args, num_cpu_cores=-1):
+    mp.freeze_support()
+    os.makedirs(args.vis_out_dir, exist_ok=True)
+
+    thresh = args.downsample_density
+    if args.mode == 'mesh':
+        pbar = tqdm(total=9)
+        pbar.set_description('read data mesh')
+        data_mesh = o3d.io.read_triangle_mesh(args.data)
+
+        vertices = np.asarray(data_mesh.vertices)
+        triangles = np.asarray(data_mesh.triangles)
+        tri_vert = vertices[triangles]
+
+        pbar.update(1)
+        pbar.set_description('sample pcd from mesh')
+        v1 = tri_vert[:, 1] - tri_vert[:, 0]
+        v2 = tri_vert[:, 2] - tri_vert[:, 0]
+        l1 = np.linalg.norm(v1, axis=-1, keepdims=True)
+        l2 = np.linalg.norm(v2, axis=-1, keepdims=True)
+        area2 = np.linalg.norm(np.cross(v1, v2), axis=-1, keepdims=True)
+        non_zero_area = (area2 > 0)[:, 0]
+        l1, l2, area2, v1, v2, tri_vert = [
+            arr[non_zero_area] for arr in [l1, l2, area2, v1, v2, tri_vert]
+        ]
+        thr = thresh * np.sqrt(l1 * l2 / area2)
+        n1 = np.floor(l1 / thr)
+        n2 = np.floor(l2 / thr)
+
+        with mp.Pool() as mp_pool:
+            new_pts = mp_pool.map(sample_single_tri,
+                                  ((n1[i, 0], n2[i, 0], v1[i:i + 1], v2[i:i + 1], tri_vert[i:i + 1, 0]) for i in
+                                   range(len(n1))), chunksize=1024)
+
+        new_pts = np.concatenate(new_pts, axis=0)
+        data_pcd = np.concatenate([vertices, new_pts], axis=0)
+
+    elif args.mode == 'pcd':
+        pbar = tqdm(total=8)
+        pbar.set_description('read data pcd')
+        data_pcd_o3d = o3d.io.read_point_cloud(args.data)
+        data_pcd = np.asarray(data_pcd_o3d.points)
+
+    pbar.update(1)
+    pbar.set_description('random shuffle pcd index')
+    shuffle_rng = np.random.default_rng()
+    shuffle_rng.shuffle(data_pcd, axis=0)
+
+    pbar.update(1)
+    pbar.set_description('downsample pcd')
+    nn_engine = skln.NearestNeighbors(n_neighbors=1, radius=thresh, algorithm='kd_tree', n_jobs=num_cpu_cores)
+    nn_engine.fit(data_pcd)
+    rnn_idxs = nn_engine.radius_neighbors(data_pcd, radius=thresh, return_distance=False)
+    mask = np.ones(data_pcd.shape[0], dtype=np.bool_)
+    for curr, idxs in enumerate(rnn_idxs):
+        if mask[curr]:
+            mask[idxs] = 0
+            mask[curr] = 1
+    data_down = data_pcd[mask]
+
+    pbar.update(1)
+    pbar.set_description('masking data pcd')
+    obs_mask_file = loadmat(f'{args.dataset_dir}/ObsMask/ObsMask{args.scan}_10.mat')
+    ObsMask, BB, Res = [obs_mask_file[attr] for attr in ['ObsMask', 'BB', 'Res']]
+    BB = BB.astype(np.float32)
+
+    patch = args.patch_size
+    inbound = ((data_down >= BB[:1] - patch) & (data_down < BB[1:] + patch * 2)).sum(axis=-1) == 3
+    data_in = data_down[inbound]
+
+    data_grid = np.around((data_in - BB[:1]) / Res).astype(np.int32)
+    grid_inbound = ((data_grid >= 0) & (data_grid < np.expand_dims(ObsMask.shape, 0))).sum(axis=-1) == 3
+    data_grid_in = data_grid[grid_inbound]
+    in_obs = ObsMask[data_grid_in[:, 0], data_grid_in[:, 1], data_grid_in[:, 2]].astype(np.bool_)
+    data_in_obs = data_in[grid_inbound][in_obs]
+
+    pbar.update(1)
+    pbar.set_description('read STL pcd')
+    stl_pcd = o3d.io.read_point_cloud(args.gt)
+    stl = np.asarray(stl_pcd.points)
+
+    pbar.update(1)
+    pbar.set_description('compute data2stl')
+    nn_engine.fit(stl)
+    dist_d2s, idx_d2s = nn_engine.kneighbors(data_in_obs, n_neighbors=1, return_distance=True)
+    max_dist = args.max_dist
+    mean_d2s = dist_d2s[dist_d2s < max_dist].mean()
+
+    pbar.update(1)
+    pbar.set_description('compute stl2data')
+    ground_plane = loadmat(f'{args.dataset_dir}/ObsMask/Plane{args.scan}.mat')['P']
+
+    stl_hom = np.concatenate([stl, np.ones_like(stl[:, :1])], -1)
+    above = (ground_plane.reshape((1, 4)) * stl_hom).sum(-1) > 0
+    stl_above = stl[above]
+
+    nn_engine.fit(data_in)
+    dist_s2d, idx_s2d = nn_engine.kneighbors(stl_above, n_neighbors=1, return_distance=True)
+    mean_s2d = dist_s2d[dist_s2d < max_dist].mean()
+
+    pbar.update(1)
+    pbar.set_description('visualize error')
+    vis_dist = args.visualize_threshold
+    R = np.array([[1, 0, 0]], dtype=np.float64)
+    G = np.array([[0, 1, 0]], dtype=np.float64)
+    B = np.array([[0, 0, 1]], dtype=np.float64)
+    W = np.array([[1, 1, 1]], dtype=np.float64)
+    data_color = np.tile(B, (data_down.shape[0], 1))
+    data_alpha = dist_d2s.clip(max=vis_dist) / vis_dist
+    data_color[np.where(inbound)[0][grid_inbound][in_obs]] = R * data_alpha + W * (1 - data_alpha)
+    data_color[np.where(inbound)[0][grid_inbound][in_obs][dist_d2s[:, 0] >= max_dist]] = G
+    write_vis_pcd(f'{args.vis_out_dir}/vis_{args.scan:03}_d2gt.ply', data_down, data_color)
+    stl_color = np.tile(B, (stl.shape[0], 1))
+    stl_alpha = dist_s2d.clip(max=vis_dist) / vis_dist
+    stl_color[np.where(above)[0]] = R * stl_alpha + W * (1 - stl_alpha)
+    stl_color[np.where(above)[0][dist_s2d[:, 0] >= max_dist]] = G
+    write_vis_pcd(f'{args.vis_out_dir}/vis_{args.scan:03}_gt2d.ply', stl, stl_color)
+
+    pbar.update(1)
+    pbar.set_description('done')
+    pbar.close()
+    over_all = (mean_d2s + mean_s2d) / 2
+    print(f'ean_d2gt: {mean_d2s}; mean_gt2d: {mean_s2d}  over_all: {over_all}; .')
+
+    pparent, stem, ext = get_path_components(args.data)
+    if args.log is None:
+        path_log = os.path.join(pparent, 'eval_result.txt')
+    else:
+        path_log = args.log
+    with open(path_log, 'a+') as fLog:
+        fLog.write(f'mean_d2gt {np.round(mean_d2s, 3)} '
+                   f'mean_gt2d {np.round(mean_s2d, 3)} '
+                   f'Over_all {np.round(over_all, 3)} '
+                   f'[{stem}] \n')
+
+    return over_all, mean_d2s, mean_s2d
+
+
+if __name__ == '__main__':
+    from glob import glob
+
+    mp.freeze_support()
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data', type=str, default='data_in.ply')
+    parser.add_argument('--gt', type=str, help='ground truth')
+    parser.add_argument('--scan', type=int, default=1)
+    parser.add_argument('--mode', type=str, default='mesh', choices=['mesh', 'pcd'])
+    parser.add_argument('--dataset_dir', type=str, default='/dataset/dtu_official/SampleSet/MVS_Data')
+    parser.add_argument('--vis_out_dir', type=str, default='.')
+    parser.add_argument('--downsample_density', type=float, default=0.2)
+    parser.add_argument('--patch_size', type=float, default=60)
+    parser.add_argument('--max_dist', type=float, default=20)
+    parser.add_argument('--visualize_threshold', type=float, default=10)
+    parser.add_argument('--log', type=str, default=None)
+    args = parser.parse_args()
+
+    base_dir = "./exp"
+
+    GT_DIR = "./gt_pcd"
+
+    scans = [24, 37, 40, 55, 63, 65, 69, 83, 97, 105, 106, 110, 114, 118, 122]
+
+    for scan in scans:
+
+        print("processing scan%d" % scan)
+
+        args.data = os.path.join(base_dir, "scan{}".format(scan), "final_%03d.ply" % scan)
+
+        if not os.path.exists(args.data):
+            continue
+
+        args.gt = os.path.join(GT_DIR, "stl%03d_total.ply" % scan)
+        args.vis_out_dir = os.path.join(base_dir, "scan{}".format(scan))
+        args.scan = scan
+        os.makedirs(args.vis_out_dir, exist_ok=True)
+
+        dist_thred1 = 1
+        dist_thred2 = 2
+
+        thresh = args.downsample_density
+
+        if args.mode == 'mesh':
+            pbar = tqdm(total=9)
+            pbar.set_description('read data mesh')
+            data_mesh = o3d.io.read_triangle_mesh(args.data)
+
+            vertices = np.asarray(data_mesh.vertices)
+            triangles = np.asarray(data_mesh.triangles)
+            tri_vert = vertices[triangles]
+
+            pbar.update(1)
+            pbar.set_description('sample pcd from mesh')
+            v1 = tri_vert[:, 1] - tri_vert[:, 0]
+            v2 = tri_vert[:, 2] - tri_vert[:, 0]
+            l1 = np.linalg.norm(v1, axis=-1, keepdims=True)
+            l2 = np.linalg.norm(v2, axis=-1, keepdims=True)
+            area2 = np.linalg.norm(np.cross(v1, v2), axis=-1, keepdims=True)
+            non_zero_area = (area2 > 0)[:, 0]
+            l1, l2, area2, v1, v2, tri_vert = [
+                arr[non_zero_area] for arr in [l1, l2, area2, v1, v2, tri_vert]
+            ]
+            thr = thresh * np.sqrt(l1 * l2 / area2)
+            n1 = np.floor(l1 / thr)
+            n2 = np.floor(l2 / thr)
+
+            with mp.Pool() as mp_pool:
+                new_pts = mp_pool.map(sample_single_tri,
+                                      ((n1[i, 0], n2[i, 0], v1[i:i + 1], v2[i:i + 1], tri_vert[i:i + 1, 0]) for i in
+                                       range(len(n1))), chunksize=1024)
+
+            new_pts = np.concatenate(new_pts, axis=0)
+            data_pcd = np.concatenate([vertices, new_pts], axis=0)
+
+        elif args.mode == 'pcd':
+            pbar = tqdm(total=8)
+            pbar.set_description('read data pcd')
+            data_pcd_o3d = o3d.io.read_point_cloud(args.data)
+            data_pcd = np.asarray(data_pcd_o3d.points)
+
+        pbar.update(1)
+        pbar.set_description('random shuffle pcd index')
+        shuffle_rng = np.random.default_rng()
+        shuffle_rng.shuffle(data_pcd, axis=0)
+
+        pbar.update(1)
+        pbar.set_description('downsample pcd')
+        nn_engine = skln.NearestNeighbors(n_neighbors=1, radius=thresh, algorithm='kd_tree', n_jobs=-1)
+        nn_engine.fit(data_pcd)
+        rnn_idxs = nn_engine.radius_neighbors(data_pcd, radius=thresh, return_distance=False)
+        mask = np.ones(data_pcd.shape[0], dtype=np.bool_)
+        for curr, idxs in enumerate(rnn_idxs):
+            if mask[curr]:
+                mask[idxs] = 0
+                mask[curr] = 1
+        data_down = data_pcd[mask]
+
+        pbar.update(1)
+        pbar.set_description('masking data pcd')
+        obs_mask_file = loadmat(f'{args.dataset_dir}/ObsMask/ObsMask{args.scan}_10.mat')
+        ObsMask, BB, Res = [obs_mask_file[attr] for attr in ['ObsMask', 'BB', 'Res']]
+        BB = BB.astype(np.float32)
+
+        patch = args.patch_size
+        inbound = ((data_down >= BB[:1] - patch) & (data_down < BB[1:] + patch * 2)).sum(axis=-1) == 3
+        data_in = data_down[inbound]
+
+        data_grid = np.around((data_in - BB[:1]) / Res).astype(np.int32)
+        grid_inbound = ((data_grid >= 0) & (data_grid < np.expand_dims(ObsMask.shape, 0))).sum(axis=-1) == 3
+        data_grid_in = data_grid[grid_inbound]
+        in_obs = ObsMask[data_grid_in[:, 0], data_grid_in[:, 1], data_grid_in[:, 2]].astype(np.bool_)
+        data_in_obs = data_in[grid_inbound][in_obs]
+
+        pbar.update(1)
+        pbar.set_description('read STL pcd')
+        stl_pcd = o3d.io.read_point_cloud(args.gt)
+        stl = np.asarray(stl_pcd.points)
+
+        pbar.update(1)
+        pbar.set_description('compute data2stl')
+        nn_engine.fit(stl)
+        dist_d2s, idx_d2s = nn_engine.kneighbors(data_in_obs, n_neighbors=1, return_distance=True)
+        max_dist = args.max_dist
+        mean_d2s = dist_d2s[dist_d2s < max_dist].mean()
+
+        precision_1 = len(dist_d2s[dist_d2s < dist_thred1]) / len(dist_d2s)
+        precision_2 = len(dist_d2s[dist_d2s < dist_thred2]) / len(dist_d2s)
+
+        pbar.update(1)
+        pbar.set_description('compute stl2data')
+        ground_plane = loadmat(f'{args.dataset_dir}/ObsMask/Plane{args.scan}.mat')['P']
+
+        stl_hom = np.concatenate([stl, np.ones_like(stl[:, :1])], -1)
+        above = (ground_plane.reshape((1, 4)) * stl_hom).sum(-1) > 0
+
+        stl_above = stl[above]
+
+        nn_engine.fit(data_in)
+        dist_s2d, idx_s2d = nn_engine.kneighbors(stl_above, n_neighbors=1, return_distance=True)
+        mean_s2d = dist_s2d[dist_s2d < max_dist].mean()
+
+        recall_1 = len(dist_s2d[dist_s2d < dist_thred1]) / len(dist_s2d)
+        recall_2 = len(dist_s2d[dist_s2d < dist_thred2]) / len(dist_s2d)
+
+        pbar.update(1)
+        pbar.set_description('visualize error')
+        vis_dist = args.visualize_threshold
+        R = np.array([[1, 0, 0]], dtype=np.float64)
+        G = np.array([[0, 1, 0]], dtype=np.float64)
+        B = np.array([[0, 0, 1]], dtype=np.float64)
+        W = np.array([[1, 1, 1]], dtype=np.float64)
+        data_color = np.tile(B, (data_down.shape[0], 1))
+        data_alpha = dist_d2s.clip(max=vis_dist) / vis_dist
+        data_color[np.where(inbound)[0][grid_inbound][in_obs]] = R * data_alpha + W * (1 - data_alpha)
+        data_color[np.where(inbound)[0][grid_inbound][in_obs][dist_d2s[:, 0] >= max_dist]] = G
+        write_vis_pcd(f'{args.vis_out_dir}/vis_{args.scan:03}_d2gt.ply', data_down, data_color)
+        stl_color = np.tile(B, (stl.shape[0], 1))
+        stl_alpha = dist_s2d.clip(max=vis_dist) / vis_dist
+        stl_color[np.where(above)[0]] = R * stl_alpha + W * (1 - stl_alpha)
+        stl_color[np.where(above)[0][dist_s2d[:, 0] >= max_dist]] = G
+        write_vis_pcd(f'{args.vis_out_dir}/vis_{args.scan:03}_gt2d.ply', stl, stl_color)
+
+        pbar.update(1)
+        pbar.set_description('done')
+        pbar.close()
+        over_all = (mean_d2s + mean_s2d) / 2
+
+        fscore_1 = 2 * precision_1 * recall_1 / (precision_1 + recall_1 + 1e-6)
+        fscore_2 = 2 * precision_2 * recall_2 / (precision_2 + recall_2 + 1e-6)
+
+        print(f'over_all: {over_all}; mean_d2gt: {mean_d2s}; mean_gt2d: {mean_s2d}.')
+        print(f'precision_1mm: {precision_1};  recall_1mm: {recall_1};  fscore_1mm: {fscore_1}')
+        print(f'precision_2mm: {precision_2};  recall_2mm: {recall_2};  fscore_2mm: {fscore_2}')
+
+        pparent, stem, ext = get_path_components(args.data)
+        if args.log is None:
+            path_log = os.path.join(pparent, 'eval_result.txt')
+        else:
+            path_log = args.log
+        with open(path_log, 'w+') as fLog:
+            fLog.write(f'over_all {np.round(over_all, 3)} '
+                       f'mean_d2gt {np.round(mean_d2s, 3)} '
+                       f'mean_gt2d {np.round(mean_s2d, 3)} \n'
+                       f'precision_1mm {np.round(precision_1, 3)} '
+                       f'recall_1mm {np.round(recall_1, 3)} '
+                       f'fscore_1mm {np.round(fscore_1, 3)} \n'
+                       f'precision_2mm {np.round(precision_2, 3)} '
+                       f'recall_2mm {np.round(recall_2, 3)} '
+                       f'fscore_2mm {np.round(fscore_2, 3)} \n'
+                       f'[{stem}] \n')

SparseNeuS_demo_v1/exp/lod0/checkpoint_trash/ckpt_285000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:763c2a4934928cc089342905ba61481d6f9efc977b9729d7fc2d3eae4f0e1f9b
+size 5310703