code pruning

SparseNeuS_demo_v1/confs/blender_general_lod1_val_new.conf

@@ -1,137 +0,0 @@
-# - 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

@@ -18,8 +18,6 @@ dataset {
   valpath = ../ # !!! where you store the validation data
   testpath = ../
 
-
-
   imgScale_train = 1.0
   imgScale_test = 1.0
   nviews = 5

SparseNeuS_demo_v1/data/blender.py

@@ -1,340 +0,0 @@
-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

@@ -1,432 +0,0 @@
-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

@@ -1,427 +0,0 @@
-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

@@ -1,427 +0,0 @@
-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

@@ -1,412 +0,0 @@
-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

@@ -1,406 +0,0 @@
-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

@@ -1,411 +0,0 @@
-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

@@ -1,480 +0,0 @@
-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

@@ -1,476 +0,0 @@
-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

@@ -1,449 +0,0 @@
-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

@@ -1,396 +0,0 @@
-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

@@ -1,396 +0,0 @@
-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

@@ -1,446 +0,0 @@
-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

@@ -1,439 +0,0 @@
-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

@@ -1,470 +0,0 @@
-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

@@ -1,395 +0,0 @@
-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

@@ -1,418 +0,0 @@
-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

@@ -1,399 +0,0 @@
-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

@@ -1,393 +0,0 @@
-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

@@ -1,395 +0,0 @@
-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

@@ -1,432 +0,0 @@
-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

@@ -1,386 +0,0 @@
-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

@@ -1,410 +0,0 @@
-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

@@ -1,411 +0,0 @@
-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

@@ -1,6 +1,6 @@
 from torch.utils.data import Dataset
-from utils.misc_utils import read_pfm
 import os
+import json
 import numpy as np
 import cv2
 from PIL import Image
@@ -9,12 +9,7 @@ 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):
     """
@@ -73,10 +68,6 @@ class BlenderPerView(Dataset):
         # 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
@@ -102,8 +93,6 @@ class BlenderPerView(Dataset):
         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!'
@@ -130,9 +119,6 @@ class BlenderPerView(Dataset):
             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,
@@ -160,11 +146,6 @@ class BlenderPerView(Dataset):
         # 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
@@ -198,9 +179,8 @@ class BlenderPerView(Dataset):
         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]
+        for image_idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_idx]
             c2w = pose @ self.blender2opencv
             self.c2ws.append(c2w)
             self.w2cs.append(np.linalg.inv(c2w))
@@ -224,7 +204,6 @@ class BlenderPerView(Dataset):
         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)
@@ -258,7 +237,6 @@ class BlenderPerView(Dataset):
 
         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
@@ -312,7 +290,6 @@ class BlenderPerView(Dataset):
             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)
@@ -360,7 +337,6 @@ class BlenderPerView(Dataset):
         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)

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data3_1.py

@@ -1,414 +0,0 @@
-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

@@ -1,465 +0,0 @@
-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

@@ -1,419 +0,0 @@
-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

@@ -1,420 +0,0 @@
-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

@@ -1,428 +0,0 @@
-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

@@ -1,420 +0,0 @@
-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

@@ -1,417 +0,0 @@
-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

@@ -1,388 +0,0 @@
-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

@@ -1,389 +0,0 @@
-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

@@ -1,395 +0,0 @@
-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

@@ -1,419 +0,0 @@
-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

@@ -1,93 +0,0 @@
-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

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

SparseNeuS_demo_v1/data/dtu/lists/train.txt

@@ -1,75 +0,0 @@
-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

@@ -1,278 +0,0 @@
-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

@@ -1,376 +0,0 @@
-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

@@ -1,6 +1,5 @@
 import numpy as np
 import torch
-import pdb
 
 
 def rigid_transform(xyz, transform):

SparseNeuS_demo_v1/evaluation/clean_mesh.py

@@ -1,283 +0,0 @@
-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

@@ -1,369 +0,0 @@
-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_runner_generic_blender_val.py

@@ -1,15 +1,12 @@
-import torch
-import torch.nn.functional as F
-from torch.utils.data import DataLoader
-import argparse
 import os
 import logging
+import argparse
 import numpy as np
-import cv2 as cv
-import trimesh
 from shutil import copyfile
+import torch
+from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
-from icecream import ic
+from rich import print
 from tqdm import tqdm
 from pyhocon import ConfigFactory
 
@@ -17,23 +14,12 @@ import sys
 sys.path.append(os.path.dirname(__file__))
 
 from models.fields import SingleVarianceNetwork
-
 from models.featurenet import FeatureNet
-
 from models.trainer_generic import GenericTrainer
-
 from models.sparse_sdf_network import SparseSdfNetwork
-
 from models.rendering_network import GeneralRenderingNetwork
-
-from datetime import datetime
-
-from data.dtu_general import MVSDatasetDtuPerView
-
-from utils.training_utils import tocuda
 from data.blender_general_narrow_all_eval_new_data import BlenderPerView
 
-from termcolor import colored
 
 from datetime import datetime
 
@@ -45,14 +31,14 @@ class Runner:
         self.device = torch.device('cuda:%d' % local_rank)
         # self.device = torch.device('cuda')
         self.num_devices = torch.cuda.device_count()
-        self.is_continue = is_continue
+        self.is_continue = is_continue or (mode == "export_mesh")
         self.is_restore = is_restore
         self.restore_lod0 = restore_lod0
         self.mode = mode
         self.model_list = []
         self.logger = logging.getLogger('exp_logger')
 
-        print(colored("detected %d GPUs" % self.num_devices, "red"))
+        print("detected %d GPUs" % self.num_devices)
 
         self.conf_path = conf_path
         self.conf = ConfigFactory.parse_file(conf_path)
@@ -63,7 +49,7 @@ class Runner:
         else:
             self.base_exp_dir = self.conf['general.base_exp_dir']
         self.conf['general.base_exp_dir'] = self.base_exp_dir # jha use this when testing
-        print(colored("base_exp_dir:  " + self.base_exp_dir, 'yellow'))
+        print("base_exp_dir: " + self.base_exp_dir)
         os.makedirs(self.base_exp_dir, exist_ok=True)
         self.iter_step = 0
         self.val_step = 0
@@ -121,12 +107,11 @@ class Runner:
                 **self.conf['model.rendering_network_lod1']).to(self.device)
         if self.mode == 'export_mesh' or self.mode == 'val':
             # base_exp_dir_to_store = os.path.join(self.base_exp_dir, '{:%Y_%m_%d_%H_%M_%S}'.format(datetime.now()))
-            print("save mesh to:", os.path.join("../", args.specific_dataset_name))
             base_exp_dir_to_store = os.path.join("../", args.specific_dataset_name) #"../gradio_tmp" # MODIFIED
         else:
             base_exp_dir_to_store = self.base_exp_dir
 
-        print(colored(f"Store in: {base_exp_dir_to_store}", "blue"))
+        print(f"Store in: {base_exp_dir_to_store}")
         # Renderer model
         self.trainer = GenericTrainer(
             self.rendering_network_outside,
@@ -149,7 +134,7 @@ class Runner:
 
         # Load checkpoint
         latest_model_name = None
-        if is_continue:
+        if self.is_continue:
             model_list_raw = os.listdir(os.path.join(self.base_exp_dir, 'checkpoints'))
             model_list = []
             for model_name in model_list_raw:
@@ -238,7 +223,7 @@ class Runner:
         epochs = int(1 + res_step // len(dataloader))
 
         self.adjust_learning_rate()
-        print(colored("starting training learning rate: {:.5f}".format(self.optimizer.param_groups[0]['lr']), "yellow"))
+        print("starting training learning rate: {:.5f}".format(self.optimizer.param_groups[0]['lr']))
 
         background_rgb = None
         if self.use_white_bkgd:
@@ -247,7 +232,7 @@ class Runner:
 
         for epoch_i in range(epochs):
 
-            print(colored("current epoch %d" % epoch_i, 'red'))
+            print("current epoch %d" % epoch_i)
             dataloader = tqdm(dataloader)
 
             for batch in dataloader:
@@ -365,8 +350,8 @@ class Runner:
                             losses_lod1['color_fine_loss'].mean() if losses_lod1 is not None else 0,
                             self.optimizer.param_groups[0]['lr']))
 
-                    print(colored('alpha_inter_ratio_lod0 = {:.4f} alpha_inter_ratio_lod1 = {:.4f}\n'.format(
-                        alpha_inter_ratio_lod0, alpha_inter_ratio_lod1), 'green'))
+                    print('alpha_inter_ratio_lod0 = {:.4f} alpha_inter_ratio_lod1 = {:.4f}\n'.format(
+                        alpha_inter_ratio_lod0, alpha_inter_ratio_lod1))
 
                     if losses_lod0 is not None:
                         # print("[TEST]: weights_sum in print", losses_lod0['weights_sum'].mean())
@@ -469,7 +454,7 @@ class Runner:
                     # 3. load the new state dict
                     network.load_state_dict(pretrained_dict)
                 except:
-                    print(colored(comment + " load fails", 'yellow'))
+                    print(comment + " load fails")
 
         checkpoint = torch.load(os.path.join(self.base_exp_dir, 'checkpoints', checkpoint_name),
                                 map_location=self.device)
@@ -497,7 +482,7 @@ class Runner:
             try:
                 self.optimizer.load_state_dict(checkpoint['optimizer'])
             except:
-                print(colored("load optimizer fails", "yellow"))
+                print("load optimizer fails")
             self.iter_step = checkpoint['iter_step']
             self.val_step = checkpoint['val_step'] if 'val_step' in checkpoint.keys() else 0
 
@@ -532,15 +517,10 @@ class Runner:
         torch.save(checkpoint,
                    os.path.join(self.base_exp_dir, 'checkpoints', 'ckpt_{:0>6d}.pth'.format(self.iter_step)))
 
-    def validate(self, idx=-1, resolution_level=-1):
+    def validate(self, resolution_level=-1):
         # validate image
-
-        ic(self.iter_step, idx)
+        print("iter_step: ", self.iter_step)
         self.logger.info('Validate begin')
-
-        if idx < 0:
-            idx = self.val_step
-            # idx = np.random.randint(len(self.val_dataset))
         self.val_step += 1
 
         try:
@@ -576,16 +556,9 @@ class Runner:
         )
 
 
-    def export_mesh(self, idx=-1, resolution_level=-1):
-        # validate image
-
-        ic(self.iter_step, idx)
+    def export_mesh(self, resolution_level=-1):
+        print("iter_step: ", self.iter_step)
         self.logger.info('Validate begin')
-        import time 
-        start1 = time.time()
-        if idx < 0:
-            idx = self.val_step
-            # idx = np.random.randint(len(self.val_dataset))
         self.val_step += 1
 
         try:
@@ -598,7 +571,6 @@ class Runner:
 
         background_rgb = None
         if self.use_white_bkgd:
-            # background_rgb = torch.ones([1, 3]).to(self.device)
             background_rgb = 1.0
 
         batch['batch_idx'] = torch.tensor([x for x in range(self.batch_size)])
@@ -609,8 +581,6 @@ class Runner:
         else:
             alpha_inter_ratio_lod0 = 1.
         alpha_inter_ratio_lod1 = self.get_alpha_inter_ratio(self.anneal_start_lod1, self.anneal_end_lod1)
-        end1 = time.time()
-        print("time for getting data", end1 - start1)
         self.trainer(
             batch,
             background_rgb=background_rgb,

SparseNeuS_demo_v1/models/patch_projector.py

@@ -208,4 +208,4 @@ def normalize(flow, h, w, clamp=None):
 
 def build_patch_offset(h_patch_size):
     offsets = torch.arange(-h_patch_size, h_patch_size + 1)
-    return torch.stack(torch.meshgrid(offsets, offsets)[::-1], dim=-1).view(1, -1, 2)  # nb_pixels_patch * 2
+    return torch.stack(torch.meshgrid(offsets, offsets, indexing="ij")[::-1], dim=-1).view(1, -1, 2)  # nb_pixels_patch * 2