Remove editor module

src/histogram_transfer/__init__.py

@@ -0,0 +1,4 @@
+from .regrain import regrain
+from .pdf_transfer_1d import pdf_transfer_1d
+from .pdf_transfer_3d import pdf_transfer_3d
+from .apply_histogram import apply_histogram

src/histogram_transfer/apply_histogram.py

@@ -0,0 +1,40 @@
+from histogram_transfer import pdf_transfer_3d
+import numpy as np
+from scipy.ndimage import zoom
+
+
+def apply_histogram(original_image, target_histogram, bin_count: int):
+    actual_predicted_histogram = target_histogram.cpu().detach().numpy().squeeze()
+
+    scale = 64 / bin_count
+    scaled_predicted_histogram = zoom(actual_predicted_histogram, scale, order=3)
+    scaled_predicted_histogram = (
+        scaled_predicted_histogram / scaled_predicted_histogram.sum()
+    )
+
+    [h, w, _] = np.array(original_image).shape
+
+    histogram = np.round(scaled_predicted_histogram * h * w).astype(int)
+
+    rgb_vectors = []
+
+    for r in range(histogram.shape[0]):
+        for g in range(histogram.shape[1]):
+            for b in range(histogram.shape[2]):
+                # Append the RGB value 'count' times to the list
+                for _ in range(histogram[r, g, b]):
+                    rgb_vectors.append([r, g, b])
+
+    rgb_vectors = np.array(rgb_vectors)
+    np.random.shuffle(rgb_vectors)
+    rgb_vectors = rgb_vectors * 256 / 64
+
+    return pdf_transfer_3d(
+        source=np.array(original_image),
+        target_flattened=rgb_vectors.transpose(),
+        relaxation=0.9,
+        bin_count=3500,
+        iterations=50,
+        smoothness=1,
+        should_regrain=True,
+    )

src/histogram_transfer/pdf_transfer_1d.py

@@ -0,0 +1,13 @@
+import numpy as np
+
+
+def pdf_transfer_1d(pX: np.ndarray, pY: np.ndarray) -> np.ndarray:
+    PX = np.cumsum(pX + np.finfo(float).eps)
+    PX /= PX[-1]
+
+    PY = np.cumsum(pY + np.finfo(float).eps)
+    PY /= PY[-1]
+
+    f = np.interp(PX, PY, np.arange(len(pX)))
+
+    return f

src/histogram_transfer/pdf_transfer_3d.py

@@ -0,0 +1,47 @@
+import numpy as np
+from utils import generate_rotation_matrices
+from histogram_transfer import pdf_transfer_1d
+from 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,
+    should_regrain: bool = True,
+):
+    [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) if should_regrain else result

src/histogram_transfer/regrain.py

@@ -0,0 +1,87 @@
+from scipy.ndimage import zoom
+import numpy as np
+
+
+NBITS = [4, 16, 32, 64, 64, 64]
+
+
+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, 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, 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:
+        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,
+            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, smoothness
+    )
+    return img_arr_out
+
+
+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)
+    last_pad_0 = lambda arr: np.concatenate((arr[1:, :], arr[-1:, :]), axis=0)
+    last_pad_1 = lambda arr: np.concatenate((arr[:, 1:], arr[:, -1:]), axis=1)
+
+    delta_x = last_pad_1(img_arr_in) - first_pad_1(img_arr_in)
+    delta_y = last_pad_0(img_arr_in) - first_pad_0(img_arr_in)
+    delta = np.sqrt((delta_x**2 + delta_y**2).sum(axis=2, keepdims=True))
+
+    psi = 256 * delta / 5
+    psi[psi > 1] = 1
+    phi = 30 * 2 ** (-level) / (1 + 10 * delta / smoothness)
+
+    phi1 = (last_pad_1(phi) + phi) / 2
+    phi2 = (last_pad_0(phi) + phi) / 2
+    phi3 = (first_pad_1(phi) + phi) / 2
+    phi4 = (first_pad_0(phi) + phi) / 2
+
+    rho = 1 / 5.0
+    for i in range(nbit):
+        den = psi + phi1 + phi2 + phi3 + phi4
+        num = (
+            np.tile(psi, [1, 1, c]) * img_arr_col
+            + np.tile(phi1, [1, 1, c])
+            * (last_pad_1(img_arr_out) - last_pad_1(img_arr_in) + img_arr_in)
+            + np.tile(phi2, [1, 1, c])
+            * (last_pad_0(img_arr_out) - last_pad_0(img_arr_in) + img_arr_in)
+            + np.tile(phi3, [1, 1, c])
+            * (first_pad_1(img_arr_out) - first_pad_1(img_arr_in) + img_arr_in)
+            + np.tile(phi4, [1, 1, c])
+            * (first_pad_0(img_arr_out) - first_pad_0(img_arr_in) + img_arr_in)
+        )
+        img_arr_out = (
+            num / np.tile(den + 1e-6, [1, 1, c]) * (1 - rho) + rho * img_arr_out
+        )
+    return img_arr_out
+
+
+def _resize_image(data, target_width, target_height):
+    return zoom(data, (target_height / data.shape[0], target_width / data.shape[1], 1))