seiza-cli 0.37.0

use anyhow::{Context, Result};
use clap::{Args, Subcommand, ValueEnum};
use seiza_stretch::{
    ColorStrategy, GhsParams, ImageShape, LayeredParams, MedianAnchorParams, StretchConfig,
    StretchModel, StretchParams,
};
use std::path::PathBuf;

#[derive(Args)]
pub(crate) struct StretchArgs {
    /// Linear mono or RGB FITS or XISF input
    input: PathBuf,
    /// Display-referred PNG, JPEG, or TIFF output
    #[arg(short, long)]
    output: PathBuf,
    /// How RGB channels share statistics and transfer curves
    #[arg(long, value_enum, default_value = "linked")]
    color_strategy: ColorStrategyArg,
    /// Maximum pooled scalar samples retained by data-driven models
    #[arg(long, default_value_t = crate::preview::MAXIMUM_SAMPLES)]
    max_analysis_samples: usize,
    #[command(subcommand)]
    model: StretchModelArgs,
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum ColorStrategyArg {
    /// One distribution and curve shared by all channels
    Linked,
    /// Independent distribution and curve for each channel
    Unlinked,
    /// Stretch Rec.709 luminance and preserve RGB chromaticity
    LuminancePreserving,
}

impl From<ColorStrategyArg> for ColorStrategy {
    fn from(value: ColorStrategyArg) -> Self {
        match value {
            ColorStrategyArg::Linked => Self::Linked,
            ColorStrategyArg::Unlinked => Self::Unlinked,
            ColorStrategyArg::LuminancePreserving => Self::LuminancePreserving,
        }
    }
}

impl ColorStrategyArg {
    fn summary(self) -> &'static str {
        match self {
            Self::Linked => "linked channels",
            Self::Unlinked => "unlinked channels",
            Self::LuminancePreserving => "luminance-preserving",
        }
    }
}

