use super::super::CudaError;
#[cfg(feature = "_cuda-libraries")]
pub(super) use crate::inference::native_model::segmentation::LSTM_LAYERS;
pub(super) use crate::inference::native_model::segmentation::{
CLASSES, CONV_KERNEL, FEATURES, HIDDEN, LEAKY_SLOPE, LINEAR, NORM_EPSILON, POOL, SINC_CHANNELS,
SINC_KERNEL, SINC_STRIDE, WINDOW_SAMPLES,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SegmentationShape {
pub batch: usize,
pub samples: usize,
pub sinc: usize,
pub pool0: usize,
pub conv1: usize,
pub pool1: usize,
pub conv2: usize,
pub frames: usize,
}
impl SegmentationShape {
pub fn new(batch: usize, samples: usize) -> Result<Self, CudaError> {
let too_short = CudaError::BufferLength {
context: "segmentation window samples",
expected: WINDOW_SAMPLES,
actual: samples,
};
if batch == 0 {
return Err(CudaError::BufferLength {
context: "segmentation batch",
expected: 1,
actual: 0,
});
}
let valid = |len: usize, kernel: usize| len.checked_sub(kernel).map(|rest| rest + 1);
let shape = (|| {
let sinc = samples.checked_sub(SINC_KERNEL)? / SINC_STRIDE + 1;
let pool0 = sinc / POOL;
let conv1 = valid(pool0, CONV_KERNEL)?;
let pool1 = conv1 / POOL;
let conv2 = valid(pool1, CONV_KERNEL)?;
let frames = conv2 / POOL;
(frames > 0).then_some(Self {
batch,
samples,
sinc,
pool0,
conv1,
pool1,
conv2,
frames,
})
})();
shape.ok_or(too_short)
}
}