- Integrate Convolutional Block Attention Module (CBAM) for improved feature focus - Implement temperature scaling and confidence-based status reporting - Automate dataset acquisition using kagglehub - Update ONNX opset to 18 and refine preprocessing validation
60 lines
1.5 KiB
Rust
60 lines
1.5 KiB
Rust
mod config;
|
|
mod error;
|
|
mod image;
|
|
mod model;
|
|
mod routes;
|
|
mod telemetry;
|
|
|
|
use anyhow::Result;
|
|
use config::Config;
|
|
use model::ModelService;
|
|
use routes::{router, AppState};
|
|
use std::sync::Arc;
|
|
use tokio::net::TcpListener;
|
|
use tracing_subscriber::EnvFilter;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
// Initialize tracing subscriber with EnvFilter
|
|
tracing_subscriber::fmt()
|
|
.with_env_filter(EnvFilter::from_default_env())
|
|
.init();
|
|
|
|
// Load configuration from environment
|
|
let config = Config::from_env()?;
|
|
|
|
// Create ModelService with calibration and wrap in Arc
|
|
let model = Arc::new(ModelService::with_calibration(
|
|
&config.model_path,
|
|
config.input_size,
|
|
config.temperature,
|
|
config.conf_threshold_high,
|
|
config.conf_threshold_low,
|
|
));
|
|
tracing::info!(
|
|
model_loaded = model.is_loaded(),
|
|
model_path = ?config.model_path,
|
|
temperature = config.temperature,
|
|
conf_high = config.conf_threshold_high,
|
|
conf_low = config.conf_threshold_low,
|
|
"Model service initialized"
|
|
);
|
|
|
|
// Set model load status metric
|
|
telemetry::model_load_status().set(if model.is_loaded() { 1.0 } else { 0.0 });
|
|
|
|
// Create AppState
|
|
let state = AppState { model };
|
|
|
|
// Bind TcpListener to host:port
|
|
let addr = format!("{}:{}", config.host, config.port);
|
|
let listener = TcpListener::bind(&addr).await?;
|
|
|
|
tracing::info!(address = %addr, "Server listening");
|
|
|
|
// Serve with Axum
|
|
axum::serve(listener, router(state)).await?;
|
|
|
|
Ok(())
|
|
}
|