sam2.rs

use anyhow::Result;
use clap::Args;
use usls::{Config, DType, Device, Scale};

#[derive(Args, Debug)]
pub struct Sam2Args {
    /// Scale: t, s, b, l
    #[arg(long, default_value = "t")]
    pub scale: String,

    /// Encoder Dtype: fp32, fp16, q4f16, etc.
    #[arg(long, default_value = "fp32")]
    pub encoder_dtype: DType,

    /// Encoder Device: cpu, cuda:0, mps, coreml, openvino:CPU, etc.
    #[arg(long, global = true, default_value = "cpu")]
    pub encoder_device: Device,

    /// Decoder Dtype: fp32, fp16, q4f16, etc.
    #[arg(long, default_value = "fp32")]
    pub decoder_dtype: DType,

    /// Decoder Device: cpu, cuda:0, mps, coreml, openvino:CPU, etc.
    #[arg(long, global = true, default_value = "cpu")]
    pub decoder_device: Device,

    /// Processor device (for pre/post processing)
    #[arg(long, global = true, default_value = "cpu")]
    pub processor_device: Device,

    /// Batch size
    #[arg(long, global = true, default_value_t = 1)]
    pub batch: usize,

    /// Min batch size (TensorRT)
    #[arg(long, global = true, default_value_t = 1)]
    pub min_batch: usize,

    /// Max batch size (TensorRT)
    #[arg(long, global = true, default_value_t = 4)]
    pub max_batch: usize,

    /// num dry run
    #[arg(long, global = true, default_value_t = 3)]
    pub num_dry_run: usize,
}

pub fn config(args: &Sam2Args) -> Result<Config> {
    let config = match args.scale.parse()? {
        Scale::T => Config::sam2_1_tiny(),
        Scale::S => Config::sam2_1_small(),
        Scale::B => Config::sam2_1_base_plus(),
        Scale::L => Config::sam2_1_large(),
        _ => anyhow::bail!("Unsupported SAM2.1 scale: {}. Try t, s, b, l.", args.scale),
    }
    .with_encoder_dtype(args.encoder_dtype)
    .with_encoder_device(args.encoder_device)
    .with_decoder_dtype(args.decoder_dtype)
    .with_decoder_device(args.decoder_device)
    .with_batch_size_min_opt_max_all(args.min_batch, args.batch, args.max_batch)
    .with_num_dry_run_all(args.num_dry_run)
    .with_image_processor_device(args.processor_device);

    Ok(config)
}