onnxruntime_ng/download/vision/
image_manipulation.rs

1//! Module defining image manipulation models available to download.
2//!
3//! See [https://github.com/onnx/models#image_manipulation](https://github.com/onnx/models#image_manipulation)
4
5use crate::download::{vision::Vision, AvailableOnnxModel, ModelUrl};
6
7/// Image Manipulation
8///
9/// > Image manipulation models use neural networks to transform input images to modified output images. Some
10/// > popular models in this category involve style transfer or enhancing images by increasing resolution.
11///
12/// Source: [https://github.com/onnx/models#image_manipulation](https://github.com/onnx/models#image_manipulation)
13#[derive(Debug, Clone)]
14pub enum ImageManipulation {
15    /// Super Resolution
16    ///
17    /// > The Super Resolution machine learning model sharpens and upscales the input image to refine the
18    /// > details and improve quality.
19    ///
20    /// Source: [https://github.com/onnx/models/tree/master/vision/super_resolution/sub_pixel_cnn_2016](https://github.com/onnx/models/tree/master/vision/super_resolution/sub_pixel_cnn_2016)
21    ///
22    /// Variant downloaded: ONNX Version 1.5 with Opset Version 10.
23    SuperResolution,
24    /// Fast Neural Style Transfer
25    ///
26    /// > This artistic style transfer model mixes the content of an image with the style of another image.
27    /// > Examples of the styles can be seen
28    /// > [in this PyTorch example](https://github.com/pytorch/examples/tree/master/fast_neural_style#models).
29    ///
30    /// Source: [https://github.com/onnx/models/tree/master/vision/style_transfer/fast_neural_style](https://github.com/onnx/models/tree/master/vision/style_transfer/fast_neural_style)
31    FastNeuralStyleTransfer(FastNeuralStyleTransferStyle),
32}
33
34/// Fast Neural Style Transfer Style
35///
36/// Source: [https://github.com/onnx/models/tree/master/vision/style_transfer/fast_neural_style](https://github.com/onnx/models/tree/master/vision/style_transfer/fast_neural_style)
37///
38/// Variant downloaded: ONNX Version 1.4 with Opset Version 9.
39#[derive(Debug, Clone)]
40pub enum FastNeuralStyleTransferStyle {
41    /// Mosaic style
42    Mosaic,
43    /// Candy style
44    Candy,
45    /// RainPrincess style
46    RainPrincess,
47    /// Udnie style
48    Udnie,
49    /// Pointilism style
50    Pointilism,
51}
52
53impl ModelUrl for ImageManipulation {
54    fn fetch_url(&self) -> &'static str {
55        match self {
56            ImageManipulation::SuperResolution => "https://github.com/onnx/models/raw/master/vision/super_resolution/sub_pixel_cnn_2016/model/super-resolution-10.onnx",
57            ImageManipulation::FastNeuralStyleTransfer(style) => style.fetch_url(),
58        }
59    }
60}
61
62impl ModelUrl for FastNeuralStyleTransferStyle {
63    fn fetch_url(&self) -> &'static str {
64        match self {
65            FastNeuralStyleTransferStyle::Mosaic => "https://github.com/onnx/models/raw/master/vision/style_transfer/fast_neural_style/model/mosaic-9.onnx",
66            FastNeuralStyleTransferStyle::Candy => "https://github.com/onnx/models/raw/master/vision/style_transfer/fast_neural_style/model/candy-9.onnx",
67            FastNeuralStyleTransferStyle::RainPrincess => "https://github.com/onnx/models/raw/master/vision/style_transfer/fast_neural_style/model/rain-princess-9.onnx",
68            FastNeuralStyleTransferStyle::Udnie => "https://github.com/onnx/models/raw/master/vision/style_transfer/fast_neural_style/model/udnie-9.onnx",
69            FastNeuralStyleTransferStyle::Pointilism => "https://github.com/onnx/models/raw/master/vision/style_transfer/fast_neural_style/model/pointilism-9.onnx",
70        }
71    }
72}
73
74impl From<ImageManipulation> for AvailableOnnxModel {
75    fn from(model: ImageManipulation) -> Self {
76        AvailableOnnxModel::Vision(Vision::ImageManipulation(model))
77    }
78}
79
80impl From<FastNeuralStyleTransferStyle> for AvailableOnnxModel {
81    fn from(style: FastNeuralStyleTransferStyle) -> Self {
82        AvailableOnnxModel::Vision(Vision::ImageManipulation(
83            ImageManipulation::FastNeuralStyleTransfer(style),
84        ))
85    }
86}
onnxruntime_ng/download/vision/image_manipulation.rs

onnxruntime_ng/download/vision/
image_manipulation.rs
