extract.rs

use clap::Parser;
use image::{GenericImage, GenericImageView, Pixel, Rgba, RgbaImage};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::ops::{Deref, DerefMut};
use std::path::PathBuf;

/// Extract individual, non-square sprites from a given spritesheet
#[derive(Parser, Serialize, Deserialize, Clone, Debug)]
#[clap(author, version = "0.7.0")]
pub struct Extract {
	/// The path to the spritesheet to extract sprites from
	#[serde(default)]
	pub input: String,
	/// The path to a folder that will contain the extracted sprites
	#[serde(default)]
	pub output: String,
}

struct PixelData {
	x: u32,
	y: u32,
}

#[derive(Debug)]
struct BlobBounds {
	left: u32,
	right: u32,
	top: u32,
	bottom: u32,
}

impl BlobBounds {
	pub fn width(&self) -> u32 {
		(self.right + 1).saturating_sub(self.left)
	}
	pub fn height(&self) -> u32 {
		(self.bottom + 1).saturating_sub(self.top)
	}
}

impl Default for BlobBounds {
	fn default() -> Self {
		Self {
			left: u32::MAX,
			right: u32::MIN,
			bottom: u32::MIN,
			top: u32::MAX,
		}
	}
}

#[derive(Default)]
struct PixelBlob {
	pixels: Vec<PixelData>,
}

impl PixelBlob {
	pub fn find_bounds(&self) -> BlobBounds {
		self.pixels
			.iter()
			.fold(BlobBounds::default(), |bounds, pix| BlobBounds {
				bottom: bounds.bottom.max(pix.y),
				top: bounds.top.min(pix.y),
				left: bounds.left.min(pix.x),
				right: bounds.right.max(pix.x),
			})
	}
}

impl Deref for PixelBlob {
	type Target = Vec<PixelData>;

	fn deref(&self) -> &Self::Target {
		&self.pixels
	}
}

impl DerefMut for PixelBlob {
	fn deref_mut(&mut self) -> &mut Self::Target {
		&mut self.pixels
	}
}

macro_rules! check_point {
	($visited: expr, $to_check: expr, $x: expr, $y: expr) => {
		if !$visited.contains(&($x, $y)) && !$to_check.contains(&($x, $y)) {
			$to_check.push(($x, $y));
		}
	};
}

impl Extract {
	pub fn run(&self, image: &impl GenericImage<Pixel = Rgba<u8>>) -> anyhow::Result<()> {
		let mut visited = HashSet::with_capacity((image.width() * image.height()) as usize);
		let mut blobs = Vec::with_capacity(1);

		let image_width = image.width();
		let image_height = image.height();

		for (x, y, pixel) in image.pixels() {
			if visited.contains(&(x, y)) {
				continue;
			}

			let alpha = pixel.channels()[3];
			if alpha == 0 {
				visited.insert((x, y));
				continue;
			}

			let mut current = PixelBlob::default();
			let mut to_check = Vec::with_capacity(4);
			to_check.push((x, y));

			while let Some((cx, cy)) = to_check.pop() {
				if visited.contains(&(cx, cy)) {
					continue;
				}
				visited.insert((cx, cy));

				let pixel = image.get_pixel(cx, cy);
				let alpha = pixel.channels()[3];
				if alpha == 0 {
					continue;
				}

				current.push(PixelData { x: cx, y: cy });

				let x = [-1, 0, 1];
				let y = [-1, 0, 1];

				for dx in x {
					for dy in y {
						let new_x = cx as i32 + dx;
						let new_y = cy as i32 + dy;

						if new_x >= 0
							&& new_y >= 0 && new_x < image_width as i32
							&& new_y < image_height as i32
						{
							check_point!(visited, to_check, new_x as u32, new_y as u32);
						}
					}
				}
			}

			if !current.is_empty() {
				blobs.push(std::mem::take(&mut current));
			}
		}

		for (idx, blob) in blobs.iter().enumerate() {
			let bounds = blob.find_bounds();
			let width = bounds.width();
			let height = bounds.height();

			if width == 0 || height == 0 {
				continue;
			}

			let sub = image.view(bounds.left, bounds.top, width, height);
			let mut new_image = RgbaImage::from_pixel(width, height, Rgba::from([0, 0, 0, 0]));
			let pixels = sub
				.pixels()
				.flat_map(|(_, _, pixel)| pixel.channels().to_vec())
				.collect::<Vec<u8>>();

			new_image.copy_from_slice(pixels.as_slice());
			new_image
				.save(PathBuf::from(self.output.as_str()).join(format!("image_{}.png", idx)))?;
		}

		Ok(())
	}
}