Spaces:
Runtime error
Runtime error
File size: 2,738 Bytes
4893ce0 |
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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
"""
Preprocessing ArkitScenes
"""
import os
import argparse
import glob
import plyfile
import numpy as np
import pandas as pd
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor
from itertools import repeat
import torch
def read_plymesh(filepath):
"""Read ply file and return it as numpy array. Returns None if emtpy."""
with open(filepath, "rb") as f:
plydata = plyfile.PlyData.read(f)
if plydata.elements:
vertices = pd.DataFrame(plydata["vertex"].data).values
faces = np.stack(plydata["face"].data["vertex_indices"], axis=0)
return vertices, faces
def face_normal(vertex, face):
v01 = vertex[face[:, 1]] - vertex[face[:, 0]]
v02 = vertex[face[:, 2]] - vertex[face[:, 0]]
vec = np.cross(v01, v02)
length = np.sqrt(np.sum(vec**2, axis=1, keepdims=True)) + 1.0e-8
nf = vec / length
area = length * 0.5
return nf, area
def vertex_normal(vertex, face):
nf, area = face_normal(vertex, face)
nf = nf * area
nv = np.zeros_like(vertex)
for i in range(face.shape[0]):
nv[face[i]] += nf[i]
length = np.sqrt(np.sum(nv**2, axis=1, keepdims=True)) + 1.0e-8
nv = nv / length
return nv
def parse_scene(scene_path, output_dir):
print(f"Parsing scene {scene_path}")
split = os.path.basename(os.path.dirname(os.path.dirname(scene_path)))
scene_id = os.path.basename(os.path.dirname(scene_path))
vertices, faces = read_plymesh(scene_path)
coords = vertices[:, :3]
colors = vertices[:, 3:6]
data_dict = dict(coord=coords, color=colors, scene_id=scene_id)
data_dict["normal"] = vertex_normal(coords, faces)
torch.save(data_dict, os.path.join(output_dir, split, f"{scene_id}.pth"))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--dataset_root",
required=True,
help="Path to the ScanNet dataset containing scene folders",
)
parser.add_argument(
"--output_root",
required=True,
help="Output path where train/val folders will be located",
)
opt = parser.parse_args()
# Create output directories
train_output_dir = os.path.join(opt.output_root, "Training")
os.makedirs(train_output_dir, exist_ok=True)
val_output_dir = os.path.join(opt.output_root, "Validation")
os.makedirs(val_output_dir, exist_ok=True)
# Load scene paths
scene_paths = sorted(glob.glob(opt.dataset_root + "/3dod/*/*/*_mesh.ply"))
# Preprocess data.
pool = ProcessPoolExecutor(max_workers=mp.cpu_count())
# pool = ProcessPoolExecutor(max_workers=1)
print("Processing scenes...")
_ = list(pool.map(parse_scene, scene_paths, repeat(opt.output_root)))
|