use anyhow::{Context, Result, bail};
use image::{ImageFormat, ImageReader};
use std::borrow::Cow;
use std::io::Cursor;
use std::path::{Path, PathBuf};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImageMimeType {
Png,
Jpeg,
Webp,
Gif,
Tiff,
Bmp,
}
impl ImageMimeType {
pub fn from_extension(ext: &str) -> Option<Self> {
match ext.to_ascii_lowercase().as_str() {
"png" => Some(Self::Png),
"jpg" | "jpeg" => Some(Self::Jpeg),
"webp" => Some(Self::Webp),
"gif" => Some(Self::Gif),
"tiff" | "tif" => Some(Self::Tiff),
"bmp" => Some(Self::Bmp),
_ => None,
}
}
}
pub struct PreparedImage<'a> {
pub bytes: Cow<'a, [u8]>,
pub mime_type: &'static str,
}
pub fn prepare_image_for_upload(bytes: &[u8], mime: ImageMimeType) -> Result<PreparedImage<'_>> {
match mime {
ImageMimeType::Png => Ok(PreparedImage {
bytes: Cow::Borrowed(bytes),
mime_type: "image/png",
}),
ImageMimeType::Jpeg => Ok(PreparedImage {
bytes: Cow::Borrowed(bytes),
mime_type: "image/jpeg",
}),
ImageMimeType::Webp => Ok(PreparedImage {
bytes: Cow::Borrowed(bytes),
mime_type: "image/webp",
}),
ImageMimeType::Gif => Ok(PreparedImage {
bytes: Cow::Borrowed(bytes),
mime_type: "image/gif",
}),
ImageMimeType::Tiff | ImageMimeType::Bmp => {
let image = ImageReader::new(Cursor::new(bytes))
.with_guessed_format()?
.decode()
.context("Failed to decode image for upload")?;
let mut encoded = Cursor::new(Vec::new());
image
.write_to(&mut encoded, ImageFormat::Png)
.context("Failed to convert image to PNG for upload")?;
Ok(PreparedImage {
bytes: Cow::Owned(encoded.into_inner()),
mime_type: "image/png",
})
}
}
}
pub enum InputSource {
Pdf(PathBuf),
SingleImage {
path: PathBuf,
mime: ImageMimeType,
},
ImageFolder {
images: Vec<(PathBuf, ImageMimeType)>,
},
}
pub fn resolve_input(path: &Path) -> Result<InputSource> {
if !path.exists() {
bail!("Not found: {}", path.display());
}
if path.is_dir() {
let mut images: Vec<(PathBuf, ImageMimeType)> = Vec::new();
let mut entries: Vec<_> = std::fs::read_dir(path)?.filter_map(|e| e.ok()).collect();
entries.sort_by_key(|e| e.file_name());
for entry in entries {
let p = entry.path();
if !p.is_file() {
continue;
}
if let Some(ext) = p.extension().and_then(|e| e.to_str())
&& let Some(mime) = ImageMimeType::from_extension(ext)
{
images.push((p, mime));
}
}
if images.is_empty() {
bail!(
"No supported image files found in directory: {}\nSupported: png, jpg, jpeg, webp, gif, tiff, tif, bmp",
path.display()
);
}
return Ok(InputSource::ImageFolder { images });
}
let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
if ext.eq_ignore_ascii_case("pdf") {
return Ok(InputSource::Pdf(path.to_path_buf()));
}
if let Some(mime) = ImageMimeType::from_extension(ext) {
return Ok(InputSource::SingleImage {
path: path.to_path_buf(),
mime,
});
}
bail!(
"Unsupported file type: .{ext}\nSupported: pdf, png, jpg, jpeg, webp, gif, tiff, tif, bmp"
);
}
#[cfg(test)]
mod tests {
use super::*;
use image::{DynamicImage, ImageFormat, Rgba, RgbaImage};
use std::fs;
use std::io::Cursor;
use tempfile::TempDir;
fn encode_test_image(format: ImageFormat) -> Vec<u8> {
let image = DynamicImage::ImageRgba8(RgbaImage::from_pixel(2, 2, Rgba([255, 0, 0, 255])));
let mut encoded = Cursor::new(Vec::new());
image.write_to(&mut encoded, format).unwrap();
encoded.into_inner()
}
#[test]
fn mime_type_from_known_extensions() {
assert_eq!(
ImageMimeType::from_extension("png"),
Some(ImageMimeType::Png)
);
assert_eq!(
ImageMimeType::from_extension("jpg"),
Some(ImageMimeType::Jpeg)
);
assert_eq!(
ImageMimeType::from_extension("jpeg"),
Some(ImageMimeType::Jpeg)
);
assert_eq!(
ImageMimeType::from_extension("JPEG"),
Some(ImageMimeType::Jpeg)
);
assert_eq!(
ImageMimeType::from_extension("webp"),
Some(ImageMimeType::Webp)
);
assert_eq!(
ImageMimeType::from_extension("gif"),
Some(ImageMimeType::Gif)
);
assert_eq!(
ImageMimeType::from_extension("tiff"),
Some(ImageMimeType::Tiff)
);
assert_eq!(
ImageMimeType::from_extension("tif"),
Some(ImageMimeType::Tiff)
);
assert_eq!(
ImageMimeType::from_extension("bmp"),
Some(ImageMimeType::Bmp)
);
}
#[test]
fn mime_type_unknown_extension() {
assert_eq!(ImageMimeType::from_extension("pdf"), None);
assert_eq!(ImageMimeType::from_extension("txt"), None);
assert_eq!(ImageMimeType::from_extension(""), None);
}
#[test]
fn prepare_png_keeps_original_bytes() {
let bytes = encode_test_image(ImageFormat::Png);
let prepared = prepare_image_for_upload(&bytes, ImageMimeType::Png).unwrap();
assert_eq!(prepared.mime_type, "image/png");
assert_eq!(prepared.bytes.as_ref(), bytes.as_slice());
}
#[test]
fn prepare_tiff_converts_to_png() {
let bytes = encode_test_image(ImageFormat::Tiff);
let prepared = prepare_image_for_upload(&bytes, ImageMimeType::Tiff).unwrap();
assert_eq!(prepared.mime_type, "image/png");
assert!(prepared.bytes.as_ref().starts_with(b"\x89PNG\r\n\x1a\n"));
}
#[test]
fn prepare_bmp_converts_to_png() {
let bytes = encode_test_image(ImageFormat::Bmp);
let prepared = prepare_image_for_upload(&bytes, ImageMimeType::Bmp).unwrap();
assert_eq!(prepared.mime_type, "image/png");
assert!(prepared.bytes.as_ref().starts_with(b"\x89PNG\r\n\x1a\n"));
}
#[test]
fn resolve_pdf_file() {
let dir = TempDir::new().unwrap();
let pdf = dir.path().join("test.pdf");
fs::write(&pdf, b"fake pdf").unwrap();
match resolve_input(&pdf).unwrap() {
InputSource::Pdf(p) => assert_eq!(p, pdf),
_ => panic!("expected Pdf variant"),
}
}
#[test]
fn resolve_single_image() {
let dir = TempDir::new().unwrap();
let img = dir.path().join("screenshot.png");
fs::write(&img, b"fake png").unwrap();
match resolve_input(&img).unwrap() {
InputSource::SingleImage { path, mime } => {
assert_eq!(path, img);
assert_eq!(mime, ImageMimeType::Png);
}
_ => panic!("expected SingleImage variant"),
}
}
#[test]
fn resolve_jpeg_variants() {
let dir = TempDir::new().unwrap();
for ext in ["jpg", "jpeg", "JPG"] {
let img = dir.path().join(format!("photo.{ext}"));
fs::write(&img, b"fake jpeg").unwrap();
match resolve_input(&img).unwrap() {
InputSource::SingleImage { mime, .. } => {
assert_eq!(mime, ImageMimeType::Jpeg);
}
_ => panic!("expected SingleImage for .{ext}"),
}
}
}
#[test]
fn resolve_image_folder() {
let dir = TempDir::new().unwrap();
fs::write(dir.path().join("001.png"), b"img1").unwrap();
fs::write(dir.path().join("002.jpg"), b"img2").unwrap();
fs::write(dir.path().join("003.webp"), b"img3").unwrap();
fs::write(dir.path().join("README.md"), b"ignore me").unwrap();
fs::write(dir.path().join(".DS_Store"), b"ignore me").unwrap();
match resolve_input(dir.path()).unwrap() {
InputSource::ImageFolder { images } => {
assert_eq!(images.len(), 3);
assert!(
images[0]
.0
.file_name()
.unwrap()
.to_str()
.unwrap()
.contains("001")
);
assert!(
images[1]
.0
.file_name()
.unwrap()
.to_str()
.unwrap()
.contains("002")
);
assert!(
images[2]
.0
.file_name()
.unwrap()
.to_str()
.unwrap()
.contains("003")
);
assert_eq!(images[0].1, ImageMimeType::Png);
assert_eq!(images[1].1, ImageMimeType::Jpeg);
assert_eq!(images[2].1, ImageMimeType::Webp);
}
_ => panic!("expected ImageFolder variant"),
}
}
#[test]
fn resolve_empty_folder_errors() {
let dir = TempDir::new().unwrap();
fs::write(dir.path().join("README.md"), b"no images").unwrap();
assert!(resolve_input(dir.path()).is_err());
}
#[test]
fn resolve_unsupported_extension_errors() {
let dir = TempDir::new().unwrap();
let file = dir.path().join("data.csv");
fs::write(&file, b"data").unwrap();
assert!(resolve_input(&file).is_err());
}
#[test]
fn resolve_nonexistent_errors() {
assert!(resolve_input(Path::new("/tmp/nonexistent_lmocr_test_file.png")).is_err());
}
}