diff --git a/src/editor/histogram_transfer/__init__.py b/src/editor/histogram_transfer/__init__.py index f039d7a..bfc8e2a 100644 --- a/src/editor/histogram_transfer/__init__.py +++ b/src/editor/histogram_transfer/__init__.py @@ -1,3 +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 diff --git a/src/editor/histogram_transfer/apply_histogram.py b/src/editor/histogram_transfer/apply_histogram.py new file mode 100644 index 0000000..6d7e9ea --- /dev/null +++ b/src/editor/histogram_transfer/apply_histogram.py @@ -0,0 +1,40 @@ +from editor.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, + )