Add an initial Segformer implementation (#1617)

* add segformer

* Make the id2label field optional.

---------

Co-authored-by: laurent <laurent.mazare@gmail.com>

candle-examples/examples/segformer/main.rs

@@ -0,0 +1,155 @@
+use candle::Device;
+use candle::Module;
+use candle_nn::VarBuilder;
+use candle_transformers::models::segformer::{
+    Config, ImageClassificationModel, SemanticSegmentationModel,
+};
+use clap::{Args, Parser, Subcommand};
+use image::Rgb;
+use imageproc::integral_image::ArrayData;
+use std::collections::HashMap;
+use std::path::PathBuf;
+
+#[derive(Parser)]
+#[clap(about, version, long_about = None)]
+struct CliArgs {
+    #[arg(long, help = "use cpu")]
+    cpu: bool,
+    #[command(subcommand)]
+    command: Commands,
+}
+#[derive(Args, Debug)]
+struct SegmentationArgs {
+    #[arg(
+        long,
+        help = "name of the huggingface hub model",
+        default_value = "nvidia/segformer-b0-finetuned-ade-512-512"
+    )]
+    model_name: String,
+    #[arg(
+        long,
+        help = "path to the label file in json format",
+        default_value = "candle-examples/examples/segformer/assets/labels.json"
+    )]
+    label_path: PathBuf,
+    #[arg(long, help = "path to for the output mask image")]
+    output_path: PathBuf,
+    #[arg(help = "path to image as input")]
+    image: PathBuf,
+}
+
+#[derive(Args, Debug)]
+struct ClassificationArgs {
+    #[arg(
+        long,
+        help = "name of the huggingface hub model",
+        default_value = "paolinox/segformer-finetuned-food101"
+    )]
+    model_name: String,
+    #[arg(help = "path to image as input")]
+    image: PathBuf,
+}
+
+#[derive(Subcommand, Debug)]
+enum Commands {
+    Segment(SegmentationArgs),
+    Classify(ClassificationArgs),
+}
+
+fn get_vb_and_config(model_name: String, device: &Device) -> anyhow::Result<(VarBuilder, Config)> {
+    println!("loading model {} via huggingface hub", model_name);
+    let api = hf_hub::api::sync::Api::new()?;
+    let api = api.model(model_name.clone());
+    let model_file = api.get("model.safetensors")?;
+    println!("model {} downloaded and loaded", model_name);
+    let vb =
+        unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], candle::DType::F32, device)? };
+    let config = std::fs::read_to_string(api.get("config.json")?)?;
+    let config: Config = serde_json::from_str(&config)?;
+    println!("{:?}", config);
+    Ok((vb, config))
+}
+
+#[derive(Debug, serde::Deserialize)]
+struct LabelItem {
+    index: u32,
+    color: String,
+}
+
+fn segmentation_task(args: SegmentationArgs, device: &Device) -> anyhow::Result<()> {
+    let label_file = std::fs::read_to_string(&args.label_path)?;
+    let label_items: Vec<LabelItem> = serde_json::from_str(&label_file)?;
+    let label_colors: HashMap<u32, Rgb<u8>> = label_items
+        .iter()
+        .map(|x| {
+            (x.index - 1, {
+                let color = x.color.trim_start_matches('#');
+                let r = u8::from_str_radix(&color[0..2], 16).unwrap();
+                let g = u8::from_str_radix(&color[2..4], 16).unwrap();
+                let b = u8::from_str_radix(&color[4..6], 16).unwrap();
+                Rgb([r, g, b])
+            })
+        })
+        .collect();
+
+    let image = candle_examples::imagenet::load_image224(args.image)?
+        .unsqueeze(0)?
+        .to_device(device)?;
+    let (vb, config) = get_vb_and_config(args.model_name, device)?;
+    let num_labels = label_items.len();
+
+    let model = SemanticSegmentationModel::new(&config, num_labels, vb)?;
+    let segmentations = model.forward(&image)?;
+
+    // generate a mask image
+    let mask = &segmentations.squeeze(0)?.argmax(0)?;
+    let (h, w) = mask.dims2()?;
+    let mask = mask.flatten_all()?.to_vec1::<u32>()?;
+    let mask = mask
+        .iter()
+        .flat_map(|x| label_colors[x].data())
+        .collect::<Vec<u8>>();
+    let mask: image::ImageBuffer<image::Rgb<u8>, Vec<u8>> =
+        image::ImageBuffer::from_raw(w as u32, h as u32, mask).unwrap();
+    // resize
+    let mask = image::DynamicImage::from(mask);
+    let mask = mask.resize_to_fill(
+        w as u32 * 4,
+        h as u32 * 4,
+        image::imageops::FilterType::CatmullRom,
+    );
+    mask.save(args.output_path.clone())?;
+    println!("mask image saved to {:?}", args.output_path);
+    Ok(())
+}
+
+fn classification_task(args: ClassificationArgs, device: &Device) -> anyhow::Result<()> {
+    let image = candle_examples::imagenet::load_image224(args.image)?
+        .unsqueeze(0)?
+        .to_device(device)?;
+    let (vb, config) = get_vb_and_config(args.model_name, device)?;
+    let num_labels = 7;
+    let model = ImageClassificationModel::new(&config, num_labels, vb)?;
+    let classification = model.forward(&image)?;
+    let classification = candle_nn::ops::softmax_last_dim(&classification)?;
+    let classification = classification.squeeze(0)?;
+    println!(
+        "classification logits {:?}",
+        classification.to_vec1::<f32>()?
+    );
+    let label_id = classification.argmax(0)?.to_scalar::<u32>()?;
+    let label_id = format!("{}", label_id);
+    println!("label: {}", config.id2label[&label_id]);
+    Ok(())
+}
+
+pub fn main() -> anyhow::Result<()> {
+    let args = CliArgs::parse();
+    let device = candle_examples::device(args.cpu)?;
+    if let Commands::Segment(args) = args.command {
+        segmentation_task(args, &device)?
+    } else if let Commands::Classify(args) = args.command {
+        classification_task(args, &device)?
+    }
+    Ok(())
+}