46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
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)
|