426 lines
13 KiB
Rust
426 lines
13 KiB
Rust
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 CreateIdentityRequest {
|
|
pub face_json_path: String,
|
|
pub identity_name: String,
|
|
pub schema: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct CreateIdentityResponse {
|
|
pub success: bool,
|
|
pub message: String,
|
|
pub identity_uuid: Option<String>,
|
|
pub identity_name: String,
|
|
pub total_vectors: Option<i32>,
|
|
pub angle_coverage: Option<Vec<String>>,
|
|
pub quality_avg: Option<f64>,
|
|
}
|
|
|
|
pub fn identity_routes() -> Router<crate::api::server::AppState> {
|
|
Router::new()
|
|
.route("/api/v1/identities", get(list_identities))
|
|
.route("/api/v1/identity", post(create_identity))
|
|
.route("/api/v1/faces/candidates", get(list_face_candidates))
|
|
}
|
|
|
|
/// 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 create_identity(
|
|
State(_state): State<crate::api::server::AppState>,
|
|
Json(req): Json<CreateIdentityRequest>,
|
|
) -> Result<Json<CreateIdentityResponse>, (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<i32>, Option<Vec<String>>, Option<f64>)> =
|
|
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(CreateIdentityResponse {
|
|
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(CreateIdentityResponse {
|
|
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,
|
|
})),
|
|
}
|
|
}
|
|
|
|
/// List all global identities
|
|
async fn list_identities(
|
|
State(_state): State<crate::api::server::AppState>,
|
|
Query(query): Query<ListIdentitiesQuery>,
|
|
) -> Result<Json<IdentityListResponse>, (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 id_table = crate::core::db::schema::table_name("identities");
|
|
|
|
let total: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {}", id_table))
|
|
.fetch_one(db.pool()).await
|
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Count error: {}", e)))?;
|
|
|
|
let sql = format!("SELECT id, uuid, name, metadata FROM {} ORDER BY id DESC LIMIT $1 OFFSET $2", id_table);
|
|
|
|
let rows: Vec<(i32, uuid::Uuid, String, Option<serde_json::Value>)> = 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<IdentityResponse> = rows
|
|
.into_iter()
|
|
.map(|r| IdentityResponse {
|
|
id: r.0,
|
|
identity_uuid: r.1.to_string().replace('-', ""),
|
|
name: r.2,
|
|
metadata: r.3,
|
|
})
|
|
.collect();
|
|
|
|
let identities_table = crate::core::db::schema::table_name("identities");
|
|
let total_identities: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {}", identities_table))
|
|
.fetch_one(db.pool()).await.unwrap_or(0);
|
|
let tmdb_identities: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {} WHERE source = 'tmdb'", identities_table))
|
|
.fetch_one(db.pool()).await.unwrap_or(0);
|
|
let auto_identities: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {} WHERE source = 'auto'", identities_table))
|
|
.fetch_one(db.pool()).await.unwrap_or(0);
|
|
|
|
Ok(Json(IdentityListResponse {
|
|
identities,
|
|
count: total,
|
|
page,
|
|
page_size,
|
|
total_identities,
|
|
tmdb_identities,
|
|
auto_identities,
|
|
}))
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct ListIdentitiesQuery {
|
|
pub page: Option<usize>,
|
|
pub page_size: Option<usize>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct FaceCandidatesQuery {
|
|
pub file_uuid: Option<String>,
|
|
pub min_confidence: Option<f64>,
|
|
pub page: Option<usize>,
|
|
pub page_size: Option<usize>,
|
|
pub limit: Option<usize>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct FaceCandidate {
|
|
pub id: i32,
|
|
pub face_id: Option<String>,
|
|
pub file_uuid: String,
|
|
pub frame_number: i32,
|
|
pub confidence: f32,
|
|
pub bbox: Option<serde_json::Value>,
|
|
pub attributes: Option<serde_json::Value>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct FaceCandidatesResponse {
|
|
pub candidates: Vec<FaceCandidate>,
|
|
pub total: i64,
|
|
pub page: usize,
|
|
pub page_size: usize,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct IdentityResponse {
|
|
pub id: i32,
|
|
pub identity_uuid: String,
|
|
pub name: String,
|
|
pub metadata: Option<serde_json::Value>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct IdentityListResponse {
|
|
pub identities: Vec<IdentityResponse>,
|
|
pub count: i64,
|
|
pub page: usize,
|
|
pub page_size: usize,
|
|
pub total_identities: i64,
|
|
pub tmdb_identities: i64,
|
|
pub auto_identities: i64,
|
|
}
|
|
|
|
async fn list_face_candidates(
|
|
Query(query): Query<FaceCandidatesQuery>,
|
|
) -> Result<Json<FaceCandidatesResponse>, (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,
|
|
jsonb_build_object('x', x, 'y', y, 'width', width, 'height', height) as bbox,
|
|
NULL::jsonb as 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>,
|
|
String,
|
|
i32,
|
|
f32,
|
|
Option<serde_json::Value>,
|
|
Option<serde_json::Value>,
|
|
),
|
|
>(&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,
|
|
jsonb_build_object('x', x, 'y', y, 'width', width, 'height', height) as bbox,
|
|
NULL::jsonb as 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>,
|
|
String,
|
|
i32,
|
|
f32,
|
|
Option<serde_json::Value>,
|
|
Option<serde_json::Value>,
|
|
),
|
|
>(&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<FaceCandidate> = 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,
|
|
}))
|
|
}
|