Spaces:
Running
on
Zero
Running
on
Zero
# https://raw.githubusercontent.com/3dlg-hcvc/omages/refs/heads/main/src/evals/fpd_eval.py | |
import json | |
import os | |
import random | |
from tqdm import tqdm | |
import glob | |
from pdb import set_trace as st | |
import trimesh | |
import sys | |
import numpy as np | |
import scipy # should be version 1.11.1 | |
import torch | |
from pathlib import Path | |
import argparse | |
# from point_e.evals.feature_extractor import PointNetClassifier, get_torch_devices | |
from feature_extractor import PointNetClassifier, get_torch_devices | |
from point_e.evals.fid_is import compute_statistics | |
from point_e.evals.fid_is import compute_inception_score | |
from point_e.evals.npz_stream import NpzStreamer | |
import numpy as np | |
def rotation_matrix(axis, angle): | |
""" | |
Returns a rotation matrix for a given axis and angle in radians. | |
:param axis: str, the axis to rotate around ('x', 'y', or 'z') | |
:param angle: float, the rotation angle in radians | |
:return: 3x3 rotation matrix | |
""" | |
if axis == 'x': | |
return np.array([[1, 0, 0], | |
[0, np.cos(angle), -np.sin(angle)], | |
[0, np.sin(angle), np.cos(angle)]]) | |
elif axis == 'y': | |
return np.array([[np.cos(angle), 0, np.sin(angle)], | |
[0, 1, 0], | |
[-np.sin(angle), 0, np.cos(angle)]]) | |
elif axis == 'z': | |
return np.array([[np.cos(angle), -np.sin(angle), 0], | |
[np.sin(angle), np.cos(angle), 0], | |
[0, 0, 1]]) | |
else: | |
raise ValueError("Axis must be 'x', 'y', or 'z'.") | |
def rotate_point_cloud(point_cloud, rotations): | |
""" | |
Rotates a point cloud along specified axes by the given angles. | |
:param point_cloud: Nx3 numpy array of points | |
:param rotations: list of tuples [(axis, angle_in_degrees), ...] | |
Example: [('x', 90), ('y', 45)] for composite rotations | |
:return: Rotated point cloud as Nx3 numpy array | |
""" | |
rotated_cloud = point_cloud.copy() | |
for axis, angle in rotations: | |
angle_rad = np.radians(angle) # Convert degrees to radians | |
R = rotation_matrix(axis, angle_rad) | |
rotated_cloud = np.dot(rotated_cloud, R.T) # Apply rotation matrix | |
return rotated_cloud | |
from functools import partial | |
# transformation dictionary | |
transformation_dict = { | |
'gso': partial(rotate_point_cloud, rotations=[('x', 0)]), # no transformation | |
'LGM_fixpose': partial(rotate_point_cloud, rotations=[('x', 90), ('z', 180)]), | |
'CRM/Animals': partial(rotate_point_cloud, rotations=[('x', 90), ('z', 180)]), | |
'Lara': partial(rotate_point_cloud, rotations=[('x', -110), ('z', 33)]), | |
'ln3diff-lite/Animals': partial(rotate_point_cloud, rotations=[('x', 90)]), | |
'One-2-3-45/Animals': partial(rotate_point_cloud, rotations=[('x', 90), ('z', 180)]), | |
'splatter-img': partial(rotate_point_cloud, rotations=[('x', -60)]), | |
# | |
'OpenLRM/Animals': partial(rotate_point_cloud, rotations=[('x', 0)]), | |
'shape-e/Animals': partial(rotate_point_cloud, rotations=[('x', 0)]), | |
'objv-gt': partial(rotate_point_cloud, rotations=[('x', 0)]), | |
# un-aligned | |
# 'ditl-fromditlPCD-fixPose-tomesh': partial(rotate_point_cloud, rotations=[('x', 0)]), | |
'GA': partial(rotate_point_cloud, rotations=[('x', 0)]), | |
'scale3d/eval/eval_nerf/Animals': partial(rotate_point_cloud, rotations=[('x', 0)]), | |
'scale3d/eval/eval_mesh/Animals': partial(rotate_point_cloud, rotations=[('x', 180), ('z', 180)]), | |
# 'ditl-fromditlPCD-fixPose-tomesh-ditxlPCD': partial(rotate_point_cloud, rotations=[('x', 0)]), | |
} | |
class PFID_evaluator(): | |
def __init__(self, devices=['cuda:0'], batch_size=256, cache_dir='~/.temp/PFID_evaluator'): | |
self.__dict__.update(locals()) | |
cache_dir = os.path.expanduser(cache_dir) | |
if not os.path.exists(cache_dir): | |
os.makedirs(cache_dir) | |
self.devices = [torch.device(d) for d in devices] | |
self.clf = PointNetClassifier(devices=self.devices, cache_dir=cache_dir, device_batch_size=self.batch_size) | |
def compute_pfid(self, pc_1, pc_2, return_feature=False): | |
# print("computing first batch activations") | |
# save clouds to npz files | |
npz_path1 = os.path.join(self.cache_dir, "temp1.npz") | |
npz_path2 = os.path.join(self.cache_dir, "temp2.npz") | |
np.savez(npz_path1, arr_0=pc_1) | |
np.savez(npz_path2, arr_0=pc_2) | |
features_1, _ = self.clf.features_and_preds(NpzStreamer(npz_path1)) | |
stats_1 = compute_statistics(features_1) | |
# print(features_1.max(), features_1.min(), features_1.mean(), features_1.std() ) | |
# print(stats_1.mu.shape, stats_1.sigma.shape) | |
features_2, _ = self.clf.features_and_preds(NpzStreamer(npz_path2)) | |
stats_2 = compute_statistics(features_2) | |
# print(features_2.max(), features_2.min(), features_2.mean(), features_2.std() ) | |
# print(stats_2.mu.shape, stats_2.sigma.shape) | |
if return_feature: | |
return features_1, features_2 | |
#PFID = stats_1.frechet_distance(stats_2) # same result as the next line | |
PFID= frechet_distance(stats_1.mu, stats_1.sigma, stats_2.mu, stats_2.sigma) | |
PKID = kernel_distance(features_1, features_2) | |
print(f"P-FID: {PFID}", f"P-KID: {PKID}") | |
return dict(PFID=PFID, PKID=PKID) | |
# from https://github.com/GaParmar/clean-fid/blob/main/cleanfid/fid.py | |
""" | |
Numpy implementation of the Frechet Distance. | |
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1) | |
and X_2 ~ N(mu_2, C_2) is | |
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)). | |
Stable version by Danica J. Sutherland. | |
Params: | |
mu1 : Numpy array containing the activations of a layer of the | |
inception net (like returned by the function 'get_predictions') | |
for generated samples. | |
mu2 : The sample mean over activations, precalculated on an | |
representative data set. | |
sigma1: The covariance matrix over activations for generated samples. | |
sigma2: The covariance matrix over activations, precalculated on an | |
representative data set. | |
""" | |
def frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6): | |
mu1 = np.atleast_1d(mu1) | |
mu2 = np.atleast_1d(mu2) | |
sigma1 = np.atleast_2d(sigma1) | |
sigma2 = np.atleast_2d(sigma2) | |
assert mu1.shape == mu2.shape, \ | |
'Training and test mean vectors have different lengths' | |
assert sigma1.shape == sigma2.shape, \ | |
'Training and test covariances have different dimensions' | |
diff = mu1 - mu2 | |
# Product might be almost singular | |
covmean, _ = scipy.linalg.sqrtm(sigma1.dot(sigma2), disp=False) | |
if not np.isfinite(covmean).all(): | |
msg = ('fid calculation produces singular product; ' | |
'adding %s to diagonal of cov estimates') % eps | |
print(msg) | |
offset = np.eye(sigma1.shape[0]) * eps | |
covmean = scipy.linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) | |
# Numerical error might give slight imaginary component | |
if np.iscomplexobj(covmean): | |
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): | |
m = np.max(np.abs(covmean.imag)) | |
raise ValueError('Imaginary component {}'.format(m)) | |
covmean = covmean.real | |
tr_covmean = np.trace(covmean) | |
return (diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean) | |
""" | |
Compute the KID score given the sets of features | |
""" | |
def kernel_distance(feats1, feats2, num_subsets=100, max_subset_size=1000): | |
n = feats1.shape[1] | |
m = min(min(feats1.shape[0], feats2.shape[0]), max_subset_size) | |
t = 0 | |
for _subset_idx in range(num_subsets): | |
x = feats2[np.random.choice(feats2.shape[0], m, replace=False)] | |
y = feats1[np.random.choice(feats1.shape[0], m, replace=False)] | |
a = (x @ x.T / n + 1) ** 3 + (y @ y.T / n + 1) ** 3 | |
b = (x @ y.T / n + 1) ** 3 | |
t += (a.sum() - np.diag(a).sum()) / (m - 1) - b.sum() * 2 / m | |
kid = t / num_subsets / m | |
return float(kid) | |
# load and calculate fid, kid, is | |
def normalize_point_clouds(pc: np.ndarray) -> np.ndarray: | |
# centroids = np.mean(pc, axis=1, keepdims=True) | |
centroids = np.mean(pc, axis=1, keepdims=True) | |
pc = pc - centroids | |
m = np.max(np.sqrt(np.sum(pc**2, axis=-1, keepdims=True)), axis=1, keepdims=True) | |
pc = pc / m | |
return pc | |
class PCDPathDataset(torch.utils.data.Dataset): | |
def __init__(self, pcd_file_path, transformation, rand_aug=False): | |
# files = sorted(glob.glob(f'{pcd_file_path}/*.ply') ) | |
objv_dataset = '/mnt/sfs-common/yslan/Dataset/Obajverse/chunk-jpeg-normal/bs_16_fixsave3/170K/512/' | |
dataset_json = os.path.join(objv_dataset, 'dataset.json') | |
with open(dataset_json, 'r') as f: | |
dataset_json = json.load(f) | |
# all_objs = dataset_json['Animals'][::3][:6250] | |
all_objs = dataset_json['Animals'][::3][1100:2200] | |
all_objs = all_objs[:600] | |
if 'GA' in pcd_file_path or 'objv-gt' in pcd_file_path: | |
files = [os.path.join(pcd_file_path, f"{obj.replace('/', '-')}_pcd_4096.ply") for obj in all_objs] | |
filter_files = [] | |
for file in files: | |
if os.path.exists(file): | |
filter_files.append(file) | |
self.files = filter_files | |
else: | |
self.files = sorted(glob.glob(f'{pcd_file_path}/*.ply') ) | |
self.transformation = transformation | |
# self.transforms = transforms | |
# self.reso=reso | |
self.rand_aug = rand_aug | |
# if rand_aug: | |
# else: | |
# self.rand_transform = None | |
def __len__(self): | |
return len(self.files) | |
def __getitem__(self, i): | |
path = self.files[i] | |
pcd = trimesh.load(path).vertices # pcu may fail sometimes | |
pcd = normalize_point_clouds(pcd[None])[0] | |
if self.transformation is not None: | |
pcd = self.transformation(pcd) | |
# if self.rand_aug: | |
# rand_rot = [('x', random.randint(0,359)), ('y', random.randint(0,359)), ('z', random.randint(0,359))] | |
# rand_transform = partial(rotate_point_cloud, rotations=rand_rot) # no transformation | |
# pcd = rand_transform(pcd) # since no canonical space | |
# try: | |
# assert pcd.shape[1]==4096 | |
# except Exception as e: | |
# print(path) | |
return pcd | |
def main(): | |
parser = argparse.ArgumentParser() | |
parser.add_argument("--cache_dir", type=str, default=None) | |
parser.add_argument("batch_1", type=str) | |
parser.add_argument("batch_2", type=str) | |
args = parser.parse_args() | |
print("creating classifier...") | |
clf = PointNetClassifier(devices=get_torch_devices(), cache_dir=args.cache_dir) | |
worker=2 | |
# force_recompute = False | |
force_recompute = True | |
feat_1_path = os.path.join(args.batch_1, 'feat.npy') | |
pred_1_path = os.path.join(args.batch_1, 'pred.npy') | |
# if not force_recompute and all(os.path.exists(path) for path in [feat_1_path, pred_1_path]): | |
# if all(os.path.exists(path) for path in [feat_1_path, pred_1_path]): | |
if not force_recompute and all(os.path.exists(path) for path in [feat_1_path, pred_1_path]): | |
# if all(os.path.exists(path) for path in [feat_1_path, pred_1_path]): | |
print("loading activations", args.batch_1) | |
features_1 = np.load(feat_1_path) | |
preds_1 = np.load(pred_1_path) | |
else: | |
print("computing activations", args.batch_1) | |
# gt_dataset = PCDPathDataset(args.batch_1, transformation_dict['gso']) | |
gt_dataset = PCDPathDataset(args.batch_1, transformation_dict['gso'], rand_aug=False) | |
# gt_dataset = PCDPathDataset(args.batch_1, None, rand_aug=True) | |
# gt | |
gt_loader = torch.utils.data.DataLoader(gt_dataset, | |
batch_size=64, | |
shuffle=False, | |
drop_last=False, | |
num_workers=worker) | |
features_1, preds_1 = clf.features_and_preds(gt_loader) | |
np.save(feat_1_path, features_1) | |
np.save(pred_1_path, preds_1) | |
feat_2_path = os.path.join(args.batch_2, 'feat.npy') | |
pred_2_path = os.path.join(args.batch_2, 'pred.npy') | |
if not force_recompute and all(os.path.exists(path) for path in [feat_2_path, pred_2_path]): | |
features_2 = np.load(feat_2_path) | |
preds_2 = np.load(pred_2_path) | |
print("loading activations", args.batch_2) | |
else: | |
gen_path_base='/mnt/sfs-common/yslan/Repo/3dgen/FID-KID-Outputdir-objv/3D-metrics-fps/' | |
method_name = str(Path(args.batch_2).relative_to(gen_path_base)) | |
# print("computing activations", args.batch_2) | |
print("computing activations", method_name) | |
# method_name = args.batch_2.split('/')[-1] | |
# st() | |
pcd_transformation = transformation_dict[method_name] | |
pred_dataset = PCDPathDataset(args.batch_2, transformation=pcd_transformation, rand_aug=False) | |
# pred_dataset = PCDPathDataset(args.batch_2, transformation=None, rand_aug=True) | |
# worker=0 | |
pred_loader = torch.utils.data.DataLoader(pred_dataset, | |
batch_size=64, | |
shuffle=False, | |
drop_last=False, | |
num_workers=worker) | |
features_2, preds_2 = clf.features_and_preds(pred_loader) | |
np.save(feat_2_path, features_2) | |
np.save(feat_2_path, preds_2) | |
print("computing statistics") | |
stats_1 = compute_statistics(features_1) | |
# print(features_1.max(), features_1.min(), features_1.mean(), features_1.std() ) | |
# print(stats_1.mu.shape, stats_1.sigma.shape) | |
stats_2 = compute_statistics(features_2) | |
# print(features_2.max(), features_2.min(), features_2.mean(), features_2.std() ) | |
# print(stats_2.mu.shape, stats_2.sigma.shape) | |
# if return_feature: | |
# return features_1, features_2 | |
#PFID = stats_1.frechet_distance(stats_2) # same result as the next line | |
PFID= frechet_distance(stats_1.mu, stats_1.sigma, stats_2.mu, stats_2.sigma) | |
PKID = kernel_distance(features_1, features_2) | |
# _, preds = clf.features_and_preds(pred_loader) | |
# print(f"P-IS: {compute_inception_score(preds)}") | |
# print(f"P-IS: {compute_inception_score(preds)}") | |
method_name = args.batch_2.split('/')[-1] | |
# print(method_name, f"P-FID: {PFID}", f"P-KID: {PKID}", f"P-IS: {compute_inception_score(preds_2)}") | |
print(method_name, f"P-FID: {PFID}", f"P-KID: {PKID}") | |
# return dict(PFID=PFID, PKID=PKID) | |
if __name__ == "__main__": | |
main() | |