janitor

janitor
git clone https://git.prime8.dev/janitor.git
Log | Files | Refs | README | LICENSE

processing.rs (2940B)

      1 use crate::audio::{create_fbank, extract_audio, resample, NUM_MEL_BINS, SAMPLE_RATE};
      2 use anyhow::{Context, Result};
      3 use byte_slice_cast::AsByteSlice;
      4 use reqwest::Client;
      5 use safetensors::{serialize, tensor::TensorView, Dtype};
      6 use serde::Deserialize;
      7 use std::sync::OnceLock;
      8 use std::{
      9     collections::HashMap,
     10     path::{Path, PathBuf},
     11 };
     12 use tokio::fs;
     13 use tokio::sync::Semaphore;
     14 use tokio::task::{block_in_place, spawn_blocking};
     15 
     16 #[derive(Debug, Clone, Copy, Deserialize)]
     17 pub enum Label {
     18     Speech,
     19     Music,
     20     Noise,
     21 }
     22 
     23 const MAX_OPEN_FILES: usize = 128;
     24 static PERMITS: Semaphore = Semaphore::const_new(MAX_OPEN_FILES);
     25 
     26 static CLIENT: OnceLock<Client> = OnceLock::new();
     27 
     28 fn get_client() -> &'static Client {
     29     CLIENT.get_or_init(Client::new)
     30 }
     31 
     32 async fn label<'a>(tensor: &'a TensorView<'a>, url: String) -> Result<Label> {
     33     let bytes = block_in_place(|| -> Result<Vec<u8>> {
     34         let tensors = HashMap::from([("fbank", tensor)]);
     35         let bytes = serialize(tensors, &None).context("Failed to serialize tensor")?;
     36         Ok(bytes)
     37     })?;
     38 
     39     let label = get_client()
     40         .post(&url)
     41         .body(bytes)
     42         .send()
     43         .await
     44         .with_context(|| format!("Failed to send request to {url}"))?
     45         .json::<Label>()
     46         .await
     47         .context("Failed to parse response")?;
     48 
     49     Ok(label)
     50 }
     51 
     52 pub async fn process(path: PathBuf, url: String) -> Result<(PathBuf, Label)> {
     53     let _permit = PERMITS.acquire().await?;
     54 
     55     let name = path
     56         .file_name()
     57         .expect("path should have a file name")
     58         .to_string_lossy()
     59         .into_owned();
     60 
     61     let buffer = fs::read(&path)
     62         .await
     63         .with_context(|| format!("Failed to read {name}"))?;
     64 
     65     let fbank = spawn_blocking(move || -> Result<Box<[[f32; NUM_MEL_BINS]]>> {
     66         let mut audio = extract_audio(buffer).context("Failed to extract audio")?;
     67         resample(&mut audio, SAMPLE_RATE);
     68         create_fbank(&mut audio).context("Failed to create filter bank")
     69     })
     70     .await??;
     71 
     72     let tensor = TensorView::new(
     73         Dtype::F32,
     74         vec![fbank.len(), NUM_MEL_BINS],
     75         fbank.as_byte_slice(),
     76     )
     77     .context("Failed to create tensor")?;
     78 
     79     let label = label(&tensor, url)
     80         .await
     81         .with_context(|| format!("Failed to label {name}"))?;
     82 
     83     Ok((path, label))
     84 }
     85 
     86 #[derive(Clone)]
     87 pub struct ResultPathOptions {
     88     pub speech_dir: Option<PathBuf>,
     89     pub music_dir: Option<PathBuf>,
     90     pub noise_dir: Option<PathBuf>,
     91 }
     92 
     93 pub fn get_result_path(path: &Path, label: &Label, options: &ResultPathOptions) -> Option<PathBuf> {
     94     let dir = match label {
     95         Label::Speech => options.speech_dir.as_ref(),
     96         Label::Music => options.music_dir.as_ref(),
     97         Label::Noise => options.noise_dir.as_ref(),
     98     }?
     99     .clone();
    100 
    101     let name = path.file_name()?.to_string_lossy().into_owned();
    102     let mut result = dir;
    103     result.push(name);
    104 
    105     Some(result)
    106 }