dspk_synapse_ml/models/unitrefine/
classifier.rs

1//! SpikeInterface UnitRefine & Bombcell automated cluster quality curation (`unitrefine/curation-v1`).
2
3use std::path::Path;
4
5use dspk_core::{ComputeTarget, DspError, DspResult};
6use dspk_synapse::UnitQualityLabel;
7use serde::{Deserialize, Serialize};
8
9use crate::manifest::TensorIoSpec;
10use crate::runtime::{
11    OnnxRuntimeSession, RuntimeTensor, validate_tensor_port,
12};
13
14pub const UNITREFINE_CURATION_MODEL_ID: &str = "unitrefine/curation-v1";
15pub const UNITREFINE_BOMBCELL_MODEL_ID: &str = "unitrefine/bombcell-metrics-v1";
16
17/// Prediction emitted by [`UnitRefineClassifier`] for a single unit cluster.
18#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
19pub struct UnitCurationResult {
20    pub label: UnitQualityLabel,
21    /// Calibrated softmax probabilities `[p_single_unit, p_multi_unit, p_noise]`.
22    pub probabilities: [f32; 3],
23}
24
25/// Pretrained ONNX unit quality curator (`unitrefine/curation-v1`).
26pub struct UnitRefineClassifier {
27    session: OnnxRuntimeSession,
28    io_spec: Option<TensorIoSpec>,
29}
30
31impl UnitRefineClassifier {
32    /// Pulls and initializes `unitrefine/curation-v1` from [`crate::runtime::pull_model`].
33    #[cfg(feature = "hub")]
34    pub fn from_hub(target: ComputeTarget) -> DspResult<Self> {
35        let (manifest, path) = crate::runtime::pull_model(UNITREFINE_CURATION_MODEL_ID)?;
36        let compute_target = target.checked().map_err(|e| DspError::ComputeError(e.to_string()))?;
37        let session = OnnxRuntimeSession::from_file(&path, compute_target)?;
38        Ok(Self {
39            session,
40            io_spec: manifest.io_spec,
41        })
42    }
43
44    /// Loads a UnitRefine `.onnx` model from disk.
45    pub fn from_onnx_file(
46        path: impl AsRef<Path>,
47        target: ComputeTarget,
48    ) -> DspResult<Self> {
49        let compute_target = target.checked().map_err(|e| DspError::ComputeError(e.to_string()))?;
50        let session = OnnxRuntimeSession::from_file(path, compute_target)?;
51        Ok(Self {
52            session,
53            io_spec: None,
54        })
55    }
56
57    /// Classifies a `[num_units, num_features]` quality metrics matrix into
58    /// `[SingleUnit, MultiUnit, Noise]` labels and softmax probabilities.
59    pub fn classify_metrics(
60        &self,
61        metrics: &[f32],
62        num_units: usize,
63        num_features: usize,
64    ) -> DspResult<Vec<UnitCurationResult>> {
65        let in_shape = vec![num_units, num_features];
66        if let Some(io) = &self.io_spec {
67            if let Some(in_port) = io.inputs.first() {
68                validate_tensor_port(in_port, &in_shape)?;
69            }
70        }
71
72        let input = RuntimeTensor::new(metrics.to_vec(), in_shape)?;
73        let output = self.session.run(input)?;
74
75        if let Some(io) = &self.io_spec {
76            if let Some(out_port) = io.outputs.first() {
77                validate_tensor_port(out_port, &output.shape)?;
78            }
79        }
80
81        let num_classes = *output.shape.last().unwrap_or(&0);
82        if num_classes != 3 {
83            return Err(DspError::InvalidConfig(format!(
84                "UnitRefineClassifier expects 3 output classes [SUA, MUA, Noise], got {num_classes}"
85            )));
86        }
87
88        let mut results = Vec::with_capacity(num_units);
89        for row in output.data.chunks_exact(3) {
90            let max_v = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
91            let e0 = (row[0] - max_v).exp();
92            let e1 = (row[1] - max_v).exp();
93            let e2 = (row[2] - max_v).exp();
94            let sum = (e0 + e1 + e2).max(1e-12);
95            let probs = [e0 / sum, e1 / sum, e2 / sum];
96
97            let label = if probs[0] >= probs[1] && probs[0] >= probs[2] {
98                UnitQualityLabel::SingleUnit
99            } else if probs[1] >= probs[2] {
100                UnitQualityLabel::MultiUnit
101            } else {
102                UnitQualityLabel::Noise
103            };
104
105            results.push(UnitCurationResult {
106                label,
107                probabilities: probs,
108            });
109        }
110        Ok(results)
111    }
112}
dspk_synapse_ml/models/unitrefine/classifier.rs

dspk_synapse_ml/models/unitrefine/
classifier.rs
