|
|
import os
|
|
|
import copy
|
|
|
import sys
|
|
|
import importlib
|
|
|
import argparse
|
|
|
import pandas as pd
|
|
|
from easydict import EasyDict as edict
|
|
|
from functools import partial
|
|
|
import numpy as np
|
|
|
import open3d as o3d
|
|
|
import utils3d
|
|
|
|
|
|
|
|
|
def _voxelize(file, sha256, output_dir):
|
|
|
mesh = o3d.io.read_triangle_mesh(os.path.join(output_dir, 'renders', sha256, 'mesh.ply'))
|
|
|
|
|
|
vertices = np.clip(np.asarray(mesh.vertices), -0.5 + 1e-6, 0.5 - 1e-6)
|
|
|
mesh.vertices = o3d.utility.Vector3dVector(vertices)
|
|
|
voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds(mesh, voxel_size=1/64, min_bound=(-0.5, -0.5, -0.5), max_bound=(0.5, 0.5, 0.5))
|
|
|
vertices = np.array([voxel.grid_index for voxel in voxel_grid.get_voxels()])
|
|
|
assert np.all(vertices >= 0) and np.all(vertices < 64), "Some vertices are out of bounds"
|
|
|
vertices = (vertices + 0.5) / 64 - 0.5
|
|
|
utils3d.io.write_ply(os.path.join(output_dir, 'voxels', f'{sha256}.ply'), vertices)
|
|
|
return {'sha256': sha256, 'voxelized': True, 'num_voxels': len(vertices)}
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
parser.add_argument('--output_dir', type=str, required=True,
|
|
|
help='Directory to save the metadata')
|
|
|
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')
|
|
|
parser.add_argument('--num_views', type=int, default=150,
|
|
|
help='Number of views to render')
|
|
|
dataset_utils.add_args(parser)
|
|
|
parser.add_argument('--rank', type=int, default=0)
|
|
|
parser.add_argument('--world_size', type=int, default=1)
|
|
|
parser.add_argument('--max_workers', type=int, default=None)
|
|
|
opt = parser.parse_args(sys.argv[2:])
|
|
|
opt = edict(vars(opt))
|
|
|
|
|
|
os.makedirs(os.path.join(opt.output_dir, 'voxels'), exist_ok=True)
|
|
|
|
|
|
|
|
|
if not os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
|
|
|
raise ValueError('metadata.csv not found')
|
|
|
metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
|
|
|
if opt.instances is None:
|
|
|
if opt.filter_low_aesthetic_score is not None:
|
|
|
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
|
|
if 'rendered' not in metadata.columns:
|
|
|
raise ValueError('metadata.csv does not have "rendered" column, please run "build_metadata.py" first')
|
|
|
metadata = metadata[metadata['rendered'] == True]
|
|
|
if 'voxelized' in metadata.columns:
|
|
|
metadata = metadata[metadata['voxelized'] == False]
|
|
|
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]
|
|
|
records = []
|
|
|
|
|
|
|
|
|
for sha256 in copy.copy(metadata['sha256'].values):
|
|
|
if os.path.exists(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply')):
|
|
|
pts = utils3d.io.read_ply(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply'))[0]
|
|
|
records.append({'sha256': sha256, 'voxelized': True, 'num_voxels': len(pts)})
|
|
|
metadata = metadata[metadata['sha256'] != sha256]
|
|
|
|
|
|
print(f'Processing {len(metadata)} objects...')
|
|
|
|
|
|
|
|
|
func = partial(_voxelize, output_dir=opt.output_dir)
|
|
|
voxelized = dataset_utils.foreach_instance(metadata, opt.output_dir, func, max_workers=opt.max_workers, desc='Voxelizing')
|
|
|
voxelized = pd.concat([voxelized, pd.DataFrame.from_records(records)])
|
|
|
voxelized.to_csv(os.path.join(opt.output_dir, f'voxelized_{opt.rank}.csv'), index=False)
|
|
|
|
