image_postprocess/processor.py

#!/usr/bin/env python3
"""
processor.py

Main pipeline for image postprocessing with an optional realistic camera-pipeline simulator.
Added support for applying 1D PNG/.npy LUTs and .cube 3D LUTs via --lut.
"""

import argparse
import os
from PIL import Image
import numpy as np
import piexif
from datetime import datetime

from .utils import (
    add_gaussian_noise,
    clahe_color_correction,
    randomized_perturbation,
    fourier_match_spectrum,
    auto_white_balance_ref,
    load_lut,
    apply_lut,
)
from .camera_pipeline import simulate_camera_pipeline


def add_fake_exif():
    """
    Generates a plausible set of fake EXIF data.
    Returns:
        bytes: The EXIF data as a byte string, ready for insertion.
    """
    now = datetime.now()
    datestamp = now.strftime("%Y:%m:%d %H:%M:%S")

    zeroth_ifd = {
        piexif.ImageIFD.Make: b"PurinCamera",
        piexif.ImageIFD.Model: b"Model420X",
        piexif.ImageIFD.Software: b"NovaImageProcessor",
        piexif.ImageIFD.DateTime: datestamp.encode('utf-8'),
    }
    exif_ifd = {
        piexif.ExifIFD.DateTimeOriginal: datestamp.encode('utf-8'),
        piexif.ExifIFD.DateTimeDigitized: datestamp.encode('utf-8'),
        piexif.ExifIFD.ExposureTime: (1, 125),  # 1/125s
        piexif.ExifIFD.FNumber: (28, 10),      # F/2.8
        piexif.ExifIFD.ISOSpeedRatings: 200,
        piexif.ExifIFD.FocalLength: (50, 1),    # 50mm
    }
    gps_ifd = {}

    exif_dict = {"0th": zeroth_ifd, "Exif": exif_ifd, "GPS": gps_ifd, "1st": {}, "thumbnail": None}
    exif_bytes = piexif.dump(exif_dict)
    return exif_bytes


def process_image(path_in, path_out, args):
    img = Image.open(path_in).convert('RGB')
    arr = np.array(img)

    # --- Auto white-balance (if enabled) ---
    if args.awb:
        if args.ref:
            try:
                ref_img_awb = Image.open(args.ref).convert('RGB')
                ref_arr_awb = np.array(ref_img_awb)
                arr = auto_white_balance_ref(arr, ref_arr_awb)
            except Exception as e:
                print(f"Warning: failed to load AWB reference '{args.ref}': {e}. Skipping AWB.")
        else:
            print("Applying AWB using grey-world assumption...")
            # Assuming auto_white_balance_ref with a None reference
            # triggers the grey-world algorithm as described.
            arr = auto_white_balance_ref(arr, None)

    # apply CLAHE color correction (contrast)
    arr = clahe_color_correction(arr, clip_limit=args.clahe_clip, tile_grid_size=(args.tile, args.tile))

    # FFT spectral matching reference (separate flag: --fft-ref)
    ref_arr_fft = None
    if args.fft_ref:
        try:
            ref_img_fft = Image.open(args.fft_ref).convert('RGB')
            ref_arr_fft = np.array(ref_img_fft)
        except Exception as e:
            print(f"Warning: failed to load FFT reference '{args.fft_ref}': {e}. Skipping FFT reference matching.")
            ref_arr_fft = None

    arr = fourier_match_spectrum(arr, ref_img_arr=ref_arr_fft, mode=args.fft_mode,
                                 alpha=args.fft_alpha, cutoff=args.cutoff,
                                 strength=args.fstrength, randomness=args.randomness,
                                 phase_perturb=args.phase_perturb, radial_smooth=args.radial_smooth,
                                 seed=args.seed)

    arr = add_gaussian_noise(arr, std_frac=args.noise_std, seed=args.seed)
    arr = randomized_perturbation(arr, magnitude_frac=args.perturb, seed=args.seed)

    # call the camera simulator if requested
    if args.sim_camera:
        arr = simulate_camera_pipeline(arr,
                                       bayer=not args.no_no_bayer,
                                       jpeg_cycles=args.jpeg_cycles,
                                       jpeg_quality_range=(args.jpeg_qmin, args.jpeg_qmax),
                                       vignette_strength=args.vignette_strength,
                                       chroma_aberr_strength=args.chroma_strength,
                                       iso_scale=args.iso_scale,
                                       read_noise_std=args.read_noise,
                                       hot_pixel_prob=args.hot_pixel_prob,
                                       banding_strength=args.banding_strength,
                                       motion_blur_kernel=args.motion_blur_kernel,
                                       seed=args.seed)

    # --- LUT application (optional) ---
    if args.lut:
        try:
            lut = load_lut(args.lut)
            # Ensure array is uint8 for LUT application
            arr_uint8 = np.clip(arr, 0, 255).astype(np.uint8)
            arr_lut = apply_lut(arr_uint8, lut, strength=args.lut_strength)
            # Ensure output is uint8
            arr = np.clip(arr_lut, 0, 255).astype(np.uint8)
        except Exception as e:
            print(f"Warning: failed to load/apply LUT '{args.lut}': {e}. Skipping LUT.")

    out_img = Image.fromarray(arr)

    # Generate fake EXIF data and save it with the image
    fake_exif_bytes = add_fake_exif()
    out_img.save(path_out, exif=fake_exif_bytes)


