evaluations/fidkid-pytorch/sample_surface_objv.py

# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.

#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from pdb import set_trace as st
import json
import math
import pytorch3d.ops
import numpy as np
import os
import argparse
import multiprocessing as mp
from multiprocessing import Pool
import trimesh
import tqdm
import torch
import nvdiffrast.torch as dr
import kaolin as kal
import glob
import ipdb

parser = argparse.ArgumentParser(description='sample surface points from mesh')
parser.add_argument(
    '--n_proc', type=int, default=8,
    help='Number of processes to run in parallel'
         '(0 means sequential execution).')
parser.add_argument(
    '--n_points', type=int, default=4096,
    help='Number of points to sample per model.')
parser.add_argument(
    '--n_views', type=int, default=100,
    help='Number of views per model.')
parser.add_argument(
    '--image_height', type=int, default=640,
    help='Depth image height.')
parser.add_argument(
    '--image_width', type=int, default=640,
    help='Depth image width.')
parser.add_argument(
    '--focal_length_x', type=float, default=640,
    help='Focal length in x direction.')
parser.add_argument(
    '--focal_length_y', type=float, default=640,
    help='Focal length in y direction.')
parser.add_argument(
    '--principal_point_x', type=float, default=320,
    help='Principal point location in x direction.')
parser.add_argument(
    '--principal_point_y', type=float, default=320,
    help='Principal point location in y direction.')
parser.add_argument("--shape_root", type=str, default='/mnt/petrelfs/caoziang/3D_generation/Checkpoint_all/diffusion_shapenet_testmodel27_omni_ablation2/ddpm_5000/test', help="path to the save resules shapenet dataset")
parser.add_argument("--save_root", type=str, default='/mnt/petrelfs/caoziang/3D_generation/Checkpoint_all/diffusion_shapenet_testmodel27_omni_ablation2/ddpm_vis_ab2surface', help="path to the split shapenet dataset")

options = parser.parse_args()

# create array for inverse mapping
coordspx2 = np.stack(np.nonzero(np.ones((options.image_height, options.image_width))), -1).astype(np.float32)
coordspx2 = coordspx2[:, ::-1]
fusion_intrisics = np.array(
    [
        [options.focal_length_x, 0, options.principal_point_x],
        [0, options.focal_length_y, options.principal_point_y],
        [0, 0, 1]
    ])
# glctx = dr.RasterizeGLContext() # EGL/egl.h: No such file or directory
glctx = dr.RasterizeCudaContext()


def CalcLinearZ(depth):
    # depth = depth * 2 - 1
    zFar = 100.0
    zNear = 0.1
    linear = zNear / (zFar - depth * (zFar - zNear)) * zFar
    return linear


def projection_cv_new(fx, fy, cx, cy, width, height, n=1.0, f=50.0):
    return np.array(
        [[-2 * fx / width, 0.0, (width - 2 * cx) / width, 0.0],
         [0.0, -2 * fy / height, (height - 2 * cy) / height, 0.0],
         [0.0, 0.0, (-f - n) / (f - n), -2.0 * f * n / (f - n)],
         [0.0, 0.0, -1.0, 0.0]])


def interpolate(attr, rast, attr_idx, rast_db=None):
    return dr.interpolate(
        attr.contiguous(), rast, attr_idx, rast_db=rast_db,
        diff_attrs=None if rast_db is None else 'all')


def render_nvdiffrast(v_pos, tris, T_bx4x4):
    # T_bx4x4 - world to cam
    proj = projection_cv_new(
        fx=options.focal_length_x, fy=options.focal_length_y, cx=options.principal_point_x,
        cy=options.principal_point_y,
        width=options.image_width, height=options.image_height, n=0.1, f=100.0)

    fix = torch.eye(4, dtype=torch.float32, device='cuda')
    fix[2, 2] = -1
    fix[1, 1] = -1
    fix[0, 0] = -1
    fix = fix.unsqueeze(0).repeat(T_bx4x4.shape[0], 1, 1)

    proj = torch.tensor(proj, dtype=torch.float32, device='cuda').unsqueeze(0).repeat(T_bx4x4.shape[0], 1, 1)
    T_world_cam_bx4x4 = torch.bmm(fix, T_bx4x4)
    mvp = torch.bmm(proj, T_world_cam_bx4x4)
    v_pos_clip = torch.matmul(
        torch.nn.functional.pad(v_pos, pad=(0, 1), mode='constant', value=1.0),
        torch.transpose(mvp, 1, 2))
    rast, db = dr.rasterize(
        glctx, torch.tensor(v_pos_clip, dtype=torch.float32, device='cuda'), tris.int(),
        (options.image_height, options.image_width))

    v_pos_cam = torch.matmul(
        torch.nn.functional.pad(v_pos, pad=(0, 1), mode='constant', value=1.0),
        torch.transpose(T_world_cam_bx4x4, 1, 2))
    gb_pos_cam, _ = interpolate(v_pos_cam, rast, tris.int())
    depth_maps = gb_pos_cam[..., 2].abs()
    return depth_maps


