diff options
Diffstat (limited to 'Scripts/gaussianize')
| -rwxr-xr-x | Scripts/gaussianize | 496 |
1 files changed, 496 insertions, 0 deletions
diff --git a/Scripts/gaussianize b/Scripts/gaussianize new file mode 100755 index 0000000..2ca915d --- /dev/null +++ b/Scripts/gaussianize @@ -0,0 +1,496 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "numpy", +# "scipy", +# "pillow", +# "openexr", +# "imath", +# "matplotlib" +# ] +# /// + +""" +Gaussianization for Histogram-Preserving Blending +Based on Burley's "On Histogram-preserving Blending for Randomized Texture Tiling" (2019) + +This implementation uses per-channel 1D histogram transformation with: +- Truncated Gaussian distribution (avoiding values outside [0,1]) +- Soft-clipping contrast operator +- Fast 1D LUT generation from input histogram +""" + +import argparse +import sys +from pathlib import Path +import numpy as np +from scipy import special +from PIL import Image +import OpenEXR +import Imath + + +def load_image(image_path: Path) -> np.ndarray: + """Load a PNG or EXR image and return as float64 array (H, W, 3) in [0,1].""" + suffix = image_path.suffix.lower() + if suffix == '.exr': + exr = OpenEXR.InputFile(str(image_path)) + header = exr.header() + dw = header['dataWindow'] + w = dw.max.x - dw.min.x + 1 + h = dw.max.y - dw.min.y + 1 + channels = [] + for ch in ('R', 'G', 'B'): + raw = exr.channel(ch, Imath.PixelType(Imath.PixelType.HALF)) + channels.append(np.frombuffer(raw, dtype=np.float16).reshape(h, w)) + return np.stack(channels, axis=-1).astype(np.float64) + else: + return np.array(Image.open(image_path).convert("RGB")).astype(np.float64) / 255.0 + + +def save_image(image: np.ndarray, output_path: Path): + """Save an RGB float image as EXR or PNG depending on the file suffix.""" + suffix = output_path.suffix.lower() + if suffix == ".exr": + image_to_save = image.astype(np.float16) + h, w, _ = image_to_save.shape + header = OpenEXR.Header(w, h) + half_chan = Imath.Channel(Imath.PixelType(Imath.PixelType.HALF)) + header['channels'] = {'R': half_chan, 'G': half_chan, 'B': half_chan} + out = OpenEXR.OutputFile(str(output_path), header) + out.writePixels({ + 'R': image_to_save[:, :, 0].tobytes(), + 'G': image_to_save[:, :, 1].tobytes(), + 'B': image_to_save[:, :, 2].tobytes(), + }) + out.close() + elif suffix == ".png": + clipped = np.clip(image, 0.0, 1.0) + Image.fromarray(np.round(clipped * 255.0).astype(np.uint8), mode="RGB").save(output_path) + else: + raise ValueError(f"Unsupported output format '{output_path.suffix}'. Use .exr or .png.") + + +class TruncatedGaussian: + """Truncated Gaussian distribution for histogram transformation.""" + + def __init__(self, sigma: float = 1.0 / 6.0): + """ + Initialize truncated Gaussian centered at 0.5 with given sigma. + Default sigma = 1/6 as recommended in Burley's paper. + """ + self.sigma = sigma + self.mu = 0.5 + + # Burley's C(sigma) is the reciprocal normalization factor required + # after truncating the Gaussian to the [0, 1] interval. + self.C = 1.0 / special.erf(1.0 / (2.0 * np.sqrt(2.0) * sigma)) + + def inverse_cdf(self, u: np.ndarray) -> np.ndarray: + """ + Inverse CDF of truncated Gaussian distribution. + Maps uniform values in [0,1] to truncated Gaussian in [0,1]. + + Equation (3) from Burley's paper: + CDF^-1_[G](u; σ) = 1/2 + sqrt(2)σ * erfinv((2u - 1) / C(σ)) + """ + u = np.clip(u, 0.0, 1.0) + result = 0.5 + np.sqrt(2.0) * self.sigma * special.erfinv((2.0 * u - 1.0) / self.C) + return np.clip(result, 0.0, 1.0) + + def cdf(self, x: np.ndarray) -> np.ndarray: + """ + CDF of truncated Gaussian distribution. + Maps truncated Gaussian values to uniform [0,1]. + """ + x = np.clip(x, 0.0, 1.0) + result = 0.5 * (1.0 + self.C * special.erf((x - 0.5) / (np.sqrt(2.0) * self.sigma))) + return np.clip(result, 0.0, 1.0) + + +def _cdf_bin_edges(histogram: np.ndarray) -> np.ndarray: + """Return normalized CDF samples on histogram bin edges.""" + histogram = np.asarray(histogram, dtype=np.float64) + total = float(histogram.sum()) + if total <= 0.0: + return np.linspace(0.0, 1.0, len(histogram) + 1) + return np.concatenate(([0.0], np.cumsum(histogram / total, dtype=np.float64))) + + +def _occupied_bin_mid_quantiles(histogram: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Return source-value centers and CDF midpoints for occupied bins.""" + histogram = np.asarray(histogram, dtype=np.float64) + n_bins = len(histogram) + cdf_edges = _cdf_bin_edges(histogram) + bin_mass = cdf_edges[1:] - cdf_edges[:-1] + occupied = bin_mass > 0.0 + + if not np.any(occupied): + centers = (np.arange(n_bins, dtype=np.float64) + 0.5) / n_bins + return centers, centers + + centers = (np.arange(n_bins, dtype=np.float64) + 0.5) / n_bins + mid_quantiles = cdf_edges[:-1] + 0.5 * bin_mass + return centers[occupied], mid_quantiles[occupied] + + +def build_gaussianization_lut(histogram: np.ndarray, lut_size: int = 4096) -> np.ndarray: + """ + Build 1D LUT for Gaussianizing a channel based on its histogram. + + Algorithm 1 from Burley's paper: + 1. Compute CDF from histogram + 2. Transform through inverse CDF of truncated Gaussian + """ + gaussian = TruncatedGaussian() + value_centers, mid_quantiles = _occupied_bin_mid_quantiles(histogram) + mapped_centers = gaussian.inverse_cdf(mid_quantiles) + + sample_positions = np.linspace(0.0, 1.0, lut_size) + return np.interp( + sample_positions, + value_centers, + mapped_centers, + left=mapped_centers[0], + right=mapped_centers[-1], + ) + + +def apply_lut(image: np.ndarray, lut: np.ndarray) -> np.ndarray: + """Apply a 1D LUT to an image channel using linear interpolation.""" + lut_size = len(lut) + coords = np.clip(image, 0.0, 1.0) * (lut_size - 1) + indices0 = np.floor(coords).astype(np.int32) + indices1 = np.minimum(indices0 + 1, lut_size - 1) + alpha = coords - indices0 + return (1.0 - alpha) * lut[indices0] + alpha * lut[indices1] + + +def _deterministic_noise(shape: tuple[int, int], channel: int) -> np.ndarray: + """Generate stable per-pixel noise in [0, 1) from pixel coordinates.""" + height, width = shape + yy, xx = np.indices((height, width), dtype=np.uint32) + state = xx * np.uint32(0x1F123BB5) ^ yy * np.uint32(0x159A55E5) ^ np.uint32(channel + 1) * np.uint32(0x2C1B3C6D) + state ^= state >> np.uint32(16) + state *= np.uint32(0x7FEB352D) + state ^= state >> np.uint32(15) + state *= np.uint32(0x846CA68B) + state ^= state >> np.uint32(16) + return state.astype(np.float64) / float(np.iinfo(np.uint32).max) + + +def dither_channel(channel: np.ndarray, quantization_step: float, channel_index: int) -> np.ndarray: + """Spread repeated quantized values across their source bucket deterministically.""" + if quantization_step <= 0.0: + return channel + noise = _deterministic_noise(channel.shape, channel_index) - 0.5 + return np.clip(channel + noise * quantization_step, 0.0, 1.0) + + +def _soft_clipping_lower_half(x_hat: np.ndarray, W_hat: float) -> np.ndarray: + """Evaluate Burley's Eq. 4 on the lower half of the domain.""" + linear_start = (2.0 - W_hat) / 4.0 + linear = (x_hat - 0.5) / W_hat + 0.5 + + if W_hat >= (2.0 / 3.0): + t = x_hat / (2.0 - W_hat) + quadratic = 8.0 * (1.0 / W_hat - 1.0) * (t ** 2) + (3.0 - 2.0 / W_hat) * t + return np.where(x_hat >= linear_start, linear, quadratic) + + quadratic_start = (2.0 - 3.0 * W_hat) / 4.0 + quadratic = ((x_hat - quadratic_start) / W_hat) ** 2 + return np.where( + x_hat >= linear_start, + linear, + np.where(x_hat >= quadratic_start, quadratic, 0.0), + ) + + +def soft_clipping_contrast(x_hat: np.ndarray, W_hat: float) -> np.ndarray: + """ + Soft-clipping contrast operator S*_[G] from Equation (4) in Burley's paper. + + This is a piecewise function that: + - Is linear in the middle half of the range + - Blends smoothly to 0 or 1 using quadratic segments at the ends + """ + if not (0.0 < W_hat <= 1.0): + raise ValueError(f"W_hat must be in (0, 1], got {W_hat}") + + x_hat = np.clip(x_hat, 0.0, 1.0) + lower_input = np.where(x_hat <= 0.5, x_hat, 1.0 - x_hat) + lower_result = _soft_clipping_lower_half(lower_input, W_hat) + result = np.where(x_hat <= 0.5, lower_result, 1.0 - lower_result) + return np.clip(result, 0.0, 1.0) + + +def gaussianize_texture( + image: np.ndarray, + verbose: bool = True, + quantization_step: float = 0.0, +) -> tuple[np.ndarray, list]: + """ + Gaussianize a texture using per-channel 1D histogram transformation. + + Returns: + - Gaussianized image + - List of inverse LUTs (one per channel) for restoration + """ + _, _, c = image.shape + if c != 3: + raise ValueError(f"Expected RGB image with 3 channels, got {c}") + + # Process each channel independently + gaussianized = np.zeros_like(image) + inverse_luts = [] + + for ch in range(3): + if verbose: + channel_name = ['R', 'G', 'B'][ch] + print(f"Processing channel {channel_name}...") + + # Break ties inside quantized source buckets before building the transport. + channel = dither_channel(image[:, :, ch], quantization_step, ch) + + # Compute histogram (using 4096 bins for better precision) + hist, _ = np.histogram(channel.flatten(), bins=4096, range=(0.0, 1.0)) + + # Build Gaussianization LUT + lut = build_gaussianization_lut(hist, lut_size=4096) + + # Apply LUT to channel + gaussianized[:, :, ch] = apply_lut(channel, lut) + + # Build inverse LUT for later restoration + inverse_lut = build_inverse_lut(hist, lut_size=4096) + inverse_luts.append(inverse_lut) + + return gaussianized, inverse_luts + + +def build_inverse_lut(original_histogram: np.ndarray, lut_size: int = 4096) -> np.ndarray: + """ + Build inverse LUT to restore original histogram from Gaussianized values. + This maps from Gaussian distribution back to original distribution. + """ + gaussian = TruncatedGaussian() + value_centers, mid_quantiles = _occupied_bin_mid_quantiles(original_histogram) + + gaussian_values = np.linspace(0.0, 1.0, lut_size) + uniform_values = gaussian.cdf(gaussian_values) + return np.interp( + uniform_values, + mid_quantiles, + value_centers, + left=value_centers[0], + right=value_centers[-1], + ) + + +def histogram_preserving_blend( + textures: list[np.ndarray], + weights: np.ndarray, + inverse_luts: list[np.ndarray] | list[list[np.ndarray]], + gamma: float = 1.0 +) -> np.ndarray: + """ + Perform histogram-preserving blend of multiple Gaussianized textures. + + Algorithm 2 from Burley's paper: + 1. Optionally exponentiate weights + 2. Linear blend + 3. Compute variance scale factor W_hat + 4. Apply soft-clipping contrast operator + 5. Apply inverse LUTs to restore original histogram + + Args: + textures: List of Gaussianized textures + weights: Blending weights (must sum to 1) + inverse_luts: Shared inverse LUTs for the source texture, or repeated copies + of the same LUT set for each texture. + gamma: Exponent for weight adjustment (Eq. 5) + """ + n_textures = len(textures) + if len(weights) != n_textures: + raise ValueError(f"Number of weights ({len(weights)}) must match number of textures ({n_textures})") + + if len(inverse_luts) == 3 and all(np.asarray(lut).ndim == 1 for lut in inverse_luts): + shared_inverse_luts = inverse_luts + else: + if len(inverse_luts) != n_textures: + raise ValueError( + "inverse_luts must be either one shared RGB LUT set or one repeated set per texture" + ) + shared_inverse_luts = inverse_luts[0] + for lut_set in inverse_luts[1:]: + if any(not np.array_equal(ref, cur) for ref, cur in zip(shared_inverse_luts, lut_set)): + raise ValueError( + "Burley's per-channel method assumes all blended tiles share the same histogram LUTs" + ) + + # Normalize weights + weights = np.array(weights, dtype=np.float64) / np.sum(weights) + + # Apply weight exponentiation if gamma != 1 (Equation 5) + if gamma != 1.0: + weights_exp = np.power(weights, gamma) + weights = weights_exp / np.sum(weights_exp) + + # Linear blend (Equation 1) + blended = np.zeros_like(textures[0]) + for tex, w in zip(textures, weights): + blended += w * tex + + # Compute variance scale factor (Equation 2) + W_hat = np.sqrt(np.sum(weights ** 2)) + + # Apply contrast restoration per channel + result = np.zeros_like(blended) + for ch in range(3): + # Apply soft-clipping contrast operator (Equation 4) + result[:, :, ch] = soft_clipping_contrast(blended[:, :, ch], W_hat) + + # Apply inverse LUT to restore the shared source histogram. + result[:, :, ch] = apply_lut(result[:, :, ch], shared_inverse_luts[ch]) + + return result + + +def verify_histogram(image_path: Path, output_path: Path): + """Generate a minimal RGB histogram verification figure.""" + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + + img = load_image(image_path) + + fig = plt.figure(figsize=(4.0, 2.0), facecolor='#bcbcbc') + plot_ax = fig.add_axes((0.0, 0.0, 1.0, 1.0)) + plot_ax.set_facecolor('#bcbcbc') + for spine in plot_ax.spines.values(): + spine.set_visible(False) + plot_ax.set_xticks([]) + plot_ax.set_yticks([]) + + kernel = np.array([1.0, 2.0, 3.0, 2.0, 1.0], dtype=np.float64) + kernel /= kernel.sum() + curve_max = 0.0 + colors = ('#ff1a1a', '#00aa22', '#003cff') + + for channel, color in enumerate(colors): + hist, edges = np.histogram(img[:, :, channel].ravel(), bins=512, range=(0.0, 1.0), density=True) + hist = np.convolve(hist, kernel, mode='same') + centers = 0.5 * (edges[:-1] + edges[1:]) + curve_max = max(curve_max, float(hist.max())) + plot_ax.plot(centers, hist, color=color, linewidth=0.9, antialiased=True) + + plot_ax.set_xlim(0.0, 1.0) + plot_ax.set_ylim(0.0, curve_max * 1.18 if curve_max > 0.0 else 1.0) + fig.savefig(output_path, dpi=180, facecolor=fig.get_facecolor(), edgecolor='none') + plt.close(fig) + print(f"Saved histogram to {output_path}") + + +def save_lut_as_image( + luts: list[np.ndarray], + output_path: Path, + width: int = 2048, + height: int = 2048, +): + """Save inverse LUTs as a 2D texture with LUT samples running across columns.""" + if width <= 0 or height <= 0: + raise ValueError(f"Invalid LUT image size {width}x{height}") + + lut_size = len(luts[0]) + if any(len(lut) != lut_size for lut in luts): + raise ValueError("All inverse LUT channels must have the same length") + + src_coords = np.linspace(0.0, 1.0, lut_size) + dst_coords = np.linspace(0.0, 1.0, width) + packed_columns = np.stack( + [np.interp(dst_coords, src_coords, lut) for lut in luts], + axis=-1, + ) + + lut_image = np.broadcast_to(packed_columns[np.newaxis, :, :], (height, width, 3)).copy() + save_image(lut_image, output_path) + + +def main(): + parser = argparse.ArgumentParser( + description="Gaussianize texture using Burley's per-channel histogram-preserving method", + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "input", + type=Path, + help="Path to the input texture (PNG or EXR)" + ) + parser.add_argument( + "-o", "--output", + type=Path, + default=None, + help="Output path. Defaults to <input>_gaussianized.exr" + ) + parser.add_argument( + "--inverse-lut", + action="store_true", + default=True, + help="Also save the inverse LUT as <input>_inverse_lut.exr" + ) + parser.add_argument( + "--verify", + action="store_true", + help="Generate histogram visualization instead of processing" + ) + parser.add_argument( + "-v", "--verbose", + action="store_true", + help="Print detailed progress information" + ) + + args = parser.parse_args() + + if not args.input.exists(): + print(f"Error: Input file '{args.input}' does not exist.", file=sys.stderr) + sys.exit(1) + + if args.verify: + # Generate histogram visualization + hist_path = args.input.with_name(args.input.stem + "_histogram.png") + verify_histogram(args.input, hist_path) + else: + # Load input image + print(f"Loading {args.input}...") + image = load_image(args.input) + quantization_step = 0.0 if args.input.suffix.lower() == ".exr" else (1.0 / 255.0) + + # Gaussianize the texture + print("Applying per-channel Gaussianization...") + gaussianized, inverse_luts = gaussianize_texture( + image, + verbose=args.verbose, + quantization_step=quantization_step, + ) + + # Determine output path + if args.output is None: + args.output = args.input.with_name(args.input.stem + "_gaussianized.exr") + + # Save Gaussianized texture + print(f"Saving Gaussianized texture to {args.output}...") + save_image(gaussianized, args.output) + + # Optionally save inverse LUT + if args.inverse_lut: + lut_path = args.input.with_name(args.input.stem + "_inverse_lut.exr") + print(f"Saving inverse LUT to {lut_path}...") + save_lut_as_image(inverse_luts, lut_path) + + print("Done!") + + +if __name__ == "__main__": + main() |