From d4589d9600b5bddd3482c9f35414d4f843255d13 Mon Sep 17 00:00:00 2001
From: gennyble <gen@nyble.dev>
Date: Mon, 9 Oct 2023 02:28:37 -0500
Subject: accept jpeg input with zune_jpeg

---
 Cargo.lock         | 34 +++++++++++++++++++++++++++++++++-
 squash/Cargo.toml  |  1 +
 squash/src/main.rs | 52 ++++++++++++++++++++++++++++++++++++++++++++++++----
 3 files changed, 82 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d04e722..4256a12 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -20,6 +20,12 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
+[[package]]
+name = "bitflags"
+version = "2.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b4682ae6287fcf752ecaabbfcc7b6f9b72aa33933dc23a554d853aea8eea8635"
+
 [[package]]
 name = "bytemuck"
 version = "1.14.0"
@@ -73,6 +79,12 @@ dependencies = [
  "miniz_oxide",
 ]
 
+[[package]]
+name = "log"
+version = "0.4.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f"
+
 [[package]]
 name = "miniz_oxide"
 version = "0.7.1"
@@ -88,7 +100,7 @@ name = "png"
 version = "0.17.10"
 source = "git+https://github.com/image-rs/image-png.git?rev=f10238a1e886b228e7da5301e5c0f5011316f2d6#f10238a1e886b228e7da5301e5c0f5011316f2d6"
 dependencies = [
- "bitflags",
+ "bitflags 1.3.2",
  "crc32fast",
  "fdeflate",
  "flate2",
@@ -118,4 +130,24 @@ dependencies = [
  "camino",
  "colorsquash",
  "png",
+ "zune-jpeg",
+]
+
+[[package]]
+name = "zune-core"
+version = "0.2.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "29ca36c2e02af0d8d7ee977542bfe33ed1c516be73d3c1faa4420af46e96ceee"
+dependencies = [
+ "bitflags 2.4.0",
+]
+
+[[package]]
+name = "zune-jpeg"
+version = "0.3.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2848e8f4f29dbdcc79910ab3abdff22bb0bacef8556f2a983b5ca950d8b4991e"
+dependencies = [
+ "log",
+ "zune-core",
 ]
diff --git a/squash/Cargo.toml b/squash/Cargo.toml
index 16c450a..1b92d38 100644
--- a/squash/Cargo.toml
+++ b/squash/Cargo.toml
@@ -13,3 +13,4 @@ camino = "1.1.6"
 # png has a change to ignore extra iCCP blocks my test image needed. it hasn't
 # been released yet, so we're using the git here. the commit we require is e4b4811
 png = { git = "https://github.com/image-rs/image-png.git", rev = "f10238a1e886b228e7da5301e5c0f5011316f2d6" }
+zune-jpeg = "0.3.17"
diff --git a/squash/src/main.rs b/squash/src/main.rs
index b6eb699..da3cfbf 100644
--- a/squash/src/main.rs
+++ b/squash/src/main.rs
@@ -1,9 +1,10 @@
 use std::{fs::File, io::BufWriter};
 
-use anyhow::bail;
+use anyhow::{anyhow, bail};
 use camino::{Utf8Path, Utf8PathBuf};
 use colorsquash::Squasher;
 use png::{ColorType, Decoder, Encoder};
+use zune_jpeg::{zune_core::colorspace::ColorSpace, JpegDecoder};
 
 fn main() -> Result<(), anyhow::Error> {
 	// I should use clap or at least getopt, but this is fine. It's 20LOC.
@@ -13,8 +14,13 @@ fn main() -> Result<(), anyhow::Error> {
 	};
 	let mut argv = std::env::args().skip(1);
 
-	let color_count: u8 = if let Some(Ok(count)) = argv.next().map(|r| r.parse()) {
-		count
+	let color_count: u8 = if let Some(Ok(count)) = argv.next().map(|r| r.parse::<usize>()) {
+		if count > 256 {
+			eprintln!("max colour count must be 256 or below");
+			std::process::exit(1);
+		} else {
+			(count - 1) as u8
+		}
 	} else {
 		usage()
 	};
@@ -31,12 +37,29 @@ fn main() -> Result<(), anyhow::Error> {
 		usage();
 	};
 
-	let mut image = get_png(input_path)?;
+	let mut image = match input_path.extension() {
+		None => {
+			eprintln!("can't determine input filetype!\nSupported input types: PNG, JPG");
+			std::process::exit(1);
+		}
+		Some("png") => get_png(input_path)?,
+		Some("jpg") | Some("jpeg") => get_jpg(input_path)?,
+		Some(ext) => {
+			eprintln!("unknown filetype '{ext}'!\nSupported input types: PNG, JPG");
+			std::process::exit(1);
+		}
+	};
 
 	let squasher = Squasher::new(color_count, &image.data);
 	let size = squasher.map_over(&mut image.data);
 	image.data.resize(size, 0);
 
+	println!(
+		"selected {} colours of max {}",
+		squasher.palette().len(),
+		color_count
+	);
+
 	// PNG Output
 	let file = File::create(output_path)?;
 	let bufw = BufWriter::new(file);
@@ -80,6 +103,27 @@ fn get_png<P: AsRef<Utf8Path>>(path: P) -> Result<Image, anyhow::Error> {
 	}
 }
 
+fn get_jpg<P: AsRef<Utf8Path>>(path: P) -> Result<Image, anyhow::Error> {
+	let content = std::fs::read(path.as_ref())?;
+	let mut dec = JpegDecoder::new(&content);
+	let pixels = dec.decode()?;
+	let info = dec
+		.info()
+		.ok_or(anyhow!("image had no info; this should be impossible"))?;
+
+	let colorspace = dec.get_output_colorspace();
+	match colorspace {
+		Some(ColorSpace::RGB) => (),
+		_ => bail!("colorspace {colorspace:?} not supported"),
+	}
+
+	Ok(Image {
+		width: info.width as usize,
+		height: info.height as usize,
+		data: pixels,
+	})
+}
+
 struct Image {
 	width: usize,
 	height: usize,
-- 
cgit 1.4.1-3-g733a5