hole_mesh_generation.py

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from sdf import stl as ss
# from stl import mesh as stl_mesh
import numpy as np
from multiprocessing.pool import ThreadPool
import itertools
from sdf import progress,mesh
from functools import partial
import torch
from dataset.kitti360.labels import labels, id2label
import open3d as o3d
from tqdm import tqdm
from wisp.ops.differential import finitediff_gradient, tetrahedron_gradient, autodiff_gradient
import argparse

num_lods = 3
# device = 'cuda1'
@torch.no_grad()
def query_sdf(coords):
        shape = coords.shape
        # pidx = grid.blas.query(torch.FloatTensor(coords.reshape(-1, 3)).to(device), grid.active_lods[lod_idx], with_parents=True)
        # mask = pidx.pidx[:,grid.active_lods[0]:].max(1).values<0
        
        # mask = pidx.pidx[:,self.grid.active_lods[1]:].max(1).values<0
        # num_samples = coords.shape[1]
        if len(shape) == 2:
                coords = torch.FloatTensor(coords[:, None]).to(device)
        dis = torch.norm(coords - origin_all,dim=-1)
        cube_idx = dis.min(-1).indices + START#[N,CUBE,3]
        grids = []
        decodes = []
        masks_all = torch.zeros([coords.shape[0]],dtype=torch.bool).to(coords.device)
        sdfs_all = torch.zeros_like(coords[:,:,0:1])
        if torch.unique(cube_idx).tolist()[-1]>1:
            yx = 1
        for i in torch.unique(cube_idx).tolist():
            data = torch.load(os.path.join(result_dir, f'grid/optimed_grid_{i}.pth')) #
            # data = torch.load(os.path.join(result_dir, f'grid/pretrained_sdf.pth')) #
            grid = data['grid']
            # grids.append(grid)
            decoder = data['decoder']
            # decodes.append(decoder)
            mask_cube = cube_idx==i
            # masks.append(mask_cube)
            points = coords[mask_cube] - origin_all[i - START]
            feats = grid['sdf'].interpolate(points, grid['sdf'].num_lods-1)
            sdfs = decoder['sdf'].forward_sdf(feats)

            query_results = grid['sdf'].blas.query(points.reshape(-1, 3), grid['sdf'].active_lods[grid['sdf'].num_lods-1], with_parents=True)
            mask = query_results.pidx[:,grid['sdf'].active_lods[0]]<0

            sdfs[mask] = 10.0 
            masks_all[mask_cube] = mask
            sdfs_all[mask_cube] = sdfs
      
        # feats = grid.interpolate(coords, lod_idx)
        # sdfs = decoder(feats) #/self.scale 
        # # sdfs[mask] = torch.abs(sdfs[mask])
        # sdfs[mask] = 10.0 
        return sdfs_all.squeeze().detach().cpu().numpy(), ~masks_all.detach().cpu().numpy()  

