File size: 2,195 Bytes
49ce4bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import json, os, sys, time, numpy as np, trimesh
from multiprocessing import Pool
sys.path.insert(0, '/root')
from caption_filter import is_single_object
root = '/root/lowpoly'
OUT = '/root/data'
os.makedirs(OUT, exist_ok=True)
paths = {}
for dirpath, _, names in os.walk(root):
for n in names:
if n.endswith('.obj'):
paths[n[:-4]] = os.path.join(dirpath, n)
print('obj files on disk:', len(paths), flush=True)
caps = json.load(open('/root/captions.json'))
uids = sorted(u for u in caps if u in paths)
print('captioned meshes:', len(uids), flush=True)
def vox(u):
try:
m = trimesh.load(paths[u], force='mesh')
v = m.vertices - m.vertices.mean(0)
s = np.abs(v).max()
m.vertices = v / (s if s > 0 else 1)
g = m.voxelized(pitch=2.0/31).matrix.astype(np.uint8)
pad = np.zeros((32,32,32), np.uint8)
sh = [min(32,d) for d in g.shape]
pad[:sh[0],:sh[1],:sh[2]] = g[:sh[0],:sh[1],:sh[2]]
if pad.sum() == 0:
return u, None
return u, np.packbits(pad.reshape(-1))
except Exception:
return u, None
t0 = time.time()
ok_uids, rows, failed = [], [], 0
with Pool(32) as p:
for i, (u, packed) in enumerate(p.imap(vox, uids, chunksize=64)):
if packed is None:
failed += 1
else:
ok_uids.append(u); rows.append(packed)
if (i+1) % 5000 == 0:
print(f' voxelized {i+1}/{len(uids)} failed {failed} {time.time()-t0:.0f}s', flush=True)
vox_arr = np.stack(rows)
np.save(f'{OUT}/voxels_packed.npy', vox_arr)
captions = [caps[u] for u in ok_uids]
keep = [bool(is_single_object(c)) for c in captions]
json.dump({'uids': ok_uids, 'captions': captions, 'keep': keep},
open(f'{OUT}/meta.json', 'w'))
occ = np.unpackbits(vox_arr, axis=1).mean(1) * 100
print(f'\ndone in {time.time()-t0:.0f}s')
print(f'meshes voxelized {len(ok_uids)}, failed/empty {failed}')
print(f'array {vox_arr.shape} {vox_arr.nbytes/2**20:.0f} MiB packed')
print(f'filter survivors {sum(keep)} = {sum(keep)/len(keep)*100:.1f}%')
print(f'surface occupancy mean {occ.mean():.2f}% median {np.median(occ):.2f}%')
print('BUILDDONE', flush=True)
|