use axum::{ body::Body, extract::{Path, Query, State}, http::{header, StatusCode}, response::{IntoResponse, Json}, routing::{get, post}, Router, }; use serde::{Deserialize, Serialize}; use std::process::Command; use crate::core::db::{Database, PostgresDb}; #[derive(Debug, Deserialize)] pub struct RegisterFromPersonRequest { pub file_uuid: String, pub person_id: String, pub identity_name: String, pub metadata: Option, } #[derive(Debug, Deserialize)] pub struct RegisterFromFaceRequest { pub face_json_path: String, pub identity_name: String, pub schema: Option, } #[derive(Debug, Serialize)] pub struct RegisterFromPersonResponse { pub success: bool, pub message: String, pub identity_id: i32, pub identity_name: String, pub person_id: String, } #[derive(Debug, Serialize)] pub struct RegisterFromFaceResponse { pub success: bool, pub message: String, pub identity_uuid: Option, pub identity_name: String, pub total_vectors: Option, pub angle_coverage: Option>, pub quality_avg: Option, } pub fn identity_routes() -> Router { Router::new() .route("/api/v1/identities/from-person", post(register_from_person)) .route("/api/v1/identities/from-face", post(register_from_face)) .route("/api/v1/identities", get(list_identities)) .route("/api/v1/faces/candidates", get(list_face_candidates)) .route( "/api/v1/identities/:identity_id/faces", get(get_identity_faces), ) .route("/api/v1/faces/:face_id/thumbnail", get(get_face_thumbnail)) } /// Register a Global Identity from face.json with multi-angle reference vectors. /// Calls select_face_reference_vectors_v2.py for automatic reference selection. async fn register_from_face( State(_state): State, Json(req): Json, ) -> Result, (StatusCode, String)> { let schema = req.schema.unwrap_or("dev".to_string()); let python_path = std::env::var("MOMENTRY_PYTHON_PATH").unwrap_or("/opt/homebrew/bin/python3.11".to_string()); let scripts_dir = std::env::var("MOMENTRY_SCRIPTS_DIR").unwrap_or_else(|_| { let mut path = std::env::current_dir().unwrap_or_default(); path.push("scripts"); path.to_string_lossy().to_string() }); let script_path = format!("{}/select_face_reference_vectors_v2.py", scripts_dir); tracing::info!( "Registering identity '{}' from face.json: {}", req.identity_name, req.face_json_path ); let output = Command::new(&python_path) .arg(&script_path) .arg("--face-json") .arg(&req.face_json_path) .arg("--identity-name") .arg(&req.identity_name) .arg("--register") .arg("--schema") .arg(&schema) .output() .map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to execute script: {}", e), ) })?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Script failed: {}", stderr), )); } let db = PostgresDb::init().await.map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {}", e), ) })?; let query = r#" SELECT uuid, reference_data->'total_references' as total, reference_data->'angles_covered' as angles, reference_data->'quality_avg' as quality FROM identities WHERE name = $1 ORDER BY created_at DESC LIMIT 1 "#; let row: Option<(String, Option, Option>, Option)> = sqlx::query_as(query) .bind(&req.identity_name) .fetch_optional(db.pool()) .await .map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), ) })?; match row { Some((uuid, total, angles, quality)) => Ok(Json(RegisterFromFaceResponse { success: true, message: format!( "Successfully registered identity '{}' with {} reference vectors", req.identity_name, total.unwrap_or(0) ), identity_uuid: Some(uuid), identity_name: req.identity_name, total_vectors: total, angle_coverage: angles, quality_avg: quality, })), None => Ok(Json(RegisterFromFaceResponse { success: true, message: format!( "Identity '{}' registered, but details not found", req.identity_name ), identity_uuid: None, identity_name: req.identity_name, total_vectors: None, angle_coverage: None, quality_avg: None, })), } } /// Register a Global Identity from a specific Person in a video. /// This creates/updates the Identity record, links the Person to the Identity, /// and updates the Person's name to match the Identity. async fn register_from_person( State(_state): State, Json(req): Json, ) -> Result, (StatusCode, String)> { let db = match PostgresDb::init().await { Ok(db) => db, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {}", e), )) } }; let mut tx = match db.pool().begin().await { Ok(tx) => tx, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Tx error: {}", e), )) } }; // 1. Check if Person exists let person_query = "SELECT id, name FROM person_identities WHERE person_id = $1 AND file_uuid = $2"; let person: Option<(i32, Option)> = match sqlx::query_as(person_query) .bind(&req.person_id) .bind(&req.file_uuid) .fetch_optional(&mut *tx) .await { Ok(p) => p, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), )) } }; let (person_db_id, _old_name) = match person { Some(p) => p, None => { return Err(( StatusCode::NOT_FOUND, format!( "Person '{}' not found in video '{}'", req.person_id, req.file_uuid ), )) } }; // 2. Check if Identity exists let identity_query = "SELECT id FROM identities WHERE name = $1"; let identity_id: Option = match sqlx::query_scalar(identity_query) .bind(&req.identity_name) .fetch_optional(&mut *tx) .await { Ok(id) => id, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), )) } }; let final_identity_id = if let Some(id) = identity_id { id } else { // Create new Identity let meta_json = req.metadata.clone().unwrap_or(serde_json::json!({})); let new_id: i32 = match sqlx::query_scalar( r#" INSERT INTO identities (name, embedding, metadata) VALUES ($1, NULLIF($2, '')::public.vector, $3) RETURNING id "#, ) .bind(&req.identity_name) .bind("".to_string()) // No embedding for now via this API .bind(&meta_json) .fetch_one(&mut *tx) .await { Ok(id) => id, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Insert identity error: {}", e), )) } }; new_id }; // 3. Create Binding // Columns: id, identity_id, identity_type, identity_value, confidence, metadata, created_at let binding_query = r#" INSERT INTO identity_bindings (identity_id, identity_type, identity_value, confidence, metadata) VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING "#; match sqlx::query(binding_query) .bind(final_identity_id) .bind("person_id") // identity_type .bind(&req.person_id) // identity_value .bind(1.0) // confidence .bind(serde_json::to_string(&serde_json::json!({"auto_updated": true})).unwrap()) .execute(&mut *tx) .await { Ok(_) => {} Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Binding error: {}", e), )) } }; // 4. Update Person Name let update_person = "UPDATE person_identities SET name = $1 WHERE id = $2"; match sqlx::query(update_person) .bind(&req.identity_name) .bind(person_db_id) .execute(&mut *tx) .await { Ok(_) => {} Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Update person error: {}", e), )) } }; match tx.commit().await { Ok(_) => {} Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Commit error: {}", e), )) } }; Ok(Json(RegisterFromPersonResponse { success: true, message: format!( "Successfully registered identity '{}' and linked to person '{}'", req.identity_name, req.person_id ), identity_id: final_identity_id, identity_name: req.identity_name, person_id: req.person_id, })) } /// List all global identities async fn list_identities( State(_state): State, Query(query): Query, ) -> Result, (StatusCode, String)> { let db = match PostgresDb::init().await { Ok(db) => db, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {}", e), )) } }; let page = query.page.unwrap_or(1); let page_size = query.page_size.unwrap_or(20); let offset = ((page - 1) as i64) * (page_size as i64); // 獲取總數 let count_sql = "SELECT COUNT(*) FROM identities"; let total: i64 = match sqlx::query_scalar(count_sql).fetch_one(db.pool()).await { Ok(count) => count, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Count error: {}", e), )) } }; let sql = "SELECT id, name, metadata FROM identities ORDER BY id DESC LIMIT $1 OFFSET $2"; let rows: Vec<(i32, String, Option)> = match sqlx::query_as(sql) .bind(page_size as i64) .bind(offset) .fetch_all(db.pool()) .await { Ok(rows) => rows, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), )) } }; let identities: Vec = rows .into_iter() .map(|r| IdentityResponse { id: r.0, name: r.1, metadata: r.2, }) .collect(); Ok(Json(IdentityListResponse { identities, count: total, page, page_size, })) } #[derive(Debug, Deserialize)] pub struct ListIdentitiesQuery { pub page: Option, pub page_size: Option, } #[derive(Debug, Serialize)] pub struct IdentityResponse { pub id: i32, pub name: String, pub metadata: Option, } #[derive(Debug, Serialize)] pub struct IdentityListResponse { pub identities: Vec, pub count: i64, pub page: usize, pub page_size: usize, } #[derive(Debug, Deserialize)] pub struct FaceCandidatesQuery { pub file_uuid: Option, pub min_confidence: Option, pub page: Option, pub page_size: Option, pub limit: Option, } #[derive(Debug, Serialize)] pub struct FaceCandidate { pub id: i32, pub face_id: Option, pub file_uuid: String, pub frame_number: i64, pub confidence: f64, pub bbox: Option, pub attributes: Option, } #[derive(Debug, Serialize)] pub struct FaceCandidatesResponse { pub candidates: Vec, pub total: i64, pub page: usize, pub page_size: usize, } #[derive(Debug, Deserialize)] pub struct IdentityFacesQuery { pub page: Option, pub page_size: Option, pub limit: Option, } #[derive(Debug, Serialize)] pub struct IdentityFace { pub id: i32, pub face_id: Option, pub file_uuid: String, pub frame_number: i64, pub confidence: f64, pub bbox: Option, pub attributes: Option, } #[derive(Debug, Serialize)] pub struct IdentityFacesResponse { pub identity_id: i32, pub faces: Vec, pub total: i64, } async fn list_face_candidates( Query(query): Query, ) -> Result, (StatusCode, String)> { let db = match PostgresDb::init().await { Ok(db) => db, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to connect to database: {}", e), )) } }; let page = query.page.unwrap_or(1); let page_size = std::cmp::min(query.page_size.unwrap_or(15), 100); let offset = (page - 1) * page_size; let min_confidence = query.min_confidence.unwrap_or(0.5); let table = crate::core::db::schema::table_name("face_detections"); let total: i64 = if let Some(file_uuid) = &query.file_uuid { let count_sql = format!( "SELECT COUNT(*) FROM {} WHERE identity_id IS NULL AND confidence >= $1 AND file_uuid = $2", table ); match sqlx::query_scalar(&count_sql) .bind(min_confidence) .bind(file_uuid) .fetch_one(db.pool()) .await { Ok(count) => count, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Count error: {}", e), )) } } } else { let count_sql = format!( "SELECT COUNT(*) FROM {} WHERE identity_id IS NULL AND confidence >= $1", table ); match sqlx::query_scalar(&count_sql) .bind(min_confidence) .fetch_one(db.pool()) .await { Ok(count) => count, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Count error: {}", e), )) } } }; let rows = if let Some(file_uuid) = &query.file_uuid { let sql = format!( "SELECT id, face_id, file_uuid, frame_number, confidence, bbox, attributes FROM {} WHERE identity_id IS NULL AND confidence >= $1 AND file_uuid = $2 ORDER BY confidence DESC LIMIT $3 OFFSET $4", table ); match sqlx::query_as::< _, ( i32, Option, String, i64, f64, Option, Option, ), >(&sql) .bind(min_confidence) .bind(file_uuid) .bind(page_size as i64) .bind(offset as i64) .fetch_all(db.pool()) .await { Ok(rows) => rows, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), )) } } } else { let sql = format!( "SELECT id, face_id, file_uuid, frame_number, confidence, bbox, attributes FROM {} WHERE identity_id IS NULL AND confidence >= $1 ORDER BY confidence DESC LIMIT $2 OFFSET $3", table ); match sqlx::query_as::< _, ( i32, Option, String, i64, f64, Option, Option, ), >(&sql) .bind(min_confidence) .bind(page_size as i64) .bind(offset as i64) .fetch_all(db.pool()) .await { Ok(rows) => rows, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), )) } } }; let candidates: Vec = rows .into_iter() .map(|r| FaceCandidate { id: r.0, face_id: r.1, file_uuid: r.2, frame_number: r.3, confidence: r.4, bbox: r.5, attributes: r.6, }) .collect(); Ok(Json(FaceCandidatesResponse { candidates, total, page, page_size, })) } async fn get_identity_faces( axum::extract::Path(identity_id): axum::extract::Path, Query(query): Query, ) -> Result, (StatusCode, String)> { let db = match PostgresDb::init().await { Ok(db) => db, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to connect to database: {}", e), )) } }; let page_size = std::cmp::min(query.page_size.unwrap_or(100), 1000); let offset = (query.page.unwrap_or(1) - 1) * page_size; let table = crate::core::db::schema::table_name("face_detections"); let count_sql = format!("SELECT COUNT(*) FROM {} WHERE identity_id = $1", table); let total: i64 = match sqlx::query_scalar(&count_sql) .bind(identity_id) .fetch_one(db.pool()) .await { Ok(count) => count, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Count error: {}", e), )) } }; let sql = format!( "SELECT id, face_id, file_uuid, frame_number, confidence, bbox, attributes FROM {} WHERE identity_id = $1 ORDER BY confidence DESC LIMIT $2 OFFSET $3", table ); let rows = match sqlx::query_as::< _, ( i32, Option, String, i64, f64, Option, Option, ), >(&sql) .bind(identity_id) .bind(page_size as i64) .bind(offset as i64) .fetch_all(db.pool()) .await { Ok(rows) => rows, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), )) } }; let faces: Vec = rows .into_iter() .map(|r| IdentityFace { id: r.0, face_id: r.1, file_uuid: r.2, frame_number: r.3, confidence: r.4, bbox: r.5, attributes: r.6, }) .collect(); Ok(Json(IdentityFacesResponse { identity_id, faces, total, })) } async fn get_face_thumbnail( Path(face_id): Path, ) -> Result { let db = match PostgresDb::init().await { Ok(db) => db, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to connect to database: {}", e), )) } }; let table_fd = crate::core::db::schema::table_name("face_detections"); let table_v = crate::core::db::schema::table_name("videos"); let sql = format!( "SELECT fd.frame_number, fd.bbox, v.file_path, v.fps FROM {} fd JOIN {} v ON fd.file_uuid = v.uuid WHERE fd.id = $1", table_fd, table_v ); let row: Option<(i64, Option, String, f64)> = match sqlx::query_as(&sql) .bind(face_id) .fetch_optional(db.pool()) .await { Ok(row) => row, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("Query error: {}", e), )) } }; let (frame_number, bbox_json, file_path, fps) = match row { Some(r) => r, None => return Err((StatusCode::NOT_FOUND, format!("Face {} not found", face_id))), }; let bbox: Bbox = match bbox_json { Some(json) => serde_json::from_value(json).unwrap_or(Bbox { x: 0, y: 0, width: 100, height: 100, }), None => Bbox { x: 0, y: 0, width: 100, height: 100, }, }; let timestamp = frame_number as f64 / fps; let crop_filter = format!("crop={}:{}:{}:{}", bbox.width, bbox.height, bbox.x, bbox.y); let output = match Command::new("ffmpeg") .args(&[ "-ss", ×tamp.to_string(), "-i", &file_path, "-vf", &crop_filter, "-frames:v", "1", "-f", "image2pipe", "-vcodec", "mjpeg", "-", ]) .output() { Ok(o) => o, Err(e) => { return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("ffmpeg error: {}", e), )) } }; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); return Err(( StatusCode::INTERNAL_SERVER_ERROR, format!("ffmpeg failed: {}", stderr), )); } let response = axum::response::Response::builder() .status(StatusCode::OK) .header(header::CONTENT_TYPE, "image/jpeg") .header(header::CACHE_CONTROL, "public, max-age=3600") .body(Body::from(output.stdout)) .map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("Response error: {}", e), ) })?; Ok(response) } #[derive(Debug, Deserialize)] struct Bbox { x: i32, y: i32, width: i32, height: i32, }