146 lines
4.9 KiB
Python
146 lines
4.9 KiB
Python
import argparse
|
|
import os
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
import trimesh
|
|
from tqdm import tqdm
|
|
|
|
from src.data import IndexField, PointCloudField, Shapes3dDataset
|
|
from src.eval import MeshEvaluator
|
|
from src.utils import load_config, load_pointcloud
|
|
|
|
np.set_printoptions(precision=4)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="MNIST toy experiment")
|
|
parser.add_argument("config", type=str, help="Path to config file.")
|
|
parser.add_argument("--no_cuda", action="store_true", default=False, help="disables CUDA training")
|
|
parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)")
|
|
parser.add_argument("--iter", type=int, metavar="S", help="the training iteration to be evaluated.")
|
|
|
|
args = parser.parse_args()
|
|
cfg = load_config(args.config, "configs/default.yaml")
|
|
use_cuda = not args.no_cuda and torch.cuda.is_available()
|
|
torch.device("cuda" if use_cuda else "cpu")
|
|
cfg["data"]["data_type"]
|
|
# Shorthands
|
|
out_dir = cfg["train"]["out_dir"]
|
|
generation_dir = os.path.join(out_dir, cfg["generation"]["generation_dir"])
|
|
|
|
if cfg["generation"].get("iter", 0) != 0:
|
|
generation_dir += "_%04d" % cfg["generation"]["iter"]
|
|
elif args.iter is not None:
|
|
generation_dir += "_%04d" % args.iter
|
|
|
|
print("Evaluate meshes under %s" % generation_dir)
|
|
|
|
out_file = os.path.join(generation_dir, "eval_meshes_full.pkl")
|
|
out_file_class = os.path.join(generation_dir, "eval_meshes.csv")
|
|
|
|
# PYTORCH VERSION > 1.0.0
|
|
assert float(torch.__version__.split(".")[-3]) > 0
|
|
|
|
pointcloud_field = PointCloudField(cfg["data"]["pointcloud_file"])
|
|
fields = {
|
|
"pointcloud": pointcloud_field,
|
|
"idx": IndexField(),
|
|
}
|
|
|
|
print("Test split: ", cfg["data"]["test_split"])
|
|
|
|
dataset_folder = cfg["data"]["path"]
|
|
dataset = Shapes3dDataset(
|
|
dataset_folder, fields, cfg["data"]["test_split"], categories=cfg["data"]["class"], cfg=cfg,
|
|
)
|
|
|
|
# Loader
|
|
test_loader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=0, shuffle=False)
|
|
|
|
# Evaluator
|
|
evaluator = MeshEvaluator(n_points=100000)
|
|
|
|
eval_dicts = []
|
|
print("Evaluating meshes...")
|
|
for _it, data in enumerate(tqdm(test_loader)):
|
|
if data is None:
|
|
print("Invalid data.")
|
|
continue
|
|
|
|
mesh_dir = os.path.join(generation_dir, "meshes")
|
|
pointcloud_dir = os.path.join(generation_dir, "pointcloud")
|
|
|
|
# Get index etc.
|
|
idx = data["idx"].item()
|
|
try:
|
|
model_dict = dataset.get_model_dict(idx)
|
|
except AttributeError:
|
|
model_dict = {"model": str(idx), "category": "n/a"}
|
|
|
|
modelname = model_dict["model"]
|
|
category_id = model_dict["category"]
|
|
|
|
try:
|
|
category_name = dataset.metadata[category_id].get("name", "n/a")
|
|
except AttributeError:
|
|
category_name = "n/a"
|
|
|
|
if category_id != "n/a":
|
|
mesh_dir = os.path.join(mesh_dir, category_id)
|
|
pointcloud_dir = os.path.join(pointcloud_dir, category_id)
|
|
|
|
# Evaluate
|
|
pointcloud_tgt = data["pointcloud"].squeeze(0).numpy()
|
|
normals_tgt = data["pointcloud.normals"].squeeze(0).numpy()
|
|
|
|
eval_dict = {
|
|
"idx": idx,
|
|
"class id": category_id,
|
|
"class name": category_name,
|
|
"modelname": modelname,
|
|
}
|
|
eval_dicts.append(eval_dict)
|
|
|
|
# Evaluate mesh
|
|
if cfg["test"]["eval_mesh"]:
|
|
mesh_file = os.path.join(mesh_dir, "%s.off" % modelname)
|
|
|
|
if os.path.exists(mesh_file):
|
|
mesh = trimesh.load(mesh_file, process=False)
|
|
eval_dict_mesh = evaluator.eval_mesh(mesh, pointcloud_tgt, normals_tgt)
|
|
for k, v in eval_dict_mesh.items():
|
|
eval_dict[k + " (mesh)"] = v
|
|
else:
|
|
print("Warning: mesh does not exist: %s" % mesh_file)
|
|
|
|
# Evaluate point cloud
|
|
if cfg["test"]["eval_pointcloud"]:
|
|
pointcloud_file = os.path.join(pointcloud_dir, "%s.ply" % modelname)
|
|
|
|
if os.path.exists(pointcloud_file):
|
|
pointcloud = load_pointcloud(pointcloud_file).astype(np.float32)
|
|
eval_dict_pcl = evaluator.eval_pointcloud(pointcloud, pointcloud_tgt)
|
|
for k, v in eval_dict_pcl.items():
|
|
eval_dict[k + " (pcl)"] = v
|
|
else:
|
|
print("Warning: pointcloud does not exist: %s" % pointcloud_file)
|
|
|
|
# Create pandas dataframe and save
|
|
eval_df = pd.DataFrame(eval_dicts)
|
|
eval_df.set_index(["idx"], inplace=True)
|
|
eval_df.to_pickle(out_file)
|
|
|
|
# Create CSV file with main statistics
|
|
eval_df_class = eval_df.groupby(by=["class name"]).mean()
|
|
eval_df_class.loc["mean"] = eval_df_class.mean()
|
|
eval_df_class.to_csv(out_file_class)
|
|
|
|
# Print results
|
|
print(eval_df_class)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|