glint_mask_tools/

cli.rs

/// Command-line interface for the glint mask generation tool.
///
/// This module provides a comprehensive CLI that dynamically generates
/// subcommands for each supported sensor type, matching the functionality
/// of the original Python implementation.
use clap::{Arg, ArgAction, ArgMatches, Command, Parser};
use lazy_static::lazy_static;
use std::path::PathBuf;
use tracing::{error, info};

use crate::{
    algorithms::ThresholdAlgorithm,
    core::{
        masker::MaskerBuilder,
        postprocessor::{CompositePostProcessor, MetashapeConverter, PixelBufferProcessor},
        sensor::SensorRegistry,
    },
    error::{GlintError, Result},
    loaders::{BigTiffLoader, ConfigurableMultiFileLoader, SingleFileLoader},
};

lazy_static! {
    static ref SENSOR_REGISTRY: SensorRegistry = SensorRegistry::from_user_config();
}

/// Glint mask generation tool for UAV and aerial imagery
#[derive(Parser)]
#[command(name = "glint-mask")]
#[command(about = "Generate glint masks for specular reflection in UAV and aerial imagery")]
#[command(version)]
pub struct Cli {
    /// Enable verbose logging
    #[arg(short, long, global = true)]
    pub verbose: bool,
}

/// Common arguments for sensor processing commands
#[derive(Debug)]
pub struct SensorArgs {
    /// Input directory containing images or single image file
    pub input: PathBuf,

    /// Output directory for generated masks
    pub output: PathBuf,

    /// Threshold values for each band (0.0-1.0)
    pub thresholds: Option<Vec<f64>>,

    /// Pixel buffer radius for mask dilation
    pub pixel_buffer: usize,
}

impl From<&ArgMatches> for SensorArgs {
    fn from(matches: &ArgMatches) -> Self {
        let input = matches
            .get_one::<String>("input")
            .map(PathBuf::from)
            .unwrap();
        let output = matches
            .get_one::<String>("output")
            .map(PathBuf::from)
            .unwrap();
        let thresholds = matches
            .get_many::<f64>("thresholds")
            .map(|v| v.copied().collect());
        let pixel_buffer = matches.get_one::<usize>("pixel_buffer").copied().unwrap();

        Self {
            input,
            output,
            thresholds,
            pixel_buffer,
        }
    }
}

/// Initialize the CLI and run the appropriate command
pub fn run() -> Result<()> {
    let mut cmd = Command::new("glint-mask")
        .about("Generate glint masks for specular reflection in UAV and aerial imagery")
        .version(env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::new("verbose")
                .short('v')
                .long("verbose")
                .action(ArgAction::SetTrue)
                .help("Enable verbose logging")
                .global(true),
        )
        .subcommand_required(true);

    for sensor in SENSOR_REGISTRY.sensors() {
        let mut subcommand = Command::new(sensor.id.as_str()).about(sensor.name.as_str());
        subcommand = subcommand
            .arg(
                Arg::new("input")
                    .help("Path to input directory or single image file")
                    .required(true),
            )
            .arg(
                Arg::new("output")
                    .help("Path to output directory for mask files")
                    .required(true),
            )
            .arg(
                Arg::new("thresholds")
                    .short('t')
                    .long("thresholds")
                    .value_delimiter(',')
                    .value_parser(clap::value_parser!(f64))
                    .help("Comma-separated threshold values for each band (e.g., 0.8,0.9,0.7)"),
            )
            .arg(
                Arg::new("pixel_buffer")
                    .short('b')
                    .long("pixel_buffer")
                    .default_value("0")
                    .value_parser(clap::value_parser!(usize))
                    .help("Pixel buffer radius to expand masks"),
            );
        cmd = cmd.subcommand(subcommand);
    }

    cmd = cmd.subcommand(Command::new("list-sensors").about("List available sensors"));
    cmd = cmd.subcommand(
        Command::new("sensor-info")
            .about("Show sensor information")
            .arg(
                Arg::new("sensor_id")
                    .help("Sensor ID to show information for")
                    .required(true),
            ),
    );

    let matches = cmd.get_matches();

    // Initialize logging
    let verbose = matches.get_flag("verbose");
    let subscriber = tracing_subscriber::fmt()
        .with_max_level(if verbose {
            tracing::Level::DEBUG
        } else {
            tracing::Level::INFO
        })
        .finish();

    tracing::subscriber::set_global_default(subscriber)
        .map_err(|e| GlintError::processing(format!("Failed to initialize logging: {}", e)))?;

    match matches.subcommand() {
        Some(("list-sensors", _)) => list_sensors(),
        Some(("sensor-info", sub_matches)) => {
            let sensor_id = sub_matches.get_one::<String>("sensor_id").unwrap();
            show_sensor_info(sensor_id)
        }
        Some((sensor_id, sub_matches)) => {
            let args = SensorArgs::from(sub_matches);
            process_sensor(sensor_id, args)
        }
        _ => unreachable!("Subcommand is required"),
    }
}

