third_party/TRELLIS.2/data_toolkit/dual_grid.py

import os
import sys
import importlib
import argparse
import pandas as pd
import numpy as np
import torch
import pickle
import o_voxel
from easydict import EasyDict as edict
from functools import partial


def _dual_grid_mesh(file, metadatum, mesh_dump_root, root):
    sha256 = metadatum['sha256']
    try:
        pack = {'sha256': sha256}
        data = None
        for res in opt.resolution:
            need_process = False

            # check if already processed
            if os.path.exists(os.path.join(root, f'dual_grid_{res}', f'{sha256}.vxz')):
                try:
                    info = o_voxel.io.read_vxz_info(os.path.join(root, f'dual_grid_{res}', f'{sha256}.vxz'))
                    pack[f'dual_grid_converted_{res}'] = True
                    pack[f'dual_grid_size_{res}'] = info['num_voxel']
                except Exception as e:
                    print(f'Error reading {sha256}.vxz: {e}')
                    need_process = True
            else:
                need_process = True

            # process mesh
            if need_process:
                if data is None:
                    with open(os.path.join(mesh_dump_root, 'mesh_dumps', f'{sha256}.pickle'), 'rb') as f:
                        dump = pickle.load(f)
                    start = 0
                    vertices = []
                    faces = []
                    for obj in dump['objects']:
                        if obj['vertices'].size == 0 or obj['faces'].size == 0:
                            continue
                        vertices.append(obj['vertices'])
                        faces.append(obj['faces'] + start)
                        start += len(obj['vertices'])
                    vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
                    faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
                    vertices_min = vertices.min(dim=0)[0]
                    vertices_max = vertices.max(dim=0)[0]
                    center = (vertices_min + vertices_max) / 2
                    scale = 0.99999 / (vertices_max - vertices_min).max()
                    vertices = (vertices - center) * scale
                    assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
                    data = {'vertices': vertices, 'faces': faces}

                voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
                    **data,
                    grid_size=res,
                    aabb=[[-0.5,-0.5,-0.5],[0.5,0.5,0.5]],
                    face_weight=1.0,
                    boundary_weight=0.2,
                    regularization_weight=1e-2,
                    timing=False,
                )
                dual_vertices = dual_vertices * res - voxel_indices
                assert torch.all(dual_vertices >= -1e-3) and torch.all(dual_vertices <= 1+1e-3), 'dual_vertices out of range'
                dual_vertices = torch.clamp(dual_vertices, 0, 1)
                dual_vertices = (dual_vertices * 255).type(torch.uint8)
                intersected = (intersected[:, 0:1] + 2 * intersected[:, 1:2] + 4 * intersected[:, 2:3]).type(torch.uint8)
                
                o_voxel.io.write_vxz(
                    os.path.join(root, f'dual_grid_{res}', f'{sha256}.vxz'),
                    voxel_indices,
                    {'vertices': dual_vertices, 'intersected': intersected},
                )

                pack[f'dual_grid_converted_{res}'] = True
                pack[f'dual_grid_size_{res}'] = len(dual_vertices)

        return pack
    except Exception as e:
        print(f'Error processing {sha256}: {e}')
        return {'sha256': sha256, 'error': str(e)}


if __name__ == '__main__':
    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')

    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str, required=True,
                        help='Directory to save the metadata')
    parser.add_argument('--mesh_dump_root', type=str, default=None,
                        help='Directory to load mesh dumps')
    parser.add_argument('--dual_grid_root', type=str, default=None,
                        help='Directory to save dual grids')
    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
                        help='Filter objects with aesthetic score lower than this value')
    parser.add_argument('--instances', type=str, default=None,
                        help='Instances to process')
    dataset_utils.add_args(parser)
    parser.add_argument('--rank', type=int, default=0)
    parser.add_argument('--resolution', type=str, default=256)
    parser.add_argument('--world_size', type=int, default=1)
    parser.add_argument('--max_workers', type=int, default=0)
    opt = parser.parse_args(sys.argv[2:])
    opt = edict(vars(opt))
    opt.resolution = [int(x) for x in opt.resolution.split(',')]
    opt.mesh_dump_root = opt.mesh_dump_root or opt.root
    opt.dual_grid_root = opt.dual_grid_root or opt.root

    for res in opt.resolution:
        os.makedirs(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'new_records'), exist_ok=True)

    # get file list
    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
        raise ValueError('metadata.csv not found')
    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
    if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')):
        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')).set_index('sha256'))
    for res in opt.resolution:
        if os.path.exists(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'metadata.csv')):
            dual_grid_metadata = pd.read_csv(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'metadata.csv')).set_index('sha256')
            dual_grid_metadata = dual_grid_metadata.rename(columns={'dual_grid_converted': f'dual_grid_converted_{res}', 'dual_grid_size': f'dual_grid_size_{res}'})
            metadata = metadata.combine_first(dual_grid_metadata)
    metadata = metadata.reset_index()
    if opt.instances is None:
        if opt.filter_low_aesthetic_score is not None:
            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
        metadata = metadata[metadata['mesh_dumped'] == True]
        mask = np.zeros(len(metadata), dtype=bool)
        for res in opt.resolution:
            if f'dual_grid_converted_{res}' in metadata.columns:
                mask |= metadata[f'dual_grid_converted_{res}'] != True
            else:
                mask[:] = True
                break
        metadata = metadata[mask]
    else:
        if os.path.exists(opt.instances):
            with open(opt.instances, 'r') as f:
                instances = f.read().splitlines()
        else:
            instances = opt.instances.split(',')
        metadata = metadata[metadata['sha256'].isin(instances)]

    start = len(metadata) * opt.rank // opt.world_size
    end = len(metadata) * (opt.rank + 1) // opt.world_size
    metadata = metadata[start:end]
    
    print(f'Processing {len(metadata)} objects...')

    # process objects
    func = partial(_dual_grid_mesh, root=opt.dual_grid_root, mesh_dump_root=opt.mesh_dump_root)
    dual_grids = dataset_utils.foreach_instance(metadata, None, func, max_workers=opt.max_workers, no_file=True, desc='Dual griding')
    if 'error' in dual_grids.columns:
        errors = dual_grids[dual_grids[f'error'].notna()]
        with open('errors.txt', 'w') as f:
            f.write('\n'.join(errors['sha256'].tolist()))
    for res in opt.resolution:
        if f'dual_grid_converted_{res}' in dual_grids.columns:
            dual_grid_metadata = dual_grids[dual_grids[f'dual_grid_converted_{res}'] == True]
            if len(dual_grid_metadata) > 0:
                dual_grid_metadata = dual_grid_metadata[['sha256', f'dual_grid_converted_{res}', f'dual_grid_size_{res}']]
                dual_grid_metadata = dual_grid_metadata.rename(columns={f'dual_grid_converted_{res}': 'dual_grid_converted', f'dual_grid_size_{res}': 'dual_grid_size'})
                dual_grid_metadata.to_csv(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'new_records', f'part_{opt.rank}.csv'), index=False)