granatpy 0.1.0__py3-none-any.whl

granatpy/__init__.py

@@ -0,0 +1,4 @@
+from .core import naturalize_rgb_image, naturalize_single_channel
+from .metrics import compute_all_metrics, compare_images
+from .visualization import overlay_naturalness_heatmap
+from .naturalizer import GraNat

granatpy/core.py

@@ -0,0 +1,327 @@
+import numpy as np
+import math
+
+__all__ = ["naturalize_rgb_image", "naturalize_single_channel"]
+
+# Constants from the original code
+EPS = 0.0001
+T1_pr = 0.3754
+N_Lap = 2041
+Lap_Offset = 1020
+N_Grad = 512
+Grad_Offset = 256
+T2_pr = [0.2421, 0.2550, 0.2474, 0.24816666]
+
+verbose = False
+
+def _calculate_psnr(nf):
+    """Calculate PSNR from naturalization factor"""
+    if 0 <= nf <= 0.934:
+        return 23.65 * math.exp(0.6 * nf) - 20.0 * math.exp(-7.508 * nf)
+    elif 0.934 < nf < 1.07:
+        return float('inf')  # > 40
+    elif 1.07 <= nf < 1.9:
+        return -11.566 * nf + 52.776
+    else:
+        return 13.06 * nf**4 - 121.4 * nf**3 + 408.5 * nf**2 - 595.5 * nf + 349
+
+def _find_t_error(p_t, N, offset, T):
+    """Calculate error for finding optimal T parameter"""
+    error = 0.0
+    for i in range(-offset, N - offset):
+        tmp = math.atan(T * i) - p_t[i + offset]
+        error += tmp * tmp
+    return error
+
+def _find_t_golden_section(data, N, offset, eps):
+    """
+    Find optimal T parameter using golden section search.
+    EXACT implementation matching the Java code.
+    """
+    left = 0.0
+    right = 1.0
+    
+    # Create p_t array (shift and rescale) - use float32 to match Java float
+    p_t = np.array([(x - 0.5) * math.pi for x in data], dtype=np.float32)
+    
+    m1 = 0.0
+    m2 = 0.0
+    
+    # While the precision is bigger than eps
+    while right - left >= eps:
+        # Move left and right by 1/3
+        m1 = left + (right - left) / 3
+        m2 = right - (right - left) / 3
+        
+        # Evaluate on m1 and m2, and move the extreme point
+        error1 = _find_t_error(p_t, N, offset, m1)
+        error2 = _find_t_error(p_t, N, offset, m2)
+        
+        if error1 <= error2:
+            right = m2
+        else:
+            left = m1
+    
+    # Return the average (this is critical - matches Java exactly)
+    return (m1 + m2) / 2
+
+def _compute_gradient_field_forward(image):
+    """
+    Compute forward gradient field matching Java implementation.
+    G_x = I(x+1, y) - I(x, y)
+    G_y = I(x, y+1) - I(x, y)
+    """
+    gx = np.zeros_like(image, dtype=np.float32)
+    gy = np.zeros_like(image, dtype=np.float32)
+    
+    # Forward differences
+    gx[:, :-1] = image[:, 1:] - image[:, :-1]
+    gy[:-1, :] = image[1:, :] - image[:-1, :]
+    
+    return gx, gy
+
+def _compute_laplacian_field(image):
+    """Compute Laplacian field using 5-point stencil."""
+    lap = np.zeros_like(image, dtype=np.float32)
+    
+    # 5-point stencil: -4*center + right + left + down + up
+    lap[1:-1, 1:-1] = (
+        image[1:-1, 2:] +      # right
+        image[1:-1, :-2] +     # left
+        image[2:, 1:-1] +      # down
+        image[:-2, 1:-1] -     # up
+        4 * image[1:-1, 1:-1]  # center
+    )
+    
+    return lap
+
+def _calculate_laplace_field_and_gradient_optimized(image, lap_cdf, grad_d):
+    """
+    Calculate laplacian field and gradient histogram.
+    
+    Key insight: Java stores grad_d[gy, gx] due to indexD[1-i]
+    """
+    # Ensure image is float32 in [0, 255] range
+    if image.max() <= 1:
+        image = (image * 255).astype(np.float32)
+    else:
+        image = image.astype(np.float32)
+
+    if verbose:
+        print(f"  Image stats: shape={image.shape}, min={image.min():.2f}, max={image.max():.2f}, mean={image.mean():.2f}")
+    
+    # Compute laplacian and gradients
+    laplace_field = _compute_laplacian_field(image)
+    gx, gy = _compute_gradient_field_forward(image)
+    
+    # Get valid regions (exclude borders)
+    valid_laplace = laplace_field[1:-1, 1:-1]
+    valid_gx = gx[1:-1, 1:-1]
+    valid_gy = gy[1:-1, 1:-1]
+
+    if verbose:
+        print(f"  Valid region: {valid_laplace.shape}, n_pixels={valid_laplace.size}")
+        print(f"  Laplacian: min={valid_laplace.min():.2f}, max={valid_laplace.max():.2f}")
+        print(f"  Gradient X: min={valid_gx.min():.2f}, max={valid_gx.max():.2f}")
+        print(f"  Gradient Y: min={valid_gy.min():.2f}, max={valid_gy.max():.2f}")
+    
+    n_pixels = valid_laplace.size
+    f = np.float32(1.0 / n_pixels)  # Use float32 to match Java
+    
+    # Vectorized histogram binning for laplacian
+    lap_bins = (valid_laplace + Lap_Offset).astype(np.int32)
+    lap_bins = np.clip(lap_bins, 0, N_Lap - 1)
+    np.add.at(lap_cdf, lap_bins.ravel(), f)
+
+    if verbose:
+        print(f"  Laplacian histogram: non-zero bins={np.count_nonzero(lap_cdf)}, sum={lap_cdf.sum():.6f}")
+    
+    # Vectorized 2D gradient histogram
+    # Java: indexD[1-i] means grad_d[gy, gx]
+    gx_bins = np.clip((Grad_Offset + valid_gx).astype(np.int32), 0, N_Grad - 1)
+    gy_bins = np.clip((Grad_Offset + valid_gy).astype(np.int32), 0, N_Grad - 1)
+    
+    gx_flat = gx_bins.ravel()
+    gy_flat = gy_bins.ravel()
+    
+    # Accumulate: grad_d[gy, gx]
+    np.add.at(grad_d, (gy_flat, gx_flat), f)
+
+    if verbose:
+        print(f"  Gradient 2D histogram: non-zero bins={np.count_nonzero(grad_d)}, sum={grad_d.sum():.6f}")
+
+def _calculate_lap_cdf(lap_cdf):
+    """Convert laplacian histogram to CDF (in-place)"""
+    np.cumsum(lap_cdf, out=lap_cdf, dtype=np.float32)
+
+def _convert_grad_2d_to_cdf(grad_d):
+    """
+    Convert 2D gradient HISTOGRAM to 2D CDF (in-place).
+    This performs cumsum in both dimensions using float32 precision.
+    """
+    # Ensure we're working with float32
+    if grad_d.dtype != np.float32:
+        grad_d = grad_d.astype(np.float32)
+    
+    # Row integration (cumsum along axis 0) - matches Java first loop
+    np.cumsum(grad_d, axis=0, out=grad_d, dtype=np.float32)
+    
+    # Column integration (cumsum along axis 1) - matches Java second loop
+    np.cumsum(grad_d, axis=1, out=grad_d, dtype=np.float32)
+
+def _calculate_grad_cdf(grad_cdf, grad_d):
+    """
+    Calculate 1D gradient CDFs from 2D CDF.
+    
+    Java code integrates (sums) the 2D CDF to get marginal CDFs:
+    - Gradx[gx] = sum over gy of grad_d[gy, gx]
+    - Grady[gy] = sum over gx of grad_d[gy, gx]
+    """
+    # Integrate the 2D CDF - use float32 for consistency
+    grad_x = np.sum(grad_d, axis=0, dtype=np.float32)
+    grad_y = np.sum(grad_d, axis=1, dtype=np.float32)
+
+    # Scale by 255.0 (not 256!) as per Java code
+    grad_x = grad_x / np.float32(255.0)
+    grad_y = grad_y / np.float32(255.0)
+    
+    # Store in output array
+    grad_cdf[:, 0] = grad_x
+    grad_cdf[:, 1] = grad_y
+
+def _find_naturalization_factor(image, t2_prior):
+    """Find optimal naturalization factor"""
+    # Initialize histograms - use float32 to match Java
+    lap_cdf = np.zeros(N_Lap, dtype=np.float32)
+    grad_d = np.zeros((N_Grad, N_Grad), dtype=np.float32)
+    
+    # Calculate laplacian and gradient histograms
+    _calculate_laplace_field_and_gradient_optimized(image, lap_cdf, grad_d)
+    
+    # Convert 2D gradient histogram to 2D CDF
+    _convert_grad_2d_to_cdf(grad_d)
+    
+    # Extract 1D CDFs from 2D CDF
+    grad_cdf = np.zeros((N_Grad, 2), dtype=np.float32)
+    _calculate_grad_cdf(grad_cdf, grad_d)
+    
+    # Convert laplacian histogram to CDF
+    _calculate_lap_cdf(lap_cdf)
+    
+    # Find T1 parameter - USE GOLDEN SECTION (critical fix!)
+    t1_x = _find_t_golden_section(grad_cdf[:, 0], N_Grad, Grad_Offset, EPS)
+    t1_y = _find_t_golden_section(grad_cdf[:, 1], N_Grad, Grad_Offset, EPS)
+    t1 = (t1_x + t1_y) / (2 * T1_pr)
+    
+    # Find T2 parameter - USE GOLDEN SECTION (critical fix!)
+    t2 = _find_t_golden_section(lap_cdf, N_Lap, Lap_Offset, EPS) / t2_prior
+    
+    # Calculate naturalization factor
+    theta = 0.5
+    nf = (1 - theta) * t1 + theta * t2
+
+    if verbose:
+        print(f"  Results: T1={t1:.6f}, T2={t2:.6f}, Nf={nf:.6f}, T2_prior={t2_prior:.4f}")
+    
+    return nf, t1, t2
+
+def _naturalize_image(image, nf):
+    """Apply naturalization to image"""
+    # Calculate mean
+    mean_original = np.mean(image)
+    
+    # Apply naturalization: (I - mean) * Nf + mean
+    result = (image - mean_original) * nf + mean_original
+    
+    # Round and clamp to [0, 255] - matches Java rounding
+    result = np.round(result)
+    result = np.clip(result, 0, 255)
+    
+    return result.astype(np.uint8)
+
+def naturalize_single_channel(image: np.ndarray, channel_idx: int = 0) -> tuple:
+    """
+    Naturalize a single image channel using its corresponding T2 prior.
+
+    The T2 prior varies per channel to account for differences in the
+    statistical properties of R, G, and B channels. For grayscale images,
+    pass channel_idx=3 to use the averaged prior.
+
+    Args:
+        image: 2D numpy array of shape (H, W) representing a single
+               channel, with pixel values in the range [0, 255].
+        channel_idx: Index indicating which channel is being processed.
+                     0=Red, 1=Green, 2=Blue, 3=Grayscale. Defaults to 0.
+
+    Returns:
+        tuple: A tuple of (result, nf, t1, t2, psnr) where:
+            result: Naturalized channel as a uint8 numpy array of shape (H, W).
+            nf:     Naturalization factor applied to this channel.
+            t1:     T1 parameter derived from the gradient CDF.
+            t2:     T2 parameter derived from the Laplacian CDF.
+            psnr:   Estimated PSNR value based on the naturalization factor.
+    """
+    # Select appropriate T2 prior - matches Java indexing
+    if channel_idx <= 2:
+        t2_prior = T2_pr[2 - channel_idx]  # R=0->2, G=1->1, B=2->0
+    else:
+        t2_prior = T2_pr[3]  # Grayscale average
+    
+    # Find optimal naturalization factor
+    nf, t1, t2 = _find_naturalization_factor(image, t2_prior)
+    
+    # Apply naturalization
+    result = _naturalize_image(image, nf)
+    
+    # Calculate PSNR
+    psnr = _calculate_psnr(nf)
+    
+    return result, nf, t1, t2, psnr
+
+def naturalize_rgb_image(image: np.ndarray) -> tuple:
+    """
+    Naturalize an RGB or grayscale image by adjusting each channel independently.
+
+    For RGB images, each channel is processed separately using its corresponding
+    T2 prior. For grayscale images, a single averaged T2 prior is used instead.
+
+    Args:
+        image: Input image as a numpy array. Either (H, W) for grayscale
+               or (H, W, 3) for RGB, with pixel values in the range [0, 255].
+
+    Returns:
+        tuple: A tuple of (result, nfs, t1s, t2s, psnrs) where:
+            result: Naturalized image as a uint8 numpy array, same shape as input.
+            nfs:    List of naturalization factors, one per channel.
+            t1s:    List of T1 parameters, one per channel.
+            t2s:    List of T2 parameters, one per channel.
+            psnrs:  List of estimated PSNR values, one per channel.
+    """
+    if len(image.shape) == 2:
+        # Single channel
+        result, nf, t1, t2, psnr = naturalize_single_channel(image, channel_idx=3)
+        return result, [nf], [t1], [t2], [psnr]
+    else:
+        # RGB image
+        results = []
+        nfs = []
+        t1s = []
+        t2s = []
+        psnrs = []
+        
+        for i in range(min(3, image.shape[2])):
+            channel = image[:, :, i]
+            result, nf, t1, t2, psnr = naturalize_single_channel(channel, channel_idx=i)
+            results.append(result)
+            nfs.append(nf)
+            t1s.append(t1)
+            t2s.append(t2)
+            psnrs.append(psnr)
+        
+        # Combine channels
+        if len(results) == 3:
+            result = np.stack(results, axis=2)
+        else:
+            result = results[0]
+            
+        return result, nfs, t1s, t2s, psnrs