/// Process images using the specified sensor configuration
fn process_sensor(sensor_id: &str, args: SensorArgs) -> Result<()> {
    info!("Processing {} images", sensor_id.to_uppercase());
    info!("Input: {}", args.input.display());
    info!("Output: {}", args.output.display());

    // Validate input and output paths
    if !args.input.exists() {
        return Err(GlintError::Io(std::io::Error::new(
            std::io::ErrorKind::NotFound,
            format!("Input path does not exist: {}", args.input.display()),
        )));
    }

    crate::utils::ensure_directory_exists(&args.output)?;

    // Get sensor configuration
    let sensor = SENSOR_REGISTRY
        .get_sensor(sensor_id)
        .ok_or_else(|| GlintError::sensor(format!("Unknown sensor: {}", sensor_id)))?
        .clone();

    // Determine thresholds
    let thresholds = if let Some(user_thresholds) = args.thresholds {
        if user_thresholds.len() != sensor.band_count() {
            return Err(GlintError::BandCountMismatch {
                expected: sensor.band_count(),
                actual: user_thresholds.len(),
            });
        }
        user_thresholds
    } else {
        sensor.default_thresholds()
    };

    info!("Using thresholds: {:?}", thresholds);
    info!("Pixel buffer: {} pixels", args.pixel_buffer);

    // Create algorithm
    let algorithm = ThresholdAlgorithm::new(thresholds)?;

    // Create post-processor pipeline
    let mut postprocessor = CompositePostProcessor::new();

    if args.pixel_buffer > 0 {
        postprocessor =
            postprocessor.add_processor(Box::new(PixelBufferProcessor::new(args.pixel_buffer)));
    }

    // Always add Metashape converter
    postprocessor = postprocessor.add_processor(Box::new(MetashapeConverter::new()));

    // Create loader based on sensor type
    let loader = create_loader(&sensor)?;

    // Build masker
    let masker = MaskerBuilder::new()
        .with_sensor(sensor.clone())
        .with_algorithm(Box::new(algorithm))
        .with_postprocessor(Box::new(postprocessor))
        .with_loader(loader)
        .build()?;

    // Process images
    let stats = masker.process_directory(&args.input, &args.output, None)?;

    // Report results
    info!("Processing complete!");
    info!("Total captures: {}", stats.total_captures);
    info!("Successful: {}", stats.successful_captures);
    info!("Failed: {}", stats.failed_captures);
    info!("Success rate: {:.1}%", stats.success_rate());

    if !stats.errors.is_empty() {
        error!("Errors encountered:");
        for error in &stats.errors {
            error!("  {}", error);
        }
    }

    if !stats.all_successful() {
        std::process::exit(1);
    }

    Ok(())
}

/// Create the appropriate image loader for a sensor
fn create_loader(
    sensor: &crate::core::sensor::Sensor,
) -> Result<Box<dyn crate::core::ImageLoader>> {
    match sensor.loader_type.as_str() {
        "single_file" => {
            if sensor.band_count() == 1 {
                Ok(Box::new(SingleFileLoader::grayscale()?))
            } else if sensor.band_count() == 3 {
                Ok(Box::new(SingleFileLoader::rgb()?))
            } else if sensor.band_count() == 4 {
                // For 4-band sensors, check the bit depth to determine the right loader
                if sensor.bit_depth == 8 {
                    // 8-bit 4-band (e.g., CIR) - create a custom 8-bit loader
                    let extensions = vec![
                        "tif".to_string(),
                        "tiff".to_string(),
                        "jpg".to_string(),
                        "jpeg".to_string(),
                        "png".to_string(),
                    ];
                    Ok(Box::new(SingleFileLoader::new(extensions, 4, 8)?))
                } else {
                    // 16-bit 4-band
                    Ok(Box::new(SingleFileLoader::tiff_16bit(4)?))
                }
            } else {
                Err(GlintError::validation(format!(
                    "Unsupported band count for single file loader: {}",
                    sensor.band_count()
                )))
            }
        }
        "multifile" => {
            // Create configurable multifile loader from sensor configuration
            Ok(Box::new(ConfigurableMultiFileLoader::from_config(
                &sensor.loader_config,
            )?))
        }
        "big_tiff" => {
            // Extract chunk size from loader config if available
            let chunk_size = sensor
                .loader_config
                .get("chunk_size")
                .and_then(|s| s.parse::<usize>().ok());

            // Extract extensions from loader config
            let extensions = sensor
                .loader_config
                .get("extensions")
                .map(|s| s.split(',').map(|ext| ext.trim().to_string()).collect())
                .unwrap_or_else(|| vec!["tif".to_string(), "tiff".to_string()]);

            Ok(Box::new(BigTiffLoader::new(
                extensions,
                sensor.band_count(),
                sensor.bit_depth,
                chunk_size,
            )?))
        }
        _ => Err(GlintError::validation(format!(
            "Unknown loader type: {}",
            sensor.loader_type
        ))),
    }
}

