hy3dshape/tools/watertight/watertight_and_sample.py

# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
# except for the third-party components listed below.
# Hunyuan 3D does not impose any additional limitations beyond what is outlined
# in the repsective licenses of these third-party components.
# Users must comply with all terms and conditions of original licenses of these third-party
# components and must ensure that the usage of the third party components adheres to
# all relevant laws and regulations.

# For avoidance of doubts, Hunyuan 3D means the large language models and
# their software and algorithms, including trained model weights, parameters (including
# optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
# fine-tuning enabling code and other elements of the foregoing made publicly available
# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.

import argparse
import igl
import numpy as np
import os
from scipy.stats import truncnorm
import trimesh

def random_sample_pointcloud(mesh, num = 30000):
    points, face_idx = mesh.sample(num, return_index=True)
    normals = mesh.face_normals[face_idx]
    rng = np.random.default_rng()
    index = rng.choice(num, num, replace=False)
    return points[index], normals[index]

def sharp_sample_pointcloud(mesh, num=16384):
    V = mesh.vertices
    N = mesh.face_normals
    VN = mesh.vertex_normals
    F = mesh.faces
    VN2 = np.ones(V.shape[0])
    for i in range(3):
        dot = np.stack((VN2[F[:,i]], np.sum(VN[F[:,i]] * N, axis=-1)), axis=-1)
        VN2[F[:,i]] = np.min(dot, axis=-1)

    sharp_mask = VN2<0.985
    # collect edge
    edge_a = np.concatenate((F[:,0],F[:,1],F[:,2]))
    edge_b = np.concatenate((F[:,1],F[:,2],F[:,0]))
    sharp_edge = ((sharp_mask[edge_a] * sharp_mask[edge_b]))
    edge_a = edge_a[sharp_edge>0]
    edge_b = edge_b[sharp_edge>0]

    sharp_verts_a = V[edge_a]
    sharp_verts_b = V[edge_b]
    sharp_verts_an = VN[edge_a]
    sharp_verts_bn = VN[edge_b]

    weights = np.linalg.norm(sharp_verts_b - sharp_verts_a, axis=-1)
    weights /= np.sum(weights)

    random_number = np.random.rand(num)
    w = np.random.rand(num,1)
    index = np.searchsorted(weights.cumsum(), random_number)
    samples = w * sharp_verts_a[index] + (1 - w) * sharp_verts_b[index]
    normals = w * sharp_verts_an[index] + (1 - w) * sharp_verts_bn[index]
    return samples, normals

def sample_sdf(mesh, random_surface, sharp_surface):
    n_volume_points = sharp_surface.shape[0] * 2
    vol_points = (np.random.rand(n_volume_points, 3) - 0.5) * 2 * 1.05

    a, b = -0.25, 0.25
    mu = 0

    # get near points (add offset on surface points)
    offset1 = truncnorm.rvs((a - mu) / 0.005, (b - mu) / 0.005, loc=mu, scale=0.005, size=(len(random_surface), 3))
    offset2 = truncnorm.rvs((a - mu) / 0.05, (b - mu) / 0.05, loc=mu, scale=0.05,  size=(len(random_surface), 3))
    random_near_points = np.concatenate([
        random_surface + offset1,
        random_surface + offset2
    ], axis=0)

    unit_num = len(sharp_surface) // 6
    sharp_near_points = np.concatenate([
        sharp_surface[:unit_num] + np.random.normal(scale=0.001, size=(unit_num, 3)),
        sharp_surface[unit_num:unit_num*2] + np.random.normal(scale=0.003, size=(unit_num,3)),
        sharp_surface[unit_num*2:unit_num*3] + np.random.normal(scale=0.06, size=(unit_num,3)),
        sharp_surface[unit_num*3:unit_num*4] + np.random.normal(scale=0.01, size=(unit_num,3)),
        sharp_surface[unit_num*4:unit_num*5] + np.random.normal(scale=0.02, size=(unit_num,3)),
        sharp_surface[unit_num*5:] + np.random.normal(scale=0.04, size=(len(sharp_surface)-5*unit_num,3))
    ], axis=0)

    np.random.shuffle(random_near_points)
    np.random.shuffle(sharp_near_points)

    sign_type = igl.SIGNED_DISTANCE_TYPE_FAST_WINDING_NUMBER
    try:
        vol_sdf, I, C = igl.signed_distance(
            vol_points.astype(np.float32), 
            mesh.vertices, mesh.faces, 
            return_normals=False,
            sign_type=sign_type)
    except:
        vol_sdf, I, C = igl.signed_distance(
            vol_points.astype(np.float32), 
            mesh.vertices, mesh.faces, 
            return_normals=False)
    try:
        random_near_sdf, I, C = igl.signed_distance(
            random_near_points.astype(np.float32), 
            mesh.vertices, mesh.faces, 
            return_normals=False,
            sign_type=sign_type)
    except:
        random_near_sdf, I, C = igl.signed_distance(
            random_near_points.astype(np.float32), 
            mesh.vertices, mesh.faces, 
            return_normals=False)
    try:
        sharp_near_sdf, I, C = igl.signed_distance(
            sharp_near_points.astype(np.float32), 
            mesh.vertices, mesh.faces, 
            return_normals=False,
            sign_type=sign_type)
    except:
        sharp_near_sdf, I, C = igl.signed_distance(
            sharp_near_points.astype(np.float32), 
            mesh.vertices, mesh.faces, 
            return_normals=False)
        
    vol_label = -vol_sdf
    random_near_label = -random_near_sdf
    sharp_near_label = -sharp_near_sdf

    data = {
        "vol_points": vol_points.astype(np.float16),
        "vol_label": vol_label.astype(np.float16),
        "random_near_points": random_near_points.astype(np.float16),
        "random_near_label": random_near_label.astype(np.float16),
        "sharp_near_points": sharp_near_points.astype(np.float16),
        "sharp_near_label": sharp_near_label.astype(np.float16)
    }
    return data

