outline/
lib.rs

1//! # Outline
2//!
3//! Image background removal with flexible mask processing options.
4//!
5//! Powered by ONNX Runtime ([`ort`](https://docs.rs/ort)) and VTracer, and works with U2-Net, BiRefNet,
6//! and other ONNX models with a compatible input/output shape.
7//!
8//! # Quick Start
9//!
10//! ```no_run
11//! use outline::Outline;
12//!
13//! let outline = Outline::new("model.onnx");
14//! let session = outline.for_image("input.png")?;
15//! let matte = session.matte();
16//!
17//! // Compose the foreground directly from the raw matte (soft edges)
18//! let foreground = matte.foreground()?;
19//! foreground.save("foreground.png")?;
20//!
21//! // Process the mask and save it
22//! let mask = matte.blur().threshold().processed()?;
23//! mask.save("mask.png")?;
24//! # Ok::<_, outline::OutlineError>(())
25//! ```
26
27#![cfg_attr(docsrs, feature(doc_cfg))]
28
29mod config;
30mod error;
31mod foreground;
32mod inference;
33mod mask;
34mod vectorizer;
35
36#[doc(inline)]
37pub use crate::config::{
38    DEFAULT_MODEL_PATH, ENV_MODEL_PATH, InferenceSettings, MaskProcessingOptions,
39};
40#[doc(inline)]
41pub use crate::error::{OutlineError, OutlineResult};
42pub use vectorizer::MaskVectorizer;
43
44#[cfg(feature = "vectorizer-vtracer")]
45#[cfg_attr(docsrs, doc(cfg(feature = "vectorizer-vtracer")))]
46#[doc(inline)]
47pub use vectorizer::vtracer::{TraceOptions, VtracerSvgVectorizer, trace_to_svg_string};
48
49use std::path::Path;
50use std::path::PathBuf;
51use std::sync::Arc;
52
53use image::imageops::FilterType;
54use image::{GrayImage, RgbImage, RgbaImage};
55
56use crate::foreground::compose_foreground;
57use crate::inference::run_matte_pipeline;
58use crate::mask::{MaskOperation, apply_operations, operations_from_options};
59
60/// Entry point for configuring and running background matting inference.
61///
62/// This is the main interface for loading an ONNX model and processing images to extract
63/// foreground subjects. Configure model path, inference settings, and default mask processing
64/// options, then call [`for_image`](Outline::for_image) to run inference on individual images.
65#[derive(Debug, Clone)]
66pub struct Outline {
67    /// Inference settings for model and image handling.
68    settings: InferenceSettings,
69    /// If nothing is specified and processing is requested, these options will be used.
70    default_mask_processing: MaskProcessingOptions,
71}
72
73impl Outline {
74    /// Create a new `Outline` instance with the given model path and default settings.
75    pub fn new(model_path: impl Into<PathBuf>) -> Self {
76        Self {
77            settings: InferenceSettings::new(model_path),
78            default_mask_processing: MaskProcessingOptions::default(),
79        }
80    }
81
82    /// Construct Outline using env var `ENV_MODEL_PATH` or fallback to `DEFAULT_MODEL_PATH`.
83    pub fn from_env_or_default() -> Self {
84        let resolved = std::env::var_os(ENV_MODEL_PATH)
85            .map(PathBuf::from)
86            .unwrap_or_else(|| PathBuf::from(DEFAULT_MODEL_PATH));
87        Self::new(resolved)
88    }
89
90    /// Try constructing Outline strictly from env var; returns error if not set.
91    pub fn try_from_env() -> OutlineResult<Self> {
92        if let Some(from_env) = std::env::var_os(ENV_MODEL_PATH) {
93            return Ok(Self::new(PathBuf::from(from_env)));
94        }
95        Err(OutlineError::Io(std::io::Error::new(
96            std::io::ErrorKind::NotFound,
97            format!(
98                "Model path not specified in env {}; set the variable to proceed",
99                ENV_MODEL_PATH
100            ),
101        )))
102    }
103
104    /// Set the filter used to resize the input image for the model.
105    pub fn with_input_resize_filter(mut self, filter: FilterType) -> Self {
106        self.settings.input_resize_filter = filter;
107        self
108    }
109
110    /// Set the filter used to resize the output matte to the original image size.
111    pub fn with_output_resize_filter(mut self, filter: FilterType) -> Self {
112        self.settings.output_resize_filter = filter;
113        self
114    }
115
116    /// Set the number of intra-op threads for the inference.
117    pub fn with_intra_threads(mut self, intra_threads: Option<usize>) -> Self {
118        self.settings.intra_threads = intra_threads;
119        self
120    }
121
122    /// Set the default mask processing options to use when none are specified.
123    pub fn with_default_mask_processing(mut self, options: MaskProcessingOptions) -> Self {
124        self.default_mask_processing = options;
125        self
126    }
127
128    /// Get a reference to the default mask processing options.
129    pub fn default_mask_processing(&self) -> &MaskProcessingOptions {
130        &self.default_mask_processing
131    }
132
133    /// Run the inference pipeline for a single image, returning the orginal image, raw matte, and processing options,
134    /// wrapped in an `InferencedMatte`.
135    pub fn for_image(&self, image_path: impl AsRef<Path>) -> OutlineResult<InferencedMatte> {
136        let (rgb, matte) = run_matte_pipeline(&self.settings, image_path.as_ref())?;
137        Ok(InferencedMatte::new(
138            rgb,
139            matte,
140            self.default_mask_processing.clone(),
141        ))
142    }
143}
144
145/// Inference result containing the original RGB image and raw matte prediction.
146///
147/// Returned by [`Outline::for_image`] after running model inference.
148///
149/// # Example
150/// ```no_run
151/// use outline::Outline;
152///
153/// let outline = Outline::new("model.onnx");
154/// let session = outline.for_image("input.png")?;
155/// let matte = session.matte();
156///
157/// // Access the original image and raw matte directly
158/// let rgb = session.rgb_image();
159/// let raw_matte = session.raw_matte();
160///
161/// // Compose the foreground from the raw matte
162/// let foreground = matte.foreground()?;
163/// foreground.save("foreground.png")?;
164/// # Ok::<_, outline::OutlineError>(())
165/// ```
166#[derive(Debug, Clone)]
167pub struct InferencedMatte {
168    rgb_image: Arc<RgbImage>,
169    raw_matte: Arc<GrayImage>,
170    default_mask_processing: MaskProcessingOptions,
171}
172
173impl InferencedMatte {
174    fn new(
175        rgb_image: RgbImage,
176        raw_matte: GrayImage,
177        default_mask_processing: MaskProcessingOptions,
178    ) -> Self {
179        Self {
180            rgb_image: Arc::new(rgb_image),
181            raw_matte: Arc::new(raw_matte),
182            default_mask_processing,
183        }
184    }
185
186    /// Get a reference to the original RGB image.
187    pub fn rgb_image(&self) -> &RgbImage {
188        self.rgb_image.as_ref()
189    }
190
191    /// Get a reference to the raw grayscale matte.
192    pub fn raw_matte(&self) -> &GrayImage {
193        self.raw_matte.as_ref()
194    }
195
196    /// Begin building a mask processing pipeline from the raw matte.
197    pub fn matte(&self) -> MatteHandle {
198        MatteHandle {
199            rgb_image: Arc::clone(&self.rgb_image),
200            raw_matte: Arc::clone(&self.raw_matte),
201            default_mask_processing: self.default_mask_processing.clone(),
202            operations: Vec::new(),
203        }
204    }
205}
206
207/// Builder for chaining mask processing operations on the raw matte.
208///
209/// The raw matte is the soft, grayscale alpha prediction from the model.
210///
211/// # Example
212/// ```no_run
213/// use outline::Outline;
214///
215/// let outline = Outline::new("model.onnx");
216/// let session = outline.for_image("input.jpg")?;
217///
218/// // Chain operations and execute them
219/// let mask = session.matte()
220///     .blur_with(6.0)          // Smooth edges
221///     .threshold_with(120)     // Convert to binary
222///     .dilate_with(5.0)        // Expand slightly
223///     .processed()?;           // Execute operations
224///
225/// mask.save("mask.png")?;
226/// # Ok::<_, outline::OutlineError>(())
227/// ```
228#[derive(Debug, Clone)]
229pub struct MatteHandle {
230    rgb_image: Arc<RgbImage>,
231    raw_matte: Arc<GrayImage>,
232    default_mask_processing: MaskProcessingOptions,
233    operations: Vec<MaskOperation>,
234}
235
236impl MatteHandle {
237    /// Get the raw grayscale matte.
238    pub fn raw(&self) -> GrayImage {
239        (*self.raw_matte).clone()
240    }
241
242    /// Consume the handle and return the raw grayscale matte.
243    pub fn into_image(self) -> GrayImage {
244        (*self.raw_matte).clone()
245    }
246
247    /// Save the raw grayscale matte to the specified path.
248    pub fn save(&self, path: impl AsRef<Path>) -> OutlineResult<()> {
249        self.raw_matte.as_ref().save(path)?;
250        Ok(())
251    }
252
253    /// Add a blur operation using the default sigma.
254    pub fn blur(mut self) -> Self {
255        let sigma = self.default_mask_processing.blur_sigma;
256        self.operations.push(MaskOperation::Blur { sigma });
257        self
258    }
259
260    /// Add a blur operation with a custom sigma.
261    pub fn blur_with(mut self, sigma: f32) -> Self {
262        self.operations.push(MaskOperation::Blur { sigma });
263        self
264    }
265
266    /// Add a threshold operation using the default mask threshold.
267    pub fn threshold(mut self) -> Self {
268        let value = self.default_mask_processing.mask_threshold;
269        self.operations.push(MaskOperation::Threshold { value });
270        self
271    }
272
273    /// Add a threshold operation with a custom value.
274    pub fn threshold_with(mut self, value: u8) -> Self {
275        self.operations.push(MaskOperation::Threshold { value });
276        self
277    }
278
279    /// Add a dilation operation using the default radius.
280    ///
281    /// **Note**: Dilation typically works best on binary masks. Consider calling
282    /// [`threshold`](MatteHandle::threshold) before `dilate` if working with a soft matte.
283    pub fn dilate(mut self) -> Self {
284        let radius = self.default_mask_processing.dilation_radius;
285        self.operations.push(MaskOperation::Dilate { radius });
286        self
287    }
288
289    /// Add a dilation operation with a custom radius.
290    ///
291    /// **Note**: Dilation typically works best on binary masks. Consider calling
292    /// [`threshold`](MatteHandle::threshold) before `dilate` if working with a soft matte.
293    pub fn dilate_with(mut self, radius: f32) -> Self {
294        self.operations.push(MaskOperation::Dilate { radius });
295        self
296    }
297
298    /// Add a hole-filling operation to the processing pipeline.
299    ///
300    /// **Note**: Hole-filling typically works best on binary masks. Consider calling
301    /// [`threshold`](MatteHandle::threshold) before `fill_holes` if working with a soft matte.
302    pub fn fill_holes(mut self) -> Self {
303        let threshold = self.default_mask_processing.mask_threshold;
304        self.operations.push(MaskOperation::FillHoles { threshold });
305        self
306    }
307
308    /// Process the raw matte with the accumulated operations and default options.
309    pub fn processed(self) -> OutlineResult<MaskHandle> {
310        self.process_with_options(None)
311    }
312
313    /// Process the raw matte with the accumulated operations and custom options.
314    pub fn processed_with(self, options: &MaskProcessingOptions) -> OutlineResult<MaskHandle> {
315        self.process_with_options(Some(options))
316    }
317
318    /// Helper function to process with options.
319    fn process_with_options(
320        mut self,
321        options: Option<&MaskProcessingOptions>,
322    ) -> OutlineResult<MaskHandle> {
323        let mut ops = std::mem::take(&mut self.operations);
324        match options {
325            Some(custom) => ops.extend(operations_from_options(custom)),
326            None if ops.is_empty() => {
327                ops.extend(operations_from_options(&self.default_mask_processing))
328            }
329            None => {}
330        }
331
332        let mask = apply_operations(self.raw_matte.as_ref(), &ops);
333        Ok(MaskHandle::new(
334            Arc::clone(&self.rgb_image),
335            mask,
336            self.default_mask_processing,
337        ))
338    }
339
340    /// Compose the RGBA foreground image from the RGB image and the raw matte.
341    pub fn foreground(&self) -> OutlineResult<ForegroundHandle> {
342        let rgba = compose_foreground(self.rgb_image.as_ref(), self.raw_matte.as_ref())?;
343        Ok(ForegroundHandle { image: rgba })
344    }
345
346    /// Trace the raw matte using the specified vectorizer and options.
347    pub fn trace<V>(&self, vectorizer: &V, options: &V::Options) -> OutlineResult<V::Output>
348    where
349        V: MaskVectorizer,
350    {
351        vectorizer.vectorize(self.raw_matte.as_ref(), options)
352    }
353}
354
355/// Processed mask image with optional further refinement and output generation.
356///
357/// Represents a concrete mask image (typically binary after thresholding) produced by executing
358/// operations from a [`MatteHandle`].
359///
360/// # Example
361/// ```no_run
362/// use outline::{Outline, VtracerSvgVectorizer, TraceOptions};
363///
364/// let outline = Outline::new("model.onnx");
365/// let session = outline.for_image("input.jpg")?;
366/// let mask = session.matte().blur().threshold().processed()?;
367///
368/// // Generate multiple outputs from the mask
369/// mask.save("mask.png")?;
370/// mask.foreground()?.save("subject.png")?;
371///
372/// let vectorizer = VtracerSvgVectorizer;
373/// let svg = mask.trace(&vectorizer, &TraceOptions::default())?;
374/// std::fs::write("outline.svg", svg)?;
375/// # Ok::<_, outline::OutlineError>(())
376/// ```
377#[derive(Debug, Clone)]
378pub struct MaskHandle {
379    rgb_image: Arc<RgbImage>,
380    mask: GrayImage,
381    default_mask_processing: MaskProcessingOptions,
382    operations: Vec<MaskOperation>,
383}
384
385impl MaskHandle {
386    fn new(
387        rgb_image: Arc<RgbImage>,
388        mask: GrayImage,
389        default_mask_processing: MaskProcessingOptions,
390    ) -> Self {
391        Self {
392            rgb_image,
393            mask,
394            default_mask_processing,
395            operations: Vec::new(),
396        }
397    }
398
399    /// Get the raw  mask.
400    pub fn raw(&self) -> GrayImage {
401        self.mask.clone()
402    }
403
404    /// Get a reference to the mask.
405    pub fn image(&self) -> &GrayImage {
406        &self.mask
407    }
408
409    /// Consume the handle and return the mask.
410    pub fn into_image(self) -> GrayImage {
411        self.mask
412    }
413
414    /// Save the mask to the specified path.
415    pub fn save(&self, path: impl AsRef<Path>) -> OutlineResult<()> {
416        self.mask.save(path)?;
417        Ok(())
418    }
419
420    /// Add a blur operation using the default sigma.
421    pub fn blur(mut self) -> Self {
422        let sigma = self.default_mask_processing.blur_sigma;
423        self.operations.push(MaskOperation::Blur { sigma });
424        self
425    }
426
427    /// Add a blur operation with a custom sigma.
428    pub fn blur_with(mut self, sigma: f32) -> Self {
429        self.operations.push(MaskOperation::Blur { sigma });
430        self
431    }
432
433    /// Add a threshold operation using the default mask threshold.
434    pub fn threshold(mut self) -> Self {
435        let value = self.default_mask_processing.mask_threshold;
436        self.operations.push(MaskOperation::Threshold { value });
437        self
438    }
439
440    /// Add a threshold operation with a custom value.
441    pub fn threshold_with(mut self, value: u8) -> Self {
442        self.operations.push(MaskOperation::Threshold { value });
443        self
444    }
445
446    /// Add a dilation operation using the default radius.
447    ///
448    /// **Note**: Dilation typically works best on binary masks. If this mask is still grayscale,
449    /// consider calling [`threshold`](MaskHandle::threshold) first.
450    pub fn dilate(mut self) -> Self {
451        let radius = self.default_mask_processing.dilation_radius;
452        self.operations.push(MaskOperation::Dilate { radius });
453        self
454    }
455
456    /// Add a dilation operation with a custom radius.
457    ///
458    /// **Note**: Dilation typically works best on binary masks. If this mask is still grayscale,
459    /// consider calling [`threshold`](MaskHandle::threshold) first.
460    pub fn dilate_with(mut self, radius: f32) -> Self {
461        self.operations.push(MaskOperation::Dilate { radius });
462        self
463    }
464
465    /// Add a hole-filling operation to the processing pipeline.
466    ///
467    /// **Note**: Hole-filling typically works best on binary masks. If this mask is still grayscale,
468    /// consider calling [`threshold`](MaskHandle::threshold) first.
469    pub fn fill_holes(mut self) -> Self {
470        let threshold = self.default_mask_processing.mask_threshold;
471        self.operations.push(MaskOperation::FillHoles { threshold });
472        self
473    }
474
475    /// Process the mask with the accumulated operations and default options.
476    pub fn processed(self) -> OutlineResult<MaskHandle> {
477        self.process_with_options(None)
478    }
479
480    /// Process the mask with the accumulated operations and custom options.
481    pub fn processed_with(self, options: &MaskProcessingOptions) -> OutlineResult<MaskHandle> {
482        self.process_with_options(Some(options))
483    }
484
485    /// Helper function to process with options.
486    fn process_with_options(
487        mut self,
488        options: Option<&MaskProcessingOptions>,
489    ) -> OutlineResult<MaskHandle> {
490        let mut ops = std::mem::take(&mut self.operations);
491        match options {
492            Some(custom) => ops.extend(operations_from_options(custom)),
493            None if ops.is_empty() => {
494                ops.extend(operations_from_options(&self.default_mask_processing))
495            }
496            None => {}
497        }
498
499        let mask = apply_operations(&self.mask, &ops);
500        Ok(MaskHandle::new(
501            self.rgb_image,
502            mask,
503            self.default_mask_processing,
504        ))
505    }
506
507    /// Compose the RGBA foreground image from the RGB image and the current mask.
508    pub fn foreground(&self) -> OutlineResult<ForegroundHandle> {
509        let rgba = compose_foreground(self.rgb_image.as_ref(), &self.mask)?;
510        Ok(ForegroundHandle { image: rgba })
511    }
512
513    /// Trace the current mask using the specified vectorizer and options.
514    pub fn trace<V>(&self, vectorizer: &V, options: &V::Options) -> OutlineResult<V::Output>
515    where
516        V: MaskVectorizer,
517    {
518        vectorizer.vectorize(&self.mask, options)
519    }
520}
521
522/// Composed RGBA foreground image with transparent background.
523///
524/// Final output produced by composing the original RGB image with a mask as the alpha channel.
525/// The mask's grayscale values map to alpha, producing smooth or hard edges depending on processing.
526/// Obtain by calling [`foreground`](MatteHandle::foreground) on a [`MatteHandle`] or [`MaskHandle`].
527///
528/// # Example
529/// ```no_run
530/// use outline::Outline;
531///
532/// let outline = Outline::new("model.onnx");
533/// let session = outline.for_image("input.jpg")?;
534///
535/// // Soft edges from raw matte
536/// let soft = session.matte().foreground()?;
537/// soft.save("soft-cutout.png")?;
538///
539/// // Hard edges from processed mask
540/// let hard = session.matte()
541///     .blur()
542///     .threshold()
543///     .processed()?
544///     .foreground()?;
545/// hard.save("hard-cutout.png")?;
546/// # Ok::<_, outline::OutlineError>(())
547/// ```
548pub struct ForegroundHandle {
549    image: RgbaImage,
550}
551
552impl ForegroundHandle {
553    /// Get a reference to the RGBA foreground image.
554    pub fn image(&self) -> &RgbaImage {
555        &self.image
556    }
557
558    /// Consume the handle and return the RGBA foreground image.
559    pub fn into_image(self) -> RgbaImage {
560        self.image
561    }
562
563    /// Save the RGBA foreground image to the specified path.
564    pub fn save(&self, path: impl AsRef<Path>) -> OutlineResult<()> {
565        self.image.save(path)?;
566        Ok(())
567    }
568}
569
570#[cfg(test)]
571mod tests {
572    use super::*;
573    use std::sync::Mutex;
574
575    // Serialize env-var tests so they don't race each other.
576    static ENV_LOCK: Mutex<()> = Mutex::new(());
577
578    mod outline_new {
579        use super::*;
580
581        #[test]
582        fn user_value_is_stored_directly() {
583            let outline = Outline::new("/explicit/model.onnx");
584            assert_eq!(
585                outline.settings.model_path,
586                PathBuf::from("/explicit/model.onnx")
587            );
588        }
589
590        #[test]
591        fn user_value_ignores_env_var() {
592            let _lock = ENV_LOCK.lock().unwrap();
593            // SAFETY: serialized by ENV_LOCK; no other thread reads this var concurrently.
594            unsafe { std::env::set_var(ENV_MODEL_PATH, "env.onnx") };
595            let outline = Outline::new("user.onnx");
596            unsafe { std::env::remove_var(ENV_MODEL_PATH) };
597            assert_eq!(outline.settings.model_path, PathBuf::from("user.onnx"));
598        }
599    }
600
601    mod outline_from_env_or_default {
602        use super::*;
603
604        #[test]
605        fn uses_env_var_when_set() {
606            let _lock = ENV_LOCK.lock().unwrap();
607            // SAFETY: serialized by ENV_LOCK; no other thread reads this var concurrently.
608            unsafe { std::env::set_var(ENV_MODEL_PATH, "/from/env.onnx") };
609            let outline = Outline::from_env_or_default();
610            unsafe { std::env::remove_var(ENV_MODEL_PATH) };
611            assert_eq!(outline.settings.model_path, PathBuf::from("/from/env.onnx"));
612        }
613
614        #[test]
615        fn falls_back_to_default_when_env_unset() {
616            let _lock = ENV_LOCK.lock().unwrap();
617            // SAFETY: serialized by ENV_LOCK; no other thread reads this var concurrently.
618            unsafe { std::env::remove_var(ENV_MODEL_PATH) };
619            let outline = Outline::from_env_or_default();
620            assert_eq!(
621                outline.settings.model_path,
622                PathBuf::from(DEFAULT_MODEL_PATH)
623            );
624        }
625    }
626
627    mod outline_try_from_env {
628        use super::*;
629
630        #[test]
631        fn succeeds_when_env_set() {
632            let _lock = ENV_LOCK.lock().unwrap();
633            // SAFETY: serialized by ENV_LOCK; no other thread reads this var concurrently.
634            unsafe { std::env::set_var(ENV_MODEL_PATH, "/from/env.onnx") };
635            let result = Outline::try_from_env();
636            unsafe { std::env::remove_var(ENV_MODEL_PATH) };
637            let outline = result.expect("should succeed when env is set");
638            assert_eq!(outline.settings.model_path, PathBuf::from("/from/env.onnx"));
639        }
640
641        #[test]
642        fn errors_when_env_unset() {
643            let _lock = ENV_LOCK.lock().unwrap();
644            // SAFETY: serialized by ENV_LOCK; no other thread reads this var concurrently.
645            unsafe { std::env::remove_var(ENV_MODEL_PATH) };
646            let result = Outline::try_from_env();
647            assert!(result.is_err());
648        }
649    }
650}
outline/lib.rs

outline/
lib.rs
