Skip to content

Commit

Permalink
add depth
Browse files Browse the repository at this point in the history
  • Loading branch information
@wangyida
wangyida committed Feb 27, 2024
1 parent 54c4e4c commit 00d5e19
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 52 deletions.
6 changes: 3 additions & 3 deletions configs/neuralangelo-hmvs-eagle-wbg.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ dataset:
img_downscale: 4 # better with more than 1000 x 1000 resolution
cam_downscale: 200 # if false, to rescale the cameras/pts positions automatically
repose: true # false
up_est_method: ground # camera/ground/z-axis
up_est_method: ground # no-change/camera/ground/z-axis
center_est_method: point # camera/point
n_test_traj_steps: 60
apply_mask: false
Expand Down Expand Up @@ -52,8 +52,8 @@ model:
xyz_encoding_config:
otype: ProgressiveBandHashGrid
n_levels: 16
n_features_per_level: 4 # 8
log2_hashmap_size: 20 # 22
n_features_per_level: 8
log2_hashmap_size: 22
base_resolution: 32
per_level_scale: 1.3195079107728942
include_xyz: true
Expand Down
8 changes: 4 additions & 4 deletions configs/neus-hmvs-eagle-wbg.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,12 +8,12 @@ dataset:
pose_path: /home/yidaw/Documents/datasets/Buildings/Building/Tile_+002_+002_AT.bin
pcd_path: /home/yidaw/Documents/datasets/Buildings/Building/building_interest.ply # Bridge/Tile_+004_+008_AT.ply
mesh_outpath: false
img_downscale: 4 # better with more than 1000 x 1000 resolution
cam_downscale: 200 # 250 # if false, to rescale the cameras/pts positions automatically
img_downscale: 2 # 4 # better with more than 1000 x 1000 resolution
cam_downscale: 200 # if false, to rescale the cameras/pts positions automatically
repose: true # false
up_est_method: ground # no-change/camera/ground/z-axis
up_est_method: ground # no-change/camera/ground/z-axis
center_est_method: point # camera/point
n_test_traj_steps: 60 # 16
n_test_traj_steps: 60
apply_mask: false
apply_depth: false
load_data_on_gpu: false
Expand Down
106 changes: 61 additions & 45 deletions datasets/hmvs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -325,8 +325,14 @@ def setup(self, config, split):
self.config.pose_path)
except:
all_c2w, intrinsics = np.load(os.path.join(self.config.root_dir,
'extrinsics.npy')), np.load(os.path.join(self.config.root_dir,
'intrinsics.npy'))
'extrinsics.npy'), allow_pickle=True), np.load(os.path.join(self.config.root_dir,
'intrinsics.npy'), allow_pickle=True)
try:
all_c2w = all_c2w[()]['FRONT']
intrinsics = intrinsics[()]['FRONT']
print(colored('Using front camera alone', 'blue'))
except:
print(colored('Using all available cameras', 'blue'))
fns = sorted(os.listdir(os.path.join(self.config.root_dir,
'images')))
fns = [os.path.join(self.config.root_dir, 'images', fn) for fn in fns]
Expand Down Expand Up @@ -371,7 +377,7 @@ def setup(self, config, split):

