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

feat: trace-level matching, health watcher/worker status, timezone config

Browse Source
This commit is contained in:
Accusys
2026-05-21 01:08:30 +08:00
parent 8ede4be159
commit bebaa743ed
60 changed files with 6110 additions and 1586 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core/processor/cut.rs
+188 -28
View File
@@ -1,5 +1,6 @@
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::process::Command;
use std::time::Duration;
use super::executor::PythonExecutor;
@@ -27,13 +28,21 @@ pub async fn process_cut(
output_path: &str,
uuid: Option<&str>,
) -> Result<CutResult> {
// Try native ffmpeg-based scene detection first
let result = try_native_cut(video_path);
if let Ok(r) = result {
let json = serde_json::to_string_pretty(&r)?;
std::fs::write(output_path, &json)
.with_context(|| format!("Failed to write {:?}", output_path))?;
return Ok(r);
}
// Fallback: Python scenedetect
tracing::warn!("[CUT] Native impl failed, falling back to Python");
let executor = PythonExecutor::new()?;
let script_path = executor.script_path("cut_processor.py");
tracing::info!("[CUT] Starting scene detection: {}", video_path);
if !script_path.exists() {
tracing::warn!("[CUT] Script not found, returning empty result");
return Ok(CutResult {
frame_count: 0,
fps: 0.0,
@@ -53,19 +62,179 @@ pub async fn process_cut(
.with_context(|| format!("Failed to run {:?}", script_path))?;
let json_str = std::fs::read_to_string(output_path).context("Failed to read CUT output")?;
let result: CutResult =
serde_json::from_str(&json_str).context("Failed to parse CUT output")?;
tracing::info!("[CUT] Result: {} scenes detected", result.scenes.len());
Ok(result)
}
// ── Native ffmpeg scene detection ─────────────────────────────────
fn try_native_cut(video_path: &str) -> Result<CutResult> {
// Step 1: Get video info (fps, frame count)
let probe = Command::new("ffprobe")
.args([
"-v",
"quiet",
"-print_format",
"json",
"-show_format",
"-show_streams",
video_path,
])
.output()
.context("Failed to run ffprobe")?;
let probe_info: serde_json::Value =
serde_json::from_slice(&probe.stdout).context("Failed to parse ffprobe output")?;
let streams = probe_info["streams"]
.as_array()
.map_or(vec![], |s| s.clone());
let video_stream = streams.iter().find(|s| s["codec_type"] == "video");
let fps = video_stream
.and_then(|s| s["r_frame_rate"].as_str().and_then(parse_fraction))
.unwrap_or(30.0);
let total_frames: u64 = video_stream
.and_then(|s| s["nb_frames"].as_str())
.and_then(|s| s.parse().ok())
.unwrap_or(0);
let duration: f64 = probe_info["format"]["duration"]
.as_str()
.and_then(|s| s.parse().ok())
.unwrap_or(0.0);
// Step 2: Use ffmpeg scene detection filter
// The `scene` filter computes the difference between consecutive frames
// and outputs when the difference exceeds the threshold (0.3 = medium sensitivity)
let scene_output = Command::new("ffprobe")
.args([
"-v",
"quiet",
"-show_entries",
"frame=pts_time",
"-of",
"compact=p=0:nk=1",
"-f",
"lavfi",
&format!("movie={},select='gt(scene\\,0.3)',showinfo", video_path),
"-show_frames",
])
.output()
.context("Failed to run ffmpeg scene detection")?;
let stderr_output = String::from_utf8_lossy(&scene_output.stderr);
let mut scene_times: Vec<f64> = Vec::new();
// Parse ffmpeg showinfo output for scene changes
// Format: [Parsed_showinfo...] pts:123.456 pts_time:123.456 ...
for line in stderr_output.lines() {
if line.contains("pts_time:") {
if let Some(pos) = line.find("pts_time:") {
let rest = &line[pos + 9..];
let time_str = rest.split_whitespace().next().unwrap_or("");
if let Ok(t) = time_str.parse::<f64>() {
scene_times.push(t);
}
}
}
}
// Step 3: Build scenes from cut points
let mut scenes: Vec<CutScene> = Vec::new();
let mut prev_time = 0.0;
let mut prev_frame: u64 = 0;
for (i, &cut_time) in scene_times.iter().enumerate() {
let end_frame = (cut_time * fps).round() as u64;
let start_frame = prev_frame;
if end_frame > start_frame {
scenes.push(CutScene {
scene_number: (i + 1) as u32,
start_frame: prev_frame,
end_frame: end_frame.saturating_sub(1),
start_time: prev_time,
end_time: cut_time - (1.0 / fps),
});
}
prev_time = cut_time;
prev_frame = end_frame;
}
// Final scene (last cut point → end of video)
if total_frames > 0 && prev_frame < total_frames {
scenes.push(CutScene {
scene_number: (scenes.len() + 1) as u32,
start_frame: prev_frame,
end_frame: total_frames.saturating_sub(1),
start_time: prev_time,
end_time: duration,
});
}
// If no scenes detected, create a single scene covering the whole video
if scenes.is_empty() && total_frames > 0 {
scenes.push(CutScene {
scene_number: 1,
start_frame: 0,
end_frame: total_frames.saturating_sub(1),
start_time: 0.0,
end_time: duration,
});
}
Ok(CutResult {
frame_count: total_frames,
fps,
scenes,
})
}
/// Parse fractional fps like "30000/1001" into f64
fn parse_fraction(s: &str) -> Option<f64> {
if let Some(pos) = s.find('/') {
let num: f64 = s[..pos].parse().ok()?;
let den: f64 = s[pos + 1..].parse().ok()?;
if den > 0.0 {
return Some(num / den);
}
}
s.parse::<f64>().ok()
}
// ── Tests ─────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_fraction() {
let r = parse_fraction("30000/1001").unwrap();
assert!((r - 29.97).abs() < 0.01);
}
#[test]
fn test_parse_fraction_int() {
let r = parse_fraction("30").unwrap();
assert!((r - 30.0).abs() < 0.01);
}
#[test]
fn test_parse_fraction_invalid() {
assert!(parse_fraction("not/a/num").is_none());
}
#[test]
fn test_parse_fraction_zero_den() {
assert!(parse_fraction("1/0").is_none());
}
#[test]
fn test_cut_result_serialization() {
let result = CutResult {
@@ -81,8 +250,9 @@ mod tests {
};
let json = serde_json::to_string(&result).unwrap();
assert!(json.contains("frame_count"));
assert!(json.contains("scene_number"));
assert!(json.contains("1"));
assert!(json.contains("fps"));
}
#[test]
@@ -90,20 +260,23 @@ mod tests {
let json = r#"{
"frame_count": 100,
"fps": 30.0,
"scenes": [
{"scene_number": 1, "start_frame": 0, "end_frame": 30, "start_time": 0.0, "end_time": 1.0},
{"scene_number": 2, "start_frame": 31, "end_frame": 60, "start_time": 1.033, "end_time": 2.0}
]
"scenes": [{
"scene_number": 1,
"start_frame": 0,
"end_frame": 30,
"start_time": 0.0,
"end_time": 1.0
}]
}"#;
let result: CutResult = serde_json::from_str(json).unwrap();
assert_eq!(result.frame_count, 100);
assert_eq!(result.scenes.len(), 2);
assert_eq!(result.scenes[1].scene_number, 2);
assert_eq!(result.scenes.len(), 1);
assert_eq!(result.scenes[0].scene_number, 1);
assert_eq!(result.scenes[0].start_frame, 0);
}
#[test]
fn test_cut_result_empty_scenes() {
fn test_empty_scenes() {
let result = CutResult {
frame_count: 0,
fps: 0.0,
@@ -111,17 +284,4 @@ mod tests {
};
assert!(result.scenes.is_empty());
}
#[test]
fn test_cut_scene_times() {
let scene = CutScene {
scene_number: 1,
start_frame: 0,
end_frame: 30,
start_time: 0.0,
end_time: 1.0,
};
assert!(scene.end_time > scene.start_time);
assert_eq!(scene.scene_number, 1);
}
}
src/core/processor/executor.rs
+13 -13
View File
@@ -109,11 +109,10 @@ pub fn validate_python_env() -> Result<()> {
tracing::warn!("Expected Python 3.11, got: {}", version.trim());
}
let scripts_dir = std::env::var("MOMENTRY_SCRIPTS_DIR")
.unwrap_or_else(|_| {
let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
manifest.join("scripts").to_string_lossy().to_string()
});
let scripts_dir = std::env::var("MOMENTRY_SCRIPTS_DIR").unwrap_or_else(|_| {
let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
manifest.join("scripts").to_string_lossy().to_string()
});
let script_path = PathBuf::from(&scripts_dir);
if !script_path.exists() {
anyhow::bail!("Scripts directory not found at {}", scripts_dir);
@@ -133,11 +132,10 @@ impl PythonExecutor {
pub fn new() -> Result<Self> {
let python_path = std::env::var("MOMENTRY_PYTHON_PATH")
.unwrap_or_else(|_| "/opt/homebrew/bin/python3.11".to_string());
let scripts_dir = std::env::var("MOMENTRY_SCRIPTS_DIR")
.unwrap_or_else(|_| {
let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
manifest.join("scripts").to_string_lossy().to_string()
});
let scripts_dir = std::env::var("MOMENTRY_SCRIPTS_DIR").unwrap_or_else(|_| {
let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
manifest.join("scripts").to_string_lossy().to_string()
});
let python_bin = PathBuf::from(&python_path);
let scripts_path = PathBuf::from(&scripts_dir);
@@ -173,7 +171,8 @@ impl PythonExecutor {
if let Some(expected_hash) = self.checksums.get(&rel_path) {
let output = std::process::Command::new("shasum")
.arg("-a").arg("256")
.arg("-a")
.arg("256")
.arg(&script_path)
.output()
.context("Failed to run shasum for integrity check")?;
@@ -235,8 +234,9 @@ impl PythonExecutor {
}
// Verify script integrity via SHA256 checksum before execution
self.verify_script_integrity(script_name)
.context("Pre-execution integrity check failed — possible version mismatch or corruption")?;
self.verify_script_integrity(script_name).context(
"Pre-execution integrity check failed — possible version mismatch or corruption",
)?;
// 標記輸出檔為處理中（add .tmp suffix）
let output_path = args.get(1).map(|p| std::path::PathBuf::from(p));
src/core/processor/heuristic_scene.rs
+68 -28
View File
@@ -44,22 +44,59 @@ pub enum CrowdSize {
/// Indoor-indicative YOLO classes (COCO labels)
const INDOOR_CLASSES: &[&str] = &[
"chair", "couch", "bed", "dining table", "toilet", "tv", "laptop",
"microwave", "oven", "refrigerator", "sink", "book", "clock",
"vase", "potted plant",
"chair",
"couch",
"bed",
"dining table",
"toilet",
"tv",
"laptop",
"microwave",
"oven",
"refrigerator",
"sink",
"book",
"clock",
"vase",
"potted plant",
];
/// Vehicle-indicative classes (person + vehicle = transport scene)
const VEHICLE_CLASSES: &[&str] = &[
"car", "truck", "bus", "train", "boat", "aeroplane", "bicycle", "motorbike",
"car",
"truck",
"bus",
"train",
"boat",
"aeroplane",
"bicycle",
"motorbike",
];
/// Outdoor-indicative YOLO classes
const OUTDOOR_CLASSES: &[&str] = &[
"car", "truck", "bus", "train", "boat", "airplane",
"traffic light", "fire hydrant", "stop sign", "parking meter",
"bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant",
"bear", "zebra", "giraffe", "tree",
"car",
"truck",
"bus",
"train",
"boat",
"airplane",
"traffic light",
"fire hydrant",
"stop sign",
"parking meter",
"bench",
"bird",
"cat",
"dog",
"horse",
"sheep",
"cow",
"elephant",
"bear",
"zebra",
"giraffe",
"tree",
];
/// Build heuristic scene metadata from disk files (yolo.json + DB face data).
@@ -113,13 +150,14 @@ pub async fn build_heuristic_scene_meta(
// Get face counts grouped by frame
let fd_table = schema::table_name("face_detections");
let face_rows: Vec<(i64, i64)> = sqlx::query_as(
&format!("SELECT frame_number, COUNT(*) as fc \
let face_rows: Vec<(i64, i64)> = sqlx::query_as(&format!(
"SELECT frame_number, COUNT(*) as fc \
FROM {} \
WHERE file_uuid = $1 AND frame_number IS NOT NULL \
GROUP BY frame_number \
ORDER BY frame_number", fd_table),
)
ORDER BY frame_number",
fd_table
))
.bind(file_uuid)
.fetch_all(pool)
.await
@@ -166,7 +204,10 @@ pub async fn build_heuristic_scene_meta(
let outdoor_ratio = outdoor_objects as f64 / frame_count.max(1) as f64;
let total_indicator = indoor_ratio + outdoor_ratio;
let (indoor_score, outdoor_score) = if total_indicator > 0.0 {
(indoor_ratio / total_indicator, outdoor_ratio / total_indicator)
(
indoor_ratio / total_indicator,
outdoor_ratio / total_indicator,
)
} else {
(0.5, 0.5)
};
@@ -187,17 +228,13 @@ pub async fn build_heuristic_scene_meta(
.map(|c| class_frame_presence.get(*c).copied().unwrap_or(0))
.sum();
let person_ratio = person_frames as f64 / frame_count.max(1) as f64;
let likely_vehicle = person_ratio > 0.5 && vehicle_frames > 0
&& outdoor_score > 0.3;
let likely_vehicle = person_ratio > 0.5 && vehicle_frames > 0 && outdoor_score > 0.3;
// Dominant objects: rank by frame presence (not total count)
let mut sorted: Vec<_> = class_frame_presence.into_iter().collect();
sorted.sort_by(|a, b| b.1.cmp(&a.1));
let dominant_objects: Vec<String> = sorted
.iter()
.take(3)
.map(|(cls, _)| cls.clone())
.collect();
let dominant_objects: Vec<String> =
sorted.iter().take(3).map(|(cls, _)| cls.clone()).collect();
segments.push(SceneSegmentMeta {
segment_index: idx as u32 + 1,
@@ -229,12 +266,15 @@ pub async fn build_heuristic_scene_meta(
/// Full pipeline entry point: reads CUT data, generates heuristic metadata, writes JSON.
/// Called from job_worker post-processing trigger.
pub async fn generate_scene_meta(db: &crate::core::db::PostgresDb, file_uuid: &str) -> Result<usize> {
pub async fn generate_scene_meta(
db: &crate::core::db::PostgresDb,
file_uuid: &str,
) -> Result<usize> {
let pool = db.pool();
// Read CUT segment boundaries from cut.json
let cut_path = Path::new(crate::core::config::OUTPUT_DIR.as_str())
.join(format!("{}.cut.json", file_uuid));
let cut_path =
Path::new(crate::core::config::OUTPUT_DIR.as_str()).join(format!("{}.cut.json", file_uuid));
let segments: Vec<(i64, i64, f64, f64)> = if cut_path.exists() {
let cut_str = tokio::fs::read_to_string(&cut_path)
.await
@@ -250,8 +290,7 @@ pub async fn generate_scene_meta(db: &crate::core::db::PostgresDb, file_uuid: &s
start_time: f64,
end_time: f64,
}
let cut: CutJson = serde_json::from_str(&cut_str)
.context("Failed to parse cut.json")?;
let cut: CutJson = serde_json::from_str(&cut_str).context("Failed to parse cut.json")?;
cut.scenes
.into_iter()
.map(|s| (s.start_frame, s.end_frame, s.start_time, s.end_time))
@@ -259,9 +298,10 @@ pub async fn generate_scene_meta(db: &crate::core::db::PostgresDb, file_uuid: &s
} else {
// Fallback: query DB for video duration, make one segment
let videos_table = schema::table_name("videos");
let (total_frames, duration): (Option<i64>, Option<f64>) = sqlx::query_as(
&format!("SELECT total_frames, duration FROM {} WHERE file_uuid = $1", videos_table),
)
let (total_frames, duration): (Option<i64>, Option<f64>) = sqlx::query_as(&format!(
"SELECT total_frames, duration FROM {} WHERE file_uuid = $1",
videos_table
))
.bind(file_uuid)
.fetch_optional(pool)
.await
src/core/processor/mod.rs
+4 -1
View File
@@ -10,6 +10,7 @@ pub mod ocr;
pub mod pose;
pub mod scene_classification;
pub mod story;
pub mod tkg;
pub mod visual_chunk;
pub mod yolo;
@@ -25,7 +26,8 @@ pub use face_recognition::{
RecognizedFaceDetection,
};
pub use heuristic_scene::{
build_heuristic_scene_meta, generate_scene_meta, CrowdSize, HeuristicSceneMeta, SceneSegmentMeta,
build_heuristic_scene_meta, generate_scene_meta, CrowdSize, HeuristicSceneMeta,
SceneSegmentMeta,
};
pub use ocr::{process_ocr, OcrFrame, OcrResult, OcrText};
pub use pose::{process_pose, Bbox, Keypoint, PersonPose, PoseFrame, PoseResult};
@@ -34,5 +36,6 @@ pub use scene_classification::{
SceneSegment,
};
pub use story::{process_story, StoryChildChunk, StoryParentChunk, StoryResult, StoryStats};
pub use tkg::{build_tkg, TkgResult};
pub use visual_chunk::{process_visual_chunk, process_visual_chunk_advanced, VisualChunkResult};
pub use yolo::{process_yolo, YoloFrame, YoloObject, YoloResult};
src/core/processor/story.rs
+51 -25
View File
@@ -106,7 +106,10 @@ pub async fn process_story(
}
// Fallback: Python script
tracing::warn!("[STORY] Native impl failed, falling back to Python: {:?}", result.err());
tracing::warn!(
"[STORY] Native impl failed, falling back to Python: {:?}",
result.err()
);
let executor = PythonExecutor::new()?;
let script_path = executor.script_path("story_processor.py");
@@ -145,7 +148,11 @@ pub async fn process_story(
// ── Native implementation ─────────────────────────────────────────
fn try_native_story(_video_path: &str, output_path: &str, _uuid: Option<&str>) -> Result<StoryResult> {
fn try_native_story(
_video_path: &str,
output_path: &str,
_uuid: Option<&str>,
) -> Result<StoryResult> {
let output_dir = Path::new(output_path).parent().unwrap_or(Path::new("."));
let basename = Path::new(output_path)
.file_stem()
@@ -160,8 +167,7 @@ fn try_native_story(_video_path: &str, output_path: &str, _uuid: Option<&str>) -
let asr_data: AsrData = if asr_path.exists() {
let content = std::fs::read_to_string(&asr_path)
.with_context(|| format!("Failed to read {:?}", asr_path))?;
serde_json::from_str(&content)
.with_context(|| format!("Failed to parse {:?}", asr_path))?
serde_json::from_str(&content).with_context(|| format!("Failed to parse {:?}", asr_path))?
} else {
AsrData { segments: vec![] }
};
@@ -169,8 +175,7 @@ fn try_native_story(_video_path: &str, output_path: &str, _uuid: Option<&str>) -
let cut_data: CutData = if cut_path.exists() {
let content = std::fs::read_to_string(&cut_path)
.with_context(|| format!("Failed to read {:?}", cut_path))?;
serde_json::from_str(&content)
.with_context(|| format!("Failed to parse {:?}", cut_path))?
serde_json::from_str(&content).with_context(|| format!("Failed to parse {:?}", cut_path))?
} else {
CutData { scenes: vec![] }
};
@@ -376,22 +381,39 @@ fn generate_narrative(texts: &[String], objects: &[String], start: f64, end: f64
let mut unique: Vec<&String> = objects.iter().collect();
unique.sort();
unique.dedup();
let objs = unique.iter().take(5).map(|s| (*s).as_str()).collect::<Vec<_>>().join(", ");
let objs = unique
.iter()
.take(5)
.map(|s| (*s).as_str())
.collect::<Vec<_>>()
.join(", ");
parts.push(format!("Visuals: {}", objs));
}
format!("[{:.0}s-{:.0}s] {}", start, end, parts.join(" | "))
}
fn generate_scene_narrative(objects: &[String], start: f64, end: f64, scene_count: usize) -> String {
fn generate_scene_narrative(
objects: &[String],
start: f64,
end: f64,
scene_count: usize,
) -> String {
let mut unique: Vec<&String> = objects.iter().collect();
unique.sort();
unique.dedup();
let top5: Vec<&String> = unique.iter().take(5).cloned().collect();
if !top5.is_empty() {
let obj_str = top5.iter().map(|s| s.as_str()).collect::<Vec<_>>().join(", ");
format!("[{:.0}s-{:.0}s] {} scenes. Visuals: {}.", start, end, scene_count, obj_str)
let obj_str = top5
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
.join(", ");
format!(
"[{:.0}s-{:.0}s] {} scenes. Visuals: {}.",
start, end, scene_count, obj_str
)
} else {
format!("[{:.0}s-{:.0}s] {} video scenes.", start, end, scene_count)
}
@@ -408,7 +430,8 @@ mod tests {
let text = generate_narrative(
&["Hello world".to_string()],
&["person".to_string()],
0.0, 5.0,
0.0,
5.0,
);
assert!(text.contains("[0s-5s]"));
assert!(text.contains("Speech:"));
@@ -576,7 +599,10 @@ mod tests {
};
assert_eq!(result.parent_chunks[0].child_chunk_ids.len(), 2);
assert!(result.child_chunks.iter().all(|c| c.parent_chunk_id.is_some()));
assert!(result
.child_chunks
.iter()
.all(|c| c.parent_chunk_id.is_some()));
assert!(result.parent_chunks[0].parent_chunk_id.is_none());
}
@@ -594,11 +620,7 @@ mod tests {
std::fs::write(&asr_path, r#"{"segments":[]}"#).unwrap();
std::fs::write(&cut_path, r#"{"scenes":[]}"#).unwrap();
let result = try_native_story(
"/dummy.mp4",
out_path.to_str().unwrap(),
None,
).unwrap();
let result = try_native_story("/dummy.mp4", out_path.to_str().unwrap(), None).unwrap();
assert_eq!(result.stats.total_child_chunks, 0);
assert_eq!(result.stats.total_parent_chunks, 0);
@@ -616,13 +638,17 @@ mod tests {
let cut_path = dir.join(format!("{}.cut.json", basename));
let out_path = dir.join(format!("{}.story.json", basename));
std::fs::write(&asr_path, r#"{
std::fs::write(
&asr_path,
r#"{
"segments": [
{"start": 0.0, "end": 2.5, "text": "Hello", "confidence": 0.95},
{"start": 2.5, "end": 5.0, "text": "World", "confidence": 0.92},
{"start": 5.0, "end": 7.5, "text": "Foo", "confidence": 0.90}
]
}"#).unwrap();
}"#,
)
.unwrap();
std::fs::write(&cut_path, r#"{
"scenes": [
@@ -631,11 +657,7 @@ mod tests {
]
}"#).unwrap();
let result = try_native_story(
"/dummy.mp4",
out_path.to_str().unwrap(),
None,
).unwrap();
let result = try_native_story("/dummy.mp4", out_path.to_str().unwrap(), None).unwrap();
assert_eq!(result.stats.asr_children, 3);
assert_eq!(result.stats.cut_children, 2);
@@ -649,7 +671,11 @@ mod tests {
for child in &result.child_chunks {
if child.source == "asr" {
assert!(child.parent_chunk_id.is_some());
assert!(child.parent_chunk_id.as_ref().unwrap().starts_with("story_asr_"));
assert!(child
.parent_chunk_id
.as_ref()
.unwrap()
.starts_with("story_asr_"));
}
}
src/core/processor/tkg.rs
+703
View File
@@ -0,0 +1,703 @@
use anyhow::{Context, Result};
use serde::Deserialize;
use sqlx::PgPool;
use std::collections::HashMap;
use std::path::Path;
use crate::core::db::postgres_db::PostgresDb;
fn t(name: &str) -> String {
let schema = std::env::var("DATABASE_SCHEMA").unwrap_or_else(|_| "dev".to_string());
if schema == "public" {
name.to_string()
} else {
format!("{}.{}", schema, name)
}
}
// ── Input data structs ────────────────────────────────────────────
#[derive(Debug, Deserialize)]
struct YoloJson {
#[serde(default)]
frames: HashMap<String, YoloFrameEntry>,
}
#[derive(Debug, Deserialize)]
struct YoloFrameEntry {
#[serde(default)]
detections: Vec<YoloDetEntry>,
#[serde(default)]
objects: Vec<YoloDetEntry>,
}
#[derive(Debug, Deserialize)]
struct YoloDetEntry {
#[serde(default)]
class_name: String,
#[serde(default)]
confidence: f64,
}
#[derive(Debug, Deserialize)]
struct AsrxJson {
#[serde(default)]
segments: Vec<AsrxSegmentEntry>,
#[serde(default)]
speaker_stats: Option<HashMap<String, AsrxSpeakerStat>>,
}
#[derive(Debug, Deserialize)]
struct AsrxSegmentEntry {
#[serde(default)]
speaker_id: String,
#[serde(default)]
start_time: f64,
#[serde(default)]
end_time: f64,
#[allow(dead_code)]
start_frame: i64,
#[allow(dead_code)]
end_frame: i64,
}
#[derive(Debug, Deserialize)]
struct AsrxSpeakerStat {
#[serde(default)]
count: i64,
}
// ── Face detection trace ──────────────────────────────────────────
#[derive(Debug, sqlx::FromRow)]
struct FaceTraceRow {
trace_id: i64,
frame_count: i64,
start_f: i64,
end_f: i64,
avg_x: Option<f64>,
avg_y: Option<f64>,
avg_w: Option<f64>,
avg_h: Option<f64>,
}
#[derive(Debug, sqlx::FromRow)]
struct FaceDetectionRow {
trace_id: i64,
frame_number: i64,
#[allow(dead_code)]
x: Option<f64>,
#[allow(dead_code)]
y: Option<f64>,
#[allow(dead_code)]
width: Option<f64>,
#[allow(dead_code)]
height: Option<f64>,
}
// ── Public API ────────────────────────────────────────────────────
pub struct TkgResult {
pub face_trace_nodes: usize,
pub object_nodes: usize,
pub speaker_nodes: usize,
pub co_occurrence_edges: usize,
pub speaker_face_edges: usize,
pub face_face_edges: usize,
}
pub async fn build_tkg(db: &PostgresDb, file_uuid: &str, output_dir: &str) -> Result<TkgResult> {
let pool = db.pool();
let n_face = build_face_trace_nodes(pool, file_uuid).await?;
let n_objects = build_yolo_object_nodes(pool, file_uuid, output_dir).await?;
let n_speakers = build_speaker_nodes(pool, file_uuid, output_dir).await?;
let e_co = build_co_occurrence_edges(pool, file_uuid, output_dir).await?;
let e_sf = build_speaker_face_edges(pool, file_uuid, output_dir).await?;
let e_ff = build_face_face_edges(pool, file_uuid).await?;
Ok(TkgResult {
face_trace_nodes: n_face,
object_nodes: n_objects,
speaker_nodes: n_speakers,
co_occurrence_edges: e_co,
speaker_face_edges: e_sf,
face_face_edges: e_ff,
})
}
// ── Node builders ─────────────────────────────────────────────────
async fn build_face_trace_nodes(pool: &PgPool, file_uuid: &str) -> Result<usize> {
let face_table = t("face_detections");
let nodes_table = t("tkg_nodes");
let rows = sqlx::query_as::<_, FaceTraceRow>(&format!(
r#"
SELECT trace_id,
COUNT(*)::bigint as frame_count,
MIN(frame_number) as start_f,
MAX(frame_number) as end_f,
AVG(x::float8) as avg_x,
AVG(y::float8) as avg_y,
AVG(width::float8) as avg_w,
AVG(height::float8) as avg_h
FROM {}
WHERE file_uuid = $1 AND trace_id IS NOT NULL
GROUP BY trace_id
ORDER BY trace_id
"#,
face_table
))
.bind(file_uuid)
.fetch_all(pool)
.await?;
let mut count = 0;
for row in &rows {
let external_id = format!("trace_{}", row.trace_id);
let label = format!("Face Trace {}", row.trace_id);
let props = serde_json::json!({
"frame_count": row.frame_count,
"start_frame": row.start_f,
"end_frame": row.end_f,
"avg_bbox": {
"x": row.avg_x.unwrap_or(0.0).round() as i64,
"y": row.avg_y.unwrap_or(0.0).round() as i64,
"width": row.avg_w.unwrap_or(0.0).round() as i64,
"height": row.avg_h.unwrap_or(0.0).round() as i64,
}
});
sqlx::query(&format!(
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 = COALESCE(EXCLUDED.properties, tkg_nodes.properties),
label = COALESCE(NULLIF(EXCLUDED.label, ''), tkg_nodes.label)
"#,
nodes_table
))
.bind("face_trace")
.bind(&external_id)
.bind(file_uuid)
.bind(&label)
.bind(serde_json::to_string(&props)?)
.execute(pool)
.await?;
count += 1;
}
Ok(count)
}
async fn build_yolo_object_nodes(
pool: &PgPool,
file_uuid: &str,
output_dir: &str,
) -> Result<usize> {
let yolo_path = Path::new(output_dir).join(format!("{}.yolo.json", file_uuid));
if !yolo_path.exists() {
return Ok(0);
}
let content = std::fs::read_to_string(&yolo_path)
.with_context(|| format!("Failed to read {:?}", yolo_path))?;
let yolo: YoloJson = serde_json::from_str(&content)
.with_context(|| format!("Failed to parse {:?}", yolo_path))?;
let mut class_counts: HashMap<String, i64> = HashMap::new();
for fdata in yolo.frames.values() {
let dets = if !fdata.detections.is_empty() {
&fdata.detections
} else {
&fdata.objects
};
for det in dets {
*class_counts.entry(det.class_name.clone()).or_insert(0) += 1;
}
}
let nodes_table = t("tkg_nodes");
let mut count = 0;
for (cls, cnt) in &class_counts {
let props = serde_json::json!({ "total_detections": cnt });
sqlx::query(&format!(
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 = COALESCE(EXCLUDED.properties, tkg_nodes.properties)
"#,
nodes_table
))
.bind("object")
.bind(cls)
.bind(file_uuid)
.bind(cls)
.bind(serde_json::to_string(&props)?)
.execute(pool)
.await?;
count += 1;
}
Ok(count)
}
async fn build_speaker_nodes(pool: &PgPool, file_uuid: &str, output_dir: &str) -> Result<usize> {
let asrx_path = Path::new(output_dir).join(format!("{}.asrx.json", file_uuid));
if !asrx_path.exists() {
return Ok(0);
}
let content = std::fs::read_to_string(&asrx_path)
.with_context(|| format!("Failed to read {:?}", asrx_path))?;
let asrx: AsrxJson = serde_json::from_str(&content)
.with_context(|| format!("Failed to parse {:?}", asrx_path))?;
let stats = asrx.speaker_stats.unwrap_or_default();
let nodes_table = t("tkg_nodes");
let mut count = 0;
for (sid, stat) in &stats {
let props = serde_json::json!({ "segment_count": stat.count });
sqlx::query(&format!(
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 = COALESCE(EXCLUDED.properties, tkg_nodes.properties)
"#,
nodes_table
))
.bind("speaker")
.bind(sid)
.bind(file_uuid)
.bind(sid)
.bind(serde_json::to_string(&props)?)
.execute(pool)
.await?;
count += 1;
}
Ok(count)
}
// ── Edge builders ─────────────────────────────────────────────────
async fn build_co_occurrence_edges(
pool: &PgPool,
file_uuid: &str,
output_dir: &str,
) -> Result<usize> {
let yolo_path = Path::new(output_dir).join(format!("{}.yolo.json", file_uuid));
if !yolo_path.exists() {
return Ok(0);
}
let content = std::fs::read_to_string(&yolo_path)?;
let yolo: YoloJson = serde_json::from_str(&content)?;
let face_table = t("face_detections");
let nodes_table = t("tkg_nodes");
let edges_table = t("tkg_edges");
let face_rows = sqlx::query_as::<_, FaceDetectionRow>(&format!(
r#"SELECT trace_id, frame_number, x, y, width, height
FROM {} WHERE file_uuid = $1 AND trace_id IS NOT NULL
ORDER BY frame_number"#,
face_table
))
.bind(file_uuid)
.fetch_all(pool)
.await?;
let mut edge_count = 0;
for face in &face_rows {
let frame_str = face.frame_number.to_string();
let yolo_frame = match yolo.frames.get(&frame_str) {
Some(f) => f,
None => continue,
};
let dets = if !yolo_frame.detections.is_empty() {
&yolo_frame.detections
} else {
&yolo_frame.objects
};
if dets.is_empty() {
continue;
}
let external_id = format!("trace_{}", face.trace_id);
let face_node: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='face_trace' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&external_id)
.fetch_optional(pool)
.await?;
let face_node_id = match face_node {
Some((id,)) => id,
None => continue,
};
for det in dets {
let obj_node: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='object' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&det.class_name)
.fetch_optional(pool)
.await?;
let obj_node_id = match obj_node {
Some((id,)) => id,
None => continue,
};
let edge_props = serde_json::json!({
"frame": face.frame_number,
"object_confidence": det.confidence,
});
if let Err(e) = sqlx::query(&format!(
r#"
INSERT INTO {} (edge_type, source_node_id, target_node_id, file_uuid, properties)
VALUES ($1, $2, $3, $4, $5::jsonb)
ON CONFLICT (file_uuid, edge_type, source_node_id, target_node_id)
DO UPDATE SET properties = COALESCE(EXCLUDED.properties, tkg_edges.properties)
"#,
edges_table
))
.bind("CO_OCCURS_WITH")
.bind(face_node_id)
.bind(obj_node_id)
.bind(file_uuid)
.bind(serde_json::to_string(&edge_props)?)
.execute(pool)
.await
{
tracing::warn!(
"[TKG] Edge insert failed (trace={}, obj={}): {}",
face.trace_id,
det.class_name,
e
);
continue;
}
edge_count += 1;
}
}
Ok(edge_count)
}
async fn build_speaker_face_edges(
pool: &PgPool,
file_uuid: &str,
output_dir: &str,
) -> Result<usize> {
let asrx_path = Path::new(output_dir).join(format!("{}.asrx.json", file_uuid));
if !asrx_path.exists() {
return Ok(0);
}
let content = std::fs::read_to_string(&asrx_path)?;
let asrx: AsrxJson = serde_json::from_str(&content)?;
if asrx.segments.is_empty() {
return Ok(0);
}
let face_table = t("face_detections");
let nodes_table = t("tkg_nodes");
let edges_table = t("tkg_edges");
let traces = sqlx::query_as::<_, (i64, i64, i64)>(&format!(
r#"SELECT trace_id, MIN(frame_number) as start_f, MAX(frame_number) as end_f
FROM {} WHERE file_uuid = $1 AND trace_id IS NOT NULL
GROUP BY trace_id"#,
face_table
))
.bind(file_uuid)
.fetch_all(pool)
.await?;
// Calculate fps from last segment
let last = asrx.segments.last().unwrap();
let fps = if last.end_time > 0.0 {
last.end_frame as f64 / last.end_time
} else {
30.0
};
let mut edge_count = 0;
for (tid, sf, ef) in &traces {
let face_ext_id = format!("trace_{}", tid);
let face_node: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='face_trace' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&face_ext_id)
.fetch_optional(pool)
.await?;
let face_node_id = match face_node {
Some((id,)) => id,
None => continue,
};
let face_start_sec = *sf as f64 / fps;
let face_end_sec = *ef as f64 / fps;
for seg in &asrx.segments {
let seg_start = seg.start_time;
let seg_end = seg.end_time;
let overlap_start = face_start_sec.max(seg_start);
let overlap_end = face_end_sec.min(seg_end);
if overlap_start >= overlap_end {
continue;
}
let overlap_dur = overlap_end - overlap_start;
let face_dur = face_end_sec - face_start_sec;
if face_dur <= 0.0 {
continue;
}
let overlap_ratio = overlap_dur / face_dur;
if overlap_ratio < 0.3 {
continue;
}
let speaker_node: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='speaker' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&seg.speaker_id)
.fetch_optional(pool)
.await?;
let speaker_node_id = match speaker_node {
Some((id,)) => id,
None => continue,
};
let edge_props = serde_json::json!({
"overlap_ratio": (overlap_ratio * 1000.0).round() / 1000.0,
"overlap_duration_s": (overlap_dur * 10.0).round() / 10.0,
"face_time_range": format!("{:.1}-{:.1}s", face_start_sec, face_end_sec),
"speaker_time_range": format!("{:.1}-{:.1}s", seg_start, seg_end),
});
sqlx::query(&format!(
r#"
INSERT INTO {} (edge_type, source_node_id, target_node_id, file_uuid, properties)
VALUES ($1, $2, $3, $4, $5::jsonb)
ON CONFLICT (file_uuid, edge_type, source_node_id, target_node_id)
DO UPDATE SET properties = COALESCE(EXCLUDED.properties, tkg_edges.properties)
"#,
edges_table
))
.bind("SPEAKS_AS")
.bind(face_node_id)
.bind(speaker_node_id)
.bind(file_uuid)
.bind(serde_json::to_string(&edge_props)?)
.execute(pool)
.await?;
edge_count += 1;
}
}
Ok(edge_count)
}
async fn build_face_face_edges(pool: &PgPool, file_uuid: &str) -> Result<usize> {
let face_table = t("face_detections");
let nodes_table = t("tkg_nodes");
let edges_table = t("tkg_edges");
let rows: Vec<(i64, i64, i64)> = sqlx::query_as(&format!(
r#"
SELECT a.trace_id AS tid_a, b.trace_id AS tid_b, a.frame_number
FROM {} a
JOIN {} b
ON a.file_uuid = b.file_uuid
AND a.frame_number = b.frame_number
AND a.trace_id < b.trace_id
WHERE a.file_uuid = $1
AND a.trace_id IS NOT NULL
AND b.trace_id IS NOT NULL
ORDER BY a.frame_number
"#,
face_table, face_table
))
.bind(file_uuid)
.fetch_all(pool)
.await?;
if rows.is_empty() {
return Ok(0);
}
// Deduplicate by pair
let mut pair_frames: HashMap<(i64, i64), Vec<i64>> = HashMap::new();
for (tid_a, tid_b, frame) in &rows {
let key = if *tid_a < *tid_b {
(*tid_a, *tid_b)
} else {
(*tid_b, *tid_a)
};
pair_frames.entry(key).or_default().push(*frame);
}
let mut edge_count = 0;
for ((tid_a, tid_b), frames) in &pair_frames {
let ext_a = format!("trace_{}", tid_a);
let ext_b = format!("trace_{}", tid_b);
let n_a: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='face_trace' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&ext_a)
.fetch_optional(pool)
.await?;
let n_b: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='face_trace' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&ext_b)
.fetch_optional(pool)
.await?;
let (n_a_id, n_b_id) = match (n_a, n_b) {
(Some((a,)), Some((b,))) => (a, b),
_ => continue,
};
let edge_props = serde_json::json!({
"first_frame": frames[0],
"frame_count": frames.len() as i64,
});
sqlx::query(&format!(
r#"
INSERT INTO {} (edge_type, source_node_id, target_node_id, file_uuid, properties)
VALUES ($1, $2, $3, $4, $5::jsonb)
ON CONFLICT (file_uuid, edge_type, source_node_id, target_node_id)
DO UPDATE SET properties = COALESCE(EXCLUDED.properties, tkg_edges.properties)
"#,
edges_table
))
.bind("CO_OCCURS_WITH")
.bind(n_a_id)
.bind(n_b_id)
.bind(file_uuid)
.bind(serde_json::to_string(&edge_props)?)
.execute(pool)
.await?;
edge_count += 1;
}
Ok(edge_count)
}
// ── Tests ─────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_yolo_json_deserialize() {
let json = r#"{
"frames": {
"1": {"time_seconds": 0.0, "detections": [{"class_name": "person", "confidence": 0.9}]},
"2": {"time_seconds": 1.0, "detections": [{"class_name": "chair", "confidence": 0.8}]}
}
}"#;
let yolo: YoloJson = serde_json::from_str(json).unwrap();
assert_eq!(yolo.frames.len(), 2);
assert_eq!(yolo.frames["1"].detections[0].class_name, "person");
}
#[test]
fn test_yolo_json_empty_frames() {
let json = r#"{"frames": {}}"#;
let yolo: YoloJson = serde_json::from_str(json).unwrap();
assert!(yolo.frames.is_empty());
}
#[test]
fn test_asrx_json_deserialize() {
let json = r#"{
"segments": [
{"speaker_id": "SPEAKER_01", "start_time": 0.0, "end_time": 2.0, "start_frame": 0, "end_frame": 60}
],
"speaker_stats": {"SPEAKER_01": {"count": 1}}
}"#;
let asrx: AsrxJson = serde_json::from_str(json).unwrap();
assert_eq!(asrx.segments.len(), 1);
assert_eq!(asrx.segments[0].speaker_id, "SPEAKER_01");
}
#[test]
fn test_asrx_json_no_stats() {
let json = r#"{"segments": []}"#;
let asrx: AsrxJson = serde_json::from_str(json).unwrap();
assert!(asrx.speaker_stats.is_none());
}
#[test]
fn test_yolo_objects_fallback() {
let json = r#"{
"frames": {
"1": {"objects": [{"class_name": "person"}]}
}
}"#;
let yolo: YoloJson = serde_json::from_str(json).unwrap();
assert_eq!(yolo.frames["1"].objects[0].class_name, "person");
assert!(yolo.frames["1"].detections.is_empty());
}
#[test]
fn test_tkg_result() {
let r = TkgResult {
face_trace_nodes: 5,
object_nodes: 10,
speaker_nodes: 3,
co_occurrence_edges: 20,
speaker_face_edges: 8,
face_face_edges: 4,
};
assert_eq!(r.face_trace_nodes, 5);
assert_eq!(r.object_nodes, 10);
assert_eq!(r.speaker_nodes, 3);
}
}