def build_argparser():
    p = argparse.ArgumentParser(description="Image postprocessing pipeline with camera simulation and LUT support")
    p.add_argument('input', help='Input image path')
    p.add_argument('output', help='Output image path')

    # AWB Options
    p.add_argument('--awb', action='store_true', help='Enable automatic white balancing. Uses grey-world if --ref is not provided.')
    p.add_argument('--ref', help='Optional reference image for auto white-balance (only used if --awb is enabled)', default=None)
    
    p.add_argument('--noise-std', type=float, default=0.02, help='Gaussian noise std fraction of 255 (0-0.1)')
    p.add_argument('--clahe-clip', type=float, default=2.0, help='CLAHE clip limit')
    p.add_argument('--tile', type=int, default=8, help='CLAHE tile grid size')
    p.add_argument('--cutoff', type=float, default=0.25, help='Fourier cutoff (0..1)')
    p.add_argument('--fstrength', type=float, default=0.9, help='Fourier blend strength (0..1)')
    p.add_argument('--randomness', type=float, default=0.05, help='Randomness for Fourier mask modulation')
    p.add_argument('--perturb', type=float, default=0.008, help='Randomized perturb magnitude fraction (0..0.05)')
    p.add_argument('--seed', type=int, default=None, help='Random seed for reproducibility')

    # FFT-matching options
    p.add_argument('--fft-ref', help='Optional reference image for FFT spectral matching', default=None)
    p.add_argument('--fft-mode', choices=('auto','ref','model'), default='auto', help='FFT mode: auto picks ref if available')
    p.add_argument('--fft-alpha', type=float, default=1.0, help='Alpha for 1/f model (spectrum slope)')
    p.add_argument('--phase-perturb', type=float, default=0.08, help='Phase perturbation strength (radians)')
    p.add_argument('--radial-smooth', type=int, default=5, help='Radial smoothing (bins) for spectrum profiles')

    # Camera-simulator options
    p.add_argument('--sim-camera', action='store_true', help='Enable camera-pipeline simulation (Bayer, CA, vignette, JPEG cycles)')
    p.add_argument('--no-no-bayer', dest='no_no_bayer', action='store_false', help='Disable Bayer/demosaic step (double negative kept for backward compat)')
    p.set_defaults(no_no_bayer=True)
    p.add_argument('--jpeg-cycles', type=int, default=1, help='Number of JPEG recompression cycles to apply')
    p.add_argument('--jpeg-qmin', type=int, default=88, help='Min JPEG quality for recompression')
    p.add_argument('--jpeg-qmax', type=int, default=96, help='Max JPEG quality for recompression')
    p.add_argument('--vignette-strength', type=float, default=0.35, help='Vignette strength (0..1)')
    p.add_argument('--chroma-strength', type=float, default=1.2, help='Chromatic aberration strength (pixels)')
    p.add_argument('--iso-scale', type=float, default=1.0, help='ISO/exposure scale for Poisson noise')
    p.add_argument('--read-noise', type=float, default=2.0, help='Read noise sigma for sensor noise')
    p.add_argument('--hot-pixel-prob', type=float, default=1e-6, help='Per-pixel probability of hot pixel')
    p.add_argument('--banding-strength', type=float, default=0.0, help='Horizontal banding amplitude (0..1)')
    p.add_argument('--motion-blur-kernel', type=int, default=1, help='Motion blur kernel size (1 = none)')

    # LUT options
    p.add_argument('--lut', type=str, default=None, help='Path to a 1D PNG (256x1) or .npy LUT, or a .cube 3D LUT')
    p.add_argument('--lut-strength', type=float, default=0.1, help='Strength to blend LUT (0.0 = no effect, 1.0 = full LUT)')

    return p


if __name__ == "__main__":
    args = build_argparser().parse_args()
    if not os.path.exists(args.input):
        print("Input not found:", args.input)
        raise SystemExit(2)
    process_image(args.input, args.output, args)
    print("Saved:", args.output)