Skip to main content

oxigdal_ml/models/
mod.rs

1//! Model management for OxiGDAL ML
2//!
3//! This module provides model loading, management, and configuration.
4
5pub mod onnx;
6pub mod tflite;
7
8#[cfg(feature = "coreml")]
9pub mod coreml;
10
11pub use onnx::{ExecutionProvider, ModelMetadata, OnnxModel, SessionConfig};
12
13#[cfg(not(feature = "tflite"))]
14pub use tflite::{Delegate, QuantizationParams, TensorDataType, TensorInfo, TfLiteConfig};
15
16#[cfg(feature = "tflite")]
17pub use tflite::{
18    Delegate, QuantizationParams, TensorDataType, TensorInfo, TfLiteConfig, TfLiteModel,
19};
20
21use crate::error::Result;
22use oxigdal_core::buffer::RasterBuffer;
23
24/// Trait for ML models
25pub trait Model: Send + Sync {
26    /// Returns the model metadata
27    fn metadata(&self) -> &ModelMetadata;
28
29    /// Predicts on a single raster buffer
30    ///
31    /// # Errors
32    /// Returns an error if prediction fails
33    fn predict(&mut self, input: &RasterBuffer) -> Result<RasterBuffer>;
34
35    /// Predicts on multiple raster buffers (batch prediction)
36    ///
37    /// # Errors
38    /// Returns an error if prediction fails
39    fn predict_batch(&mut self, inputs: &[RasterBuffer]) -> Result<Vec<RasterBuffer>>;
40
41    /// Returns the expected input shape (channels, height, width)
42    fn input_shape(&self) -> (usize, usize, usize);
43
44    /// Returns the expected output shape (channels, height, width)
45    fn output_shape(&self) -> (usize, usize, usize);
46}