spann3r/tools/render_dtu.py

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)