diff --git a/src/core/processor/tkg.rs b/src/core/processor/tkg.rs index 0c5fb82..2663b58 100644 --- a/src/core/processor/tkg.rs +++ b/src/core/processor/tkg.rs @@ -5,8 +5,8 @@ use std::collections::HashMap; use std::path::Path; use crate::core::db::postgres_db::PostgresDb; -use crate::core::progress::{TkgPhase, TkgProgress, TkgStats, publish_tkg_progress}; use crate::core::db::redis_client::RedisClient; +use crate::core::progress::{publish_tkg_progress, TkgPhase, TkgProgress, TkgStats}; use std::sync::Arc; /// Face point data retrieved from Qdrant _faces collection @@ -36,22 +36,72 @@ async fn scroll_face_points(file_uuid: &str) -> Vec { let points = match qdrant.scroll_all_points("_faces", filter, 500).await { Ok(pts) => pts, Err(e) => { - tracing::warn!("[TKG-Qdrant] Failed to scroll _faces for {}: {}", file_uuid, e); + tracing::error!( + "[TKG-Qdrant] Failed to scroll _faces for {}: {}", + file_uuid, + e + ); return vec![]; } }; - points.iter().filter_map(|p| { - let payload = &p["payload"]; - let trace_id = payload["trace_id"].as_i64().filter(|&t| t > 0)?; - let frame = payload["frame"].as_i64()?; - let bbox = &payload["bbox"]; - let x = bbox["x"].as_f64().unwrap_or(0.0); - let y = bbox["y"].as_f64().unwrap_or(0.0); - let w = bbox["width"].as_f64().unwrap_or(0.0); - let h = bbox["height"].as_f64().unwrap_or(0.0); - Some(FacePoint { trace_id, frame, x, y, w, h }) - }).collect() + points + .iter() + .filter_map(|p| { + let payload = &p["payload"]; + let trace_id = payload["trace_id"].as_i64().filter(|&t| t > 0)?; + let frame = payload["frame"].as_i64()?; + let bbox = &payload["bbox"]; + let x = bbox["x"].as_f64().unwrap_or(0.0); + let y = bbox["y"].as_f64().unwrap_or(0.0); + let w = bbox["width"].as_f64().unwrap_or(0.0); + let h = bbox["height"].as_f64().unwrap_or(0.0); + Some(FacePoint { + trace_id, + frame, + x, + y, + w, + h, + }) + }) + .collect() +} + +/// Build frame-to-face-points index for O(1) lookup +fn build_frame_index<'a>(points: &'a [FacePoint]) -> HashMap> { + let mut index: HashMap> = HashMap::new(); + for fp in points { + index.entry(fp.frame).or_default().push(fp); + } + index +} + +/// Find matching trace_id for a face bbox using in-memory frame index +fn find_trace_in_frame( + frame_index: &HashMap>, + frame: i64, + x: f64, + y: f64, + w: f64, + h: f64, +) -> Option { + let points = frame_index.get(&frame)?; + let tcx = x + w / 2.0; + let tcy = y + h / 2.0; + + let mut best: Option<(i64, f64)> = None; + for fp in points { + let cx = fp.x + fp.w / 2.0; + let cy = fp.y + fp.h / 2.0; + let dist = ((cx - tcx).powi(2) + (cy - tcy).powi(2)).sqrt(); + match best { + Some((_, best_dist)) if dist < best_dist => best = Some((fp.trace_id, dist)), + None => best = Some((fp.trace_id, dist)), + _ => {} + } + } + best.map(|(tid, _)| tid) } fn t(name: &str) -> String { @@ -91,7 +141,10 @@ async fn populate_face_detections_from_face_json( {"key": "trace_id", "match": {"value": 1}} ] }); - let traced_points = match qdrant.scroll_points("_faces", traced_filter.clone(), 1, None).await { + let traced_points = match qdrant + .scroll_points("_faces", traced_filter.clone(), 1, None) + .await + { Ok((pts, _)) => pts.len(), Err(_) => 0, }; @@ -102,9 +155,9 @@ async fn populate_face_detections_from_face_json( file_uuid ); progress.update_phase(TkgPhase::FaceTracing, 1.0, "Faces already traced"); - if let Some(r) = redis.as_ref() { - publish_tkg_progress(&r, file_uuid, &progress).await; - } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(&r, file_uuid, &progress).await; + } return Ok(()); } @@ -124,7 +177,8 @@ async fn populate_face_detections_from_face_json( progress.stats.total_faces = all_face_points.len() as i64; if all_face_points.is_empty() { - let face_json_path = std::path::Path::new(output_dir).join(format!("{}.face.json", file_uuid)); + let face_json_path = + std::path::Path::new(output_dir).join(format!("{}.face.json", file_uuid)); if !face_json_path.exists() { info!( "[TKG-Phase0] No faces in Qdrant and no face.json for {}, skipping", @@ -161,13 +215,20 @@ async fn populate_face_detections_from_face_json( match result { Ok(()) => { - let new_traced = match qdrant.scroll_points("_faces", traced_filter.clone(), 1, None).await { + let new_traced = match qdrant + .scroll_points("_faces", traced_filter.clone(), 1, None) + .await + { Ok((pts, _)) => pts.len(), Err(_) => 0, }; - let traced_count = match qdrant.scroll_all_points("_faces", traced_filter.clone(), 500).await { + let traced_count = match qdrant + .scroll_all_points("_faces", traced_filter.clone(), 500) + .await + { Ok(pts) => { - let mut traces: std::collections::HashSet = std::collections::HashSet::new(); + let mut traces: std::collections::HashSet = + std::collections::HashSet::new(); for p in &pts { if let Some(tid) = p["payload"]["trace_id"].as_i64() { if tid > 0 { @@ -181,9 +242,11 @@ async fn populate_face_detections_from_face_json( }; progress.stats.traced_faces = new_traced as i64; progress.stats.total_traces = traced_count; - progress.update_phase(TkgPhase::FaceTracing, 1.0, &format!( - "Face tracing complete: {} traces", traced_count - )); + progress.update_phase( + TkgPhase::FaceTracing, + 1.0, + &format!("Face tracing complete: {} traces", traced_count), + ); if let Some(r) = redis.as_ref() { publish_tkg_progress(&r, file_uuid, &progress).await; } @@ -590,7 +653,7 @@ pub async fn build_tkg( output_dir: &str, redis: Option>, ) -> Result { - use crate::core::progress::{TkgPhase, TkgProgress, publish_tkg_progress}; + use crate::core::progress::{publish_tkg_progress, TkgPhase, TkgProgress}; let pool = db.pool(); let mut progress = TkgProgress::new(file_uuid); @@ -604,16 +667,22 @@ pub async fn build_tkg( // Phase 0: Populate trace_id from face.json progress.update_phase(TkgPhase::FaceTracing, 0.0, "Tracing faces..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } - if let Err(e) = populate_face_detections_from_face_json(pool, output_dir, file_uuid, redis.clone()).await { + if let Err(e) = + populate_face_detections_from_face_json(pool, output_dir, file_uuid, redis.clone()).await + { tracing::warn!( "[TKG-Phase0] populate_face_detections failed: {} (continuing)", e ); } progress.update_phase(TkgPhase::FaceTracing, 1.0, "Face tracing complete"); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let pose_data = load_face_pose_data(output_dir, file_uuid) .map_err(|e| { @@ -627,126 +696,292 @@ pub async fn build_tkg( output_dir ); - // ── Node builders ───────────────────────────────────────────────── - progress.update_phase(TkgPhase::FaceTrackNodes, 0.0, "Building face track nodes..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } - let n_face = build_face_track_nodes(pool, file_uuid, &pose_data).await?; - progress.stats.face_track_nodes = n_face as i64; - progress.update_phase(TkgPhase::FaceTrackNodes, 1.0, &format!("{} face tracks", n_face)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + // Fetch face points from Qdrant ONCE and share across all builders + let face_points = scroll_face_points(file_uuid).await; + tracing::info!( + "[TKG] Fetched {} face points from Qdrant for {}", + face_points.len(), + file_uuid + ); + let frame_index = build_frame_index(&face_points); - progress.update_phase(TkgPhase::GazeTrackNodes, 0.0, "Building gaze track nodes..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } - let n_gaze = build_gaze_track_nodes(pool, file_uuid, &pose_data).await?; + // ── Node builders ───────────────────────────────────────────────── + progress.update_phase( + TkgPhase::FaceTrackNodes, + 0.0, + "Building face track nodes...", + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } + let n_face = build_face_track_nodes(pool, file_uuid, &pose_data, &face_points).await?; + progress.stats.face_track_nodes = n_face as i64; + progress.update_phase( + TkgPhase::FaceTrackNodes, + 1.0, + &format!("{} face tracks", n_face), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } + + progress.update_phase( + TkgPhase::GazeTrackNodes, + 0.0, + "Building gaze track nodes...", + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } + let n_gaze = build_gaze_track_nodes(pool, file_uuid, &pose_data, &face_points).await?; progress.stats.gaze_track_nodes = n_gaze as i64; - progress.update_phase(TkgPhase::GazeTrackNodes, 1.0, &format!("{} gaze tracks", n_gaze)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::GazeTrackNodes, + 1.0, + &format!("{} gaze tracks", n_gaze), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } progress.update_phase(TkgPhase::LipTrackNodes, 0.0, "Building lip track nodes..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } - let n_lip = build_lip_track_nodes(pool, file_uuid, output_dir, &pose_data).await?; + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } + let n_lip = build_lip_track_nodes(pool, file_uuid, output_dir, &pose_data, &frame_index).await?; progress.stats.lip_track_nodes = n_lip as i64; - progress.update_phase(TkgPhase::LipTrackNodes, 1.0, &format!("{} lip tracks", n_lip)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::LipTrackNodes, + 1.0, + &format!("{} lip tracks", n_lip), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } - progress.update_phase(TkgPhase::TextRegionNodes, 0.0, "Building text region nodes..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::TextRegionNodes, + 0.0, + "Building text region nodes...", + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let n_text = build_text_region_nodes(pool, file_uuid).await?; progress.stats.text_region_nodes = n_text as i64; - progress.update_phase(TkgPhase::TextRegionNodes, 1.0, &format!("{} text regions", n_text)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::TextRegionNodes, + 1.0, + &format!("{} text regions", n_text), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } - progress.update_phase(TkgPhase::AppearanceNodes, 0.0, "Building appearance nodes..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } - let n_appearance = build_appearance_trace_nodes(pool, file_uuid, output_dir, &pose_data).await?; + progress.update_phase( + TkgPhase::AppearanceNodes, + 0.0, + "Building appearance nodes...", + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } + let n_appearance = + build_appearance_trace_nodes(pool, file_uuid, output_dir, &pose_data, &frame_index).await?; progress.stats.appearance_nodes = n_appearance as i64; - progress.update_phase(TkgPhase::AppearanceNodes, 1.0, &format!("{} appearances", n_appearance)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::AppearanceNodes, + 1.0, + &format!("{} appearances", n_appearance), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } progress.update_phase(TkgPhase::AccessoryNodes, 0.0, "Building accessory nodes..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let n_accessories = build_accessory_nodes(pool, file_uuid, output_dir).await?; progress.stats.accessory_nodes = n_accessories as i64; - progress.update_phase(TkgPhase::AccessoryNodes, 1.0, &format!("{} accessories", n_accessories)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::AccessoryNodes, + 1.0, + &format!("{} accessories", n_accessories), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } progress.update_phase(TkgPhase::ObjectNodes, 0.0, "Building object nodes..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let n_objects = build_yolo_object_nodes(pool, file_uuid, output_dir).await?; progress.stats.object_nodes = n_objects as i64; - progress.update_phase(TkgPhase::ObjectNodes, 1.0, &format!("{} objects", n_objects)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::ObjectNodes, + 1.0, + &format!("{} objects", n_objects), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } progress.update_phase(TkgPhase::HandNodes, 0.0, "Building hand nodes..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let n_hands = build_hand_nodes(pool, file_uuid, output_dir).await?; progress.stats.hand_nodes = n_hands as i64; progress.update_phase(TkgPhase::HandNodes, 1.0, &format!("{} hands", n_hands)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } progress.update_phase(TkgPhase::SpeakerNodes, 0.0, "Building speaker nodes..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let n_speakers = build_speaker_nodes(pool, file_uuid, output_dir).await?; progress.stats.speaker_nodes = n_speakers as i64; - progress.update_phase(TkgPhase::SpeakerNodes, 1.0, &format!("{} speakers", n_speakers)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::SpeakerNodes, + 1.0, + &format!("{} speakers", n_speakers), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } // ── Edge builders ───────────────────────────────────────────────── - progress.update_phase(TkgPhase::CoOccurrenceEdges, 0.0, "Building co-occurrence edges..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::CoOccurrenceEdges, + 0.0, + "Building co-occurrence edges...", + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let e_co = build_co_occurrence_edges(pool, file_uuid, output_dir).await?; progress.stats.co_occurrence_edges = e_co as i64; - progress.update_phase(TkgPhase::CoOccurrenceEdges, 1.0, &format!("{} co-occurrence", e_co)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::CoOccurrenceEdges, + 1.0, + &format!("{} co-occurrence", e_co), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } - progress.update_phase(TkgPhase::SpeakerFaceEdges, 0.0, "Building speaker-face edges..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::SpeakerFaceEdges, + 0.0, + "Building speaker-face edges...", + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let e_sf = build_speaker_face_edges(pool, file_uuid, output_dir).await?; progress.stats.speaker_face_edges = e_sf as i64; - progress.update_phase(TkgPhase::SpeakerFaceEdges, 1.0, &format!("{} speaker-face", e_sf)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::SpeakerFaceEdges, + 1.0, + &format!("{} speaker-face", e_sf), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } progress.update_phase(TkgPhase::FaceFaceEdges, 0.0, "Building face-face edges..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } - let e_ff = build_face_face_edges(pool, file_uuid, &pose_data).await?; + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } + let e_ff = build_face_face_edges(pool, file_uuid, &pose_data, &face_points).await?; progress.stats.face_face_edges = e_ff as i64; progress.update_phase(TkgPhase::FaceFaceEdges, 1.0, &format!("{} face-face", e_ff)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } - progress.update_phase(TkgPhase::MutualGazeEdges, 0.0, "Building mutual gaze edges..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } - let e_mg = build_mutual_gaze_edges(pool, file_uuid, &pose_data).await?; + progress.update_phase( + TkgPhase::MutualGazeEdges, + 0.0, + "Building mutual gaze edges...", + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } + let e_mg = build_mutual_gaze_edges(pool, file_uuid, &pose_data, &face_points).await?; progress.stats.mutual_gaze_edges = e_mg as i64; - progress.update_phase(TkgPhase::MutualGazeEdges, 1.0, &format!("{} mutual gaze", e_mg)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::MutualGazeEdges, + 1.0, + &format!("{} mutual gaze", e_mg), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } progress.update_phase(TkgPhase::LipSyncEdges, 0.0, "Building lip sync edges..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let e_ls = build_lip_sync_edges(pool, file_uuid, output_dir, &pose_data).await?; progress.stats.lip_sync_edges = e_ls as i64; progress.update_phase(TkgPhase::LipSyncEdges, 1.0, &format!("{} lip sync", e_ls)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } - progress.update_phase(TkgPhase::HasAppearanceEdges, 0.0, "Building has-appearance edges..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::HasAppearanceEdges, + 0.0, + "Building has-appearance edges...", + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let e_ha = build_has_appearance_edges(pool, file_uuid).await?; progress.stats.has_appearance_edges = e_ha as i64; - progress.update_phase(TkgPhase::HasAppearanceEdges, 1.0, &format!("{} has-appearance", e_ha)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::HasAppearanceEdges, + 1.0, + &format!("{} has-appearance", e_ha), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } progress.update_phase(TkgPhase::WearsEdges, 0.0, "Building wears edges..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let e_w = build_wears_edges(pool, file_uuid).await?; progress.stats.wears_edges = e_w as i64; progress.update_phase(TkgPhase::WearsEdges, 1.0, &format!("{} wears", e_w)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } - progress.update_phase(TkgPhase::HandObjectEdges, 0.0, "Building hand-object edges..."); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::HandObjectEdges, + 0.0, + "Building hand-object edges...", + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } let e_ho = build_hand_object_edges(pool, file_uuid, output_dir).await?; progress.stats.hand_object_edges = e_ho as i64; - progress.update_phase(TkgPhase::HandObjectEdges, 1.0, &format!("{} hand-object", e_ho)); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + progress.update_phase( + TkgPhase::HandObjectEdges, + 1.0, + &format!("{} hand-object", e_ho), + ); + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } // Calculate totals progress.stats.total_nodes = progress.stats.face_track_nodes @@ -769,9 +1004,15 @@ pub async fn build_tkg( + progress.stats.hand_object_edges; progress.mark_completed(); - if let Some(r) = redis.as_ref() { publish_tkg_progress(r, file_uuid, &progress).await; } + if let Some(r) = redis.as_ref() { + publish_tkg_progress(r, file_uuid, &progress).await; + } - tracing::info!("[TKG] Complete: {} nodes, {} edges", progress.stats.total_nodes, progress.stats.total_edges); + tracing::info!( + "[TKG] Complete: {} nodes, {} edges", + progress.stats.total_nodes, + progress.stats.total_edges + ); Ok(TkgResult { face_track_nodes: n_face, @@ -800,19 +1041,21 @@ async fn build_face_track_nodes( pool: &PgPool, file_uuid: &str, pose_data: &[FacePose], + face_points: &[FacePoint], ) -> Result { - build_face_track_nodes_from_qdrant(pool, file_uuid, pose_data).await + build_face_track_nodes_from_qdrant(pool, file_uuid, pose_data, face_points).await } async fn build_face_track_nodes_from_qdrant( pool: &PgPool, file_uuid: &str, pose_data: &[FacePose], + face_points: &[FacePoint], ) -> Result { let nodes_table = t("tkg_nodes"); - // Get face points from Qdrant - let points = scroll_face_points(file_uuid).await; + // Use shared face points (passed from build_tkg) + let points = face_points; if points.is_empty() { tracing::info!("[TKG] No traced faces in Qdrant for {}", file_uuid); return Ok(0); @@ -832,7 +1075,7 @@ async fn build_face_track_nodes_from_qdrant( } let mut traces: HashMap = HashMap::new(); - for p in &points { + for p in points { let entry = traces.entry(p.trace_id).or_insert_with(|| TraceAgg { frame_count: 0, start_f: i64::MAX, @@ -865,9 +1108,7 @@ async fn build_face_track_nodes_from_qdrant( let mut pose_count = 0i64; for (frame, x, y, w, h) in &agg.frames { - if let Some((yaw, pitch, roll)) = - get_pose_for_face(*frame, *x, *y, *w, *h, pose_data) - { + if let Some((yaw, pitch, roll)) = get_pose_for_face(*frame, *x, *y, *w, *h, pose_data) { yaw_sum += yaw; pitch_sum += pitch; roll_sum += roll; @@ -1390,7 +1631,9 @@ async fn build_speaker_face_edges_from_pg( let points = scroll_face_points(file_uuid).await; let mut trace_ranges: HashMap = HashMap::new(); for p in &points { - let entry = trace_ranges.entry(p.trace_id).or_insert((i64::MAX, i64::MIN)); + let entry = trace_ranges + .entry(p.trace_id) + .or_insert((i64::MAX, i64::MIN)); entry.0 = entry.0.min(p.frame); entry.1 = entry.1.max(p.frame); } @@ -1501,27 +1744,29 @@ async fn build_face_face_edges( pool: &PgPool, file_uuid: &str, pose_data: &[FacePose], + face_points: &[FacePoint], ) -> Result { - build_face_face_edges_from_qdrant(pool, file_uuid, pose_data).await + build_face_face_edges_from_qdrant(pool, file_uuid, pose_data, face_points).await } async fn build_face_face_edges_from_qdrant( pool: &PgPool, file_uuid: &str, pose_data: &[FacePose], + face_points: &[FacePoint], ) -> Result { let nodes_table = t("tkg_nodes"); let edges_table = t("tkg_edges"); - // Get face points from Qdrant - let points = scroll_face_points(file_uuid).await; + // Use shared face points (passed from build_tkg) + let points = face_points; if points.is_empty() { return Ok(0); } // Group by frame: frame → Vec<(trace_id, x, y, w, h)> let mut frame_traces: HashMap> = HashMap::new(); - for p in &points { + for p in points { frame_traces .entry(p.frame) .or_default() @@ -1538,7 +1783,11 @@ async fn build_face_face_edges_from_qdrant( for j in (i + 1)..traces.len() { let (tid_a, _, _, _, _) = traces[i]; let (tid_b, _, _, _, _) = traces[j]; - let key = if tid_a < tid_b { (tid_a, tid_b) } else { (tid_b, tid_a) }; + let key = if tid_a < tid_b { + (tid_a, tid_b) + } else { + (tid_b, tid_a) + }; cooccurring_pairs.entry(key).or_default().push(*frame); } } @@ -1546,7 +1795,7 @@ async fn build_face_face_edges_from_qdrant( // Build frame_map from Qdrant points let mut frame_map: HashMap<(i64, i64), (f64, f64, f64, f64)> = HashMap::new(); - for p in &points { + for p in points { frame_map.insert((p.trace_id, p.frame), (p.x, p.y, p.w, p.h)); } @@ -1561,15 +1810,20 @@ async fn build_face_face_edges_from_qdrant( (Some(&(xa, ya, wa, ha)), Some(&(xb, yb, wb, hb))) => { get_pose_for_face(*frame, xa, ya, wa, ha, pose_data) .and_then(|(yaw_a, _, _)| { - get_pose_for_face(*frame, xb, yb, wb, hb, pose_data).map(|(yaw_b, _, _)| { - detect_mutual_gaze(xa, wa, yaw_a, xb, wb, yaw_b, 0.05) - }) + get_pose_for_face(*frame, xb, yb, wb, hb, pose_data).map( + |(yaw_b, _, _)| { + detect_mutual_gaze(xa, wa, yaw_a, xb, wb, yaw_b, 0.05) + }, + ) }) .unwrap_or(false) } _ => false, }; - pair_frames.entry((*tid_a, *tid_b)).or_default().push((*frame, gaze)); + pair_frames + .entry((*tid_a, *tid_b)) + .or_default() + .push((*frame, gaze)); } } @@ -1689,143 +1943,29 @@ async fn build_gaze_track_nodes( pool: &PgPool, file_uuid: &str, pose_data: &[FacePose], + face_points: &[FacePoint], ) -> Result { - // Try face.json first - let count = build_gaze_track_nodes_from_face_json(pool, file_uuid, pose_data).await?; - if count > 0 { - return Ok(count); - } - - // Fallback to Qdrant - tracing::info!("[TKG-Phase2.5] No face.json gaze data, falling back to Qdrant"); - build_gaze_track_nodes_from_qdrant(pool, file_uuid, pose_data).await -} - -async fn build_gaze_track_nodes_from_face_json( - pool: &PgPool, - file_uuid: &str, - pose_data: &[FacePose], -) -> Result { - let face_json_path = - Path::new(&*crate::core::config::OUTPUT_DIR).join(format!("{}.face.json", file_uuid)); - - if !face_json_path.exists() { - tracing::info!("[TKG-Phase2.5] No face.json for gaze_track"); - return Ok(0); - } - - let content = std::fs::read_to_string(&face_json_path)?; - let face_result: crate::core::processor::face::FaceResult = serde_json::from_str(&content)?; - - // Group faces by trace_id (assuming trace_id = 1 for all faces in face.json) - let mut frames_data: Vec<(u64, f64, f64, f64, f64)> = vec![]; - - for frame in &face_result.frames { - for face in &frame.faces { - frames_data.push(( - frame.frame, - face.x as f64, - face.y as f64, - face.width as f64, - face.height as f64, - )); - } - } - - if frames_data.is_empty() { - return Ok(0); - } - - // Compute gaze stats for trace_id = 1 - let trace_id = 1_i64; - let external_id = format!("gaze_{}", trace_id); - - let mut frame_count = 0i64; - let mut first_frame = i64::MAX; - let mut last_frame = i64::MIN; - let mut yaw_sum = 0.0f64; - let mut pitch_sum = 0.0f64; - let mut roll_sum = 0.0f64; - let mut gaze_dir_counts: HashMap<&str, i64> = HashMap::new(); - - for (frame, x, y, w, h) in &frames_data { - if let Some((yaw, pitch, roll)) = - get_pose_for_face(*frame as i64, *x, *y, *w, *h, pose_data) - { - frame_count += 1; - first_frame = first_frame.min(*frame as i64); - last_frame = last_frame.max(*frame as i64); - yaw_sum += yaw; - pitch_sum += pitch; - roll_sum += roll; - - let gaze_dir = GazeDirection::from_yaw_pitch(yaw, pitch); - *gaze_dir_counts.entry(gaze_dir.as_str()).or_default() += 1; - } - } - - if frame_count == 0 { - return Ok(0); - } - - let avg_yaw = yaw_sum / frame_count as f64; - let avg_pitch = pitch_sum / frame_count as f64; - let avg_roll = roll_sum / frame_count as f64; - let dominant_gaze = gaze_dir_counts - .iter() - .max_by_key(|(_, &c)| c) - .map(|(&d, _)| d) - .unwrap_or("unknown"); - let (gaze_dx, gaze_dy) = compute_gaze_vector(avg_yaw, avg_pitch); - - let props = serde_json::json!({ - "trace_id": trace_id, - "frame_count": frame_count, - "start_frame": first_frame, - "end_frame": last_frame, - "avg_yaw": (avg_yaw * 1000.0).round() / 1000.0, - "avg_pitch": (avg_pitch * 1000.0).round() / 1000.0, - "avg_roll": (avg_roll * 1000.0).round() / 1000.0, - "dominant_gaze": dominant_gaze, - "gaze_dx": (gaze_dx * 1000.0).round() / 1000.0, - "gaze_dy": (gaze_dy * 1000.0).round() / 1000.0, - }); - - let nodes_table = t("tkg_nodes"); - sqlx::query(&format!( - "INSERT INTO {} (file_uuid, external_id, label, node_type, properties, created_at) \ - VALUES ($1, $2, $3, 'gaze_track', $4, NOW()) \ - ON CONFLICT (file_uuid, node_type, external_id) \ - DO UPDATE SET properties = COALESCE(EXCLUDED.properties, tkg_nodes.properties)", - nodes_table - )) - .bind(file_uuid) - .bind(&external_id) - .bind(&format!("Gaze Trace {}", trace_id)) - .bind(&props) - .execute(pool) - .await?; - - tracing::info!("[TKG-Phase2.5] Built {} gaze_track node from face.json", 1); - Ok(1) + // Use Qdrant for proper trace_id grouping (face.json path hardcodes trace_id=1 for all faces) + build_gaze_track_nodes_from_qdrant(pool, file_uuid, pose_data, face_points).await } async fn build_gaze_track_nodes_from_qdrant( pool: &PgPool, file_uuid: &str, pose_data: &[FacePose], + face_points: &[FacePoint], ) -> Result { let nodes_table = t("tkg_nodes"); - // Get face points from Qdrant - let points = scroll_face_points(file_uuid).await; + // Use shared face points (passed from build_tkg) + let points = face_points; if points.is_empty() { return Ok(0); } // Group by trace_id let mut trace_frames: HashMap> = HashMap::new(); - for p in &points { + for p in points { trace_frames .entry(p.trace_id) .or_default() @@ -1943,19 +2083,20 @@ async fn build_mutual_gaze_edges( pool: &PgPool, file_uuid: &str, pose_data: &[FacePose], + face_points: &[FacePoint], ) -> Result { let nodes_table = t("tkg_nodes"); let edges_table = t("tkg_edges"); - // Get face points from Qdrant - let points = scroll_face_points(file_uuid).await; + // Use shared face points (passed from build_tkg) + let points = face_points; if points.is_empty() { return Ok(0); } // Group by frame: frame → Vec<(trace_id, x, y, w, h)> let mut frame_traces: HashMap> = HashMap::new(); - for p in &points { + for p in points { frame_traces .entry(p.frame) .or_default() @@ -2101,127 +2242,9 @@ async fn build_lip_track_nodes( file_uuid: &str, output_dir: &str, pose_data: &[FacePose], + frame_index: &HashMap>, ) -> Result { - // Try face.json first - let count = build_lip_track_nodes_from_face_json(pool, file_uuid, pose_data).await?; - if count > 0 { - return Ok(count); - } - - // Fallback to PostgreSQL - tracing::info!("[TKG-Phase2.5] No face.json lip data, falling back to PostgreSQL"); - build_lip_track_nodes_from_pg(pool, file_uuid, output_dir, pose_data).await -} - -async fn build_lip_track_nodes_from_face_json( - pool: &PgPool, - file_uuid: &str, - pose_data: &[FacePose], -) -> Result { - let face_json_path = - Path::new(&*crate::core::config::OUTPUT_DIR).join(format!("{}.face.json", file_uuid)); - - if !face_json_path.exists() { - tracing::info!("[TKG-Phase2.5] No face.json for lip_track"); - return Ok(0); - } - - let content = std::fs::read_to_string(&face_json_path)?; - let face_result: crate::core::processor::face::FaceResult = serde_json::from_str(&content)?; - - // Group faces by trace_id (assuming trace_id = 1 for all faces in face.json) - let mut frames_data: Vec<(u64, f64, f64, f64, f64, Option)> = vec![]; - - for frame in &face_result.frames { - for face in &frame.faces { - frames_data.push(( - frame.frame, - face.x as f64, - face.y as f64, - face.width as f64, - face.height as f64, - face.landmarks.clone(), - )); - } - } - - if frames_data.is_empty() { - return Ok(0); - } - - // Compute lip stats for trace_id = 1 - let trace_id = 1_i64; - let external_id = format!("lip_{}", trace_id); - - let mut frame_count = 0i64; - let mut first_frame = i64::MAX; - let mut last_frame = i64::MIN; - let mut lip_area_sum = 0.0f64; - let mut lip_openness_sum = 0.0f64; - let mut speaking_frames = 0i64; - - for (frame, x, y, w, h, landmarks) in &frames_data { - if let Some((yaw, pitch, roll)) = - get_pose_for_face(*frame as i64, *x, *y, *w, *h, pose_data) - { - frame_count += 1; - first_frame = first_frame.min(*frame as i64); - last_frame = last_frame.max(*frame as i64); - - // Compute lip area and openness from landmarks - let lip_area = compute_lip_area(landmarks.as_ref()); - let lip_openness = if lip_area > 0.0 { - lip_area / (w * h) - } else { - 0.0 - }; - - lip_area_sum += lip_area; - lip_openness_sum += lip_openness; - - // Speaking detection (lip openness > threshold) - if lip_openness > 0.02 { - speaking_frames += 1; - } - } - } - - if frame_count == 0 { - return Ok(0); - } - - let avg_lip_area = lip_area_sum / frame_count as f64; - let avg_lip_openness = lip_openness_sum / frame_count as f64; - let speaking_ratio = speaking_frames as f64 / frame_count as f64; - - let props = serde_json::json!({ - "trace_id": trace_id, - "frame_count": frame_count, - "start_frame": first_frame, - "end_frame": last_frame, - "avg_lip_area": (avg_lip_area * 1000.0).round() / 1000.0, - "avg_lip_openness": (avg_lip_openness * 1000.0).round() / 1000.0, - "speaking_frames": speaking_frames, - "speaking_ratio": (speaking_ratio * 100.0).round() / 100.0, - }); - - let nodes_table = t("tkg_nodes"); - sqlx::query(&format!( - "INSERT INTO {} (file_uuid, external_id, label, node_type, properties, created_at) \ - VALUES ($1, $2, $3, 'lip_track', $4, NOW()) \ - ON CONFLICT (file_uuid, node_type, external_id) \ - DO UPDATE SET properties = COALESCE(EXCLUDED.properties, tkg_nodes.properties)", - nodes_table - )) - .bind(file_uuid) - .bind(&external_id) - .bind(&format!("Lip Trace {}", trace_id)) - .bind(&props) - .execute(pool) - .await?; - - tracing::info!("[TKG-Phase2.5] Built {} lip_track node from face.json", 1); - Ok(1) + build_lip_track_nodes_from_pg(pool, file_uuid, output_dir, pose_data, frame_index).await } async fn build_lip_track_nodes_from_pg( @@ -2229,6 +2252,7 @@ async fn build_lip_track_nodes_from_pg( file_uuid: &str, output_dir: &str, pose_data: &[FacePose], + frame_index: &HashMap>, ) -> Result { let nodes_table = t("tkg_nodes"); @@ -2259,12 +2283,11 @@ async fn build_lip_track_nodes_from_pg( let w = face.get("width").and_then(|v| v.as_f64()).unwrap_or(0.0); let h = face.get("height").and_then(|v| v.as_f64()).unwrap_or(0.0); - // Get trace_id for this face - let trace_id = - match get_trace_for_face(pool, file_uuid, frame_num, x, y, w, h).await { - Some(tid) => tid, - None => continue, - }; + // Get trace_id for this face via in-memory lookup (no HTTP call) + let trace_id = match find_trace_in_frame(frame_index, frame_num, x, y, w, h) { + Some(tid) => tid, + None => continue, + }; // Extract lip landmarks let lips = face.get("lips"); @@ -2413,63 +2436,6 @@ fn compute_lip_area(landmarks: Option<&serde_json::Value>) -> f64 { } } -async fn get_trace_for_face( - _pool: &PgPool, - file_uuid: &str, - frame: i64, - x: f64, - y: f64, - w: f64, - h: f64, -) -> Option { - use crate::core::db::qdrant_db::QdrantDb; - use serde_json::json; - - // Get face points from Qdrant for this frame - let qdrant = QdrantDb::new(); - let filter = json!({ - "must": [ - {"key": "file_uuid", "match": {"value": file_uuid}}, - {"key": "frame", "match": {"value": frame}}, - {"key": "trace_id", "match": {"value": 1}} - ] - }); - - let points = match qdrant.scroll_points("_faces", filter, 100, None).await { - Ok((pts, _)) => pts, - Err(_) => return None, - }; - - if points.is_empty() { - return None; - } - - // Find closest by bbox center distance - let mut best: Option<(i64, f64)> = None; - let tcx = x + w / 2.0; - let tcy = y + h / 2.0; - - for point in &points { - let payload = &point["payload"]; - let tid = payload["trace_id"].as_i64().filter(|&t| t > 0)?; - let bbox = &payload["bbox"]; - let fx = bbox["x"].as_f64().unwrap_or(0.0); - let fy = bbox["y"].as_f64().unwrap_or(0.0); - let fw = bbox["width"].as_f64().unwrap_or(0.0); - let fh = bbox["height"].as_f64().unwrap_or(0.0); - - let cx = fx + fw / 2.0; - let cy = fy + fh / 2.0; - let dist = ((cx - tcx).powi(2) + (cy - tcy).powi(2)).sqrt(); - match best { - Some((_, best_dist)) if dist < best_dist => best = Some((tid, dist)), - None => best = Some((tid, dist)), - _ => {} - } - } - best.map(|(tid, _)| tid) -} - // ── Text/Sentence Trace Nodes ───────────────────────────────────── async fn build_text_region_nodes(pool: &PgPool, file_uuid: &str) -> Result { @@ -2693,6 +2659,7 @@ async fn build_appearance_trace_nodes( file_uuid: &str, output_dir: &str, pose_data: &[FacePose], + frame_index: &HashMap>, ) -> Result { let path = Path::new(output_dir).join(format!("{}.appearance.json", file_uuid)); if !path.exists() { @@ -2707,8 +2674,8 @@ async fn build_appearance_trace_nodes( let mut count = 0; if let Some(frames) = data.get("frames").and_then(|v| v.as_array()) { - // Group by person_id across frames - let mut person_data: HashMap> = HashMap::new(); + // Group by face trace_id (appearance.json has no person_id, match via nose keypoint → Qdrant) + let mut trace_data: HashMap> = HashMap::new(); for frame_entry in frames { let frame_num = frame_entry @@ -2717,37 +2684,39 @@ async fn build_appearance_trace_nodes( .unwrap_or(0); if let Some(persons) = frame_entry.get("persons").and_then(|v| v.as_array()) { for person in persons { - let pid = person - .get("person_id") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - person_data - .entry(pid) - .or_default() - .push((frame_num, person.clone())); + // Get nose keypoint as face position proxy + let nose = person + .get("keypoints") + .and_then(|kps| kps.as_array()) + .and_then(|kps| { + kps.iter() + .find(|kp| kp.get("name").and_then(|n| n.as_str()) == Some("nose")) + }); + if let Some(nose) = nose { + let nx = nose.get("x").and_then(|v| v.as_f64()).unwrap_or(0.0); + let ny = nose.get("y").and_then(|v| v.as_f64()).unwrap_or(0.0); + // Find closest face trace in Qdrant for this frame (in-memory lookup) + if let Some(trace_id) = + find_trace_in_frame(frame_index, frame_num, nx, ny, 1.0, 1.0) + { + trace_data + .entry(trace_id) + .or_default() + .push((frame_num, person.clone())); + } + } } } } - for (pid, frames) in &person_data { - let external_id = format!("appearance_{}", pid); + for (trace_id, frames) in &trace_data { + let external_id = format!("appearance_{}", trace_id); let frame_count = frames.len() as i64; let first_frame = frames.iter().map(|(f, _)| *f).min().unwrap_or(0); let last_frame = frames.iter().map(|(f, _)| *f).max().unwrap_or(0); - // Extract color features from first frame + // Get bbox info from first frame let first_person = &frames[0].1; - let dominant_colors = first_person - .get("dominant_colors") - .cloned() - .unwrap_or(serde_json::json!([])); - let body_parts = first_person - .get("body_parts") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - - // Get bbox info let bbox = first_person.get("bbox"); let bbox_x = bbox .and_then(|b| b.get("x").and_then(|v| v.as_i64())) @@ -2762,32 +2731,11 @@ async fn build_appearance_trace_nodes( .and_then(|b| b.get("height").and_then(|v| v.as_i64())) .unwrap_or(0); - // Match to face trace_id via bbox overlap - let matched_trace = match_trace_by_bbox( - pool, - file_uuid, - first_frame, - bbox_x as f64, - bbox_y as f64, - bbox_w as f64, - bbox_h as f64, - ) - .await; - - let body_part_names: Vec = body_parts - .iter() - .filter_map(|bp| bp.get("name").and_then(|n| n.as_str().map(String::from))) - .collect(); - let props = serde_json::json!({ - "person_id": pid, - "trace_id": matched_trace, + "trace_id": trace_id, "frame_count": frame_count, "start_frame": first_frame, "end_frame": last_frame, - "dominant_colors": dominant_colors, - "body_parts": body_parts, - "body_part_names": body_part_names, "avg_bbox": { "x": bbox_x, "y": bbox_y, @@ -2800,16 +2748,16 @@ async fn build_appearance_trace_nodes( r#" INSERT INTO {} (node_type, external_id, file_uuid, label, properties) VALUES ($1, $2, $3, $4, $5::jsonb) -ON CONFLICT (file_uuid, node_type, external_id) - DO UPDATE SET - properties = EXCLUDED.properties - "#, + ON CONFLICT (file_uuid, node_type, external_id) + DO UPDATE SET + properties = EXCLUDED.properties + "#, nodes_table )) .bind("appearance_trace") .bind(&external_id) .bind(file_uuid) - .bind(&format!("Appearance Trace {}", pid)) + .bind(&format!("Appearance Trace {}", trace_id)) .bind(serde_json::to_string(&props)?) .execute(pool) .await?; @@ -3117,18 +3065,6 @@ async fn build_hand_object_edges( Ok(count) } -async fn match_trace_by_bbox( - pool: &PgPool, - file_uuid: &str, - frame: i64, - x: f64, - y: f64, - w: f64, - h: f64, -) -> Option { - get_trace_for_face(pool, file_uuid, frame, x, y, w, h).await -} - // ── TKG Bridge: Representative Frame ────────────────────────────── #[derive(Debug, Serialize)] @@ -3190,7 +3126,10 @@ pub async fn query_auto_representative_frame( {"key": "identity_id", "exists": true} ] }); - let points = qdrant.scroll_all_points("_faces", face_filter, 500).await.unwrap_or_default(); + let points = qdrant + .scroll_all_points("_faces", face_filter, 500) + .await + .unwrap_or_default(); // Count appearances per identity_id let mut identity_counts: HashMap = HashMap::new(); @@ -3245,7 +3184,10 @@ pub async fn query_auto_representative_frame( {"key": "identity_id", "exists": true} ] }); - let points = qdrant.scroll_all_points("_faces", face_filter, 100).await.unwrap_or_default(); + let points = qdrant + .scroll_all_points("_faces", face_filter, 100) + .await + .unwrap_or_default(); let frame_number = frame_number.unwrap_or_else(|| { if !points.is_empty() {