huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-15 10:26:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add Resize to onnx ops (#2946)

Browse Source 
* added resize to candle-onnx, not currently working

* changed unreachable to bail, and bailed when both scales and sizes are set

* cleanup and added other unused options for this op

* cleanup

* fixed image loading to make output work

* cleanup and removed unused variables

* removed path path creation code, and changed unwrap to ?
This commit is contained in:
Kyle Birnbaum
2025-05-09 22:05:03 -07:00
committed by GitHub
parent 3d05f5cf3d
commit 36508a2c93
2 changed files with 118 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
candle-examples/examples/onnx/main.rs
View File

@ -5,12 +5,14 @@ extern crate intel_mkl_src;
extern crate accelerate_src;
use candle::{IndexOp, D};
use candle_examples::save_image;
use clap::{Parser, ValueEnum};
#[derive(Clone, Copy, Debug, ValueEnum)]
enum Which {
SqueezeNet,
EfficientNet,
EsrGan,
}
#[derive(Parser)]
@ -28,10 +30,21 @@ struct Args {
pub fn main() -> anyhow::Result<()> {
let args = Args::parse();
let image = candle_examples::imagenet::load_image224(args.image)?;
let image = match args.which {
Which::SqueezeNet | Which::EfficientNet => {
candle_examples::imagenet::load_image224(&args.image)?
}
Which::EsrGan => candle_examples::imagenet::load_image_with_std_mean(
&args.image,
128,
&[0.0f32, 0.0, 0.0],
&[1.0f32, 1.0, 1.0],
)?,
};
let image = match args.which {
Which::SqueezeNet => image,
Which::EfficientNet => image.permute((1, 2, 0))?,
Which::EsrGan => image,
};
println!("loaded image {image:?}");
@ -45,6 +58,9 @@ pub fn main() -> anyhow::Result<()> {
Which::EfficientNet => hf_hub::api::sync::Api::new()?
.model("onnx/EfficientNet-Lite4".into())
.get("efficientnet-lite4-11.onnx")?,
Which::EsrGan => hf_hub::api::sync::Api::new()?
.model("qualcomm/Real-ESRGAN-x4plus".into())
.get("Real-ESRGAN-x4plus.onnx")?,
},
};
@ -57,21 +73,40 @@ pub fn main() -> anyhow::Result<()> {
let prs = match args.which {
Which::SqueezeNet => candle_nn::ops::softmax(&output, D::Minus1)?,
Which::EfficientNet => output,
Which::EsrGan => output,
};
let prs = prs.i(0)?.to_vec1::<f32>()?;
// Sort the predictions and take the top 5
let mut top: Vec<_> = prs.iter().enumerate().collect();
top.sort_by(|a, b| b.1.partial_cmp(a.1).unwrap());
let top = top.into_iter().take(5).collect::<Vec<_>>();
match args.which {
Which::EfficientNet | Which::SqueezeNet => {
let prs = prs.i(0)?.to_vec1::<f32>()?;
// Print the top predictions
for &(i, p) in &top {
println!(
"{:50}: {:.2}%",
candle_examples::imagenet::CLASSES[i],
p * 100.0
);
// Sort the predictions and take the top 5
let mut top: Vec<_> = prs.iter().enumerate().collect();
top.sort_by(|a, b| b.1.partial_cmp(a.1).unwrap());
let top = top.into_iter().take(5).collect::<Vec<_>>();
// Print the top predictions
for &(i, p) in &top {
println!(
"{:50}: {:.2}%",
candle_examples::imagenet::CLASSES[i],
p * 100.0
);
}
}
Which::EsrGan => {
let max_pixel_val = candle::Tensor::try_from(255.0f32)?
.to_device(prs.device())?
.broadcast_as(prs.shape())?;
let out = (prs * max_pixel_val)?.i(0)?.to_dtype(candle::DType::U8)?;
let pb = std::path::PathBuf::from(args.image);
let input_file_name = pb.file_name().unwrap();
let mut output_file_name = std::ffi::OsString::from("super_");
output_file_name.push(input_file_name);
save_image(&out, output_file_name)?;
}
}
Ok(())