@ torch.no_grad()
def extract_mesh(idx):
    WORKERS = 1 #multiprocessing.cpu_count()
    SAMPLES = 2 ** 29 #32
    BATCH_SIZE = 64
   
    x0,y0,z0 = -1,-1,-0.1
    x1,y1,z1 = 1.2,1.3,0.2
    verbose = True
    sparse=False #True

    volume = (x1 - x0) * (y1 - y0) * (z1 - z0)
    step = (volume / SAMPLES) ** (1 / 3)
    dx = dy = dz = step
    dz = 0.3*step
    # dx = 3*step
    # dy = (1.5*step)//1
  
    X = np.arange(x0, x1, dx)
    Y = np.arange(y0, y1, dy)
    Z = np.arange(z0, z1, dz)

    s = BATCH_SIZE
    Xs = [X[i:i+s+1] for i in range(0, len(X), s)]
    Ys = [Y[i:i+s+1] for i in range(0, len(Y), s)]
    Zs = [Z[i:i+s+1] for i in range(0, len(Z), s)]

    batches = list(itertools.product(Xs, Ys, Zs))
    num_batches = len(batches)
    num_samples = sum(len(xs) * len(ys) * len(zs)
        for xs, ys, zs in batches)
    print('%d samples in %d batches with %d workers' %
                (num_samples, num_batches, WORKERS))
    points = []
    skipped = empty = nonempty = 0
    bar = progress.Bar(num_batches, enabled=verbose)
    pool = ThreadPool(WORKERS)
    f = partial(mesh._worker, query_sdf, sparse=sparse)
    for result in pool.imap(f, batches):
        bar.increment(1)
        if result is None:
            skipped += 1
        elif len(result) == 0:
            empty += 1
        else:
            nonempty += 1
            points.extend(result)
    bar.done()
    # ss.write_binary_stl(f'01_loam_map_init/mesh/mesh_fenge{idx}.stl', points)
    
    ss.write_binary_stl(STLPATH, points)
    
    #-------------------------------------------------------------
    if SEMANTIC or RGB:
        mymesh= o3d.io.read_triangle_mesh(STLPATH)
        vertices = torch.FloatTensor( np.asarray(mymesh.vertices) )[:,None].to(device)
        # triangles = np.asarray(mymesh.triangles) 
        # vertices = torch.FloatTensor(np.stack(points))[:,None].to(device)

        sem_batch = 50000
        semantic = torch.zeros(vertices.shape[0])
        color = torch.zeros([vertices.shape[0],3])
        pidxs = torch.zeros(vertices.shape[0])
        for i in tqdm(range(0, vertices.shape[0], sem_batch), desc='get semantic labels'):
                next_ind = min(i+sem_batch, vertices.shape[0])
                v = vertices[i:next_ind]
                dis = torch.norm(v - origin_all,dim=-1)
                cube_idx = dis.min(-1).indices + START#[N,CUBE,3]
                sem = torch.zeros(v.shape[0], dtype=torch.int64).to(device)
                rgb = torch.zeros([v.shape[0],3], dtype=torch.float32).to(device)
                pidx = torch.zeros(v.shape[0], dtype=torch.int64).to(device)
                masks_all = torch.zeros([v.shape[0]],dtype=torch.bool).to(device)
                # semantic_lists = torch.zeros_like(vertices[:,:,0:1])
                if torch.unique(cube_idx).tolist()[-1]>1:
                    yx = 1
                for j in torch.unique(cube_idx).tolist():
                    data = torch.load(os.path.join(result_dir, f'grid/optimed_grid_{j}.pth')) #optimed_grid_{i}.pth
                    grid = data['grid']
                    # grids.append(grid)
                    decoder = data['decoder']
                    mask_cube = cube_idx==j
                    # masks.append(mask_cube)
                    p_g = v[mask_cube] - origin_all[j - START]
                    if RGB:
                        feats = grid['sdf'].interpolate(p_g.reshape(-1, 3), grid['sdf'].num_lods-1)
                        radiance_feats = grid['rgb'].interpolate(p_g.reshape(-1, 3),  grid['rgb'].num_lods-1)
                        # rgb_g = decoder['rgb'](radiance_feats)
                        grad = autodiff_gradient(p_g.reshape(-1, 3), sdf)
                        grad = grad / grad.norm(2, 1)[:, None]
                        rgb_g = decoder['rgb'](radiance_feats=radiance_feats, view_dirs=-grad.reshape(-1,3))
                        rgb_g = torch.sigmoid(rgb_g)
                        rgb[mask_cube] = rgb_g
                    elif SEMANTIC:
                        # feats = grid.interpolate(p_g.reshape(-1, 3), 2)
                        # _, sem_g = decoder(feats)
                        # radiance_feats = grid['feature'].interpolate(p_g.reshape(-1, 3),  grid['feature'].num_lods-1)
                        # feats = decoder['feature'](radiance_feats)
                        ## grid
                        # feats = grid['semantic'].interpolate(p_g.reshape(-1, 3),  grid['semantic'].num_lods-1)
                        # sem_g = decoder['semantic'](feats)
                        # geo
                        feats = grid['sdf'].interpolate(p_g.reshape(-1, 3),  grid['sdf'].num_lods-1)
                        _, sem_g = decoder['sdf'](feats)
                        sem_g = torch.argmax(torch.nn.functional.softmax(sem_g.squeeze(), dim=-1), -1)
                        sem[mask_cube] = sem_g
                semantic[i:next_ind] = sem.cpu()
                color[i:next_ind] = rgb.cpu()
                # pidxs[i:next_ind] = pidx.cpu()
        
        if SEMANTIC:
            save_file = savedir+f'/mesh/mesh_semantic_{START}_{CUBE}.ply'
            v_colors = np.vstack([id2label[semID].color for semID in semantic.tolist()]) / 255.0
        elif RGB:
            save_file = savedir+f'/mesh/mesh_rgb_{START}_{CUBE}.ply'
            v_colors =color.detach().numpy()
        mymesh.vertex_colors = o3d.utility.Vector3dVector(v_colors.reshape(-1,3)) 
        o3d.io.write_triangle_mesh(save_file, mymesh)
    return points

