1 files changed, 69 insertions, 0 deletions
diff --git a/similarity.py b/similarity.py
new file mode 100644
index 0000000..fec83d6
--- /dev/null
+++ b/similarity.py
@@ -0,0 +1,69 @@
+from gaussian import Gaussian
+from utils import quat2rot
+from aabb import process_aabb, intersect, min_interval
+import numpy as np
+
+def similarity(gaussian: Gaussian, id1: int, id2: int):
+    geometry_sim = geometry_similarity(gaussian, id1, id2)
+    color_sim = color_similarity(gaussian.sh[id1], gaussian.sh[id2])
+    color_max = np.max(np.abs(gaussian.sh))
+    # return geometry_sim * geometry_coeff + color_sim * color_coeff
+    return geometry_sim, color_sim / (color_max * color_max * 48)
+    
+def color_similarity(c1: np.ndarray, c2: np.ndarray):
+    return np.sum(c1 * c2)
+
+def geometry_similarity(gaussian: Gaussian, id1: int, id2: int):
+    def geometry_similarity_using_overlap(gaussian: Gaussian, id1: int, id2: int):
+        aabb1 = process_aabb(gaussian.positions[id1], gaussian.scales[id1], gaussian.rotations[id1], scale_factor=2)
+        aabb2 = process_aabb(gaussian.positions[id2], gaussian.scales[id2], gaussian.rotations[id2], scale_factor=2)
+        # 在两个 aabb 里 stratified sample 点，对两个 Gaussian 的乘积做积分
+        # 然后除以总sample数作为重叠程度的 metric
+
+        intersected_aabb = intersect(aabb1, aabb2)
+
+        if intersected_aabb is None:
+            return 0
+
+        x_stride = (intersected_aabb[1] - intersected_aabb[0]) / (4. + 1e-3)
+        y_stride = (intersected_aabb[3] - intersected_aabb[2]) / (4. + 1e-3)
+        z_stride = (intersected_aabb[5] - intersected_aabb[4]) / (4. + 1e-3)
+        
+        score = 0
+        for i in range(int(np.floor((intersected_aabb[1] - intersected_aabb[0]) / x_stride))):
+            x = intersected_aabb[0] + i * x_stride
+            for j in range(int(np.floor((intersected_aabb[3] - intersected_aabb[2]) / y_stride))):
+                y = intersected_aabb[2] + j * y_stride
+                for k in range(int(np.floor((intersected_aabb[5] - intersected_aabb[4]) / z_stride))):
+                    z = intersected_aabb[4] + k * z_stride
+                    pos = np.array([x, y, z])
+                    score += gaussian.calc(pos, id1, 10) * gaussian.calc(pos, id2, 10) * gaussian.opacity[id1] * gaussian.opacity[id2]
+
+        scale1 = gaussian.scales[id1][0] * gaussian.scales[id1][1] * gaussian.scales[id1][2]
+        scale2 = gaussian.scales[id2][0] * gaussian.scales[id2][1] * gaussian.scales[id2][2]
+
+        # print(f"calculating pair ({id1},{id2}), used {64} samples.")
+        return score / (scale1 * scale2)
+
+    def geometry_similarity_using_position(gaussian: Gaussian, id1: int, id2: int):
+        def position_similarity(p1: np.ndarray, p2: np.ndarray):
+            return np.linalg.norm(p1 - p2)
+        
+        def scale_similarity(s1: np.ndarray, s2: np.ndarray):
+            return np.linalg.norm(s1 - s2)
+
+        def rotation_similarity(q1: np.ndarray, q2: np.ndarray):
+            return 1 - np.dot(q1, q2)
+            # r1 = quat2rot(q1)
+            # r2 = quat2rot(q2)
+            # eigenvalues, eigenvectors = np.linalg.eig(r1 @ r2.T)
+            # return np.linalg.norm(eigenvectors @ np.diag(np.log(eigenvalues)) @ np.linalg.inv(eigenvectors)) 
+        position_coeff = 1.
+        scale_coeff = 1.
+        rotation_coeff = 1.
+        return position_similarity(gaussian.positions[id1], gaussian.positions[id2]) * position_coeff + \
+            scale_similarity(gaussian.scales[id1], gaussian.scales[id2]) * scale_coeff + \
+            rotation_similarity(gaussian.rotations[id1], gaussian.rotations[id2]) * rotation_coeff
+
+    return geometry_similarity_using_overlap(gaussian, id1, id2)
+