granatpy/metrics.py

@@ -0,0 +1,99 @@
+import numpy as np
+from skimage.metrics import peak_signal_noise_ratio, structural_similarity, mean_squared_error, normalized_root_mse
+import imageio.v3 as imageio
+from typing import Dict, List
+from .core import naturalize_rgb_image
+
+
+__all__ = ["compute_all_metrics", "compare_images"]
+
+
+def _compute_nf(image: np.ndarray) -> float:
+    """Compute the mean naturalness factor across RGB channels for an image."""
+    _, nfs, _, _, _ = naturalize_rgb_image(image)
+    return float(np.mean(nfs))
+
+
+def compute_all_metrics(real: np.ndarray, synthetic: np.ndarray, verbose = False) -> Dict[str, float]:
+    """
+    Compute all image quality metrics between a reference and a synthetic image.
+
+    Args:
+        real:       Reference (real) image as a numpy array.
+        synthetic:  Comparison (synthetic) image as a numpy array.
+        verbose:    Whether to show status logs.
+
+    Returns:
+        Dictionary containing MSE, PSNR, SSIM, NRMSE, and dNf values.
+        dNf is the delta naturalness factor (synthetic Nf minus real Nf).
+    """
+    if real.shape != synthetic.shape:
+        raise ValueError(
+            f"Images must have the same shape. Got {real.shape} and {synthetic.shape}."
+        )
+
+    if verbose: print("Computing naturalness for first image...")
+    nf_real = _compute_nf(real)
+    if verbose: print("Computing naturalness for second image...")
+    nf_synthetic = _compute_nf(synthetic)
+    
+    if verbose: print("Calculating remaining metrics...")
+    results = {
+        "MSE":   mean_squared_error(real, synthetic) / (255 ** 2),
+        "PSNR":  peak_signal_noise_ratio(real, synthetic),
+        "SSIM":  structural_similarity(real, synthetic, channel_axis=-1),
+        "NRMSE": normalized_root_mse(real, synthetic),
+        "Nf_real":      nf_real,
+        "Nf_synthetic": nf_synthetic,
+        "dNf":          nf_synthetic - nf_real,
+    }
+
+    for key, value in results.items():
+        if key == "PSNR":
+            print(f"  {key:<12} {value:.3f} dB")
+        else:
+            print(f"  {key:<12} {value:.4f}")
+
+    return results
+
+
+def compare_images(real: str, images: List[str], verbose = False) -> Dict[str, Dict[str, float]]:
+    """
+    Compare a reference image against a list of synthetic images,
+    printing a summary and returning all metrics.
+
+    Args:
+        real:   Path to a reference (real) image.
+        images: List of image paths.
+        verbose: Whether to show status logs
+
+    Returns:
+        Dictionary mapping each image name to its computed metrics.
+    """
+    results = {}
+    if verbose: print(f"Reading image {real}")
+    real_img = imageio.imread(real)
+
+    for img_str in images:
+        if verbose: print(f"Reading image {img_str}")
+        img = imageio.imread(img_str)
+
+        if img.shape != real_img.shape:
+            print(f"Skipping '{img_str}': shape mismatch ({img.shape} vs {real_img.shape}).")
+            continue
+
+        metrics = compute_all_metrics(real_img, img, verbose)
+        results[img_str] = metrics
+
+        print(
+            f"{img_str}: "
+            f"PSNR={metrics['PSNR']:.2f} dB, "
+            f"MSE={metrics['MSE']:.6f}, "
+            f"SSIM={metrics['SSIM']:.4f}, "
+            f"NRMSE={metrics['NRMSE']:.4f}, "
+            f"Nf_real={metrics['Nf_real']:.4f}, "
+            f"Nf_synthetic={metrics['Nf_synthetic']:.4f}, "
+            f"dNf={metrics['dNf']:+.4f}"
+        )
+
+    return results

