Refactor pdf transfer

editor/histogram_transfer/__init__.py

@@ -1,2 +1,3 @@
 from .regrain import regrain
 from .pdf_transfer_1d import pdf_transfer_1d
+from .pdf_transfer_3d import pdf_transfer_3d

editor/histogram_transfer/pdf_transfer_3d.py

@@ -0,0 +1,46 @@
+import numpy as np
+from editor.utils import generate_rotation_matrices
+from editor.histogram_transfer import pdf_transfer_1d
+from editor.histogram_transfer import regrain
+
+
+EPSILON = 1e-6
+
+
+def pdf_transfer_3d(
+    source: np.ndarray,
+    target_flattened: np.ndarray,
+    relaxation: float = 1,
+    bin_count: int = 1000,
+    iterations: int = 25,
+    smoothness: float = 1,
+):
+    [h, w, c] = source.shape
+    source_flattened = source.reshape(-1, c).transpose()
+
+    rotation_matrices = generate_rotation_matrices(iterations)
+    for i, rotation in enumerate(rotation_matrices, start=1):
+        D0R = rotation @ source_flattened
+        D1R = rotation @ target_flattened
+        D0R_ = np.zeros_like(source_flattened)
+
+        for i in range(rotation.shape[0]):
+            datamin = min(np.min(D0R[i, :]), np.min(D1R[i, :])) - EPSILON
+            datamax = max(np.max(D0R[i, :]), np.max(D1R[i, :])) + EPSILON
+            u = np.linspace(datamin, datamax, bin_count)
+
+            p0R, _ = np.histogram(D0R[i, :], bins=u, density=True)
+            p1R, _ = np.histogram(D1R[i, :], bins=u, density=True)
+
+            f = pdf_transfer_1d(p0R, p1R)
+            mapped_values = (
+                np.interp(D0R[i, :], u[:-1], f) * (datamax - datamin) / (bin_count - 1)
+                + datamin
+            )
+            D0R_[i, :] = mapped_values
+
+        source_flattened = source_flattened + relaxation * (rotation.T @ (D0R_ - D0R))
+        source_flattened.clip(0, 255, out=source_flattened)
+
+    result = source_flattened.astype(np.uint8).transpose().reshape(h, w, c)
+    return regrain(source, result, smoothness=smoothness)