janitor

model.rs (1912B)

---

 use anyhow::Result;
 use serde::Serialize;
 use std::path::Path;
 use tch::{autocast, no_grad, CModule, Device, Tensor};
 
 #[derive(Debug, Clone, Copy, Serialize)]
 pub enum Label {
     Speech,
     Music,
     Noise,
 }
 
 pub struct Model {
     model: CModule,
 }
 
 impl Model {
     pub fn new<T>(model_path: T) -> Result<Model>
     where
         T: AsRef<Path>,
     {
         let device = Device::cuda_if_available();
         println!("Running model on {:?}", device);
         Ok(Self {
             model: CModule::load_on_device(model_path, device)?,
         })
     }
 
     pub fn label(&self, tensor: &Tensor) -> Result<Box<[Label]>> {
         let raw_output = Vec::try_from(
             no_grad(|| autocast(true, || tensor.apply(&self.model)))
                 .f_sigmoid()?
                 .flatten(0, -1),
         )?;
 
         let labels = raw_output
             .chunks(527)
             .map(|chunk| {
                 if chunk.len() != 527 {
                     return Ok(Label::Noise);
                 }
                 let output: [f32; 527] = chunk
                     .try_into()
                     .map_err(|_| anyhow::anyhow!("Failed to convert chunk to fixed-size array"))?;
 
                 let results = [
                     (Label::Speech, output[0]),
                     (Label::Music, output[137]),
                     (Label::Noise, output[513]),
                 ];
 
                 let label = if results[0].1 < 0.5 && results[1].1 < 0.5 {
                     Label::Noise
                 } else {
                     results
                         .iter()
                         .max_by(|a, b| a.1.total_cmp(&b.1))
                         .map(|(label, _)| *label)
                         .unwrap_or(Label::Noise)
                 };
 
                 Ok::<_, anyhow::Error>(label)
             })
             .collect::<Result<Vec<_>>>()?
             .into_boxed_slice();
 
         Ok(labels)
     }
 }