granatpy/naturalizer.py

@@ -0,0 +1,144 @@
+import numpy as np
+import imageio.v3 as imageio
+import matplotlib.pyplot as plt
+from .core import naturalize_rgb_image
+from .visualization import overlay_naturalness_heatmap
+
+
+__all__ = ["GraNat"]
+
+
+class GraNat:
+    """
+    A convenience wrapper around the GraNatPy naturalization pipeline.
+
+    This class provides an interface to the core naturalization functions,
+    plus loading, saving and image visualization without managing intermediate results manually.
+
+    Typical usage::
+
+        # From a file path
+        result = GraNat.load_image("photo.tif").naturalize()
+        # From an existing numpy array
+        result = GraNat(img_array).naturalize()
+        # Chained with visualization
+        GraNat.load_image("photo.tif").naturalize().plot_heatmap()
+
+    Attributes:
+        image:   The input image as a uint8 numpy array of shape (H, W, 3)
+                 or (H, W) for grayscale.
+        result:  The naturalized image as a uint8 numpy array. None until
+                 naturalize() has been called.
+        nfs:     List of naturalization factors, one per channel. None until
+                 naturalize() has been called.
+        t1s:     List of T1 parameters, one per channel. None until
+                 naturalize() has been called.
+        t2s:     List of T2 parameters, one per channel. None until
+                 naturalize() has been called.
+        psnrs:   List of estimated PSNR values, one per channel. None until
+                 naturalize() has been called.
+    """
+
+    def __init__(self, image: np.ndarray):
+
+        self.image = image
+        self.result = None
+        self.nfs = None
+        self.t1s = None
+        self.t2s = None
+        self.psnrs = None
+
+    @classmethod
+    def load_image(cls, path: str) -> "GraNat":
+        """
+        Load an image from disk and return a GraNat instance.
+        Supports any format readable by imageio (TIFF, PNG, JPEG, etc.).
+
+        Args:
+            path: Path to the image file.
+
+        Returns:
+            A new GraNat instance initialised with the loaded image.
+
+        Raises:
+            FileNotFoundError: If no file exists at the given path.
+            OSError: If the file exists but cannot be read as an image.
+
+        Example::
+
+            naturalizer = GraNat.load_image("photo.tif")
+        """
+        return cls(imageio.imread(path))
+
+    def naturalize(self, show = False) -> "GraNat":
+        """
+        Run the naturalization pipeline on the loaded image.
+
+        Processes each channel independently and stores the results in the
+        instance attributes. Returns self to allow method chaining.
+
+        Args:
+            show: Whether to directly plot the image with matplotlib.
+        Returns:
+            self, to allow chaining e.g. GraNat.load_image(...).naturalize()
+        Raises:
+            ValueError: If the image array is not a valid 2D or 3D numpy array.
+        Example::
+
+            naturalizer = GraNat.load_image("photo.tif").naturalize()
+            print(naturalizer.nfs)
+        """
+        self.result, self.nfs, self.t1s, self.t2s, self.psnrs = naturalize_rgb_image(self.image)
+        if show:
+            plt.imshow(self.result)
+        return self
+
+    def plot_heatmap(self, **kwargs):
+        """
+        Overlay a naturalness factor heatmap on the original image.
+
+        Divides the image into a grid of cells and computes the average naturalization factor per cell,
+        then renders a semi-transparent heatmap overlay. Keyword arguments are forwarded directly to
+        overlay_naturalness_heatmap().
+
+        Common keyword arguments:
+            grid_rows (int):  Number of rows in the grid. Default 10.
+            grid_cols (int):  Number of columns in the grid. Default 15.
+            cmap (str):       Matplotlib colormap name. Default 'YlOrRd'.
+            alpha (float):    Heatmap opacity in [0, 1]. Default 0.5.
+            figsize (tuple):  Figure size as (width, height). Default (12, 8).
+
+        Returns:
+            A tuple of (fig, ax, heatmap_data) where heatmap_data is a 2D
+            numpy array of shape (grid_rows, grid_cols) containing the
+            average naturalization factor per cell.
+
+        Raises:
+            ValueError: If the image has not been loaded correctly.
+
+        Example::
+
+            fig, ax, heatmap = GraNat.load_image("photo.tif").plot_heatmap(grid_rows=20)
+        """
+        return overlay_naturalness_heatmap(self.image, **kwargs)
+
+    def save(self, path: str) -> "GraNat":
+        """
+        Save the naturalized image to disk.
+
+        Args:
+            path: Output file path. The format is inferred from the file
+                extension (e.g. '.png', '.tif', '.jpg').
+        Returns:
+            self, to allow chaining.
+        Raises:
+            RuntimeError: If naturalize() has not been called yet.
+            OSError: If the file cannot be written to the given path.
+        Example::
+
+            GraNat.load_image("photo.tif").naturalize().save("output.png")
+        """
+        if self.result is None:
+            raise RuntimeError("No result to save. Call naturalize() first.")
+        imageio.imwrite(path, self.result)
+        return self