mesh_dir = os.path.join(self.config.root_dir, 'mesh_exp')
os.makedirs(mesh_dir, exist_ok=True)
mesh_poisson_path = os.path.join(mesh_dir, 'layout_mesh.ply')
mesh_poisson_path = os.path.join(mesh_dir, 'layout_mesh_ps.ply')
# Poisson surface on top of given GT points
if not os.path.exists(mesh_poisson_path):
print(colored(
Expand All @@ -384,8 +390,9 @@ def setup(self, config, split):
print('run Poisson surface reconstruction')
with o3d.utility.VerbosityContextManager(
o3d.utility.VerbosityLevel.Debug) as cm:
# Poisson
mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
pcd, depth=9)
pcd, depth=10, linear_fit=True)
o3d.io.write_triangle_mesh(mesh_poisson_path, mesh)
# Open3D mesh
mesh_o3d = o3d.t.geometry.TriangleMesh.from_legacy(o3d.io.read_triangle_mesh(mesh_poisson_path))
Expand Down Expand Up @@ -467,24 +474,26 @@ def setup(self, config, split):
img = img.to(
self.rank
) if self.config.load_data_on_gpu else img.cpu()
vis_mask = torch.ones_like(img[..., 0], device=img.device)
# NOTE: Lixiang rear camera should be ignored
# NOTE: visual masks
vis_mask_path = os.path.join(
vis_mask_dir, 'mask' + fns[i].split("/")[-1][6:-3] + 'png')
if not os.path.exists(vis_mask_path):
if not os.path.exists(vis_mask_path) and os.path.exists(vis_mask_path.replace(f"{vis_mask_dir}",f"{vis_mask_ori_dir}")):
os.makedirs(os.path.dirname(vis_mask_path), exist_ok=True)
vis_mask_path = vis_mask_path.replace(
f"{vis_mask_dir}",
f"{vis_mask_ori_dir}")
vis_mask = Image.open(vis_mask_path)
if vis_mask.size[0] != w or vis_mask.size[1] != h:
vis_mask = vis_mask.resize(img_wh, Image.NEAREST)
vis_mask_path = vis_mask_path.replace(
f"{vis_mask_ori_dir}",
f"{vis_mask_dir}")
vis_mask.save(vis_mask_path)
vis_mask = TF.to_tensor(vis_mask)[0]
if os.path.exists(vis_mask_path):
vis_mask = Image.open(vis_mask_path)
if vis_mask.size[0] != w or vis_mask.size[1] != h:
vis_mask = vis_mask.resize(img_wh, Image.NEAREST)
vis_mask_path = vis_mask_path.replace(
f"{vis_mask_ori_dir}",
f"{vis_mask_dir}")
vis_mask.save(vis_mask_path)
vis_mask = TF.to_tensor(vis_mask)[0]
else:
vis_mask = torch.ones_like(img[..., 0], device=img.device)
# The unusual
"""
if (i + 1) % 6 == 0:
Expand All @@ -509,35 +518,6 @@ def setup(self, config, split):
dynamic_mask.save(dynamic_mask_path)
vis_mask *= (1 - TF.to_tensor(dynamic_mask)[0])

# NOTE: foreground masks
if self.apply_mask:
mask_path = os.path.join(
mask_dir, fns[i].split("/")[-1][:-3] + 'png')
# mask_paths = list(filter(os.path.exists, mask_paths))
# assert len(mask_paths) == 1
if not os.path.exists(mask_path):
os.makedirs(os.path.join(
self.config.root_dir,
f"sky_mask_{self.config.img_downscale}"),
exist_ok=True)
mask_path = os.path.join(
mask_ori_dir,
fns[i].split("/")[-1][:-3] + 'png')
mask = Image.open(mask_path) # (H, W, 1)
if mask.size[0] != w or mask.size[1] != h:
mask = mask.resize(img_wh, Image.NEAREST)
mask.save(
os.path.join(
self.config.root_dir,
f"sky_mask_{self.config.img_downscale}",
fns[i].split("/")[-1][:-3] + 'png'))
mask = TF.to_tensor(mask)[0] # TF.pil_to_tensor does not rescale the input PIL mask
mask_fg = torch.ones_like(mask, device=img.device)
mask_fg[mask == 1] = 0 # check whether we use 1 or 255
mask = mask_fg
else:
mask = torch.ones_like(img[..., 0], device=img.device)

depth_folder = 'lidar_depth'
if self.apply_depth and os.path.exists(fns[i].replace(f"{img_folder}", f"/{depth_folder}")):
depth_format = '.npy' # 'tif' 'npy' 'pth' 'png'
Expand Down Expand Up @@ -619,12 +599,48 @@ def setup(self, config, split):
depth_rast = Image.fromarray(rays_rast['t_hit'].numpy())
depth_rast = TF.pil_to_tensor(depth_rast).permute(
1, 2, 0) / self.config.cam_downscale
inf_mask = (depth_rast == float("Inf"))
depth_rast[depth_rast == float("Inf")] = 0
depth_rast = depth_rast.to(
self.rank
) if self.config.load_data_on_gpu else depth_rast.cpu()
depth_rast_mask = (depth_rast > 0.0).to(bool) # trim points outside the contraction box off

# NOTE: foreground masks
if self.apply_mask:
mask_path = os.path.join(
mask_dir, fns[i].split("/")[-1][:-3] + 'png')
# mask_paths = list(filter(os.path.exists, mask_paths))
# assert len(mask_paths) == 1
if not os.path.exists(mask_path):
os.makedirs(os.path.join(
self.config.root_dir,
f"sky_mask_{self.config.img_downscale}"),
exist_ok=True)
mask_path = os.path.join(
mask_ori_dir,
fns[i].split("/")[-1][:-3] + 'png')
if os.path.exists(mask_path):
mask = Image.open(mask_path) # (H, W, 1)
if mask.size[0] != w or mask.size[1] != h:
mask = mask.resize(img_wh, Image.NEAREST)
mask.save(
os.path.join(
self.config.root_dir,
f"sky_mask_{self.config.img_downscale}",
fns[i].split("/")[-1][:-3] + 'png'))
mask = TF.to_tensor(mask)[0] # TF.pil_to_tensor does not rescale the input PIL mask
mask_fg = torch.ones_like(mask, device=img.device)
mask_fg[mask == 1] = 0 # check whether we use 1 or 255
mask = mask_fg
elif inf_mask is not None:
mask = (1 - inf_mask.to(int)).to(bool)
else:
print(colored('Foreground mask not available, need to disable apply_mask', 'red'))
else:
mask = torch.ones_like(img[..., 0], device=img.device)


all_vis_masks.append(vis_mask)
all_fg_masks.append(mask) # (h, w)
all_images.append(img)
Expand All @@ -635,7 +651,7 @@ def setup(self, config, split):
o3d.io.write_point_cloud(os.path.join(mesh_dir, 'layout_depth_clt.ply'), pts_clt)

directions = torch.stack(directions, dim=0)
all_c2w = torch.tensor(all_c2w)[:, :3]
all_c2w = torch.tensor(all_c2w)[:, :3].float()
all_c2w[:, :, 1:3] *= -1. # COLMAP => OpenGL

if self.config.repose:
Expand Down Expand Up @@ -756,7 +772,7 @@ def setup(self, config, split):
exit(1)
"""

self.all_c2w = self.all_c2w.float().to(self.rank)
self.all_c2w = self.all_c2w.to(self.rank)
if self.config.load_data_on_gpu:
self.all_images = self.all_images.to(self.rank)
self.all_fg_masks = self.all_fg_masks.to(self.rank)
Expand Down

0 comments on commit 00d5e19

Please sign in to comment.