use axum::{
extract::{Multipart, State},
http::StatusCode,
response::Json,
routing::{get, post},
Router,
};
use image::GenericImageView;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tower_http::cors::{Any, CorsLayer};
use tower_http::limit::RequestBodyLimitLayer;
use crate::api::Ocr;
use crate::core::output::{format_hocr, format_tsv, to_json_output};
#[derive(Debug, Clone)]
pub struct ServerConfig {
pub host: String,
pub port: u16,
pub max_upload_size_mb: usize,
}
impl Default for ServerConfig {
fn default() -> Self {
Self {
host: "127.0.0.1".to_string(),
port: 8080,
max_upload_size_mb: 20,
}
}
}
struct AppState {
ocr: Arc<Ocr>,
}
#[derive(Debug, Deserialize)]
struct OcrParams {
#[serde(default = "default_lang")]
lang: String,
#[serde(default = "default_format")]
format: String,
#[serde(default)]
preprocess: bool,
#[serde(default = "default_engine")]
engine: String,
}
fn default_lang() -> String {
"en".to_string()
}
fn default_format() -> String {
"text".to_string()
}
fn default_engine() -> String {
"pattern".to_string()
}
#[derive(Debug, Serialize)]
struct OcrResponse {
success: bool,
text: String,
confidence: f64,
language: String,
words_count: usize,
lines_count: usize,
#[serde(skip_serializing_if = "Option::is_none")]
json_output: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
hocr: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
tsv: Option<String>,
processing_time_ms: u64,
}
#[derive(Debug, Serialize)]
struct ErrorResponse {
success: bool,
error: String,
}
#[derive(Debug, Serialize)]
struct HealthResponse {
status: String,
version: String,
supported_languages: Vec<String>,
}
pub async fn run_server(config: ServerConfig) -> crate::utils::Result<()> {
let ocr = Ocr::new()?;
ocr.initialize().await?;
let state = Arc::new(AppState { ocr: Arc::new(ocr) });
let cors = CorsLayer::new()
.allow_origin(Any)
.allow_methods(Any)
.allow_headers(Any);
let app = Router::new()
.route("/health", get(health_check))
.route("/languages", get(list_languages))
.route("/recognize", post(recognize_text))
.layer(cors)
.layer(RequestBodyLimitLayer::new(
config.max_upload_size_mb * 1024 * 1024,
))
.with_state(state);
let addr = format!("{}:{}", config.host, config.port);
tracing::info!("Starting OCR API server on http://{}", addr);
let listener = tokio::net::TcpListener::bind(&addr).await?;
axum::serve(listener, app).await?;
Ok(())
}
async fn health_check(State(state): State<Arc<AppState>>) -> Json<HealthResponse> {
let metadata = state.ocr.get_metadata();
Json(HealthResponse {
status: "ok".to_string(),
version: metadata.version.clone(),
supported_languages: metadata.supported_languages.clone(),
})
}
async fn list_languages(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
let metadata = state.ocr.get_metadata();
Json(serde_json::json!({
"languages": metadata.supported_languages.clone(),
}))
}
async fn recognize_text(
State(state): State<Arc<AppState>>,
mut multipart: Multipart,
) -> std::result::Result<Json<OcrResponse>, (StatusCode, Json<ErrorResponse>)> {
let start = std::time::Instant::now();
let mut image_data: Option<Vec<u8>> = None;
let mut params = OcrParams {
lang: "en".to_string(),
format: "text".to_string(),
preprocess: false,
engine: "pattern".to_string(),
};
while let Ok(Some(field)) = multipart.next_field().await {
let name = field.name().unwrap_or("").to_string();
match name.as_str() {
"image" => {
image_data = Some(field.bytes().await.unwrap_or_default().to_vec());
}
"lang" => {
params.lang = field.text().await.unwrap_or_else(|_| "en".to_string());
}
"format" => {
params.format = field.text().await.unwrap_or_else(|_| "text".to_string());
}
"preprocess" => {
let val = field.text().await.unwrap_or_else(|_| "false".to_string());
params.preprocess = val == "true";
}
"engine" => {
params.engine = field.text().await.unwrap_or_else(|_| "pattern".to_string());
}
_ => {}
}
}
let image_data = match image_data {
Some(data) if !data.is_empty() => data,
_ => {
return Err((
StatusCode::BAD_REQUEST,
Json(ErrorResponse {
success: false,
error: "No image file provided. Use field name 'image'.".to_string(),
}),
));
}
};
let img = match image::load_from_memory(&image_data) {
Ok(img) => img,
Err(e) => {
return Err((
StatusCode::BAD_REQUEST,
Json(ErrorResponse {
success: false,
error: format!("Invalid image: {}", e),
}),
));
}
};
let gray = img.to_luma8();
let (width, height) = gray.dimensions();
let raw_pixels: Vec<u8> = gray.pixels().map(|p| p[0]).collect();
let mut config = crate::core::config::OcrConfig::default();
config.recognition.language = params.lang.clone();
config.recognition.engine = match params.engine.as_str() {
"lstm" => crate::core::config::RecognitionEngine::LSTM,
"hybrid" => crate::core::config::RecognitionEngine::Hybrid,
_ => crate::core::config::RecognitionEngine::PatternMatching,
};
config.image_processing.enable_preprocessing = params.preprocess;
config.image_processing.enable_binarization = params.preprocess;
config.image_processing.enable_noise_reduction = params.preprocess;
config.image_processing.enable_contrast_enhancement = params.preprocess;
config.image_processing.enable_deskewing = params.preprocess;
let ocr = match crate::api::Ocr::with_config(config) {
Ok(o) => o,
Err(e) => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse {
success: false,
error: format!("Failed to create OCR engine: {}", e),
}),
));
}
};
if let Err(e) = ocr.initialize().await {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse {
success: false,
error: format!("Failed to initialize OCR engine: {}", e),
}),
));
}
let result = match ocr.recognize_text(&raw_pixels, width, height).await {
Ok(r) => r,
Err(e) => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse {
success: false,
error: format!("OCR failed: {}", e),
}),
));
}
};
let elapsed = start.elapsed().as_millis() as u64;
let words_count = result.words.len();
let lines_count = result.lines.len();
let (text, json_output, hocr, tsv) = match params.format.to_lowercase().as_str() {
"json" => {
let json_val = to_json_output(&result);
(
result.text.clone(),
Some(serde_json::to_value(&json_val).unwrap_or_default()),
None,
None,
)
}
"hocr" | "html" => {
let hocr_str = format_hocr(&result).unwrap_or_default();
(result.text.clone(), None, Some(hocr_str), None)
}
"tsv" => {
let tsv_str = format_tsv(&result).unwrap_or_default();
(result.text.clone(), None, None, Some(tsv_str))
}
_ => (result.text.clone(), None, None, None),
};
Ok(Json(OcrResponse {
success: true,
text,
confidence: result.confidence as f64,
language: params.lang,
words_count,
lines_count,
json_output,
hocr,
tsv,
processing_time_ms: elapsed,
}))
}