import os
import cv2
import copy
import trimesh
import pyrender
import numpy as np
from tqdm import tqdm
import open3d as o3d
import matplotlib.pyplot as plt
def load_cam_mvsnet(file, interval_scale=1):
""" read camera txt file """
cam = np.zeros((2, 4, 4))
words = file.read().split()
# read extrinsic
for i in range(0, 4):
for j in range(0, 4):
extrinsic_index = 4 * i + j + 1
cam[0][i][j] = words[extrinsic_index]
# read intrinsic
for i in range(0, 3):
for j in range(0, 3):
intrinsic_index = 3 * i + j + 18
cam[1][i][j] = words[intrinsic_index]
if len(words) == 29:
cam[1][3][0] = words[27]
cam[1][3][1] = float(words[28]) * interval_scale
cam[1][3][2] = 192
cam[1][3][3] = cam[1][3][0] + cam[1][3][1] * cam[1][3][2]
elif len(words) == 30:
cam[1][3][0] = words[27]
cam[1][3][1] = float(words[28]) * interval_scale
cam[1][3][2] = words[29]
cam[1][3][3] = cam[1][3][0] + cam[1][3][1] * cam[1][3][2]
elif len(words) == 31:
cam[1][3][0] = words[27]
cam[1][3][1] = float(words[28]) * interval_scale
cam[1][3][2] = words[29]
cam[1][3][3] = words[30]
else:
cam[1][3][0] = 0
cam[1][3][1] = 0
cam[1][3][2] = 0
cam[1][3][3] = 0
extrinsic = cam[0].astype(np.float32)
intrinsic = cam[1].astype(np.float32)
return intrinsic, extrinsic
def render_depth_maps(mesh, poses, K, H, W, near=0.01, far=5.0):
"""
:param mesh: Mesh to be rendered
:param poses: list of camera poses (c2w under OpenGL convention)
:param K: camera intrinsics [3, 3]
:param W: width of image plane
:param H: height of image plane
:param near: near clip
:param far: far clip
:return: list of rendered depth images [H, W]
"""
mesh = pyrender.Mesh.from_trimesh(mesh)
scene = pyrender.Scene()
scene.add(mesh)
camera = pyrender.IntrinsicsCamera(fx=K[0, 0], fy=K[1, 1], cx=K[0, 2], cy=K[1, 2], znear=near, zfar=far)
camera_node = pyrender.Node(camera=camera, matrix=np.eye(4))
scene.add_node(camera_node)
renderer = pyrender.OffscreenRenderer(W, H)
render_flags = pyrender.RenderFlags.OFFSCREEN | pyrender.RenderFlags.DEPTH_ONLY
depth_maps = []
for pose in poses:
scene.set_pose(camera_node, pose)
depth = renderer.render(scene, render_flags)
depth_maps.append(depth)
return depth_maps
def render_dtu_scenes(path_to_scan, method='furu'):
path_to_cameras = os.path.join(path_to_scan, 'cams')
path_to_images = os.path.join(path_to_scan, 'images')
scan_id = int(''.join(filter(str.isdigit, os.path.basename(path_to_scan))))
if method is not None:
path_to_depths = os.path.join(path_to_scan, f'depths_{method}')
path_to_mesh = os.path.join(path_to_scan, f'{method}{scan_id:03d}_l3_surf_11_trim_8.ply')
else:
path_to_depths = os.path.join(path_to_scan, 'depths')
path_to_mesh = os.path.join(path_to_scan, f'{scan_id:03d}_pcd.ply')
#path_to_mesh = os.path.join(path_to_scan, 'stl001_total.ply')
if not os.path.exists(path_to_depths):
os.makedirs(path_to_depths)
mesh = trimesh.load_mesh(path_to_mesh)
frames = sorted(os.listdir(path_to_images))
img = cv2.imread(os.path.join(path_to_images, frames[0]))
H, W, _ = img.shape
for i, frame in enumerate(frames):
campath = os.path.join(path_to_cameras, frame.replace('.jpg', '_cam.txt'))
print(campath)
cur_intrinsics, camera_pose = load_cam_mvsnet(open(campath, 'r'))
camera_pose = np.linalg.inv(camera_pose)
camera_pose[:, 1:3] *= -1.0
print(cur_intrinsics)
depth = render_depth_maps(mesh, [camera_pose], cur_intrinsics, H, W, near=0.01, far=5000.)[0]
# plt.imshow(depth)
# plt.show()
# Save depth map
#cv2.imwrite(os.path.join(path_to_depths, frame.replace('.jpg', '.png')), depth)
# depth_16bit = (depth).astype(np.uint16) # Scale to millimeters
# # Save depth map as 16-bit PNG
# depth_filename = os.path.join(path_to_depths, frame.replace('.jpg', '.png'))
# cv2.imwrite(depth_filename, depth_16bit)
# depth = depth.astype(np.float32)
depth_filename = os.path.join(path_to_depths, frame.replace('.jpg', '.npy'))
np.save(depth_filename, depth)
def get_dtu_mask(path_to_scan, method='furu'):
if method is not None:
path_to_depths = os.path.join(path_to_scan, f'depths_{method}')
path_to_masks = os.path.join(path_to_scan, f'masks_{method}')
else:
path_to_depths = os.path.join(path_to_scan, 'depths')
path_to_masks = os.path.join(path_to_scan, 'masks')
if not os.path.exists(path_to_masks):
os.makedirs(path_to_masks)
frames = sorted(os.listdir(path_to_depths))
for i, frame in enumerate(frames):
depth_filename = os.path.join(path_to_depths, frame)
depth = np.load(depth_filename)
mask_path = os.path.join(path_to_masks, frame.replace('.npy', '.png'))
mask = np.ones_like(depth) * 255
mask[depth == 0] = 0
mask[depth>900] = 0
cv2.imwrite(mask_path, mask)
def get_mesh_from_ply(path_to_scan, depth=9, density_thresh=0.1):
scan_id = int(''.join(filter(str.isdigit, os.path.basename(path_to_scan))))
path_to_ply = os.path.join(path_to_scan, f'stl{scan_id:03d}_total.ply')
pcd = o3d.io.read_point_cloud(path_to_ply)
mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
pcd, depth=depth)
vertices_to_remove = densities < np.quantile(densities, density_thresh)
new_mesh = copy.deepcopy(mesh)
new_mesh = copy.deepcopy(mesh)
new_mesh.remove_vertices_by_mask(vertices_to_remove)
# save mesh
path_to_mesh = os.path.join(path_to_scan, f'{scan_id:03d}_pcd.ply')
o3d.io.write_triangle_mesh(path_to_mesh, new_mesh)
path_to_dtu = './data/dtu_test'
scans = sorted(os.listdir(path_to_dtu))
for scan in tqdm(scans):
print(f"Processing {scan}")
path_to_scan = os.path.join(path_to_dtu, scan)
get_mesh_from_ply(path_to_scan)
render_dtu_scenes(path_to_scan, method=None)
#get_dtu_mask(path_to_scan, None)