/// List all available sensors
fn list_sensors() -> Result<()> {
    println!("Available sensors:");
    println!();

    for sensor in SENSOR_REGISTRY.sensors() {
        println!("  {} - {}", sensor.id, sensor.name);
        println!(
            "    Bands: {} ({}-bit)",
            sensor.band_count(),
            sensor.bit_depth
        );
        println!("    Band names: {}", sensor.band_names().join(", "));
        if let Some(desc) = &sensor.description {
            println!("    Description: {}", desc);
        }
        println!();
    }

    Ok(())
}

/// Show detailed information about a specific sensor
fn show_sensor_info(sensor_id: &str) -> Result<()> {
    let sensor = SENSOR_REGISTRY
        .get_sensor(sensor_id)
        .ok_or_else(|| GlintError::sensor(format!("Unknown sensor: {}", sensor_id)))?;

    println!("Sensor: {} ({})", sensor.name, sensor.id);
    println!("Bit depth: {}", sensor.bit_depth);
    println!("Loader type: {}", sensor.loader_type);

    if let Some(desc) = &sensor.description {
        println!("Description: {}", desc);
    }

    println!();
    println!("Bands ({}):", sensor.band_count());
    for (i, band) in sensor.bands.iter().enumerate() {
        print!(
            "  {}: {} (threshold: {:.3}",
            i + 1,
            band.name,
            band.default_threshold
        );
        if let Some(wavelength) = band.wavelength {
            print!(", wavelength: {} nm", wavelength);
        }
        println!(")");
        if let Some(desc) = &band.description {
            println!("      {}", desc);
        }
    }

    println!();
    println!("Default thresholds: {:?}", sensor.default_thresholds());

    if !sensor.loader_config.is_empty() {
        println!();
        println!("Loader configuration:");
        for (key, value) in &sensor.loader_config {
            println!("  {}: {}", key, value);
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cli_parse() {
        let mut cmd = Command::new("glint-mask");
        for sensor in SENSOR_REGISTRY.sensors() {
            cmd = cmd.subcommand(Command::new(sensor.id.as_str()));
        }
        let matches = cmd.try_get_matches_from(vec!["glint-mask", "rgb"]);
        assert!(matches.is_ok());
        let matches = matches.unwrap();
        assert_eq!(matches.subcommand_name(), Some("rgb"));
    }

    #[test]
    fn test_sensor_args_validation() {
        // This would be more comprehensive in a real test suite
        let mut cmd = Command::new("test");
        let mut subcommand = Command::new("rgb");
        subcommand = subcommand
            .arg(Arg::new("input").required(true))
            .arg(Arg::new("output").required(true))
            .arg(
                Arg::new("thresholds")
                    .short('t')
                    .long("thresholds")
                    .value_delimiter(',')
                    .value_parser(clap::value_parser!(f64)),
            )
            .arg(
                Arg::new("pixel_buffer")
                    .short('b')
                    .long("pixel_buffer")
                    .default_value("0")
                    .value_parser(clap::value_parser!(usize)),
            );
        cmd = cmd.subcommand(subcommand);
        let matches = cmd.get_matches_from(vec![
            "test",
            "rgb",
            "test_input",
            "test_output",
            "-t",
            "0.8,0.9,0.7",
            "-b",
            "5",
        ]);

        let sub_matches = matches.subcommand_matches("rgb").unwrap();
        let args = SensorArgs::from(sub_matches);

        assert_eq!(args.pixel_buffer, 5);
        assert_eq!(args.thresholds.as_ref().unwrap().len(), 3);
    }

    #[test]
    fn test_create_loader() {
        let rgb_sensor = SENSOR_REGISTRY.get_sensor("rgb").unwrap();
        let loader = create_loader(rgb_sensor);
        assert!(loader.is_ok());

        let p4ms_sensor = SENSOR_REGISTRY.get_sensor("p4ms").unwrap();
        let loader = create_loader(p4ms_sensor);
        assert!(loader.is_ok());

        let msre_sensor = SENSOR_REGISTRY.get_sensor("msre").unwrap();
        let loader = create_loader(msre_sensor);
        assert!(loader.is_ok());
    }
}