Skip to main content

max / audiofiles

11.2 KB · 296 lines History Blame Raw
1 //! The labelled corpus: folder-to-tag mapping, and the scratch vault built from it.
2 //!
3 //! Two modes need the same thing before they can do anything: every file under
4 //! `training/` imported, analysed, and tagged from the folder it came in, so
5 //! `sample_features` holds a vector next to a label. `afcl` exports that state as
6 //! a layer; `layer-eval` cross-validates over it. The step is expensive enough
7 //! (decode plus full analysis per file) that it is worth sharing, and subtle
8 //! enough (a fresh vault every run, a fatal unmapped folder) that two copies
9 //! would drift.
10
11 use std::collections::BTreeMap;
12 use std::path::{Path, PathBuf};
13 use std::time::Instant;
14
15 use audiofiles_core::analysis::{self, config::AnalysisConfig};
16 use audiofiles_core::db::Database;
17 use audiofiles_core::starter_rules::DRUM_CLASSES;
18 use audiofiles_core::store::SampleStore;
19 use audiofiles_core::tags;
20 use rayon::prelude::*;
21
22 /// Corpus subdirectory holding the labelled one-shots, one folder per class.
23 pub(crate) const TRAINING_SUBDIR: &str = "training";
24
25 /// Map a corpus folder name to its canonical tag.
26 ///
27 /// The folder names and [`DRUM_CLASSES`] are the same taxonomy, so this resolves
28 /// against that table rather than carrying a second copy of it. The only mismatch
29 /// is punctuation: the table labels the class `hi-hat` and the corpus folder is
30 /// `hihat`, so both sides are compared with `-` and spaces removed.
31 pub(crate) fn tag_for_folder(folder: &str) -> Option<&'static str> {
32 fn squash(s: &str) -> String {
33 s.chars()
34 .filter(|c| !matches!(c, '-' | ' ' | '_'))
35 .flat_map(char::to_lowercase)
36 .collect()
37 }
38 let want = squash(folder);
39 DRUM_CLASSES
40 .iter()
41 .find(|c| squash(c.label) == want)
42 .map(|c| c.tag)
43 }
44
45 /// The short class label for a canonical tag (`instrument.drum.kick` -> `kick`).
46 ///
47 /// Report tables are unreadable at full tag width once there is a column per
48 /// class, and the confusion matrix is square in the number of classes.
49 pub(crate) fn label_for_tag(tag: &str) -> &str {
50 DRUM_CLASSES
51 .iter()
52 .find(|c| c.tag == tag)
53 .map_or(tag, |c| c.label)
54 }
55
56 /// Collect the labelled files as (path, tag) pairs, grouped for a stable report.
57 pub(crate) fn collect_labelled(
58 training: &Path,
59 ) -> Result<BTreeMap<&'static str, Vec<PathBuf>>, String> {
60 let entries = std::fs::read_dir(training)
61 .map_err(|e| format!("cannot read {}: {e}", training.display()))?;
62
63 let mut by_tag: BTreeMap<&'static str, Vec<PathBuf>> = BTreeMap::new();
64 let mut unmapped = Vec::new();
65 for entry in entries.flatten() {
66 if !entry.path().is_dir() {
67 continue;
68 }
69 let folder = entry.file_name().to_string_lossy().to_string();
70 let Some(tag) = tag_for_folder(&folder) else {
71 unmapped.push(folder);
72 continue;
73 };
74 let mut files: Vec<PathBuf> = std::fs::read_dir(entry.path())
75 .map_err(|e| format!("cannot read {}: {e}", entry.path().display()))?
76 .flatten()
77 .map(|f| f.path())
78 .filter(|p| p.is_file())
79 .collect();
80 // Deterministic order so two runs over the same corpus produce the same
81 // layer, which is what makes the checked-in artifact reviewable.
82 files.sort();
83 by_tag.entry(tag).or_default().extend(files);
84 }
85
86 if !unmapped.is_empty() {
87 // Loud rather than silent: a folder nobody mapped is a class silently
88 // missing from the shipped layer, which reads downstream as the
89 // classifier being bad at that class rather than never having seen it.
90 unmapped.sort();
91 return Err(format!(
92 "no tag mapping for corpus folder(s): {}. Add them to DRUM_CLASSES or move them out of {}",
93 unmapped.join(", "),
94 training.display()
95 ));
96 }
97 if by_tag.is_empty() {
98 return Err(format!(
99 "no labelled class folders under {}",
100 training.display()
101 ));
102 }
103 Ok(by_tag)
104 }
105
106 /// A scratch vault holding the analysed, tagged corpus.
107 ///
108 /// Just the database. It used to carry the per-class file lists and the analysed
109 /// count as well, for the `afcl` generator's manifest; that generator went with
110 /// the bundled layer (2026-08-08, wiki `af-likeness-web`) and both meters read
111 /// the corpus back out of `db` rather than off this struct.
112 pub(crate) struct LabelledVault {
113 pub(crate) db: Database,
114 }
115
116 /// Import, analyse and tag every labelled file into a fresh vault at `vault`.
117 ///
118 /// Prints its own progress: both callers are long-running terminal modes and
119 /// this is the slow part of each.
120 pub(crate) fn build_vault(
121 corpus: &Path,
122 vault: &Path,
123 config: &AnalysisConfig,
124 ) -> Result<LabelledVault, String> {
125 let by_tag = collect_labelled(&corpus.join(TRAINING_SUBDIR))?;
126
127 let total: usize = by_tag.values().map(Vec::len).sum();
128 println!(
129 " {} labelled file(s) across {} class(es):",
130 total,
131 by_tag.len()
132 );
133 for (tag, files) in &by_tag {
134 println!(" {:<28} {:>5}", tag, files.len());
135 }
136 println!();
137
138 // A fresh vault every run. What is built from it is a pure function of the
139 // corpus, and leftovers from a previous run would silently widen it.
140 if vault.exists() {
141 std::fs::remove_dir_all(vault)
142 .map_err(|e| format!("could not clear scratch vault {}: {e}", vault.display()))?;
143 }
144 let samples_dir = vault.join("samples");
145 std::fs::create_dir_all(&samples_dir)
146 .map_err(|e| format!("could not create {}: {e}", samples_dir.display()))?;
147
148 let db = Database::open(vault.join("audiofiles.db"))
149 .map_err(|e| format!("Database::open failed: {e}"))?;
150 let store =
151 SampleStore::new(&samples_dir).map_err(|e| format!("SampleStore::new failed: {e}"))?;
152
153 // Import and tag. No VFS nodes: nothing here runs a UI query, and neither
154 // the export nor the eval needs one; both read `sample_features` joined to
155 // `tags`.
156 let start = Instant::now();
157 let mut to_analyze: Vec<(String, PathBuf)> = Vec::with_capacity(total);
158 let mut import_failures = 0usize;
159 for (tag, files) in &by_tag {
160 for path in files {
161 match store.import(path, &db) {
162 Ok(hash) => {
163 if let Err(e) = tags::add_tag(&db, &hash, tag) {
164 eprintln!(" tag {tag} on {}: {e}", path.display());
165 import_failures += 1;
166 continue;
167 }
168 to_analyze.push((hash, path.clone()));
169 }
170 Err(e) => {
171 eprintln!(" import {}: {e}", path.display());
172 import_failures += 1;
173 }
174 }
175 }
176 }
177 println!(
178 " imported {} file(s) in {:.1}s{}",
179 to_analyze.len(),
180 start.elapsed().as_secs_f64(),
181 if import_failures > 0 {
182 format!(", {import_failures} failed")
183 } else {
184 String::new()
185 }
186 );
187
188 // Analyse everything. The ingest benchmark caps this because analysis is the
189 // expensive stage and it only needs a sample; here every vector is the
190 // payload, so there is no budget to apply.
191 //
192 // Parallel, but still deterministic: `analyze_sample` is a pure function of
193 // the file, and a rayon `collect` into a Vec restores input order, so the
194 // batch written to the DB is the same sequence a serial run would write.
195 // That matters because the exported layer is a checked-in artifact.
196 let start = Instant::now();
197 let results: Vec<_> = to_analyze
198 .par_iter()
199 .filter_map(|(hash, path)| analysis::analyze_sample(hash, path, config).ok())
200 .collect();
201 let analyzed = results.len();
202 analysis::save_analysis_batch(&db, &results)
203 .map_err(|e| format!("save_analysis_batch failed: {e}"))?;
204 println!(
205 " analysed {} file(s) in {:.1}s{}",
206 analyzed,
207 start.elapsed().as_secs_f64(),
208 if analyzed < to_analyze.len() {
209 format!(", {} failed to analyse", to_analyze.len() - analyzed)
210 } else {
211 String::new()
212 }
213 );
214
215 Ok(LabelledVault { db })
216 }
217
218 #[cfg(test)]
219 mod tests {
220 use super::*;
221
222 #[test]
223 fn folder_names_map_to_canonical_tags() {
224 assert_eq!(tag_for_folder("kick"), Some("instrument.drum.kick"));
225 assert_eq!(tag_for_folder("snare"), Some("instrument.drum.snare"));
226 assert_eq!(tag_for_folder("cymbal"), Some("instrument.drum.cymbal"));
227 assert_eq!(tag_for_folder("clap"), Some("instrument.drum.clap"));
228 assert_eq!(tag_for_folder("tom"), Some("instrument.drum.tom"));
229 assert_eq!(tag_for_folder("percussion"), Some("instrument.percussion"));
230 }
231
232 #[test]
233 fn hihat_folder_matches_the_hi_hat_label() {
234 // The one place the corpus and DRUM_CLASSES disagree on spelling.
235 assert_eq!(tag_for_folder("hihat"), Some("instrument.drum.hihat"));
236 assert_eq!(tag_for_folder("hi-hat"), Some("instrument.drum.hihat"));
237 assert_eq!(tag_for_folder("Hi Hat"), Some("instrument.drum.hihat"));
238 }
239
240 #[test]
241 fn unknown_folder_has_no_tag() {
242 assert_eq!(tag_for_folder("bass"), None);
243 assert_eq!(tag_for_folder(""), None);
244 }
245
246 #[test]
247 fn every_drum_class_is_reachable_from_some_folder_name() {
248 // Guards the mapping against a DRUM_CLASSES entry whose label stops
249 // resolving; without this, a renamed label silently drops a class.
250 for class in DRUM_CLASSES {
251 assert_eq!(
252 tag_for_folder(class.label),
253 Some(class.tag),
254 "class {} no longer resolves from its own label",
255 class.label
256 );
257 }
258 }
259
260 #[test]
261 fn labels_round_trip_from_their_tag() {
262 for class in DRUM_CLASSES {
263 assert_eq!(label_for_tag(class.tag), class.label);
264 }
265 // An unknown tag prints as itself rather than vanishing.
266 assert_eq!(label_for_tag("instrument.bass"), "instrument.bass");
267 }
268
269 #[test]
270 fn collect_labelled_rejects_an_unmapped_folder() {
271 let dir = tempfile::tempdir().unwrap();
272 let training = dir.path().join(TRAINING_SUBDIR);
273 std::fs::create_dir_all(training.join("kick")).unwrap();
274 std::fs::create_dir_all(training.join("didgeridoo")).unwrap();
275 let err = collect_labelled(&training).unwrap_err();
276 assert!(err.contains("didgeridoo"), "{err}");
277 }
278
279 #[test]
280 fn collect_labelled_groups_files_under_their_tag() {
281 let dir = tempfile::tempdir().unwrap();
282 let training = dir.path().join(TRAINING_SUBDIR);
283 std::fs::create_dir_all(training.join("kick")).unwrap();
284 std::fs::create_dir_all(training.join("snare")).unwrap();
285 std::fs::write(training.join("kick/b.wav"), b"x").unwrap();
286 std::fs::write(training.join("kick/a.wav"), b"x").unwrap();
287 std::fs::write(training.join("snare/c.wav"), b"x").unwrap();
288
289 let got = collect_labelled(&training).unwrap();
290 assert_eq!(got["instrument.drum.kick"].len(), 2);
291 assert_eq!(got["instrument.drum.snare"].len(), 1);
292 // Sorted, so the artifact is reproducible across runs.
293 assert!(got["instrument.drum.kick"][0].ends_with("a.wav"));
294 }
295 }
296