mdeville/remove_background
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Compare commits

base: mdeville:25a4a204b9af8d04a323ce59c44a5583c1ad9e9b
Branches Tags
mdeville:master
..
compare: mdeville:31d84c99d74a33d980d5e3e9faeb714e0907d913
Branches Tags
mdeville:master

No commits in common. "25a4a204b9af8d04a323ce59c44a5583c1ad9e9b" and "31d84c99d74a33d980d5e3e9faeb714e0907d913" have entirely different histories.
25a4a204b9 ... 31d84c99d7

12 changed files with 106 additions and 288 deletions
Show stats Expand all files Collapse all files

6
Cargo.lock generated
View file

@ -2750,7 +2750,6 @@ checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a"
name = "remove_background"
version = "0.1.0"
dependencies = [
"anyhow",
"clap",
"hex",
"image",
@ -2759,7 +2758,6 @@ dependencies = [
"reqwest",
"sha2",
"show-image",
"thiserror 2.0.18",
"tracing",
"tracing-subscriber",
]
@ -2872,9 +2870,9 @@ dependencies = [
[[package]]
name = "rustls"
version = "0.23.39"
version = "0.23.38"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7c2c118cb077cca2822033836dfb1b975355dfb784b5e8da48f7b6c5db74e60e"
checksum = "69f9466fb2c14ea04357e91413efb882e2a6d4a406e625449bc0a5d360d53a21"
dependencies = [
"aws-lc-rs",
"once_cell",

8
Cargo.toml
View file

@ -3,20 +3,14 @@
name = "remove_background"
version = "0.1.0"
[features]
default = []
cuda = ["ort/cuda"]
[dependencies]
anyhow = "1"
clap = { version = "4", features = ["derive"] }
hex = "0.4"
image = "0.25"
ndarray = "0.17"
ort = { version = "=2.0.0-rc.12" }
ort = "=2.0.0-rc.12"
reqwest = { version = "0.13", features = ["blocking"] }
sha2 = "0.11"
show-image = { version = "0.14", features = ["image"] }
thiserror = "2"
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }

6
flake.lock
View file

@ -62,11 +62,11 @@
"nixpkgs": "nixpkgs_2"
},
"locked": {
"lastModified": 1776914043,
"narHash": "sha256-qug5r56yW1qOsjSI99l3Jm15JNT9CvS2otkXNRNtrPI=",
"lastModified": 1776827647,
"narHash": "sha256-sYixYhp5V8jCajO8TRorE4fzs7IkL4MZdfLTKgkPQBk=",
"owner": "oxalica",
"repo": "rust-overlay",
"rev": "2d35c4358d7de3a0e606a6e8b27925d981c01cc3",
"rev": "40e6ccc06e1245a4837cbbd6bdda64e21cc67379",
"type": "github"
},
"original": {

28
flake.nix
View file

@ -18,10 +18,7 @@
system:
let
overlays = [ (import rust-overlay) ];
pkgs = import nixpkgs {
inherit system overlays;
config.allowUnfree = true;
};
pkgs = import nixpkgs { inherit system overlays; };
stdenv = pkgs.stdenvAdapters.useMoldLinker pkgs.clangStdenv;
rustToolchain = pkgs.rust-bin.stable.latest.default.override {
extensions = [ "rust-src" ];
@ -35,10 +32,10 @@
];
xorgBuildInputs = with pkgs; [
libx11
libxcursor
libxi
libxrandr
xorg.libx11
xorg.libxcursor
xorg.libxi
xorg.libxrandr
];
waylandBuildInputs = with pkgs; [
@ -52,20 +49,7 @@
openssl
];
# CUDA runtime libraries required by ONNX Runtime CUDA EP.
cudaBuildInputs = with pkgs.cudaPackages; [
cuda_cudart
libcublas
libcurand
libcufft
cudnn
];
buildInputs =
xorgBuildInputs
++ waylandBuildInputs
++ graphicsBuildInputs
++ cudaBuildInputs;
buildInputs = xorgBuildInputs ++ waylandBuildInputs ++ graphicsBuildInputs;
mkShell = pkgs.mkShell.override {
stdenv = stdenv;

48
src/error.rs
View file

@ -1,48 +0,0 @@
use {std::path::PathBuf, thiserror::Error};
pub type AppResult<T> = Result<T, AppError>;
#[derive(Debug, Error)]
pub enum AppError {
#[error(transparent)]
Io(#[from] std::io::Error),
#[error(transparent)]
EnvVar(#[from] std::env::VarError),
#[error(transparent)]
Image(#[from] image::ImageError),
#[error(transparent)]
Reqwest(#[from] reqwest::Error),
#[error(transparent)]
Ort(#[from] ort::Error),
#[error("failed to configure ONNX session builder: {message}")]
OrtSessionBuilder { message: String },
#[error(transparent)]
NdarrayShape(#[from] ndarray::ShapeError),
#[error("no input data provided via stdin")]
NoStdinInput,
#[error("download failed: HTTP {status}")]
DownloadHttpStatus { status: reqwest::StatusCode },
#[error("custom model path does not exist: {path}")]
CustomModelPathMissing { path: PathBuf },
#[error("model hash verification failed: {path}")]
ModelHashVerificationFailed { path: PathBuf },
#[error("model not found in cache and offline mode is enabled: {cache_path}")]
OfflineModelMissing { cache_path: PathBuf },
#[error("expected 4D output tensor [N, C, H, W], got shape {shape:?}")]
UnexpectedTensorShape { shape: Vec<i64> },
#[error("expected batch size 1, got {batch_size}")]
UnexpectedBatchSize { batch_size: usize },
}

1
src/lib.rs
View file

@ -1,4 +1,3 @@
pub mod error;
pub mod model;
pub mod postprocessing;
pub mod sessions;

86
src/main.rs
View file

@ -1,15 +1,14 @@
use {
anyhow::{Context, Result},
clap::{Parser, Subcommand},
image::{GenericImageView, ImageReader},
remove_background::{
error::AppError,
model::Model,
postprocessing::{apply_mask, create_side_by_side},
sessions::init_session,
},
show_image::{AsImageView, create_window, event},
std::{
error::Error,
fs,
io::{Cursor, Read, Write, stdin, stdout},
path::PathBuf,
@ -69,11 +68,9 @@ struct Args {
}
#[show_image::main]
fn main() -> Result<()> {
fn main() -> Result<(), Box<dyn Error>> {
fmt()
.with_env_filter(
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
)
.with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")))
.with_writer(std::io::stderr)
.init();
@ -86,8 +83,7 @@ fn main() -> Result<()> {
"Starting remove_background"
);
let mut session = init_session(args.model_path.as_deref(), args.model, args.offline)
.context("failed to initialize inference session")?;
let mut session = init_session(args.model_path.as_deref(), args.model, args.offline)?;
match args.command {
Command::Single {
@ -97,65 +93,44 @@ fn main() -> Result<()> {
} => {
let img = if let Some(input_file) = input_file {
debug!(path = %input_file.display(), "Reading input image from file");
ImageReader::open(&input_file)
.with_context(|| {
format!("failed to open input image {}", input_file.display())
})?
.decode()
.with_context(|| {
format!("failed to decode input image {}", input_file.display())
})?
ImageReader::open(input_file)?.decode()?
} else {
debug!("Reading input image from stdin");
let mut bytes = Vec::new();
if stdin().lock().read_to_end(&mut bytes)? == 0 {
return Err(AppError::NoStdinInput.into());
return Err("No input data provided via stdin".into());
}
ImageReader::new(Cursor::new(bytes))
.with_guessed_format()?
.decode()
.context("failed to decode input image from stdin")?
.decode()?
};
let (img_width, img_height) = img.dimensions();
info!(width = img_width, height = img_height, "Loaded image");
let mask = session
.predict(&img)
.context("failed to predict segmentation mask")?;
let result_rgba = apply_mask(&img, &mask).context("failed to apply mask to image")?;
let mask = session.predict(&img)?;
let result_rgba = apply_mask(&img, &mask)?;
if debug {
debug_mode(&img, &result_rgba)?;
} else {
}
if let Some(output_file) = output_file {
result_rgba.save(&output_file).with_context(|| {
format!("failed to save output image {}", output_file.display())
})?;
result_rgba.save(&output_file)?;
info!(path = %output_file.display(), "Wrote output image to file");
} else {
debug!("Writing output image to stdout (PNG)");
let mut buffer = Cursor::new(Vec::new());
result_rgba
.write_to(&mut buffer, image::ImageFormat::Png)
.context("failed to encode PNG for stdout output")?;
result_rgba.write_to(&mut buffer, image::ImageFormat::Png)?;
let mut stdout = stdout().lock();
stdout
.write_all(&buffer.into_inner())
.context("failed to write PNG bytes to stdout")?;
}
stdout.write_all(&buffer.into_inner())?;
}
}
Command::Batch {
input_directory,
output_directory,
} => {
fs::create_dir_all(&output_directory).with_context(|| {
format!(
"failed to create output directory {}",
output_directory.display()
)
})?;
fs::create_dir_all(&output_directory)?;
info!(
input = %input_directory.display(),
output = %output_directory.display(),
@ -164,12 +139,7 @@ fn main() -> Result<()> {
let mut processed: usize = 0;
let mut failed: usize = 0;
for entry in fs::read_dir(&input_directory).with_context(|| {
format!(
"failed to read input directory {}",
input_directory.display()
)
})? {
for entry in fs::read_dir(&input_directory)? {
let entry = entry?;
let path = entry.path();
let span = tracing::info_span!("batch_item", path = %path.display());
@ -191,16 +161,10 @@ fn main() -> Result<()> {
}
};
let mask = session
.predict(&img)
.with_context(|| format!("failed to predict mask for {}", path.display()))?;
let result_rgba = apply_mask(&img, &mask)
.with_context(|| format!("failed to apply mask for {}", path.display()))?;
let mut output_path = output_directory.join(path.file_name().unwrap());
output_path.set_extension("png");
result_rgba.save(&output_path).with_context(|| {
format!("failed to save processed image {}", output_path.display())
})?;
let mask = session.predict(&img)?;
let result_rgba = apply_mask(&img, &mask)?;
let output_path = output_directory.join(path.file_name().unwrap());
result_rgba.save(&output_path)?;
processed += 1;
info!(output = %output_path.display(), "Processed image saved");
}
@ -213,10 +177,12 @@ fn main() -> Result<()> {
Ok(())
}
fn debug_mode(img: &image::DynamicImage, result_rgba: &image::RgbaImage) -> Result<()> {
fn debug_mode(
img: &image::DynamicImage,
result_rgba: &image::RgbaImage,
) -> Result<(), Box<dyn Error>> {
info!("Creating side-by-side comparison");
let composite = create_side_by_side(img, result_rgba)
.context("failed to create side-by-side debug image")?;
let composite = create_side_by_side(img, result_rgba)?;
let composite_dynamic = image::DynamicImage::ImageRgba8(composite);
let (comp_width, comp_height) = composite_dynamic.dimensions();
@ -228,7 +194,7 @@ fn debug_mode(img: &image::DynamicImage, result_rgba: &image::RgbaImage) -> Resu
"comparison",
composite_dynamic
.as_image_view()
.map_err(|e| anyhow::anyhow!(e.to_string()))?,
.map_err(|e| e.to_string())?,
)?;
info!(

130
src/model.rs
View file

@ -1,8 +1,12 @@
use {
clap::ValueEnum,
ort::session::Session,
ort::{
execution_providers::CUDAExecutionProvider,
session::{Session, builder::GraphOptimizationLevel},
},
sha2::{Digest, Sha256},
std::{
error::Error,
fs,
io::Read,
path::{Path, PathBuf},
@ -11,11 +15,6 @@ use {
tracing::{debug, info, warn},
};
use crate::error::{AppError, AppResult};
#[cfg(feature = "cuda")]
use ort::{builder::GraphOptimizationLevel, ep::CUDAExecutionProvider};
/// CLI-facing model selector. Concrete session metadata (URL, checksum,
/// preprocessing params) lives on the `Session` trait impls in `sessions/`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
@ -26,7 +25,7 @@ pub enum Model {
}
/// Get the cache directory for models
fn get_cache_dir() -> AppResult<PathBuf> {
fn get_cache_dir() -> Result<PathBuf, Box<dyn Error>> {
let home = std::env::var("HOME").or_else(|_| std::env::var("USERPROFILE"))?;
let cache_dir = Path::new(&home)
.join(".cache")
@ -36,7 +35,7 @@ fn get_cache_dir() -> AppResult<PathBuf> {
Ok(cache_dir)
}
fn download_file(url: &str, dest: &Path) -> AppResult<()> {
fn download_file(url: &str, dest: &Path) -> Result<(), Box<dyn Error>> {
info!(%url, dest = %dest.display(), "Downloading model");
let start = Instant::now();
@ -47,9 +46,7 @@ fn download_file(url: &str, dest: &Path) -> AppResult<()> {
let mut response = client.get(url).send()?;
if !response.status().is_success() {
return Err(AppError::DownloadHttpStatus {
status: response.status(),
});
return Err(format!("Failed to download model: HTTP {}", response.status()).into());
}
let mut file = std::io::BufWriter::new(fs::File::create(dest)?);
@ -71,12 +68,7 @@ fn download_file(url: &str, dest: &Path) -> AppResult<()> {
// Only emit every 10% to avoid log spam while still giving feedback.
if progress >= last_reported_pct + 10 {
last_reported_pct = progress - (progress % 10);
debug!(
percent = last_reported_pct,
downloaded,
total = total_size,
"Download progress"
);
debug!(percent = last_reported_pct, downloaded, total = total_size, "Download progress");
}
}
}
@ -90,7 +82,7 @@ fn download_file(url: &str, dest: &Path) -> AppResult<()> {
Ok(())
}
fn verify_hash(file_path: &Path, expected_hash: &str) -> AppResult<bool> {
fn verify_hash(file_path: &Path, expected_hash: &str) -> Result<bool, Box<dyn Error>> {
const BUF_SIZE: usize = 8192;
let mut file = fs::File::open(file_path)?;
@ -119,11 +111,11 @@ pub fn get_model_path(
sha256: Option<&str>,
custom_path: Option<&str>,
offline: bool,
) -> AppResult<PathBuf> {
) -> Result<PathBuf, Box<dyn Error>> {
if let Some(path) = custom_path {
let model_path = PathBuf::from(path);
if !model_path.exists() {
return Err(AppError::CustomModelPathMissing { path: model_path });
return Err(format!("Custom model path does not exist: {}", path).into());
}
info!(%path, "Using custom model");
return Ok(model_path);
@ -142,9 +134,9 @@ pub fn get_model_path(
debug!("Cached model hash OK");
} else {
warn!(path = %model_path.display(), "Cached model hash verification failed");
return Err(AppError::ModelHashVerificationFailed {
path: model_path.clone(),
});
return Err(
"Model hash verification failed. Try deleting the cached model.".into(),
);
}
}
@ -152,9 +144,11 @@ pub fn get_model_path(
}
if offline {
return Err(AppError::OfflineModelMissing {
cache_path: model_path,
});
return Err(format!(
"Model not found in cache and offline mode is enabled. Cache path: {}",
model_path.display()
)
.into());
}
info!("Model not found in cache, downloading");
@ -166,40 +160,23 @@ pub fn get_model_path(
debug!("Downloaded model hash OK");
} else {
fs::remove_file(&model_path)?;
return Err(AppError::ModelHashVerificationFailed { path: model_path });
return Err("Downloaded model hash verification failed".into());
}
}
Ok(model_path)
}
/// Create an ONNX Runtime session from a model path.
///
/// Uses the CUDA execution provider when built with `--features cuda`; otherwise runs on CPU.
pub fn create_session(model_path: &Path) -> AppResult<Session> {
#[cfg(feature = "cuda")]
/// Create an ONNX Runtime session from a model path with CUDA (falls back to CPU).
pub fn create_session(model_path: &Path) -> Result<Session, Box<dyn Error>> {
info!(path = %model_path.display(), "Loading model into ONNX Runtime with CUDA backend");
#[cfg(not(feature = "cuda"))]
info!(path = %model_path.display(), "Loading model into ONNX Runtime with CPU backend");
let start = Instant::now();
let mut builder = Session::builder()?;
#[cfg(feature = "cuda")]
let builder = builder
.with_execution_providers([CUDAExecutionProvider::default().build()])
.map_err(|err| AppError::OrtSessionBuilder {
message: err.to_string(),
})?
.with_optimization_level(GraphOptimizationLevel::Level3)
.map_err(|err| AppError::OrtSessionBuilder {
message: err.to_string(),
})?
.with_intra_threads(4)
.map_err(|err| AppError::OrtSessionBuilder {
message: err.to_string(),
})?;
let session = builder.commit_from_file(model_path)?;
let session = Session::builder()?
.with_execution_providers([CUDAExecutionProvider::default().build()])?
.with_optimization_level(GraphOptimizationLevel::Level3)?
.with_intra_threads(4)?
.commit_from_file(model_path)?;
info!(
elapsed_ms = start.elapsed().as_millis() as u64,
@ -208,54 +185,3 @@ pub fn create_session(model_path: &Path) -> AppResult<Session> {
Ok(session)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn custom_model_path_missing_returns_typed_error() {
let missing_path = PathBuf::from("/definitely/not/a/real/model.onnx");
let result = get_model_path(
"unused",
"https://example.invalid/model.onnx",
None,
Some(missing_path.to_str().expect("valid utf-8 path")),
false,
);
match result {
Err(AppError::CustomModelPathMissing { path }) => assert_eq!(path, missing_path),
other => panic!("unexpected result: {other:?}"),
}
}
#[test]
fn offline_missing_model_returns_typed_error() {
let unique_name = format!(
"test-model-{}",
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time before epoch")
.as_nanos()
);
let result = get_model_path(
&unique_name,
"https://example.invalid/model.onnx",
None,
None,
true,
);
match result {
Err(AppError::OfflineModelMissing { cache_path }) => {
let expected = format!("{unique_name}.onnx");
assert_eq!(
cache_path.file_name().and_then(|f| f.to_str()),
Some(expected.as_str())
);
}
other => panic!("unexpected result: {other:?}"),
}
}
}

15
src/postprocessing.rs
View file

@ -1,15 +1,14 @@
use {
image::{DynamicImage, GenericImageView, GrayImage, Rgba, RgbaImage, imageops::FilterType},
std::error::Error,
tracing::debug,
};
use crate::error::AppResult;
/// Compose `original` with `mask` as the alpha channel and return an RGBA image.
///
/// The mask is expected to already be grayscale. If its dimensions differ from
/// the original, it is resized with LANCZOS3.
pub fn apply_mask(original: &DynamicImage, mask: &GrayImage) -> AppResult<RgbaImage> {
pub fn apply_mask(original: &DynamicImage, mask: &GrayImage) -> Result<RgbaImage, Box<dyn Error>> {
let (orig_width, orig_height) = original.dimensions();
let (mask_width, mask_height) = mask.dimensions();
@ -17,7 +16,10 @@ pub fn apply_mask(original: &DynamicImage, mask: &GrayImage) -> AppResult<RgbaIm
if mask_width != orig_width || mask_height != orig_height {
debug!(
mask_width,
mask_height, orig_width, orig_height, "Resizing mask to match original image"
mask_height,
orig_width,
orig_height,
"Resizing mask to match original image"
);
std::borrow::Cow::Owned(image::imageops::resize(
mask,
@ -51,7 +53,10 @@ pub fn apply_mask(original: &DynamicImage, mask: &GrayImage) -> AppResult<RgbaIm
}
/// Create a side-by-side comparison image
pub fn create_side_by_side(original: &DynamicImage, result: &RgbaImage) -> AppResult<RgbaImage> {
pub fn create_side_by_side(
original: &DynamicImage,
result: &RgbaImage,
) -> Result<RgbaImage, Box<dyn Error>> {
let (width, height) = original.dimensions();
let mut composite = RgbaImage::new(width * 2, height);

13
src/sessions/birefnet_lite.rs
View file

@ -1,9 +1,8 @@
use {ort::session::Session as OrtSession, std::path::Path};
use ort::session::Session as OrtSession;
use std::error::Error;
use std::path::Path;
use crate::{
error::AppResult,
model::{create_session, get_model_path},
};
use crate::model::{create_session, get_model_path};
use super::Session;
@ -13,12 +12,12 @@ pub struct BiRefNetLiteSession {
}
impl BiRefNetLiteSession {
pub fn new(offline: bool) -> AppResult<Self> {
pub fn new(offline: bool) -> Result<Self, Box<dyn Error>> {
let path = get_model_path(Self::name(), Self::url(), Self::sha256(), None, offline)?;
Self::from_model_path(&path)
}
pub fn from_model_path(path: &Path) -> AppResult<Self> {
pub fn from_model_path(path: &Path) -> Result<Self, Box<dyn Error>> {
let inner_session = create_session(path)?;
let input_name = inner_session.inputs()[0].name().to_string();
Ok(Self {

13
src/sessions/bria.rs
View file

@ -1,9 +1,8 @@
use {ort::session::Session as OrtSession, std::path::Path};
use ort::session::Session as OrtSession;
use std::error::Error;
use std::path::Path;
use crate::{
error::AppResult,
model::{create_session, get_model_path},
};
use crate::model::{create_session, get_model_path};
use super::Session;
@ -13,12 +12,12 @@ pub struct BriaSession {
}
impl BriaSession {
pub fn new(offline: bool) -> AppResult<Self> {
pub fn new(offline: bool) -> Result<Self, Box<dyn Error>> {
let path = get_model_path(Self::name(), Self::url(), Self::sha256(), None, offline)?;
Self::from_model_path(&path)
}
pub fn from_model_path(path: &Path) -> AppResult<Self> {
pub fn from_model_path(path: &Path) -> Result<Self, Box<dyn Error>> {
let inner_session = create_session(path)?;
let input_name = inner_session.inputs()[0].name().to_string();
Ok(Self {

30
src/sessions/mod.rs
View file

@ -2,14 +2,11 @@ use {
image::{DynamicImage, GenericImageView, GrayImage, Luma, imageops::FilterType},
ndarray::{Array4, IntoDimension},
ort::{session::Session as OrtSession, value::Tensor},
std::time::Instant,
std::{error::Error, time::Instant},
tracing::{debug, info},
};
use crate::{
error::{AppError, AppResult},
model::Model,
};
use crate::model::Model;
mod birefnet_lite;
mod bria;
@ -65,7 +62,7 @@ pub trait Session {
/// Port of rembg's `BaseSession.normalize`: resize with LANCZOS,
/// scale into `[0, 1]` by dividing by the max pixel value, then apply
/// channel-wise mean/std.
fn normalize(&self, img: &DynamicImage) -> AppResult<Array4<f32>> {
fn normalize(&self, img: &DynamicImage) -> Result<Array4<f32>, Box<dyn Error>> {
let (w, h) = self.input_size();
let resized = img.resize_exact(w, h, FilterType::Lanczos3).to_rgb8();
let (width, height) = resized.dimensions();
@ -107,15 +104,12 @@ pub trait Session {
/// 3. optional sigmoid (birefnet)
/// 4. min/max normalize into `[0, 1]`
/// 5. scale to `u8`, resize to original image dimensions
fn predict(&mut self, img: &DynamicImage) -> AppResult<GrayImage> {
fn predict(&mut self, img: &DynamicImage) -> Result<GrayImage, Box<dyn Error>> {
let (orig_w, orig_h) = img.dimensions();
let preprocess_start = Instant::now();
let input = self.normalize(img)?;
debug!(
elapsed_ms = preprocess_start.elapsed().as_millis() as u64,
"Preprocessing complete"
);
debug!(elapsed_ms = preprocess_start.elapsed().as_millis() as u64, "Preprocessing complete");
let apply_sigmoid = self.apply_sigmoid();
let input_name = self.input_name().to_string();
@ -133,9 +127,11 @@ pub trait Session {
let (shape, data) = output.try_extract_tensor::<f32>()?;
if shape.len() != 4 {
return Err(AppError::UnexpectedTensorShape {
shape: shape.iter().copied().collect(),
});
return Err(format!(
"Expected 4D output tensor [N, C, H, W], got shape {:?}",
shape
)
.into());
}
let (n, _c, h, w) = (
shape[0] as usize,
@ -144,7 +140,7 @@ pub trait Session {
shape[3] as usize,
);
if n != 1 {
return Err(AppError::UnexpectedBatchSize { batch_size: n });
return Err(format!("Expected batch size 1, got {}", n).into());
}
let view = ndarray::ArrayView::from_shape(
@ -204,11 +200,11 @@ pub fn init_session(
custom_model_path: Option<&str>,
model: Model,
offline: bool,
) -> AppResult<Box<dyn Session>> {
) -> Result<Box<dyn Session>, Box<dyn Error>> {
Ok(if let Some(custom_path) = custom_model_path {
let path = std::path::PathBuf::from(custom_path);
if !path.exists() {
return Err(AppError::CustomModelPathMissing { path });
return Err(format!("Custom model path does not exist: {}", custom_path).into());
}
info!(path = %custom_path, "Using custom model");
match model {