From 1aff9a300a3a130625e64a4eb7faa266bc0a969b Mon Sep 17 00:00:00 2001 From: York <57304851+wlyh514@users.noreply.github.com> Date: Fri, 3 May 2024 05:24:56 +0800 Subject: [PATCH] Add support for conversion from python bytes (#79) * Allow conversion from python bytes * Update function name * Add convert_pixels_to_svg python function * Update README.md * Update README.md --- cmdapp/src/python.rs | 155 +++++++++++++++++++++++++++++++++++-- cmdapp/vtracer/README.md | 12 ++- cmdapp/vtracer/__init__.py | 3 +- cmdapp/vtracer/vtracer.pyi | 32 ++++++++ 4 files changed, 194 insertions(+), 8 deletions(-) diff --git a/cmdapp/src/python.rs b/cmdapp/src/python.rs index d3cf7a0..cb83b97 100644 --- a/cmdapp/src/python.rs +++ b/cmdapp/src/python.rs @@ -1,5 +1,7 @@ use crate::*; -use pyo3::prelude::*; +use image::{io::Reader, ImageFormat}; +use pyo3::{exceptions::PyException, prelude::*}; +use std::io::{BufReader, Cursor}; use std::path::PathBuf; use visioncortex::PathSimplifyMode; @@ -23,6 +25,148 @@ fn convert_image_to_svg_py( 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, + img_format: Option<&str>, // Format of the image (e.g. 'jpg', 'png'... A full list of supported formats can be found [here](https://docs.rs/image/latest/image/enum.ImageFormat.html)). If not provided, the image format will be guessed based on its contents. + colormode: Option<&str>, // "color" or "binary" + hierarchical: Option<&str>, // "stacked" or "cutout" + mode: Option<&str>, // "polygon", "spline", "none" + filter_speckle: Option, // default: 4 + color_precision: Option, // default: 6 + layer_difference: Option, // default: 16 + corner_threshold: Option, // default: 60 + length_threshold: Option, // in [3.5, 10] default: 4.0 + max_iterations: Option, // default: 10 + splice_threshold: Option, // default: 45 + path_precision: Option, // default: 8 +) -> PyResult { + 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>, // "color" or "binary" + hierarchical: Option<&str>, // "stacked" or "cutout" + mode: Option<&str>, // "polygon", "spline", "none" + filter_speckle: Option, // default: 4 + color_precision: Option, // default: 6 + layer_difference: Option, // default: 16 + corner_threshold: Option, // default: 60 + length_threshold: Option, // in [3.5, 10] default: 4.0 + max_iterations: Option, // default: 10 + splice_threshold: Option, // default: 45 + path_precision: Option, // default: 8 +) -> PyResult { + 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 = 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, + color_precision: Option, + layer_difference: Option, + corner_threshold: Option, + length_threshold: Option, + max_iterations: Option, + splice_threshold: Option, + path_precision: Option, +) -> Config { // TODO: enforce color mode with an enum so that we only // accept the strings 'color' or 'binary' let color_mode = match colormode.unwrap_or("color") { @@ -52,7 +196,7 @@ fn convert_image_to_svg_py( let splice_threshold = splice_threshold.unwrap_or(45); let max_iterations = max_iterations.unwrap_or(10); - let config = Config { + Config { color_mode, hierarchical, filter_speckle, @@ -65,15 +209,14 @@ fn convert_image_to_svg_py( splice_threshold, path_precision, ..Default::default() - }; - - convert_image_to_svg(&input_path, &output_path, config).unwrap(); - Ok(()) + } } /// A Python module implemented in Rust. #[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(()) } diff --git a/cmdapp/vtracer/README.md b/cmdapp/vtracer/README.md index 78873dd..e6beea0 100644 --- a/cmdapp/vtracer/README.md +++ b/cmdapp/vtracer/README.md @@ -52,7 +52,17 @@ vtracer.convert_image_to_svg_py(inp, out) # Single-color example. Good for line art, and much faster than full color: vtracer.convert_image_to_svg_py(inp, out, colormode='binary') -# All the bells & whistles +# Convert from raw image bytes +input_img_bytes: bytes = get_bytes() # e.g. reading bytes from a file or a HTTP request body +svg_str: str = vtracer.convert_raw_image_to_svg(input_img_bytes, img_format = 'jpg') + +# Convert from RGBA image pixels +from PIL import Image +img = Image.open(input_path).convert('RGBA') +pixels: list[tuple[int, int, int, int]] = list(img.getdata()) +svg_str: str = vtracer.convert_pixels_to_svg(pixels) + +# All the bells & whistles, also applicable to convert_raw_image_to_svg and convert_pixels_to_svg. vtracer.convert_image_to_svg_py(inp, out, colormode = 'color', # ["color"] or "binary" diff --git a/cmdapp/vtracer/__init__.py b/cmdapp/vtracer/__init__.py index 3ccb2aa..624f5b6 100644 --- a/cmdapp/vtracer/__init__.py +++ b/cmdapp/vtracer/__init__.py @@ -1 +1,2 @@ -from .vtracer import convert_image_to_svg_py \ No newline at end of file +from .vtracer import (convert_image_to_svg_py, convert_pixels_to_svg, + convert_raw_image_to_svg) diff --git a/cmdapp/vtracer/vtracer.pyi b/cmdapp/vtracer/vtracer.pyi index adc02ac..25fac77 100644 --- a/cmdapp/vtracer/vtracer.pyi +++ b/cmdapp/vtracer/vtracer.pyi @@ -15,3 +15,35 @@ def convert_image_to_svg_py(image_path: str, path_precision: Optional[int] = None, # default: 8 ) -> None: ... + +def convert_raw_image_to_svg(img_bytes: bytes, + img_format: Optional[str] = None, # Format of the image (e.g. 'jpg', 'png'... A full list of supported formats can be found [here](https://docs.rs/image/latest/image/enum.ImageFormat.html)). If not provided, the image format will be guessed based on its contents. + colormode: Optional[str] = None, # ["color"] or "binary" + hierarchical: Optional[str] = None, # ["stacked"] or "cutout" + mode: Optional[str] = None, # ["spline"], "polygon", "none" + filter_speckle: Optional[int] = None, # default: 4 + color_precision: Optional[int] = None, # default: 6 + layer_difference: Optional[int] = None, # default: 16 + corner_threshold: Optional[int] = None, # default: 60 + length_threshold: Optional[float] = None, # in [3.5, 10] default: 4.0 + max_iterations: Optional[int] = None, # default: 10 + splice_threshold: Optional[int] = None, # default: 45 + path_precision: Optional[int] = None, # default: 8 + ) -> str: + ... + +def convert_pixels_to_svg(rgba_pixels: list[tuple[int, int, int, int]], + size: tuple[int, int], + colormode: Optional[str] = None, # ["color"] or "binary" + hierarchical: Optional[str] = None, # ["stacked"] or "cutout" + mode: Optional[str] = None, # ["spline"], "polygon", "none" + filter_speckle: Optional[int] = None, # default: 4 + color_precision: Optional[int] = None, # default: 6 + layer_difference: Optional[int] = None, # default: 16 + corner_threshold: Optional[int] = None, # default: 60 + length_threshold: Optional[float] = None, # in [3.5, 10] default: 4.0 + max_iterations: Optional[int] = None, # default: 10 + splice_threshold: Optional[int] = None, # default: 45 + path_precision: Optional[int] = None, # default: 8 + ) -> str: + ... \ No newline at end of file