use crate::error::{Error, Result};
use crate::labels::parse_labels;
use crate::types::{LabelFormat, LocationScore, Prediction};
use ndarray::Array2;
use ort::session::Session;
use ort::value::Value;
use std::sync::{Arc, Mutex};
#[must_use]
#[allow(clippy::cast_precision_loss)]
pub const fn calculate_week(month: u32, day: u32) -> f32 {
let weeks_from_months = (month - 1) * 4;
let week_in_month = (day - 1) / 7 + 1;
(weeks_from_months + week_in_month) as f32
}
pub fn validate_coordinates(latitude: f32, longitude: f32) -> Result<()> {
if !(-90.0..=90.0).contains(&latitude) {
return Err(Error::InvalidCoordinates {
latitude,
longitude,
reason: format!("latitude must be in range [-90, 90], got {latitude}"),
});
}
if !(-180.0..=180.0).contains(&longitude) {
return Err(Error::InvalidCoordinates {
latitude,
longitude,
reason: format!("longitude must be in range [-180, 180], got {longitude}"),
});
}
Ok(())
}
pub fn validate_date(month: u32, day: u32) -> Result<()> {
if !(1..=12).contains(&month) {
return Err(Error::InvalidDate {
month,
day,
reason: format!("month must be in range [1, 12], got {month}"),
});
}
if !(1..=31).contains(&day) {
return Err(Error::InvalidDate {
month,
day,
reason: format!("day must be in range [1, 31], got {day}"),
});
}
Ok(())
}
#[derive(Debug)]
enum Labels {
Path(String),
InMemory(Vec<String>),
}
#[derive(Debug)]
pub struct RangeFilterBuilder {
model_path: Option<String>,
labels: Option<Labels>,
execution_providers: Vec<ort::execution_providers::ExecutionProviderDispatch>,
threshold: f32,
}
impl Default for RangeFilterBuilder {
fn default() -> Self {
Self::new()
}
}
impl RangeFilterBuilder {
#[must_use]
pub const fn new() -> Self {
Self {
model_path: None,
labels: None,
execution_providers: Vec::new(),
threshold: 0.01,
}
}
#[must_use]
pub fn model_path(mut self, path: impl Into<String>) -> Self {
self.model_path = Some(path.into());
self
}
#[must_use]
pub fn labels_path(mut self, path: impl Into<String>) -> Self {
self.labels = Some(Labels::Path(path.into()));
self
}
#[must_use]
pub fn labels(mut self, labels: Vec<String>) -> Self {
self.labels = Some(Labels::InMemory(labels));
self
}
#[must_use]
pub fn from_classifier_labels(mut self, labels: &[String]) -> Self {
self.labels = Some(Labels::InMemory(labels.to_vec()));
self
}
#[must_use]
pub fn execution_provider(
mut self,
provider: impl Into<ort::execution_providers::ExecutionProviderDispatch>,
) -> Self {
self.execution_providers.push(provider.into());
self
}
#[must_use]
pub const fn threshold(mut self, threshold: f32) -> Self {
self.threshold = threshold;
self
}
pub fn build(self) -> Result<RangeFilter> {
let model_path = self.model_path.ok_or(Error::ModelPathRequired)?;
let labels_source = self.labels.ok_or(Error::LabelsRequired)?;
let labels = match labels_source {
Labels::Path(path) => {
let content = std::fs::read_to_string(&path).map_err(|e| Error::LabelLoad {
path: path.clone(),
reason: e.to_string(),
})?;
parse_labels(&content, LabelFormat::Text)?
}
Labels::InMemory(labels) => labels,
};
let mut session_builder = Session::builder().map_err(Error::ModelLoad)?;
if !self.execution_providers.is_empty() {
session_builder = session_builder
.with_execution_providers(self.execution_providers)
.map_err(Error::ModelLoad)?;
}
let session = session_builder
.commit_from_file(&model_path)
.map_err(Error::ModelLoad)?;
let output_shapes = extract_output_shapes(&session)?;
if output_shapes.len() != 1 {
return Err(Error::ModelDetection {
reason: format!("meta model expects 1 output, got {}", output_shapes.len()),
});
}
let expected = extract_last_dim(&output_shapes[0])?;
if labels.len() != expected {
return Err(Error::LabelCount {
expected,
got: labels.len(),
});
}
Ok(RangeFilter {
inner: Arc::new(RangeFilterInner {
session: Mutex::new(session),
labels,
threshold: self.threshold,
}),
})
}
}
fn extract_output_shapes(session: &Session) -> Result<Vec<Vec<i64>>> {
session
.outputs
.iter()
.map(|output| {
let shape = output
.output_type
.tensor_shape()
.ok_or_else(|| Error::ModelDetection {
reason: "output is not a tensor".to_string(),
})?;
Ok(shape.iter().copied().collect())
})
.collect()
}
fn extract_last_dim(shape: &[i64]) -> Result<usize> {
let value = shape.last().copied().ok_or_else(|| Error::ModelDetection {
reason: "empty output shape".to_string(),
})?;
usize::try_from(value).map_err(|_| Error::ModelDetection {
reason: format!("invalid dimension: {value}"),
})
}
fn filter_batch_predictions_impl(
predictions_batch: Vec<Vec<crate::types::Prediction>>,
location_scores: &[LocationScore],
threshold: f32,
rerank: bool,
) -> Vec<Vec<crate::types::Prediction>> {
predictions_batch
.into_iter()
.map(|preds| filter_predictions_impl(&preds, location_scores, threshold, rerank))
.collect()
}
fn filter_predictions_impl(
predictions: &[Prediction],
location_scores: &[LocationScore],
threshold: f32,
rerank: bool,
) -> Vec<Prediction> {
let location_map: std::collections::HashMap<&str, f32> = location_scores
.iter()
.map(|score| (score.species.as_str(), score.score))
.collect();
let mut filtered: Vec<Prediction> = predictions
.iter()
.filter_map(|pred| {
let location_score = location_map.get(pred.species.as_str()).copied();
match location_score {
Some(score) if score >= threshold => {
let confidence = if rerank {
pred.confidence * score
} else {
pred.confidence
};
Some(Prediction {
species: pred.species.clone(),
confidence,
index: pred.index,
})
}
Some(_) => {
None
}
None => {
Some(Prediction {
species: pred.species.clone(),
confidence: pred.confidence,
index: pred.index,
})
}
}
})
.collect();
if rerank {
filtered.sort_unstable_by(|a, b| b.confidence.total_cmp(&a.confidence));
}
filtered
}
struct RangeFilterInner {
session: Mutex<Session>,
labels: Vec<String>,
threshold: f32,
}
#[derive(Clone)]
pub struct RangeFilter {
inner: Arc<RangeFilterInner>,
}
impl std::fmt::Debug for RangeFilter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RangeFilter")
.field("labels_count", &self.inner.labels.len())
.field("threshold", &self.inner.threshold)
.finish_non_exhaustive()
}
}
impl RangeFilter {
#[must_use]
pub const fn builder() -> RangeFilterBuilder {
RangeFilterBuilder::new()
}
#[allow(clippy::significant_drop_tightening)]
pub fn predict(
&self,
latitude: f32,
longitude: f32,
month: u32,
day: u32,
) -> Result<Vec<LocationScore>> {
validate_coordinates(latitude, longitude)?;
validate_date(month, day)?;
let week = calculate_week(month, day);
let input_data = vec![latitude, longitude, week];
let input_array = Array2::from_shape_vec((1, 3), input_data).map_err(|e| {
Error::RangeFilterInference(format!("failed to create input array: {e}"))
})?;
let input_value = Value::from_array(input_array).map_err(|e| {
Error::RangeFilterInference(format!("failed to create input tensor: {e}"))
})?;
let mut session = self
.inner
.session
.lock()
.map_err(|e| Error::RangeFilterInference(format!("session lock poisoned: {e}")))?;
let outputs = session
.run(ort::inputs![input_value])
.map_err(|e| Error::RangeFilterInference(e.to_string()))?;
let tensor = outputs.values().next().ok_or_else(|| {
Error::RangeFilterInference("model returned no output tensors".to_string())
})?;
let (_, data) = tensor
.try_extract_tensor::<f32>()
.map_err(|e| Error::RangeFilterInference(e.to_string()))?;
let mut scores: Vec<LocationScore> = data
.iter()
.enumerate()
.filter_map(|(i, &score)| {
if score >= self.inner.threshold && i < self.inner.labels.len() {
Some(LocationScore {
species: self.inner.labels[i].clone(),
score,
index: i,
})
} else {
None
}
})
.collect();
scores.sort_unstable_by(|a, b| b.score.total_cmp(&a.score));
Ok(scores)
}
#[must_use]
pub fn filter_predictions(
&self,
predictions: &[Prediction],
location_scores: &[LocationScore],
rerank: bool,
) -> Vec<Prediction> {
filter_predictions_impl(predictions, location_scores, self.inner.threshold, rerank)
}
#[must_use]
pub fn filter_batch_predictions(
&self,
predictions_batch: Vec<Vec<crate::types::Prediction>>,
location_scores: &[LocationScore],
rerank: bool,
) -> Vec<Vec<crate::types::Prediction>> {
filter_batch_predictions_impl(
predictions_batch,
location_scores,
self.inner.threshold,
rerank,
)
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
#![allow(clippy::float_cmp)]
use super::*;
#[test]
fn test_calculate_week_january_first() {
let week = calculate_week(1, 1);
assert_eq!(week, 1.0);
}
#[test]
fn test_calculate_week_january_eighth() {
let week = calculate_week(1, 8);
assert_eq!(week, 2.0);
}
#[test]
fn test_calculate_week_february_first() {
let week = calculate_week(2, 1);
assert_eq!(week, 5.0);
}
#[test]
fn test_calculate_week_december_last() {
let week = calculate_week(12, 31);
assert_eq!(week, 49.0);
}
#[test]
fn test_validate_coordinates_valid() {
assert!(validate_coordinates(45.0, -122.0).is_ok());
assert!(validate_coordinates(0.0, 0.0).is_ok());
assert!(validate_coordinates(-90.0, -180.0).is_ok());
assert!(validate_coordinates(90.0, 180.0).is_ok());
}
#[test]
fn test_validate_coordinates_invalid_latitude() {
let result = validate_coordinates(91.0, 0.0);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
Error::InvalidCoordinates { .. }
));
}
#[test]
fn test_validate_coordinates_invalid_longitude() {
let result = validate_coordinates(0.0, 181.0);
assert!(result.is_err());
}
#[test]
fn test_validate_date_valid() {
assert!(validate_date(1, 1).is_ok());
assert!(validate_date(6, 15).is_ok());
assert!(validate_date(12, 31).is_ok());
}
#[test]
fn test_validate_date_invalid_month_zero() {
let result = validate_date(0, 1);
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), Error::InvalidDate { .. }));
}
#[test]
fn test_validate_date_invalid_month_thirteen() {
let result = validate_date(13, 1);
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), Error::InvalidDate { .. }));
}
#[test]
fn test_validate_date_invalid_day_zero() {
let result = validate_date(1, 0);
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), Error::InvalidDate { .. }));
}
#[test]
fn test_validate_date_invalid_day_thirty_two() {
let result = validate_date(1, 32);
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), Error::InvalidDate { .. }));
}
#[test]
fn test_range_filter_builder_missing_model_path() {
let result = RangeFilter::builder().build();
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), Error::ModelPathRequired));
}
#[test]
fn test_range_filter_builder_missing_labels() {
let result = RangeFilter::builder().model_path("/tmp/model.onnx").build();
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), Error::LabelsRequired));
}
#[test]
fn test_filter_predictions_above_threshold() {
let predictions = vec![
Prediction {
species: "Species A".to_string(),
confidence: 0.8,
index: 0,
},
Prediction {
species: "Species B".to_string(),
confidence: 0.3,
index: 1,
},
Prediction {
species: "Species C".to_string(),
confidence: 0.05,
index: 2,
},
];
let location_scores = vec![
LocationScore {
species: "Species A".to_string(),
score: 0.9,
index: 0,
},
LocationScore {
species: "Species B".to_string(),
score: 0.02,
index: 1,
},
LocationScore {
species: "Species C".to_string(),
score: 0.5,
index: 2,
},
];
let threshold = 0.03;
let rerank = false;
let filtered = filter_predictions_impl(&predictions, &location_scores, threshold, rerank);
assert_eq!(filtered.len(), 2);
assert_eq!(filtered[0].species, "Species A");
assert_eq!(filtered[1].species, "Species C");
}
#[test]
fn test_filter_predictions_with_rerank() {
let predictions = vec![
Prediction {
species: "Species A".to_string(),
confidence: 0.9, index: 0,
},
Prediction {
species: "Species B".to_string(),
confidence: 0.8, index: 1,
},
Prediction {
species: "Species C".to_string(),
confidence: 0.7, index: 2,
},
];
let location_scores = vec![
LocationScore {
species: "Species A".to_string(),
score: 0.5, index: 0,
},
LocationScore {
species: "Species B".to_string(),
score: 0.9, index: 1,
},
LocationScore {
species: "Species C".to_string(),
score: 0.6, index: 2,
},
];
let threshold = 0.03;
let rerank = true;
let filtered = filter_predictions_impl(&predictions, &location_scores, threshold, rerank);
assert_eq!(filtered.len(), 3);
assert_eq!(filtered[0].species, "Species B");
assert_eq!(filtered[1].species, "Species A");
assert_eq!(filtered[2].species, "Species C");
assert!((filtered[0].confidence - 0.72).abs() < 0.001);
assert!((filtered[1].confidence - 0.45).abs() < 0.001);
assert!((filtered[2].confidence - 0.42).abs() < 0.001);
}
#[test]
fn test_filter_predictions_species_not_in_meta_model() {
let predictions = vec![
Prediction {
species: "Species A".to_string(),
confidence: 0.8,
index: 0,
},
Prediction {
species: "Species B".to_string(),
confidence: 0.7,
index: 1,
},
Prediction {
species: "Species D".to_string(), confidence: 0.9,
index: 3,
},
];
let location_scores = vec![
LocationScore {
species: "Species A".to_string(),
score: 0.9,
index: 0,
},
LocationScore {
species: "Species C".to_string(), score: 0.8,
index: 2,
},
];
let threshold = 0.03;
let rerank = false;
let filtered = filter_predictions_impl(&predictions, &location_scores, threshold, rerank);
assert_eq!(filtered.len(), 3);
assert_eq!(filtered[0].species, "Species A");
assert_eq!(filtered[0].confidence, 0.8);
assert_eq!(filtered[1].species, "Species B");
assert_eq!(filtered[1].confidence, 0.7);
assert_eq!(filtered[2].species, "Species D");
assert_eq!(filtered[2].confidence, 0.9);
}
#[test]
fn test_builder_from_classifier_labels() {
let labels = vec!["Species A".to_string(), "Species B".to_string()];
let builder = RangeFilterBuilder::new().from_classifier_labels(&labels);
assert!(matches!(builder.labels, Some(Labels::InMemory(_))));
}
#[test]
fn test_filter_batch_predictions() {
use crate::types::Prediction;
let batch1 = vec![Prediction {
species: "Species A".to_string(),
confidence: 0.8,
index: 0,
}];
let batch2 = vec![Prediction {
species: "Species B".to_string(),
confidence: 0.6,
index: 1,
}];
let predictions_batch = vec![batch1, batch2];
let location_scores = vec![
LocationScore {
species: "Species A".to_string(),
score: 0.9,
index: 0,
},
LocationScore {
species: "Species B".to_string(),
score: 0.05,
index: 1,
},
];
let threshold = 0.1;
let results =
filter_batch_predictions_impl(predictions_batch, &location_scores, threshold, false);
assert_eq!(results.len(), 2);
assert_eq!(results[0].len(), 1); assert_eq!(results[1].len(), 0); }
}