use crate::*;
use image::{io::Reader, ImageFormat};
use pyo3::{exceptions::PyException, prelude::*};
use std::io::{BufReader, Cursor};
use std::path::PathBuf;
use visioncortex::PathSimplifyMode;
#[pyfunction]
fn convert_image_to_svg_py(
image_path: &str,
out_path: &str,
colormode: Option<&str>, hierarchical: Option<&str>, mode: Option<&str>, filter_speckle: Option<usize>, color_precision: Option<i32>, layer_difference: Option<i32>, corner_threshold: Option<i32>, length_threshold: Option<f64>, max_iterations: Option<usize>, splice_threshold: Option<i32>, path_precision: Option<u32>, ) -> PyResult<()> {
let input_path = PathBuf::from(image_path);
let output_path = PathBuf::from(out_path);
let config = construct_config(
colormode,
hierarchical,
mode,
filter_speckle,
color_precision,
layer_difference,
corner_threshold,
length_threshold,
max_iterations,
splice_threshold,
path_precision,
);
convert_image_to_svg(&input_path, &output_path, config).unwrap();
Ok(())
}
#[pyfunction]
fn convert_raw_image_to_svg(
img_bytes: Vec<u8>,
img_format: Option<&str>, colormode: Option<&str>, hierarchical: Option<&str>, mode: Option<&str>, filter_speckle: Option<usize>, color_precision: Option<i32>, layer_difference: Option<i32>, corner_threshold: Option<i32>, length_threshold: Option<f64>, max_iterations: Option<usize>, splice_threshold: Option<i32>, path_precision: Option<u32>, ) -> PyResult<String> {
let config = construct_config(
colormode,
hierarchical,
mode,
filter_speckle,
color_precision,
layer_difference,
corner_threshold,
length_threshold,
max_iterations,
splice_threshold,
path_precision,
);
let mut img_reader = Reader::new(BufReader::new(Cursor::new(img_bytes)));
let img_format = img_format.and_then(|ext_name| ImageFormat::from_extension(ext_name));
let img = match img_format {
Some(img_format) => {
img_reader.set_format(img_format);
img_reader.decode()
}
None => img_reader
.with_guessed_format()
.map_err(|_| PyException::new_err("Unrecognized image format. "))?
.decode(),
};
let img = match img {
Ok(img) => img.to_rgba8(),
Err(_) => return Err(PyException::new_err("Failed to decode img_bytes. ")),
};
let (width, height) = (img.width() as usize, img.height() as usize);
let img = ColorImage {
pixels: img.as_raw().to_vec(),
width,
height,
};
let svg =
convert(img, config).map_err(|_| PyException::new_err("Failed to convert the image. "))?;
Ok(format!("{}", svg))
}
#[pyfunction]
fn convert_pixels_to_svg(
rgba_pixels: Vec<(u8, u8, u8, u8)>,
size: (usize, usize),
colormode: Option<&str>, hierarchical: Option<&str>, mode: Option<&str>, filter_speckle: Option<usize>, color_precision: Option<i32>, layer_difference: Option<i32>, corner_threshold: Option<i32>, length_threshold: Option<f64>, max_iterations: Option<usize>, splice_threshold: Option<i32>, path_precision: Option<u32>, ) -> PyResult<String> {
let expected_pixel_count = size.0 * size.1;
if rgba_pixels.len() != expected_pixel_count {
return Err(PyException::new_err(format!(
"Length of rgba_pixels does not match given image size. Expected {} ({} * {}), got {}. ",
expected_pixel_count,
size.0,
size.1,
rgba_pixels.len()
)));
}
let config = construct_config(
colormode,
hierarchical,
mode,
filter_speckle,
color_precision,
layer_difference,
corner_threshold,
length_threshold,
max_iterations,
splice_threshold,
path_precision,
);
let mut flat_pixels: Vec<u8> = vec![];
for (r, g, b, a) in rgba_pixels {
flat_pixels.push(r);
flat_pixels.push(g);
flat_pixels.push(b);
flat_pixels.push(a);
}
let mut img = ColorImage::new();
img.pixels = flat_pixels;
(img.width, img.height) = size;
let svg =
convert(img, config).map_err(|_| PyException::new_err("Failed to convert the image. "))?;
Ok(format!("{}", svg))
}
fn construct_config(
colormode: Option<&str>,
hierarchical: Option<&str>,
mode: Option<&str>,
filter_speckle: Option<usize>,
color_precision: Option<i32>,
layer_difference: Option<i32>,
corner_threshold: Option<i32>,
length_threshold: Option<f64>,
max_iterations: Option<usize>,
splice_threshold: Option<i32>,
path_precision: Option<u32>,
) -> Config {
let color_mode = match colormode.unwrap_or("color") {
"color" => ColorMode::Color,
"binary" => ColorMode::Binary,
_ => ColorMode::Color,
};
let hierarchical = match hierarchical.unwrap_or("stacked") {
"stacked" => Hierarchical::Stacked,
"cutout" => Hierarchical::Cutout,
_ => Hierarchical::Stacked,
};
let mode = match mode.unwrap_or("spline") {
"spline" => PathSimplifyMode::Spline,
"polygon" => PathSimplifyMode::Polygon,
"none" => PathSimplifyMode::None,
_ => PathSimplifyMode::Spline,
};
let filter_speckle = filter_speckle.unwrap_or(4);
let color_precision = color_precision.unwrap_or(6);
let layer_difference = layer_difference.unwrap_or(16);
let corner_threshold = corner_threshold.unwrap_or(60);
let length_threshold = length_threshold.unwrap_or(4.0);
let splice_threshold = splice_threshold.unwrap_or(45);
let max_iterations = max_iterations.unwrap_or(10);
Config {
color_mode,
hierarchical,
filter_speckle,
color_precision,
layer_difference,
mode,
corner_threshold,
length_threshold,
max_iterations,
splice_threshold,
path_precision,
..Default::default()
}
}
#[pymodule]
fn vtracer(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(convert_image_to_svg_py, m)?)?;
m.add_function(wrap_pyfunction!(convert_raw_image_to_svg, m)?)?;
m.add_function(wrap_pyfunction!(convert_pixels_to_svg, m)?)?;
Ok(())
}