def as_mesh(scene_or_mesh):
    """
    Convert a possible scene to a mesh.

    If conversion occurs, the returned mesh has only vertex and face data.
    """
    if isinstance(scene_or_mesh, trimesh.Scene):
        if len(scene_or_mesh.geometry) == 0:
            mesh = None  # empty scene
        else:
            # we lose texture information here
            mesh = trimesh.util.concatenate(
                tuple(
                    trimesh.Trimesh(vertices=g.vertices, faces=g.faces)
                    for g in scene_or_mesh.geometry.values()))
    else:
        assert (isinstance(scene_or_mesh, trimesh.Trimesh))
        mesh = scene_or_mesh
    return mesh


def render(mesh_v, mesh_f, Rs):
    """
    Render the given mesh using the generated views.

    :param base_mesh: mesh to render
    :type base_mesh: mesh.Mesh
    :param Rs: rotation matrices
    :type Rs: [numpy.ndarray]
    :return: depth maps
    :rtype: numpy.ndarray
    """
    T_bx4x4 = torch.zeros((options.n_views, 4, 4), dtype=torch.float32, device='cuda')
    T_bx4x4[:, 3, 3] = 1
    T_bx4x4[:, 2, 3] = 1
    T_bx4x4[:, :3, :3] = torch.tensor(Rs, dtype=torch.float32, device='cuda')
    depthmaps = render_nvdiffrast(
        mesh_v,
        mesh_f, T_bx4x4)
    return depthmaps


def get_points():
    """
    :param n_points: number of points
    :type n_points: int
    :return: list of points
    :rtype: numpy.ndarray
    """

    rnd = 1.
    points = []
    offset = 2. / options.n_views
    increment = math.pi * (3. - math.sqrt(5.))

    for i in range(options.n_views):
        y = ((i * offset) - 1) + (offset / 2)
        r = math.sqrt(1 - pow(y, 2))

        phi = ((i + rnd) % options.n_views) * increment

        x = math.cos(phi) * r
        z = math.sin(phi) * r

        points.append([x, y, z])
    return np.array(points)


def get_views(semi_sphere=False):
    """
    Generate a set of views to generate depth maps from.

    :param n_views: number of views per axis
    :type n_views: int
    :return: rotation matrices
    :rtype: [numpy.ndarray]
    """

    Rs = []
    points = get_points()
    if semi_sphere:
        points[:, 2] = -np.abs(points[:, 2]) - 0.1

    for i in range(points.shape[0]):
        longitude = - math.atan2(points[i, 0], points[i, 1])
        latitude = math.atan2(points[i, 2], math.sqrt(points[i, 0] ** 2 + points[i, 1] ** 2))

        R_x = np.array(
            [[1, 0, 0],
             [0, math.cos(latitude), -math.sin(latitude)],
             [0, math.sin(latitude), math.cos(latitude)]])
        R_y = np.array(
            [[math.cos(longitude), 0, math.sin(longitude)],
             [0, 1, 0],
             [-math.sin(longitude), 0, math.cos(longitude)]])
        R = R_x @ R_y
        Rs.append(R)

    return Rs


