diff --git a/src/api/trace_agent_api.rs b/src/api/trace_agent_api.rs index 21536f6..67cbf91 100644 --- a/src/api/trace_agent_api.rs +++ b/src/api/trace_agent_api.rs @@ -1,7 +1,8 @@ use axum::{ + body::Body, extract::{Path, Query, State}, - http::StatusCode, - response::Json, + http::{header, StatusCode}, + response::{IntoResponse, Json, Response}, routing::{get, post}, Router, }; @@ -20,6 +21,10 @@ pub fn trace_agent_routes() -> Router { "/api/v1/file/:file_uuid/trace/:trace_id/representative-face", get(get_representative_face), ) + .route( + "/api/v1/file/:file_uuid/trace/:trace_id/thumbnail", + get(get_trace_thumbnail), + ) } #[derive(Debug, Deserialize)] @@ -360,99 +365,84 @@ struct RepFaceResponse { representative: RepFaceResult, } -async fn get_representative_face( - State(state): State, - Path((file_uuid, trace_id)): Path<(String, i32)>, -) -> Result, (StatusCode, Json)> { - use crate::core::db::schema; +struct RepFaceSelection { + frame: i64, + x: i32, + y: i32, + w: i32, + h: i32, + conf: f64, + blur: f64, + score: f64, + video_path: String, + fps: f64, + face_count: i64, +} +async fn select_rep_face( + pool: &sqlx::PgPool, + file_uuid: &str, + trace_id: i32, + err_fn: F, +) -> Result +where + F: Fn(anyhow::Error) -> T, +{ + use crate::core::db::schema; let fd_table = schema::table_name("face_detections"); let video_table = schema::table_name("videos"); - // Get fps for timestamp calculation let fps: f64 = sqlx::query_scalar(&format!( - "SELECT COALESCE(fps, 25.0) FROM {} WHERE file_uuid = $1", - video_table + "SELECT COALESCE(fps, 25.0) FROM {} WHERE file_uuid = $1", video_table )) - .bind(&file_uuid) - .fetch_optional(state.db.pool()) + .bind(file_uuid) + .fetch_optional(pool) .await - .map_err(|e| { - (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()}))) - })? + .map_err(|e| err_fn(anyhow::anyhow!("{}", e)))? .unwrap_or(25.0); - // Get face count for this trace let face_count: (i64,) = sqlx::query_as(&format!( - "SELECT COUNT(*) FROM {} WHERE file_uuid = $1 AND trace_id = $2", + "SELECT COUNT(*) FROM {} WHERE file_uuid = $1 AND trace_id = $2", fd_table + )) + .bind(file_uuid) + .bind(trace_id) + .fetch_one(pool) + .await + .map_err(|e| err_fn(anyhow::anyhow!("{}", e)))?; + + struct Candidate { frame: i64, x: i32, y: i32, w: i32, h: i32, conf: f64, score: f64 } + + let rows = sqlx::query_as::<_, (i64, i32, i32, i32, i32, f64)>(&format!( + "SELECT frame_number::bigint, x, y, width, height, confidence::float8 \ + FROM {} WHERE file_uuid = $1 AND trace_id = $2 AND confidence > 0.7 \ + AND ((metadata->>'qc_ok')::boolean IS NULL OR (metadata->>'qc_ok')::boolean = true) \ + ORDER BY (width::float8 * height::float8) * confidence::float8 DESC LIMIT 10", fd_table )) - .bind(&file_uuid) - .bind(trace_id) - .fetch_one(state.db.pool()) + .bind(file_uuid).bind(trace_id) + .fetch_all(pool) .await - .map_err(|e| { - (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()}))) - })?; - - // Stage 1: SQL - top 10 candidates by area * confidence - #[derive(Debug)] - struct FaceCandidate { - frame: i64, - x: i32, - y: i32, - w: i32, - h: i32, - conf: f64, - score: f64, - } - - let rows = sqlx::query_as::<_, (i64, i32, i32, i32, i32, f64)>( - &format!( - "SELECT frame_number::bigint, x, y, width, height, confidence::float8 \ - FROM {} WHERE file_uuid = $1 AND trace_id = $2 AND confidence > 0.7 \ - AND ((metadata->>'qc_ok')::boolean IS NULL OR (metadata->>'qc_ok')::boolean = true) \ - ORDER BY (width::float8 * height::float8) * confidence::float8 DESC \ - LIMIT 10", - fd_table - ) - ) - .bind(&file_uuid) - .bind(trace_id) - .fetch_all(state.db.pool()) - .await - .map_err(|e| { - (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()}))) - })?; + .map_err(|e| err_fn(anyhow::anyhow!("{}", e)))?; if rows.is_empty() { - return Err((StatusCode::NOT_FOUND, Json(serde_json::json!({ - "error": "No suitable face found for this trace" - })))); + return Err(err_fn(anyhow::anyhow!("No suitable face found"))); } - let candidates: Vec = rows - .into_iter() + let candidates: Vec = rows.into_iter() .map(|(frame, x, y, w, h, conf)| { let score = (w as f64 * h as f64) * conf; - FaceCandidate { frame, x, y, w, h, conf, score } + Candidate { frame, x, y, w, h, conf, score } }) .collect(); - // Stage 2: FFmpeg blurdetect on candidates - let video_path = sqlx::query_scalar::<_, String>(&format!( - "SELECT file_path FROM {} WHERE file_uuid = $1", - video_table + let video_path: String = sqlx::query_scalar(&format!( + "SELECT file_path FROM {} WHERE file_uuid = $1", video_table )) - .bind(&file_uuid) - .fetch_optional(state.db.pool()) + .bind(file_uuid) + .fetch_optional(pool) .await - .map_err(|e| { - (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()}))) - })? - .ok_or_else(|| { - (StatusCode::NOT_FOUND, Json(serde_json::json!({"error": "Video not found"}))) - })?; + .map_err(|e| err_fn(anyhow::anyhow!("{}", e)))? + .ok_or_else(|| err_fn(anyhow::anyhow!("Video not found")))?; let mut best = candidates[0].frame; let mut best_blur = f64::MAX; @@ -460,28 +450,17 @@ async fn get_representative_face( for (i, c) in candidates.iter().enumerate() { let seek = c.frame as f64 / fps; - let output = tokio::process::Command::new("ffmpeg") - .args([ - "-ss", &format!("{:.2}", seek), - "-i", &video_path, - "-vframes", "1", - "-vf", &format!("crop={}:{}:{}:{},blurdetect", c.w, c.h, c.x, c.y), - "-f", "null", - "-", - ]) - .output() - .await; - - if let Ok(o) = output { - let stderr = String::from_utf8_lossy(&o.stderr); + if let Ok(output) = tokio::process::Command::new("ffmpeg") + .args(["-ss", &format!("{:.2}", seek), "-i", &video_path, + "-vframes", "1", "-vf", &format!("crop={}:{}:{}:{},blurdetect", c.w, c.h, c.x, c.y), + "-f", "null", "-"]) + .output().await + { + let stderr = String::from_utf8_lossy(&output.stderr); for line in stderr.lines() { if let Some(blur_str) = line.split("blur mean: ").nth(1) { if let Ok(blur) = blur_str.trim().parse::() { - if blur < best_blur { - best_blur = blur; - best = c.frame; - best_idx = i; - } + if blur < best_blur { best_blur = blur; best = c.frame; best_idx = i; } } } } @@ -489,24 +468,77 @@ async fn get_representative_face( } let chosen = &candidates[best_idx]; + Ok(RepFaceSelection { + frame: chosen.frame, x: chosen.x, y: chosen.y, w: chosen.w, h: chosen.h, + conf: chosen.conf, blur: best_blur, score: chosen.score, + video_path, fps, face_count: face_count.0, + }) +} + +async fn get_representative_face( + State(state): State, + Path((file_uuid, trace_id)): Path<(String, i32)>, +) -> Result, (StatusCode, Json)> { + let sel = select_rep_face(state.db.pool(), &file_uuid, trace_id, |e| { + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()}))) + }).await?; Ok(Json(RepFaceResponse { success: true, file_uuid, trace_id, - face_count: face_count.0, + face_count: sel.face_count, representative: RepFaceResult { - frame_number: chosen.frame, - timestamp_secs: chosen.frame as f64 / fps, - bbox: RepFaceBbox { - x: chosen.x, - y: chosen.y, - width: chosen.w, - height: chosen.h, - }, - confidence: chosen.conf, - quality_score: chosen.score, - blur_score: if best_blur == f64::MAX { 0.0 } else { best_blur }, + frame_number: sel.frame, + timestamp_secs: sel.frame as f64 / sel.fps, + bbox: RepFaceBbox { x: sel.x, y: sel.y, width: sel.w, height: sel.h }, + confidence: sel.conf, + quality_score: sel.score, + blur_score: sel.blur, }, })) } + +async fn get_trace_thumbnail( + State(state): State, + Path((file_uuid, trace_id)): Path<(String, i32)>, +) -> Result)> { + let sel = select_rep_face(state.db.pool(), &file_uuid, trace_id, |e| { + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()}))) + }).await?; + + let seek = sel.frame as f64 / sel.fps; + let tmp = std::env::temp_dir().join(format!("trace_{}_{}.jpg", file_uuid, trace_id)); + + let status = tokio::process::Command::new("ffmpeg") + .args([ + "-ss", &format!("{:.2}", seek), + "-i", &sel.video_path, + "-vframes", "1", + "-vf", &format!("crop={}:{}:{}:{},scale=320:320", sel.w, sel.h, sel.x, sel.y), + "-q:v", "2", + "-y", &tmp.to_string_lossy().to_string(), + ]) + .output() + .await + .map_err(|e| { + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()}))) + })?; + + if !status.status.success() { + return Err((StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": "FFmpeg failed"})))); + } + + let bytes = tokio::fs::read(&tmp).await.map_err(|e| { + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()}))) + })?; + + let _ = tokio::fs::remove_file(&tmp).await; + + Ok(Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "image/jpeg") + .header(header::CACHE_CONTROL, "public, max-age=86400") + .body(Body::from(bytes)) + .unwrap()) +}