admin/momentry_core
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'main' of http://192.168.110.200:3000/admin/momentry_core

Browse Source
This commit is contained in:
M5Max128
2026-05-22 08:51:08 +08:00
parent 701e71463d deb9516796
commit e1619c724a
6 changed files with 1837 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

459
src/api/trace_agent_api.rs
View File

@@ -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,
};
@@ -16,6 +17,22 @@ pub fn trace_agent_routes() -> Router<crate::api::types::AppState> {
"/api/v1/file/:file_uuid/trace/:trace_id/faces",
get(list_trace_faces),
)
.route(
"/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),
)
.route(
"/api/v1/file/:file_uuid/identities/:identity_uuid_a/co-occur-with/:identity_uuid_b",
get(get_cooccurrence),
)
.route(
"/api/v1/file/:file_uuid/tkg/rebuild",
post(rebuild_tkg),
)
}
#[derive(Debug, Deserialize)]
@@ -328,3 +345,441 @@ async fn list_trace_faces(
faces,
}))
}
#[derive(Debug, Serialize)]
struct RepFaceBbox {
x: i32,
y: i32,
width: i32,
height: i32,
}
#[derive(Debug, Serialize)]
struct RepFaceResult {
frame_number: i64,
timestamp_secs: f64,
bbox: RepFaceBbox,
confidence: f64,
quality_score: f64,
blur_score: f64,
}
#[derive(Debug, Serialize)]
struct RepFaceResponse {
success: bool,
file_uuid: String,
trace_id: i32,
face_count: i64,
representative: RepFaceResult,
}
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<F, T>(
pool: &sqlx::PgPool,
file_uuid: &str,
trace_id: i32,
err_fn: F,
) -> Result<RepFaceSelection, T>
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");
let fps: f64 = sqlx::query_scalar(&format!(
"SELECT COALESCE(fps, 25.0) FROM {} WHERE file_uuid = $1", video_table
))
.bind(file_uuid)
.fetch_optional(pool)
.await
.map_err(|e| err_fn(anyhow::anyhow!("{}", e)))?
.unwrap_or(25.0);
let face_count: (i64,) = sqlx::query_as(&format!(
"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_all(pool)
.await
.map_err(|e| err_fn(anyhow::anyhow!("{}", e)))?;
if rows.is_empty() {
return Err(err_fn(anyhow::anyhow!("No suitable face found")));
}
let candidates: Vec<Candidate> = rows.into_iter()
.map(|(frame, x, y, w, h, conf)| {
let score = (w as f64 * h as f64) * conf;
Candidate { frame, x, y, w, h, conf, score }
})
.collect();
let video_path: String = sqlx::query_scalar(&format!(
"SELECT file_path FROM {} WHERE file_uuid = $1", video_table
))
.bind(file_uuid)
.fetch_optional(pool)
.await
.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;
let mut best_idx = 0usize;
for (i, c) in candidates.iter().enumerate() {
let seek = c.frame as f64 / fps;
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::<f64>() {
if blur < best_blur { best_blur = blur; best = c.frame; best_idx = i; }
}
}
}
}
}
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<crate::api::types::AppState>,
Path((file_uuid, trace_id)): Path<(String, i32)>,
) -> Result<Json<RepFaceResponse>, (StatusCode, Json<serde_json::Value>)> {
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: sel.face_count,
representative: RepFaceResult {
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<crate::api::types::AppState>,
Path((file_uuid, trace_id)): Path<(String, i32)>,
) -> Result<Response, (StatusCode, Json<serde_json::Value>)> {
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())
}
#[derive(Debug, Serialize)]
struct CoOccurIdentity {
identity_uuid: String,
name: String,
trace_id: i32,
}
#[derive(Debug, Serialize)]
struct CoOccurRepFace {
frame_number: i64,
bbox: RepFaceBbox,
confidence: f64,
thumbnail_url: String,
}
#[derive(Debug, Serialize)]
struct CoOccurrence {
frame_number: i64,
timestamp_secs: f64,
total_cooccurrence_frames: i64,
representative_face_a: Option<CoOccurRepFace>,
representative_face_b: Option<CoOccurRepFace>,
}
#[derive(Debug, Serialize)]
struct CoOccurResponse {
success: bool,
file_uuid: String,
identity_a: CoOccurIdentity,
identity_b: CoOccurIdentity,
first_cooccurrence: CoOccurrence,
}
async fn get_cooccurrence(
State(state): State<crate::api::types::AppState>,
Path((file_uuid, identity_uuid_a, identity_uuid_b)): Path<(String, String, String)>,
) -> Result<Json<CoOccurResponse>, (StatusCode, Json<serde_json::Value>)> {
use crate::core::db::schema;
let id_table = schema::table_name("identities");
let fd_table = schema::table_name("face_detections");
// Stage 1: Get identity names and IDs
let id_a = sqlx::query_as::<_, (i32, String)>(&format!(
"SELECT id, name FROM {} WHERE uuid::text = $1 OR REPLACE(uuid::text, '-', '') = $1",
id_table
))
.bind(&identity_uuid_a)
.fetch_optional(state.db.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": "Identity A not found"})))
})?;
let id_b = sqlx::query_as::<_, (i32, String)>(&format!(
"SELECT id, name FROM {} WHERE uuid::text = $1 OR REPLACE(uuid::text, '-', '') = $1",
id_table
))
.bind(&identity_uuid_b)
.fetch_optional(state.db.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": "Identity B not found"})))
})?;
// Stage 2: Find first frame where both identity_ids appear
let cooccur: Option<(i64,)> = sqlx::query_as(
&format!(
"SELECT MIN(fd.frame_number)::bigint FROM {} fd \
WHERE fd.file_uuid = $1 AND fd.identity_id = $2 \
AND fd.frame_number IN ( \
SELECT frame_number FROM {} \
WHERE file_uuid = $1 AND identity_id = $3 \
)",
fd_table, fd_table
)
)
.bind(&file_uuid)
.bind(id_a.0)
.bind(id_b.0)
.fetch_optional(state.db.pool())
.await
.map_err(|e| {
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()})))
})?;
let (first_frame,) = cooccur.ok_or_else(|| {
(StatusCode::NOT_FOUND, Json(serde_json::json!({"error": "These two identities never appear together in this file"})))
})?;
// Get fps for timestamp
let video_table = schema::table_name("videos");
let fps: f64 = sqlx::query_scalar(&format!(
"SELECT COALESCE(fps, 25.0) FROM {} WHERE file_uuid = $1", video_table
))
.bind(&file_uuid)
.fetch_optional(state.db.pool())
.await
.map_err(|e| {
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()})))
})?
.unwrap_or(25.0);
// Stage 3: Get trace_ids for both at this frame
let trace_a: Option<(i32,)> = sqlx::query_as(
&format!("SELECT trace_id FROM {} WHERE file_uuid = $1 AND frame_number = $2 AND identity_id = $3 AND trace_id IS NOT NULL LIMIT 1", fd_table)
)
.bind(&file_uuid).bind(first_frame).bind(id_a.0)
.fetch_optional(state.db.pool()).await
.map_err(|e| {
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()})))
})?;
let trace_b: Option<(i32,)> = sqlx::query_as(
&format!("SELECT trace_id FROM {} WHERE file_uuid = $1 AND frame_number = $2 AND identity_id = $3 AND trace_id IS NOT NULL LIMIT 1", fd_table)
)
.bind(&file_uuid).bind(first_frame).bind(id_b.0)
.fetch_optional(state.db.pool()).await
.map_err(|e| {
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()})))
})?;
// Stage 4: Get representative faces for both traces (reusing select_rep_face)
let rep_a = if let Some((tid,)) = trace_a {
select_rep_face(state.db.pool(), &file_uuid, tid, |e| {
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()})))
}).await.ok().map(|sel| CoOccurRepFace {
frame_number: sel.frame,
bbox: RepFaceBbox { x: sel.x, y: sel.y, width: sel.w, height: sel.h },
confidence: sel.conf,
thumbnail_url: format!("/api/v1/file/{}/trace/{}/thumbnail", file_uuid, tid),
})
} else { None };
let rep_b = if let Some((tid,)) = trace_b {
select_rep_face(state.db.pool(), &file_uuid, tid, |e| {
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": e.to_string()})))
}).await.ok().map(|sel| CoOccurRepFace {
frame_number: sel.frame,
bbox: RepFaceBbox { x: sel.x, y: sel.y, width: sel.w, height: sel.h },
confidence: sel.conf,
thumbnail_url: format!("/api/v1/file/{}/trace/{}/thumbnail", file_uuid, tid),
})
} else { None };
// Total co-occurrence frames (from TKG if available, otherwise from face_detections)
let total_cooccurrence_frames: i64 = sqlx::query_scalar(
&format!(
"SELECT COUNT(DISTINCT fd.frame_number)::bigint FROM {} fd \
WHERE fd.file_uuid = $1 AND fd.identity_id = $2 \
AND fd.frame_number IN ( \
SELECT frame_number FROM {} \
WHERE file_uuid = $1 AND identity_id = $3 \
)",
fd_table, fd_table
)
)
.bind(&file_uuid).bind(id_a.0).bind(id_b.0)
.fetch_one(state.db.pool()).await
.unwrap_or(0);
Ok(Json(CoOccurResponse {
success: true,
file_uuid,
identity_a: CoOccurIdentity {
identity_uuid: identity_uuid_a,
name: id_a.1,
trace_id: trace_a.map(|t| t.0).unwrap_or(0),
},
identity_b: CoOccurIdentity {
identity_uuid: identity_uuid_b,
name: id_b.1,
trace_id: trace_b.map(|t| t.0).unwrap_or(0),
},
first_cooccurrence: CoOccurrence {
frame_number: first_frame,
timestamp_secs: first_frame as f64 / fps,
total_cooccurrence_frames,
representative_face_a: rep_a,
representative_face_b: rep_b,
},
}))
}
use crate::core::config::OUTPUT_DIR;
#[derive(Serialize)]
struct TkgRebuildResponse {
success: bool,
file_uuid: String,
result: Option<serde_json::Value>,
error: Option<String>,
}
async fn rebuild_tkg(
State(state): State<crate::api::types::AppState>,
Path(file_uuid): Path<String>,
) -> Json<TkgRebuildResponse> {
let result = crate::core::processor::tkg::build_tkg(
&state.db,
&file_uuid,
&OUTPUT_DIR,
)
.await;
match result {
Ok(r) => Json(TkgRebuildResponse {
success: true,
file_uuid,
result: Some(serde_json::json!({
"face_trace_nodes": r.face_trace_nodes,
"object_nodes": r.object_nodes,
"speaker_nodes": r.speaker_nodes,
"co_occurrence_edges": r.co_occurrence_edges,
"speaker_face_edges": r.speaker_face_edges,
"face_face_edges": r.face_face_edges,
})),
error: None,
}),
Err(e) => Json(TkgRebuildResponse {
success: false,
file_uuid,
result: None,
error: Some(e.to_string()),
}),
}
}