#[derive(Subcommand)]
enum StretchModelArgs {
    /// Clamp already display-referred samples to [0, 1]
    Identity,
    /// Affinely map an explicit black and white point
    Linear {
        #[arg(long)]
        black: f64,
        #[arg(long)]
        white: f64,
    },
    /// Apply asinh after an explicit black/white mapping
    Asinh {
        #[arg(long)]
        black: f64,
        #[arg(long)]
        white: f64,
        #[arg(long, default_value_t = 10.0)]
        strength: f64,
    },
    /// Resolve black/white percentiles and apply asinh
    PercentileAsinh {
        #[arg(long, default_value_t = 0.01)]
        black_percentile: f64,
        #[arg(long, default_value_t = 0.995)]
        white_percentile: f64,
        #[arg(long, default_value_t = 10.0)]
        strength: f64,
    },
    /// Apply explicit PixInsight/N.I.N.A.-family MTF parameters
    Mtf {
        #[arg(long)]
        shadows: f64,
        #[arg(long)]
        midtone: f64,
        #[arg(long)]
        highlights: f64,
    },
    /// Apply a manual Generalized Hyperbolic Stretch
    Ghs {
        /// Logarithmic stretch factor S, from 0 through 20
        #[arg(long, default_value_t = 1.0)]
        stretch_factor: f64,
        /// Local intensity b, from -5 through 15
        #[arg(long, default_value_t = 0.0, allow_negative_numbers = true)]
        local_intensity: f64,
        /// Symmetry point SP, from 0 through 1
        #[arg(long, default_value_t = 0.0)]
        symmetry_point: f64,
        /// Shadow-protection boundary LP, no greater than SP
        #[arg(long, default_value_t = 0.0)]
        protect_shadows: f64,
        /// Highlight-protection boundary HP, no less than SP
        #[arg(long, default_value_t = 1.0)]
        protect_highlights: f64,
        /// Explicit input black point
        #[arg(long, default_value_t = 0.0, allow_negative_numbers = true)]
        black: f64,
        /// Explicit input white point
        #[arg(long, default_value_t = 1.0, allow_negative_numbers = true)]
        white: f64,
    },
    /// Resolve the existing median/MAD Auto-MTF model
    AutoMtf {
        #[arg(long, default_value_t = 0.2)]
        target_median: f64,
        #[arg(long, default_value_t = -2.8, allow_negative_numbers = true)]
        shadows_clip: f64,
    },
    /// Place the median exactly on a target with a closed-form curve from a
    /// robust black point
    MedianAnchor {
        /// Where the background median lands, between 0 and 1
        #[arg(long, default_value_t = 0.25)]
        target_median: f64,
        /// Black point in robust sigmas below the median
        #[arg(long, default_value_t = 2.7)]
        black_point_sigma: f64,
        /// Use the image minimum as the black point and clip no black
        #[arg(long)]
        no_black_clip: bool,
        /// Lift the tones above the target, from 0 through 0.9
        #[arg(long, default_value_t = 0.0)]
        curves_boost: f64,
        /// Compress highlights above the knee, from 0 through 1
        #[arg(long, default_value_t = 0.0)]
        hdr_amount: f64,
        /// Where highlight compression starts, between 0 and 1
        #[arg(long, default_value_t = 0.75)]
        hdr_knee: f64,
        /// Rescale so the brightest sample reaches 1
        #[arg(long)]
        normalize: bool,
    },
    /// Stretch luminance with a median anchor curve, treating large and fine
    /// structure apart
    Layered {
        /// Where the background median lands, between 0 and 1
        #[arg(long, default_value_t = 0.15)]
        target_background: f64,
        /// How close the black point sits to the median, from 0 through 1
        #[arg(long, default_value_t = 0.5)]
        aggressiveness: f64,
        /// How far star cores are compressed, from 0 through 1
        #[arg(long, default_value_t = 0.0)]
        dynamic_range_compression: f64,
        /// Size in pixels that splits large structure from detail, 8 through 1024
        #[arg(long, default_value_t = 256)]
        scale_separation: u32,
        /// How much of the background's slope detail on bright structure
        /// keeps, from 0 through 1
        #[arg(long, default_value_t = 0.0)]
        contrast_recovery: f64,
        /// Saturation lift, from 0 through 1
        #[arg(long, default_value_t = 0.0)]
        saturation_amount: f64,
        /// How much the lift favours weakly coloured pixels, from 0 through 1
        #[arg(long, default_value_t = 0.5)]
        saturation_boost: f64,
        /// Lift saturation evenly instead of scaling it by each pixel's value
        #[arg(long)]
        no_lightness_mask: bool,
    },
}

impl StretchModelArgs {
    fn model(&self) -> StretchModel {
        match *self {
            Self::Identity => StretchModel::Identity,
            Self::Linear { black, white } => StretchModel::Linear { black, white },
            Self::Asinh {
                black,
                white,
                strength,
            } => StretchModel::Asinh {
                black,
                white,
                strength,
            },
            Self::PercentileAsinh {
                black_percentile,
                white_percentile,
                strength,
            } => StretchModel::PercentileAsinh {
                black_percentile,
                white_percentile,
                strength,
            },
            Self::Mtf {
                shadows,
                midtone,
                highlights,
            } => StretchModel::Mtf {
                shadows,
                midtone,
                highlights,
            },
            Self::Ghs {
                stretch_factor,
                local_intensity,
                symmetry_point,
                protect_shadows,
                protect_highlights,
                black,
                white,
            } => StretchModel::Ghs(GhsParams {
                stretch_factor,
                local_intensity,
                symmetry_point,
                protect_shadows,
                protect_highlights,
                black,
                white,
            }),
            Self::AutoMtf {
                target_median,
                shadows_clip,
            } => StretchModel::AutoMtf(StretchParams {
                target_median,
                shadows_clip,
            }),
            Self::MedianAnchor {
                target_median,
                black_point_sigma,
                no_black_clip,
                curves_boost,
                hdr_amount,
                hdr_knee,
                normalize,
            } => StretchModel::MedianAnchor(MedianAnchorParams {
                target_median,
                black_point_sigma: (!no_black_clip).then_some(black_point_sigma),
                curves_boost,
                hdr_amount,
                hdr_knee,
                normalize,
            }),
            Self::Layered {
                target_background,
                aggressiveness,
                dynamic_range_compression,
                scale_separation,
                contrast_recovery,
                saturation_amount,
                saturation_boost,
                no_lightness_mask,
            } => StretchModel::Layered(LayeredParams {
                target_background,
                aggressiveness,
                dynamic_range_compression,
                scale_separation,
                contrast_recovery,
                saturation_amount,
                saturation_boost,
                lightness_mask: !no_lightness_mask,
            }),
        }
    }

