atlas.rs

use crate::format::load_image;
use clap::Parser;
use etagere::{AllocId, AtlasAllocator, Rectangle, Size};
use image::{image_dimensions, GenericImage, Rgba, RgbaImage};
use rayon::prelude::{IntoParallelIterator, ParallelIterator};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt::Display;
use std::fs::File;
use std::path::PathBuf;

fn default_max_size() -> usize {
	2048
}

/// Given a set of images, create a single atlas image and metadata file containing all of the original
/// set
#[derive(Parser, Serialize, Deserialize, Clone, Debug)]
#[clap(author, version = "0.8.1")]
pub struct Atlas {
	/// A pattern evaluating to one or more image files
	#[serde(default)]
	pub glob: String,
	/// The path to use when writing the texture atlas
	#[serde(default)]
	pub output: String,
	/// The maximum width of the output texture
	#[serde(default = "default_max_size")]
	#[clap(short = 'x', long = "max_width")]
	pub max_frame_width: usize,
	/// The maximum height of the output texture
	#[serde(default = "default_max_size")]
	#[clap(short = 'y', long = "max_height")]
	pub max_frame_height: usize,
	/// Set a fixed size for each sprite to occupy, creating a fixed grid spritesheet. Sprites must
	/// be padded to the correct size before specifying an area
	#[clap(short = 'a', long = "area")]
	pub area: Option<usize>,
}

#[derive(Copy, Clone, Hash, Eq, PartialEq)]
struct AtlasIdent {
	page: usize,
	id: u32,
}

struct SliceData {
	path: PathBuf,
}

struct AtlasBuilder {
	wh: (usize, usize),
	pages: Vec<AtlasAllocator>,
}

impl AtlasBuilder {
	pub fn new(wh: (usize, usize)) -> Self {
		Self {
			wh,
			pages: Vec::with_capacity(1),
		}
	}

	pub fn insert(
		&mut self,
		width: u32,
		height: u32,
		path: PathBuf,
	) -> Option<(AtlasIdent, SliceData)> {
		for (idx, page) in self.pages.iter_mut().enumerate() {
			if let Some(alloc) = page.allocate(Size::new(width as i32, height as i32)) {
				return Some((
					AtlasIdent {
						page: idx,
						id: alloc.id.serialize(),
					},
					SliceData { path },
				));
			}
		}

		let mut new_page = AtlasAllocator::new(Size::new(self.wh.0 as i32, self.wh.1 as i32));
		if let Some(alloc) = new_page.allocate(Size::new(width as i32, height as i32)) {
			let idx = self.pages.len();
			self.pages.push(new_page);

			Some((
				AtlasIdent {
					page: idx,
					id: alloc.id.serialize(),
				},
				SliceData { path },
			))
		} else {
			None
		}
	}

	pub fn get(&self, ident: AtlasIdent) -> Option<Rectangle> {
		self.pages
			.get(ident.page)
			.map(|page| page.get(AllocId::deserialize(ident.id)))
	}
}

impl Atlas {
	pub fn run(&self) -> anyhow::Result<()> {
		let pattern = glob::glob(self.glob.as_str())?;
		let mut builder = AtlasBuilder::new((self.max_frame_width, self.max_frame_height));

		let page_content_map: HashMap<usize, Vec<(PathBuf, Rectangle)>> = pattern
			.into_iter()
			.filter_map(Result::ok)
			.flat_map(|path| {
				if let Some(fixed) = &self.area {
					Ok((*fixed as u32, *fixed as u32, path))
				} else {
					image_dimensions(&path).map(|(w, h)| (w, h, path))
				}
			})
			.flat_map(|(w, h, path)| builder.insert(w, h, path))
			.collect::<Vec<(AtlasIdent, SliceData)>>()
			.into_iter()
			.flat_map(|(ident, slice)| {
				builder
					.get(ident)
					.map(|rect| (ident.page, slice.path, rect))
			})
			.fold(HashMap::default(), |mut map, (page, path, rect)| {
				let entry = map.entry(page).or_default();
				entry.push((path, rect));
				map
			});

		page_content_map.into_par_iter().for_each(|(page, items)| {
			if let Err(err) = write_atlas(
				self.output.clone(),
				(self.max_frame_width as u32, self.max_frame_height as u32),
				page,
				items,
			) {
				log::error!("{}", err);
			}
		});

		Ok(())
	}
}

#[derive(Serialize, Deserialize)]
struct IndexArea {
	left: usize,
	right: usize,
	top: usize,
	bottom: usize,
}

#[derive(Serialize, Deserialize)]
struct IndexEntry {
	name: String,
	area: IndexArea,
}

fn write_atlas(
	output_path: impl Display,
	image_wh: (u32, u32),
	page_num: usize,
	items: Vec<(PathBuf, Rectangle)>,
) -> anyhow::Result<()> {
	let mut new_image = RgbaImage::from_pixel(image_wh.0, image_wh.1, Rgba::from([0, 0, 0, 0]));
	let mut metadata_index = Vec::with_capacity(items.len());
	for (source, spacing) in items {
		let base_name = source
			.file_stem()
			.ok_or(anyhow::Error::msg("Path had no stem"))?
			.to_str()
			.ok_or(anyhow::Error::msg("Path was not unicode"))?
			.to_string();

		let source_image = load_image(
			source
				.to_str()
				.ok_or(anyhow::Error::msg("Path was wrong"))?,
			None,
		)?;

		metadata_index.push(IndexEntry {
			name: base_name,
			area: IndexArea {
				left: spacing.min.x as usize,
				right: spacing.max.x as usize,
				bottom: spacing.max.y as usize,
				top: spacing.min.y as usize,
			},
		});

		new_image.copy_from(&source_image, spacing.min.x as u32, spacing.min.y as u32)?;
	}

	new_image.save(format!("{}_{}.png", output_path, page_num))?;
	let new_file = File::create(format!("{}_{}.json", output_path, page_num))?;
	serde_json::to_writer_pretty(new_file, &metadata_index)?;

	Ok(())
}