diff --git a/eval_pose.py b/eval_pose.py
new file mode 100644
index 0000000000000000000000000000000000000000..467a4d07ccc72670c11eb20f733d6c847810142b
--- /dev/null
+++ b/eval_pose.py
@@ -0,0 +1,308 @@
+import data
+import numpy as np
+import math
+from scipy import spatial
+from scipy.linalg import logm
+from numpy import linalg as LA
+import open3d as o3d
+import argparse
+import os
+import sys
+
+
+def read_diameter(object_name):
+ filename = f'Generated_Worlds_/Generated/{class_name}/diameter.txt'
+ with open(filename) as f:
+ diameter_in_cm = float(f.readline())
+ return diameter_in_cm * 0.01
+
+
+def transform_pts_Rt(pts, R, t):
+ """Applies a rigid transformation to 3D points.
+ :param pts: nx3 ndarray with 3D points.
+ :param R: 3x3 ndarray with a rotation matrix.
+ :param t: 3x1 ndarray with a translation vector.
+ :return: nx3 ndarray with transformed 3D points.
+ """
+ assert (pts.shape[1] == 3)
+ pts_t = R.dot(pts.T) + t.reshape((3, 1))
+ return pts_t.T
+
+
+def project_pts(pts, K, R, t):
+ """Projects 3D points.
+ :param pts: nx3 ndarray with the 3D points.
+ :param K: 3x3 ndarray with an intrinsic camera matrix.
+ :param R: 3x3 ndarray with a rotation matrix.
+ :param t: 3x1 ndarray with a translation vector.
+ :return: nx2 ndarray with 2D image coordinates of the projections.
+ """
+ assert (pts.shape[1] == 3)
+ P = K.dot(np.hstack((R, t)))
+ pts_h = np.hstack((pts, np.ones((pts.shape[0], 1))))
+ pts_im = P.dot(pts_h.T)
+ pts_im /= pts_im[2, :]
+ return pts_im[:2, :].T
+
+
+def add(R_est, t_est, R_gt, t_gt, pts):
+ """Average Distance of Model Points for objects with no indistinguishable
+ views - by Hinterstoisser et al. (ACCV'12).
+ :param R_est: 3x3 ndarray with the estimated rotation matrix.
+ :param t_est: 3x1 ndarray with the estimated translation vector.
+ :param R_gt: 3x3 ndarray with the ground-truth rotation matrix.
+ :param t_gt: 3x1 ndarray with the ground-truth translation vector.
+ :param pts: nx3 ndarray with 3D model points.
+ :return: The calculated error.
+ """
+ pts_est = transform_pts_Rt(pts, R_est, t_est)
+ pts_gt = transform_pts_Rt(pts, R_gt, t_gt)
+ e = np.linalg.norm(pts_est - pts_gt, axis=1).mean()
+ return e
+
+
+def adi(R_est, t_est, R_gt, t_gt, pts):
+ """Average Distance of Model Points for objects with indistinguishable views
+ - by Hinterstoisser et al. (ACCV'12).
+ :param R_est: 3x3 ndarray with the estimated rotation matrix.
+ :param t_est: 3x1 ndarray with the estimated translation vector.
+ :param R_gt: 3x3 ndarray with the ground-truth rotation matrix.
+ :param t_gt: 3x1 ndarray with the ground-truth translation vector.
+ :param pts: nx3 ndarray with 3D model points.
+ :return: The calculated error.
+ """
+ pts_est = transform_pts_Rt(pts, R_est, t_est)
+ pts_gt = transform_pts_Rt(pts, R_gt, t_gt)
+
+ # Calculate distances to the nearest neighbors from vertices in the
+ # ground-truth pose to vertices in the estimated pose.
+ nn_index = spatial.cKDTree(pts_est)
+ nn_dists, _ = nn_index.query(pts_gt, k=1)
+
+ e = nn_dists.mean()
+ return e
+
+
+def re(R_est, R_gt):
+ """Rotational Error.
+ :param R_est: 3x3 ndarray with the estimated rotation matrix.
+ :param R_gt: 3x3 ndarray with the ground-truth rotation matrix.
+ :return: The calculated error.
+ """
+ error_cos = float(0.5 * (np.trace(R_est.dot(np.linalg.inv(R_gt))) - 1.0))
+
+ # Avoid invalid values due to numerical errors.
+ error_cos = min(1.0, max(-1.0, error_cos))
+
+ error = math.acos(error_cos)
+ error = 180.0 * error / np.pi # Convert [rad] to [deg].
+ return error
+
+
+def te(t_est, t_gt):
+ """Translational Error.
+ :param t_est: 3x1 ndarray with the estimated translation vector.
+ :param t_gt: 3x1 ndarray with the ground-truth translation vector.
+ :return: The calculated error.
+ """
+ assert (t_est.size == t_gt.size == 3)
+ error = np.linalg.norm(t_gt - t_est)
+ return error
+
+
+def proj(R_est, t_est, R_gt, t_gt, K, pts):
+ """Average distance of projections of object model vertices [px]
+ - by Brachmann et al. (CVPR'16).
+ :param R_est: 3x3 ndarray with the estimated rotation matrix.
+ :param t_est: 3x1 ndarray with the estimated translation vector.
+ :param R_gt: 3x3 ndarray with the ground-truth rotation matrix.
+ :param t_gt: 3x1 ndarray with the ground-truth translation vector.
+ :param K: 3x3 ndarray with an intrinsic camera matrix.
+ :param pts: nx3 ndarray with 3D model points.
+ :return: The calculated error.
+ """
+ proj_est = project_pts(pts, K, R_est, t_est)
+ proj_gt = project_pts(pts, K, R_gt, t_gt)
+ e = np.linalg.norm(proj_est - proj_gt, axis=1).mean()
+ return e
+
+
+def compute_add_score(pts3d, diameter, R_pred, t_pred, R_gt, t_gt, percentage=0.1):
+ # R_gt, t_gt = pose_gt
+ # R_pred, t_pred = pose_pred
+ count = R_gt.shape[0]
+ mean_distances = np.zeros((count,), dtype=np.float32)
+ for i in range(count):
+ pts_xformed_gt = R_gt[i] * pts3d.transpose() + t_gt[i]
+ pts_xformed_pred = R_pred[i] * pts3d.transpose() + t_pred[i]
+ distance = np.linalg.norm(pts_xformed_gt - pts_xformed_pred, axis=0)
+ mean_distances[i] = np.mean(distance)
+
+ threshold = diameter * percentage
+ score = (mean_distances < threshold).sum() / count
+ return score
+
+
+def compute_adds_score(pts3d, diameter, pose_gt, pose_pred, percentage=0.1):
+ R_gt, t_gt = pose_gt
+ R_pred, t_pred = pose_pred
+
+ count = R_gt.shape[0]
+ mean_distances = np.zeros((count,), dtype=np.float32)
+ for i in range(count):
+ if np.isnan(np.sum(t_pred[i])):
+ mean_distances[i] = np.inf
+ continue
+ pts_xformed_gt = R_gt[i] * pts3d.transpose() + t_gt[i]
+ pts_xformed_pred = R_pred[i] * pts3d.transpose() + t_pred[i]
+ kdt = spatial.KDTree(pts_xformed_gt.transpose(), metric='euclidean')
+ distance, _ = kdt.query(pts_xformed_pred.transpose(), k=1)
+ mean_distances[i] = np.mean(distance)
+ threshold = diameter * percentage
+ score = (mean_distances < threshold).sum() / count
+ return score
+
+
+def compute_pose_error(diameter, pose_gt, pose_pred):
+ R_gt, t_gt = pose_gt
+ R_pred, t_pred = pose_pred
+
+ count = R_gt.shape[0]
+ R_err = np.zeros(count)
+ t_err = np.zeros(count)
+ for i in range(count):
+ if np.isnan(np.sum(t_pred[i])):
+ continue
+ r_err = logm(np.dot(R_pred[i].transpose(), R_gt[i])) / 2
+ R_err[i] = LA.norm(r_err, 'fro')
+ t_err[i] = LA.norm(t_pred[i] - t_gt[i])
+ return np.median(R_err) * 180 / np.pi, np.median(t_err) / diameter
+
+
+def cou_mask(mask_est, mask_gt):
+ """Complement over Union of 2D binary masks.
+ :param mask_est: hxw ndarray with the estimated mask.
+ :param mask_gt: hxw ndarray with the ground-truth mask.
+ :return: The calculated error.
+ """
+ mask_est_bool = mask_est.astype(bool)
+ mask_gt_bool = mask_gt.astype(bool)
+
+ inter = np.logical_and(mask_gt_bool, mask_est_bool)
+ union = np.logical_or(mask_gt_bool, mask_est_bool)
+
+ union_count = float(union.sum())
+ if union_count > 0:
+ e = 1.0 - inter.sum() / union_count
+ else:
+ e = 1.0
+ return e
+
+
+def find_nn_idx(pc_src, pc_target):
+ """
+ pc_src: (N1, 3) array
+ pc_target: (N2, 3) array
+ """
+ dist_sq = -np.dot(pc_src, pc_target.T) * 2 + \
+ np.square(np.linalg.norm(pc_src, axis=-1, keepdims=True)) + \
+ np.square(np.linalg.norm(pc_target.T, axis=0, keepdims=True))
+ idx_min = np.argmin(dist_sq, axis=0)
+ return idx_min
+
+
+def add_metric(pose_pred, pose_targets, obj_points, diameter, symm=False, percentage=0.1, gpu=False):
+
+ diam = diameter * percentage
+ model_pred = np.dot(obj_points, pose_pred[:, :3].T) + pose_pred[:, 3]
+ model_targets = np.dot(obj_points, pose_targets[:, :3].T) + pose_targets[:, 3]
+
+ if symm:
+ # if gpu:
+ # idxs = nn_utils.find_nearest_point_idx(model_pred, model_targets)
+ # else:
+ idxs = find_nn_idx(model_pred, model_targets)
+ mean_dist = np.mean(np.linalg.norm(model_pred[idxs] - model_targets, 2, 1))
+ else:
+ mean_dist = np.mean(np.linalg.norm(model_pred - model_targets, axis=-1))
+
+ return mean_dist < diam, (mean_dist, diam)
+
+
+def eval_pose(r_est, t_est, r_gt, t_gt, pc, k, diameter, sym=False):
+ add_res = add(r_est, t_est, r_gt, t_gt, pc)
+ is_add = add(r_est, t_est, r_gt, t_gt, pc) < diameter * 0.1
+ proj_res = proj(r_est, t_est, r_gt, t_gt, k, pc)
+ is_proj = proj(r_est, t_est, r_gt, t_gt, k, pc) < 5
+ adi_res = 0
+ is_adi = False
+ if sym:
+ add_res = adi(r_est, t_est, r_gt, t_gt, pc)
+ is_add = adi(r_est, t_est, r_gt, t_gt, pc) < diameter * 0.1
+
+ return add_res, is_add, proj_res, is_proj, adi_res, is_adi
+
+
+if __name__ == '__main__':
+ ap = argparse.ArgumentParser()
+ ap.add_argument("-cls_name", "--class_name", type=str,
+ default='kiwi1',
+ help="[apple2, apricot, banana1, kiwi1, lemon2, orange2, peach1, pear2]")
+ ap.add_argument("-sym", "--symmetry", type=bool,
+ default=False)
+
+ args = vars(ap.parse_args())
+
+ class_name = args["class_name"]
+ symmetry = args["symmetry"]
+
+ basePath = os.path.dirname(
+ os.path.realpath(__file__)) + '/Generated_Worlds_/Generated_Worlds_Evaluating/' + class_name
+
+ images_ls, labels_ls, mask_ls, choice_ls = data.getAllValData(class_name)
+ print(len(images_ls))
+
+ pc_path = f'/{basePath}/Models/{class_name}.ply'
+ pcd = o3d.io.read_point_cloud(pc_path)
+ pc = np.asarray(pcd.points)
+
+ add_res_ls = []
+ proj_res_ls = []
+ count_add = 0
+ count_iadd = 0
+ count_proj = 0
+
+
+
+ for idx in range(len(images_ls)):
+ # ============== Loading Pose ===============
+ pose_est = np.load(f'{basePath}/Pose_prediction/{idx}.npy')
+ r_est = pose_est[:3, :3]
+ t_est = pose_est[:3, 3]
+
+ pose_gt = np.load(f'{basePath}/Pose_transformed/{idx}.npy')
+ r_gt = np.array(pose_gt[0:3, 0:3], dtype='float64')
+ t_gt = np.array(pose_gt[0:3, 3], dtype='float64').reshape(3, 1)
+
+ # ============== Evaluation ===============
+ k = np.array([[543.25272224, 0., 320.25],
+ [0., 724.33696299, 240.33333333],
+ [0., 0., 1.]])
+ diameter = read_diameter(class_name)
+
+ add_res, is_add, proj_res, is_proj, adi_res, is_adi = eval_pose(r_est, t_est, r_gt, t_gt, pc, k, diameter, symmetry)
+
+ if is_add:
+ count_add += 1
+
+ if is_proj:
+ count_proj += 1
+ if is_adi:
+ count_iadd += 1
+
+ print(f"ADD_Res: {count_add / len(images_ls)}")
+ print(f"ADI_Res: {count_iadd / len(images_ls)}")
+ print(f"Proj_Res: {count_proj / len(images_ls)}")
+ print(f"======================")
+
+ print("Done")