#!/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)