def fusion(depthmaps, Rs):
    """
    Fuse the rendered depth maps.

    :param depthmaps: depth maps
    :type depthmaps: numpy.ndarray
    :param Rs: rotation matrices corresponding to views
    :type Rs: [numpy.ndarray]
    :return: (T)SDF
    :rtype: numpy.ndarray
    """

    # sample points inside mask
    sample_per_view = options.n_points // options.n_views
    sample_bxn = torch.zeros((options.n_views, sample_per_view), device='cuda', dtype=torch.long)
    for i in range(len(Rs)):
        mask = depthmaps[i] > 0
        valid_idx = torch.nonzero(mask.reshape(-1)).squeeze(-1)
        idx = list(range(valid_idx.shape[0]))
        np.random.shuffle(idx)
        idx = idx[:sample_per_view]
        sample_bxn[i] = torch.tensor(valid_idx[idx])

    depthmaps = torch.gather(depthmaps.reshape(options.n_views, -1), 1, sample_bxn)

    inv_Ks_bx3x3 = torch.tensor(np.linalg.inv(fusion_intrisics), dtype=torch.float32, device='cuda').unsqueeze(
        0).repeat(options.n_views, 1, 1)
    T_bx4x4 = torch.zeros((options.n_views, 4, 4), dtype=torch.float32, device='cuda')
    T_bx4x4[:, 3, 3] = 1
    T_bx4x4[:, 2, 3] = 1
    T_bx4x4[:, :3, :3] = torch.tensor(Rs, dtype=torch.float32, device='cuda')
    inv_T_bx4x4 = torch.inverse(T_bx4x4)

    tf_coords_bxpx2 = torch.tensor(coordspx2.copy(), dtype=torch.float32, device='cuda').unsqueeze(0).repeat(
        options.n_views, 1, 1)
    tf_coords_bxpx2 = torch.gather(tf_coords_bxpx2, 1, sample_bxn.unsqueeze(-1).repeat(1, 1, 2))

    tf_coords_bxpx3 = torch.cat([tf_coords_bxpx2, torch.ones_like(tf_coords_bxpx2[..., :1])], -1)
    tf_coords_bxpx3 *= depthmaps.reshape(options.n_views, -1, 1)
    tf_cam_bxpx3 = torch.bmm(inv_Ks_bx3x3, tf_coords_bxpx3.transpose(1, 2)).transpose(1, 2)
    tf_cam_bxpx4 = torch.cat([tf_cam_bxpx3, torch.ones_like(tf_cam_bxpx3[..., :1])], -1)
    tf_world_bxpx3 = torch.bmm(inv_T_bx4x4, tf_cam_bxpx4.transpose(1, 2)).transpose(1, 2)[..., :3]

    return tf_world_bxpx3.reshape(-1, 3)


def normalize(vertices, faces, normalized_scale=0.9, rotate_x=False, center_xyz=False):
    vertices = vertices.cuda()

    if center_xyz: # some mesh's center is not origin
        vertices = vertices - vertices.mean(0, keepdim=True)

    if rotate_x: # rotate along x axis for 90 degrees to match the two coordiantes
        rot_mat = torch.eye(n=3, device='cuda')
        theta = np.pi / 90  # degree
        rot_mat[1,1] = np.cos(theta)
        rot_mat[2,2] = np.cos(theta)
        rot_mat[1,2] =-np.sin(theta)
        rot_mat[2,1] = np.sin(theta)
        # ipdb.set_trace()
        vertices = rot_mat @ vertices.transpose(0,1)
        vertices = vertices.transpose(0,1)

    scale = (vertices.max(dim=0)[0] - vertices.min(dim=0)[0]).max()
    mesh_v1 = vertices / scale * normalized_scale
    mesh_f1 = faces.long().cuda()
    return mesh_v1, mesh_f1