def sdf(coords):
    shape = coords.shape
    if shape[0] == 0:
        return dict(sdf=torch.zeros_like(coords)[...,0:1])

    feats = grid['sdf'].interpolate(coords.reshape(-1, 3), grid['sdf'].num_lods-1)
    # mask = feats.sum(-1)==0
    sdfs = decoder['sdf'].forward(feats) 
    return sdfs
def get_output(grid, decoder, coords):
    shape = coords.shape
    semantics = None
    lod_idx = num_lods - 1
    if len(shape) == 2:
        coords = coords[:, None]
    num_samples = coords.shape[1]

    # TODO(ttakikawa): this should return [batch, ns, f] but it returns [batch, f]
    feats = grid['sdf'].interpolate(coords.reshape(-1, 3), lod_idx)
    # mask = feats.sum(-1)==0
    sdfs, semantics = decoder['sdf'](feats)
    semantics = semantics.reshape(-1,num_samples,semantics.shape[-1])
    sdfs = sdfs.reshape(-1,num_samples,1)
    if len(shape) == 2:
        sdfs = sdfs[:,0]
    return sdfs, semantics


#==================================================================
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('--START', type=int, default=0, 
                                help='first volume index of the mesh')
parser.add_argument('--CUBE', type=int, default=2, 
                                help='volumes number of the mesh')

args = parser.parse_args()

#==================================================================
# origins = []
START = args.START #25
CUBE = args.CUBE #1 #36
print(f'extracting mesh of CUBE_{CUBE}')
FILE = f"optimed_grid_{START}.pth" #"with_sem.pth" # 
SEMANTIC = 1
RGB = 0
result_dir = 'semantic_result_kitti/round2'
data = torch.load(os.path.join(result_dir, f'grid/optimed_grid_0.pth'))
grid = data['grid']
decoder = data['decoder']
savedir = result_dir
STLPATH= os.path.join(savedir, f'mesh/mesh_test_{START}_{CUBE}.stl') #mesh_grid_{CUBE}.stl
origin_all = []
scale = torch.load(os.path.join(result_dir, f'grid/{FILE}'))['scale']
world_origin = torch.load(os.path.join(result_dir, 'grid/optimed_grid_0.pth'))['origin']/scale
if START==0:
    origin0 = torch.load(os.path.join(result_dir, 'grid/optimed_grid_0.pth'))['origin']/scale
else:
    origin0 = torch.load(os.path.join(result_dir, f'grid/{FILE}'))['origin']/scale 
origin_all.append(origin0-world_origin)
device = scale.device
idx = 0
# for idx in range(1,CUBE):
#     origin = torch.load(os.path.join(result_dir, f'grid/optimed_grid_{idx}.pth'))['origin']/scale
#     origin_all.append(origin-origin0)
for idx in range(START+1,START+CUBE):
    origin = torch.load(os.path.join(result_dir, f'grid/optimed_grid_{idx}.pth'))['origin']/scale
    origin_all.append(origin-world_origin)
origin_all = torch.stack(origin_all)

points = extract_mesh(idx)
mesh= o3d.io.read_triangle_mesh(os.path.join(result_dir, f'mesh/mesh_test_{START}_{CUBE}.stl'))
model = torch.load(os.path.join(result_dir, f'grid/optimed_grid_0.pth'))
vertices = np.asarray(mesh.vertices)

sem_batch = 65536
semantic = torch.zeros(vertices.shape[0])
for i in tqdm(range(0, vertices.shape[0], sem_batch), desc='get semantic labels'):
    next_ind = min(i+sem_batch, vertices.shape[0])
    _, sem = get_output(model['grid'], model['decoder'], torch.from_numpy(vertices[i:next_ind]).float().cuda())
    sem = torch.argmax(torch.nn.functional.softmax(sem.squeeze(), dim=-1), -1)
    semantic[i:next_ind] = sem.cpu()

v_colors = np.vstack([id2label[semID].color for semID in semantic.tolist()])
mesh.vertex_colors = o3d.utility.Vector3dVector(v_colors/255.0)
save_file = os.path.join(result_dir, f'mesh/semantic_mesh_{START}_{CUBE}.ply')
o3d.io.write_triangle_mesh(save_file, mesh)