def SampleMesh(V, F):
    mesh = trimesh.Trimesh(vertices=V, faces=F)

    area = mesh.area
    sample_num = 499712//4

    random_surface, random_normal = random_sample_pointcloud(mesh, num=sample_num)
    random_sharp_surface, sharp_normal = sharp_sample_pointcloud(mesh, num=sample_num)

    #save_surface
    surface = np.concatenate((random_surface, random_normal), axis = 1).astype(np.float16)
    sharp_surface = np.concatenate((random_sharp_surface, sharp_normal), axis=1).astype(np.float16)

    surface_data = {
        "random_surface": surface,
        "sharp_surface": sharp_surface,
    }

    sdf_data = sample_sdf(mesh, random_surface, random_sharp_surface)
    return surface_data, sdf_data

def normalize_to_unit_box(V):
    """
    Normalize the vertices V to fit inside a unit bounding box [0,1]^3.
    V: (n,3) numpy array of vertex positions.
    Returns: normalized V
    """
    V_min = V.min(axis=0)
    V_max = V.max(axis=0)
    scale = (V_max - V_min).max() * 1.01
    V_normalized = (V - V_min) / scale
    return V_normalized

# Given: V (n x 3 array of vertices), F (m x 3 array of faces)
# Parameters epsilon/grid_res
def Watertight(V, F, epsilon = 2.0/256, grid_res = 256):
    # Compute bounding box
    min_corner = V.min(axis=0)
    max_corner = V.max(axis=0)
    padding = 0.05 * (max_corner - min_corner)
    min_corner -= padding
    max_corner += padding

    # Create a uniform grid
    x = np.linspace(min_corner[0], max_corner[0], grid_res)
    y = np.linspace(min_corner[1], max_corner[1], grid_res)
    z = np.linspace(min_corner[2], max_corner[2], grid_res)
    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
    grid_points = np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T

    # Compute SDF at grid points using igl.signed_distance with pseudo normals
    sdf, _, _ = igl.signed_distance(
        grid_points, V, F, sign_type=igl.SIGNED_DISTANCE_TYPE_PSEUDONORMAL
    )
 
    # igl.marching_cubes returns (vertices, faces)
    mc_verts, mc_faces = igl.marching_cubes(epsilon - np.abs(sdf), grid_points, grid_res, grid_res, grid_res, 0.0)

    # mc_verts: (k x 3) array of vertices of the epsilon contour
    # mc_faces: (l x 3) array of faces of the epsilon contour
    return mc_verts, mc_faces

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Process an OBJ file and output surface and SDF data.')
    parser.add_argument('--input_obj', type=str, help='Path to the input OBJ file')
    parser.add_argument('--output_prefix', type=str, default=None, 
        help='Base name for output files (default: input OBJ filename without extension)')
    args = parser.parse_args()

    input_obj = args.input_obj
    name = args.output_prefix

    V, F = igl.read_triangle_mesh(input_obj)
    V = normalize_to_unit_box(V)

    mc_verts, mc_faces = Watertight(V, F)
    surface_data, sdf_data = SampleMesh(mc_verts, mc_faces)

    parent_folder = os.path.dirname(args.output_prefix)
    os.makedirs(parent_folder, exist_ok=True)
    export_surface = f'{name}_surface.npz'
    np.savez(export_surface, **surface_data)
    export_sdf = f'{name}_sdf.npz'
    np.savez(export_sdf, **sdf_data)
    igl.write_obj(f'{name}_watertight.obj', mc_verts, mc_faces)