def sample_surface_pts(path):
    # ipdb.set_trace()
    try:
        mesh_path, output_pth, debug = path
        # mesh = kal.io.obj.import_mesh(mesh_path)
        # ipdb.set_trace()
        mesh = trimesh.load(mesh_path) # fail to load ply?
        #ipdb.set_trace()
        if mesh.vertices.shape[0] == 0:
            return

        mesh_v = torch.Tensor(mesh.vertices)
        # mesh_v, mesh_f = normalize(mesh_v, torch.Tensor(mesh.faces), normalized_scale=0.9, rotate_x=True)
        # mesh_v, mesh_f = normalize(mesh_v, torch.Tensor(mesh.faces), normalized_scale=1.0, rotate_x=True, center_xyz=True)
        mesh_v, mesh_f = normalize(mesh_v, torch.Tensor(mesh.faces), normalized_scale=1.0, rotate_x=False, center_xyz=True)

        # generate camera matrices
        # Rs = get_views()
        # Rs = get_views(semi_sphere=True)
        Rs = get_views(semi_sphere=False)
        # get depth images
        depths = render(mesh_v, mesh_f, Rs)
        # project to world space
        try:
            pcd = fusion(depths, Rs)
        except:
            return

        # fps again

        pcd, fps_idx = pytorch3d.ops.sample_farthest_points(
            # torch.from_numpy(pcd).unsqueeze(0).cuda(), K=4096,
            # torch.from_numpy(pcd).unsqueeze(0).cuda(), K=4000,
            pcd.unsqueeze(0).cuda(), K=4096,
            random_start_point=True)  # B self.latent_num
        pcd = pcd[0]

        pcd = pcd.cpu().numpy()

        #np.savez(output_pth, pcd=pcd)
        #ipdb.set_trace()
        #if debug:
        pcd = trimesh.points.PointCloud(pcd)
        pcd.export(output_pth.replace('.npz', '.obj'))
    except Exception as e:
        # print('error')
        print(e, flush=True)


if __name__ == '__main__':
    mp.set_start_method('spawn')

    shapenet_root = options.shape_root
    save_root = options.save_root

    debug = True

    cmds = [] 

    # ! for baseline samples

    objv_dataset = '/mnt/sfs-common/yslan/Dataset/Obajverse/chunk-jpeg-normal/bs_16_fixsave3/170K/512/'
    dataset_json = os.path.join(objv_dataset, 'dataset.json')
    with open(dataset_json, 'r') as f:
        dataset_json = json.load(f)

    # all_objs = dataset_json['Animals'][::3][:6250]
    all_objs = dataset_json['Animals'][::3][1100:2200]
    all_objs = all_objs[:600]

    os.makedirs(save_root, exist_ok=True)


    for obj_folder in tqdm.tqdm(all_objs):
        obj_folder = '/'.join(obj_folder.split('/')[1:])
        for idx in [0]:

            if 'CRM' in shapenet_root:
                # ipdb.set_trace()
                mesh_path = glob.glob(os.path.join(shapenet_root, obj_folder, f'{idx}', '*.obj'))[0]
            elif 'Lara' in shapenet_root:
                try:
                    mesh_path = glob.glob(os.path.join(shapenet_root, '/'.join(obj_folder.split('/')[:-1]), f'{idx}.jpg', '*.obj'))[0]
                except:
                    continue
            
            elif 'scale3d' in shapenet_root:
                if 'eval_nerf' in shapenet_root:
                    mesh_path = glob.glob(os.path.join(shapenet_root, '/'.join(obj_folder.split('/')[:-1]), '1', 'mesh', '*.ply'))[0]
                elif 'eval_mesh' in shapenet_root:
                    mesh_path = glob.glob(os.path.join(shapenet_root, '/'.join(obj_folder.split('/')[:-1]), '1', 'mesh', '*.obj'))[0]

                # files.append(os.path.join(path, '/'.join(obj_folder.split('/')[:-1]), '1', f'{idx}.png'))

            elif 'LRM' in shapenet_root:
                mesh_path = glob.glob(os.path.join(shapenet_root, '/'.join(obj_folder.split('/')[:-1]), '0', '*.ply'))[0]

            elif 'ln3diff' in shapenet_root:
                mesh_path = glob.glob(os.path.join(shapenet_root, obj_folder, '0', 'mesh.obj'))[0]
            else:
                if os.path.exists((os.path.join(shapenet_root, obj_folder, f'{idx}/mesh.obj'))):
                    mesh_path = os.path.join(shapenet_root, obj_folder, f'{idx}/mesh.obj')
                else:
                    mesh_path = os.path.join(shapenet_root, obj_folder, f'{idx}/mesh.ply')

            save_name = '-'.join(obj_folder.split('/'))
            cmds += [(mesh_path, os.path.join(save_root, f'{save_name}_pcd_4096.ply'), debug)]


    if options.n_proc == 0:
        for filepath in tqdm.tqdm(cmds):
            sample_surface_pts(filepath)
    else:
        with Pool(options.n_proc) as p:
            list(tqdm.tqdm(p.imap(sample_surface_pts, cmds), total=len(cmds)))