    fn summary(&self) -> String {
        match *self {
            Self::Identity => "identity clamp to [0, 1]".to_string(),
            Self::Linear { black, white } => format!("linear black {black}, white {white}"),
            Self::Asinh {
                black,
                white,
                strength,
            } => format!("asinh black {black}, white {white}, strength {strength}"),
            Self::PercentileAsinh {
                black_percentile,
                white_percentile,
                strength,
            } => format!(
                "percentile asinh black {black_percentile}, white {white_percentile}, strength {strength}"
            ),
            Self::Mtf {
                shadows,
                midtone,
                highlights,
            } => format!("MTF shadows {shadows}, midtone {midtone}, highlights {highlights}"),
            Self::Ghs {
                stretch_factor,
                local_intensity,
                symmetry_point,
                protect_shadows,
                protect_highlights,
                black,
                white,
            } => format!(
                "generalized hyperbolic stretch factor {stretch_factor}, local intensity {local_intensity}, symmetry {symmetry_point}, shadow protection {protect_shadows}, highlight protection {protect_highlights}, black {black}, white {white}"
            ),
            Self::AutoMtf {
                target_median,
                shadows_clip,
            } => format!("auto MTF target median {target_median}, shadows clip {shadows_clip}"),
            Self::MedianAnchor {
                target_median,
                black_point_sigma,
                no_black_clip,
                curves_boost,
                hdr_amount,
                hdr_knee,
                normalize,
            } => format!(
                "median anchor target median {target_median}, black point {}, curves boost {curves_boost}, highlight compression {hdr_amount} above {hdr_knee}{}",
                if no_black_clip {
                    "at the minimum".to_string()
                } else {
                    format!("{black_point_sigma} sigma below the median")
                },
                if normalize { ", normalized" } else { "" }
            ),
            Self::Layered {
                target_background,
                aggressiveness,
                dynamic_range_compression,
                scale_separation,
                contrast_recovery,
                saturation_amount,
                saturation_boost,
                no_lightness_mask,
            } => format!(
                "layered target background {target_background}, aggressiveness {aggressiveness}, dynamic range compression {dynamic_range_compression}, scale separation {scale_separation} px, contrast recovery {contrast_recovery}, saturation {saturation_amount} (boost {saturation_boost}{})",
                if no_lightness_mask {
                    ", no lightness mask"
                } else {
                    ""
                }
            ),
        }
    }
}

pub(crate) fn run(args: StretchArgs) -> Result<()> {
    crate::provenance::validate_path_roles([
        ("stretch input".into(), args.input.as_path()),
        ("stretch output".into(), args.output.as_path()),
    ])?;
    let frame = crate::common::open_frame(&args.input, "stretch input")?;
    let config = StretchConfig {
        model: args.model.model(),
        color_strategy: args.color_strategy.into(),
        max_analysis_samples: args.max_analysis_samples,
    };
    let plan = config
        .resolve_for(&frame.image.data, frame.image.channels)
        .context("could not resolve stretch model")?;
    let shape = ImageShape::new(frame.image.width, frame.image.height, frame.image.channels);
    let pixels = plan
        .apply_u8_image(&frame.image.data, shape)
        .context("could not apply stretch model")?;
    crate::preview::write_display_image(
        &args.output,
        frame.image.width,
        frame.image.height,
        frame.image.channels,
        pixels,
    )?;
    crate::common::wrote(
        &args.output,
        format_args!(
            "{}, {}, display-referred",
            args.model.summary(),
            args.color_strategy.summary(),
        ),
    );
    Ok(())
}