mirror of
https://github.com/PrimitiveAnything/PrimitiveAnything.git
synced 2025-09-18 05:22:48 +08:00
114 lines
4.2 KiB
Python
114 lines
4.2 KiB
Python
import argparse
|
|
from collections import defaultdict
|
|
import glob
|
|
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
import trimesh
|
|
import point_cloud_utils as pcu
|
|
from tqdm import tqdm
|
|
|
|
|
|
import numpy as np
|
|
|
|
def voxelize(points, voxel_size=0.1):
|
|
"""
|
|
Converts a set of 3D points to a voxel grid representation.
|
|
Points are quantized to the nearest voxel center.
|
|
|
|
Parameters:
|
|
- points: (N, 3) numpy array of 3D points
|
|
- voxel_size: the size of each voxel (adjust this depending on your data)
|
|
|
|
Returns:
|
|
- voxels: Set of unique voxel coordinates
|
|
"""
|
|
# Quantize the points to the nearest voxel
|
|
quantized_points = np.floor(points / voxel_size).astype(int)
|
|
# Use a set to get unique voxel coordinates
|
|
voxels = set(map(tuple, quantized_points))
|
|
return voxels
|
|
|
|
def calculate_iou(model_vox, target_vox, voxel_size=0.1):
|
|
"""
|
|
Calculate the IoU (Intersection over Union) between two point clouds.
|
|
|
|
Parameters:
|
|
- model_vox: (N, 3) numpy array of the first point cloud
|
|
- target_vox: (M, 3) numpy array of the second point cloud
|
|
- voxel_size: Size of the voxels (default is 0.1)
|
|
|
|
Returns:
|
|
- iou: Intersection over Union (IoU) score
|
|
"""
|
|
# Voxelize both point clouds
|
|
model_voxels = voxelize(model_vox, voxel_size)
|
|
target_voxels = voxelize(target_vox, voxel_size)
|
|
|
|
# Calculate intersection and union
|
|
intersection = len(model_voxels.intersection(target_voxels))
|
|
union = len(model_voxels.union(target_voxels))
|
|
|
|
# Compute IoU
|
|
iou = intersection / union if union > 0 else 0.0
|
|
return iou
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_dir", type=str, default="./results/infer/PointClouds")
|
|
parser.add_argument("--target_dir", type=str, default="./data/test_pc")
|
|
parser.add_argument("--detail", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
if not os.path.exists(args.input_dir) or not os.path.exists(args.target_dir):
|
|
print("Invalid input!")
|
|
exit(1)
|
|
|
|
model_prefix = os.path.join(args.input_dir, "*.ply")
|
|
model_path_list = sorted(list(glob.glob(model_prefix)))
|
|
|
|
distance_json = {}
|
|
distance_list = []
|
|
emu_distance_list = []
|
|
hausdorff_distance_list = []
|
|
voxel_iou_list = []
|
|
for model_path in tqdm(model_path_list):
|
|
model_name = os.path.basename(model_path)
|
|
target_path = os.path.join(args.target_dir, model_name)
|
|
if not os.path.exists(target_path):
|
|
print(f"{target_path}: not found!")
|
|
exit(1)
|
|
|
|
model_pc = np.array(trimesh.load(model_path).vertices)
|
|
target_pc = np.array(trimesh.load(target_path).vertices)
|
|
|
|
distance = pcu.chamfer_distance(model_pc, target_pc)
|
|
model_pc_downsampled = model_pc[np.random.choice(model_pc.shape[0], 1000, replace=False)]
|
|
target_pc_downsampled = target_pc[np.random.choice(target_pc.shape[0], 1000, replace=False)]
|
|
emu_distance, _ = pcu.earth_movers_distance(model_pc_downsampled, target_pc_downsampled)
|
|
hausdorff_distance = pcu.hausdorff_distance(model_pc, target_pc)
|
|
|
|
iou = calculate_iou(model_pc, target_pc, voxel_size=1/32.)
|
|
|
|
distance_list.append(distance)
|
|
emu_distance_list.append(emu_distance)
|
|
hausdorff_distance_list.append(hausdorff_distance)
|
|
voxel_iou_list.append(iou)
|
|
model_id = os.path.splitext(model_name)[0]
|
|
distance_json[model_id] = distance
|
|
|
|
print(f"{model_id}: chamfer distance: {distance:.3f}, earth movers distance: {emu_distance:.3f}, hausdorff distance: {hausdorff_distance:.3f}, voxel IoU: {iou:.3f}")
|
|
|
|
distance_json["mean"] = np.mean(distance_list)
|
|
distance_json["mean_emu"] = np.mean(emu_distance_list)
|
|
distance_json["mean_hausdorff"] = np.mean(hausdorff_distance_list)
|
|
distance_json["mean_voxel_iou"] = np.mean(voxel_iou_list)
|
|
print(f"mean chamfer distance: {np.mean(distance_list)}")
|
|
print(f"mean earth movers distance: {np.mean(emu_distance_list)}")
|
|
print(f"mean hausdorff distance: {np.mean(hausdorff_distance_list)}")
|
|
print(f"mean voxel IoU: {np.mean(voxel_iou_list)}")
|
|
with open(os.path.join(args.input_dir, "distance.json"), "w") as json_file:
|
|
json.dump(distance_json, json_file, indent=4)
|