use log::{debug, warn};
use std::path::{Path, PathBuf};
use crate::tensorflowlite_c;
pub const DEFAULT_TFLITEC_PATH: &str = if cfg!(windows) {
"tensorflowlite_c.dll"
} else if cfg!(target_os = "macos") {
"libtensorflowlite_c.dylib"
} else {
"libtensorflowlite_c.so"
};
pub const DEFAULT_TFLITECPP_PATH: &str = if cfg!(windows) {
"tensorflow-lite.dll"
} else if cfg!(target_os = "macos") {
"libtensorflow-lite.dylib"
} else {
"libtensorflow-lite.so"
};
pub fn discover() -> Result<tensorflowlite_c, libloading::Error> {
discover_with_path().map(|(lib, _)| lib)
}
pub fn discover_with_path() -> Result<(tensorflowlite_c, PathBuf), libloading::Error> {
if let Ok(path) = std::env::var("TFLITE_LIBRARY_PATH") {
debug!("TFLITE_LIBRARY_PATH={path}");
let lib = load(&path)?;
return Ok((lib, resolve_path(&path)));
}
if let Some((lib, path)) = try_vendored_with_path() {
return Ok((lib, resolve_path(&path)));
}
if cfg!(target_os = "linux") {
for version in (1..50).rev() {
for patch in (0..10).rev() {
let path = format!("{DEFAULT_TFLITECPP_PATH}.2.{version}.{patch}");
if let Ok(lib) = load(&path) {
debug!("Found TFLite library: {path}");
return Ok((lib, resolve_path(&path)));
}
}
}
}
if let Ok(lib) = load(DEFAULT_TFLITEC_PATH) {
debug!("Found TFLite library: {DEFAULT_TFLITEC_PATH}");
return Ok((lib, resolve_path(DEFAULT_TFLITEC_PATH)));
}
debug!("Trying fallback: {DEFAULT_TFLITECPP_PATH}");
let lib = load(DEFAULT_TFLITECPP_PATH)?;
Ok((lib, resolve_path(DEFAULT_TFLITECPP_PATH)))
}
fn resolve_path(path: impl AsRef<Path>) -> PathBuf {
let p = path.as_ref();
if p.is_file() {
std::fs::canonicalize(p).unwrap_or_else(|_| p.to_path_buf())
} else {
p.to_path_buf()
}
}
pub fn load(path: impl AsRef<Path>) -> Result<tensorflowlite_c, libloading::Error> {
unsafe { tensorflowlite_c::new(path.as_ref().as_os_str()) }
}
fn try_vendored_with_path() -> Option<(tensorflowlite_c, PathBuf)> {
let dir = option_env!("EDGEFIRST_TFLITE_VENDORED_DIR")?;
let lib_name = if cfg!(windows) {
"tensorflowlite_c.dll"
} else if cfg!(target_os = "macos") {
"libtensorflowlite_c.dylib"
} else {
"libtensorflowlite_c.so"
};
let path = std::path::Path::new(dir).join(lib_name);
if !path.exists() {
warn!(
"Vendored TFLite library not found at {} \
(built with `vendored` feature but library is missing at runtime — \
copy the library alongside your binary for deployment)",
path.display()
);
return None;
}
debug!("Trying vendored library: {}", path.display());
match load(&path) {
Ok(lib) => {
debug!("Loaded vendored TFLite library: {}", path.display());
Some((lib, path))
}
Err(e) => {
warn!(
"Vendored TFLite library at {} failed to load: {e}",
path.display()
);
None
}
}
}