yolo-set 0.1.0

use std::collections::HashMap;
use std::io;
use std::path::Path;

use crate::model::data_yaml::DataYaml;
use crate::model::merge_mode::MergeMode;
use crate::yolo::dir::{create_work_dir, merge_image_and_label};
use crate::yolo::dataset::{find_dataset, register_dataset_label};
use crate::yolo::label::{key_list_sorted_by_value, update_label_txt};

pub fn yolo_merge(src_folder: &Path, master_dataset: &Path, mode: MergeMode) -> io::Result<()> {
	let (work_folder, _work_folder_cleanup) = create_work_dir(src_folder, mode)?;
	yolo_merge_into(&work_folder, master_dataset)?;
	Ok(())
}

pub fn yolo_merge_into(src_folder: &Path, master_dataset: &Path) -> io::Result<()> {
	let dataset_list = find_dataset(src_folder).expect("输入的源文件夹无法正常打开");
	let data_yaml = DataYaml::read_from(&master_dataset.join("data.yaml")).unwrap_or(DataYaml::new());
	let mut label_map: HashMap<String, usize> = data_yaml.names.clone().into_iter().zip(0..).collect();
	for dataset in &dataset_list {
		println!("正在处理：{:?}", dataset);
		let transformed_label_idx_list = register_dataset_label(dataset, &mut label_map);
		update_label_txt(dataset, &transformed_label_idx_list).expect("更新数据集标签序号失败");
		merge_image_and_label(dataset, master_dataset).expect("合并图片与标签失败");
	}
	let label_list = key_list_sorted_by_value(label_map);
	DataYaml::new()
		.set_names(label_list)
		.write_to(&master_dataset.join("data.yaml"))?;
	Ok(())
}