wip

editor/histogram_transfer/regrain.py

@@ -2,42 +2,48 @@ from scipy.ndimage import zoom
 import numpy as np
 
 
-SMOOTHNESS = 1
 NBITS = [4, 16, 32, 64, 64, 64]
 
 
-def regrain(img_arr_in, img_arr_col):
+def regrain(img_arr_in, img_arr_col, smoothness: float = 1):
     """keep gradient of img_arr_in and color of img_arr_col."""
 
     img_arr_in = img_arr_in / 255.0
     img_arr_col = img_arr_col / 255.0
     img_arr_out = np.array(img_arr_in)
-    img_arr_out = _regrain_rec(img_arr_out, img_arr_in, img_arr_col, 0)
+    img_arr_out = _regrain_rec(
+        img_arr_out, img_arr_in, img_arr_col, NBITS, 0, smoothness
+    )
     img_arr_out[img_arr_out < 0] = 0
     img_arr_out[img_arr_out > 1] = 1
     img_arr_out = (255.0 * img_arr_out).astype("uint8")
     return img_arr_out
 
 
-def _regrain_rec(img_arr_out, img_arr_in, img_arr_col, level):
-    """direct translation of matlab code."""
-
+def _regrain_rec(img_arr_out, img_arr_in, img_arr_col, nbits, level, smoothness):
     [h, w, _] = img_arr_in.shape
     h2 = (h + 1) // 2
     w2 = (w + 1) // 2
-    if len(NBITS) > 1 and h2 > 20 and w2 > 20:
+    if len(nbits) > 1 and h2 > 20 and w2 > 20:
         resize_arr_in = _resize_image(img_arr_in, w2, h2)
         resize_arr_col = _resize_image(img_arr_col, w2, h2)
         resize_arr_out = _resize_image(img_arr_out, w2, h2)
         resize_arr_out = _regrain_rec(
-            resize_arr_out, resize_arr_in, resize_arr_col, NBITS[1:], level + 1
+            resize_arr_out,
+            resize_arr_in,
+            resize_arr_col,
+            nbits[1:],
+            level + 1,
+            smoothness,
         )
         img_arr_out = _resize_image(resize_arr_out, w, h)
-    img_arr_out = _solve(img_arr_out, img_arr_in, img_arr_col, NBITS[0], level)
+    img_arr_out = _solve(
+        img_arr_out, img_arr_in, img_arr_col, nbits[0], level, smoothness
+    )
     return img_arr_out
 
 
-def _solve(img_arr_out, img_arr_in, img_arr_col, nbit, level, eps=1e-6):
+def _solve(img_arr_out, img_arr_in, img_arr_col, nbit, level, smoothness):
     [width, height, c] = img_arr_in.shape
     first_pad_0 = lambda arr: np.concatenate((arr[:1, :], arr[:-1, :]), axis=0)
     first_pad_1 = lambda arr: np.concatenate((arr[:, :1], arr[:, :-1]), axis=1)
@@ -50,7 +56,7 @@ def _solve(img_arr_out, img_arr_in, img_arr_col, nbit, level, eps=1e-6):
 
     psi = 256 * delta / 5
     psi[psi > 1] = 1
-    phi = 30 * 2 ** (-level) / (1 + 10 * delta / SMOOTHNESS)
+    phi = 30 * 2 ** (-level) / (1 + 10 * delta / smoothness)
 
     phi1 = (last_pad_1(phi) + phi) / 2
     phi2 = (last_pad_0(phi) + phi) / 2
@@ -72,7 +78,7 @@ def _solve(img_arr_out, img_arr_in, img_arr_col, nbit, level, eps=1e-6):
             * (first_pad_0(img_arr_out) - first_pad_0(img_arr_in) + img_arr_in)
         )
         img_arr_out = (
-            num / np.tile(den + eps, [1, 1, c]) * (1 - rho) + rho * img_arr_out
+            num / np.tile(den + 1e-6, [1, 1, c]) * (1 - rho) + rho * img_arr_out
         )
     return img_arr_out
 

editor/utils/__init__.py

@@ -4,3 +4,4 @@ from .apply_pixel_shader import apply_pixel_shader
 from .get_colour_lut import get_colour_lut
 from .compute_histogram import compute_histogram
 from .kldiv import kldiv
+from .generate_rotation_matrices import generate_rotation_matrices

editor/utils/generate_rotation_matrices.py

@@ -0,0 +1,66 @@
+from random import shuffle
+from typing import List, Tuple
+import numpy as np
+from functools import lru_cache
+from numpy.typing import NDArray
+
+
+@lru_cache
+def generate_rotation_matrices(count: int) -> List[NDArray[np.float64]]:
+    axes = fibonacci_sphere(count)
+    shuffle(axes)
+    angles = np.linspace(0, 2 * np.pi, count, endpoint=False)
+    matrices = [_rotation_matrix(axis, angle) for axis, angle in zip(axes, angles)]
+    for matrix in matrices:
+        _check_rotation_matrix(matrix)
+    return matrices
+
+
+def fibonacci_sphere(samples: int) -> List[Tuple[float, float, float]]:
+    points = []
+    phi = np.pi * (3.0 - np.sqrt(5.0))  # Golden angle in radians
+    for i in range(samples):
+        y = 1 - (i / float(samples - 1)) * 2  # y goes from 1 to -1
+        radius = np.sqrt(1 - y * y)  # radius at y
+
+        theta = phi * i  # golden angle increment
+
+        x = np.cos(theta) * radius
+        z = np.sin(theta) * radius
+
+        points.append([x, y, z])
+    return points
+
+
+def _rotation_matrix(
+    axis: Tuple[float, float, float], theta: float
+) -> NDArray[np.float64]:
+    axis = np.asarray(axis)
+    axis = axis / np.sqrt(np.dot(axis, axis))
+    a = np.cos(theta / 2.0)
+    b, c, d = -axis * np.sin(theta / 2.0)
+
+    aa, bb, cc, dd = a * a, b * b, c * c, d * d
+    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
+    return np.array(
+        [
+            [aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
+            [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
+            [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc],
+        ]
+    )
+
+
+def _check_rotation_matrix(R: NDArray[np.float64]):
+    # Check if the matrix is square
+    if R.shape != (3, 3):
+        raise ValueError("Matrix must be 3x3.")
+
+    # Check orthogonality: R.T * R should be close to the identity matrix
+    I = np.eye(3)
+    if not np.allclose(np.dot(R.T, R), I):
+        raise ValueError("allclose")
+
+    # Check determinant: Should be +1
+    if not np.isclose(np.linalg.det(R), 1.0):
+        raise ValueError(f"det {np.linalg.det(R)}")