refactor: remove face embedding architecture - single Qdrant _faces collection
- Delete FaceEmbeddingDb module (face_embedding_db.rs) - Stub match_faces_iterative, generate_seed_embeddings, tmdb_match_handler - Remove sync_trace_embeddings, populate_face_embeddings_to_qdrant - Remove embedding from face.json output (face_processor.py) - Remove embedding from PG UPDATE (store_traced_faces.py) - Remove workspace traces staging (checkin.rs, qdrant_workspace.rs) - Fix tests: add pose_angle to Face, hand_nodes to TkgResult Disabled functions (need reimplement with _faces): - match_faces_iterative (identity agent) - generate_seed_embeddings (TMDb seeds) - tmdb_match_handler (TMDb matching) - cluster_face_embeddings, search_similar_faces - merge_traces_within_cuts
This commit is contained in:
Accusys
@@ -103,6 +103,7 @@ mod tests {
|
||||
confidence: 0.95,
|
||||
embedding: Some(vec![0.1, 0.2, 0.3]),
|
||||
landmarks: Some(serde_json::json!([[10.0, 20.0], [30.0, 40.0]])),
|
||||
pose_angle: None,
|
||||
attributes: Some(FaceAttributes {
|
||||
age: Some(30),
|
||||
gender: Some("male".to_string()),
|
||||
@@ -174,6 +175,7 @@ mod tests {
|
||||
confidence: 0.5,
|
||||
embedding: None,
|
||||
landmarks: None,
|
||||
pose_angle: None,
|
||||
attributes: None,
|
||||
};
|
||||
assert!(face.confidence >= 0.0 && face.confidence <= 1.0);
|
||||
@@ -190,6 +192,7 @@ mod tests {
|
||||
confidence: 0.95,
|
||||
embedding: Some(vec![0.1; 512]),
|
||||
landmarks: None,
|
||||
pose_angle: None,
|
||||
attributes: Some(FaceAttributes {
|
||||
age: Some(35),
|
||||
gender: Some("male".to_string()),
|
||||
|
||||
@@ -1,96 +0,0 @@
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::Duration;
|
||||
|
||||
use super::executor::PythonExecutor;
|
||||
|
||||
const MEDIAPIPE_TIMEOUT: Duration = Duration::from_secs(7200);
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipeResult {
|
||||
pub frame_count: u64,
|
||||
pub fps: f64,
|
||||
pub frames: Vec<MediaPipeFrame>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipeFrame {
|
||||
pub frame: u64,
|
||||
pub timestamp: f64,
|
||||
pub persons: Vec<MediaPipePerson>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipePerson {
|
||||
pub person_id: u64,
|
||||
pub pose: Option<MediaPipePose>,
|
||||
pub left_hand: Option<MediaPipeHand>,
|
||||
pub right_hand: Option<MediaPipeHand>,
|
||||
pub face_mesh: Option<MediaPipeFaceMesh>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipePose {
|
||||
pub landmarks: Vec<Vec<f64>>,
|
||||
pub keypoints_33: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipeHand {
|
||||
pub landmarks: Vec<Vec<f64>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipeFaceMesh {
|
||||
pub landmarks: Vec<Vec<f64>>,
|
||||
}
|
||||
|
||||
pub async fn process_mediapipe(
|
||||
video_path: &str,
|
||||
output_path: &str,
|
||||
uuid: Option<&str>,
|
||||
) -> Result<MediaPipeResult> {
|
||||
// If mediapipe.json already exists (written by face_processor), skip
|
||||
if std::path::Path::new(output_path).exists() {
|
||||
let json_str = std::fs::read_to_string(output_path).context("Failed to read MEDIAPIPE output")?;
|
||||
let result: MediaPipeResult =
|
||||
serde_json::from_str(&json_str).context("Failed to parse MEDIAPIPE output")?;
|
||||
tracing::info!("[MEDIAPIPE] Skipping (already exists): {} frames", result.frames.len());
|
||||
return Ok(result);
|
||||
}
|
||||
let executor = PythonExecutor::new()?;
|
||||
let script_name = "mediapipe_processor_v1.11.py";
|
||||
let script_path = executor.script_path(script_name);
|
||||
|
||||
tracing::info!("[MEDIAPIPE] Starting MediaPipe Holistic: {}", video_path);
|
||||
|
||||
if !script_path.exists() {
|
||||
tracing::warn!("[MEDIAPIPE] Script not found, returning empty result");
|
||||
return Ok(MediaPipeResult {
|
||||
frame_count: 0,
|
||||
fps: 0.0,
|
||||
frames: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
executor
|
||||
.run(
|
||||
script_name,
|
||||
&[video_path, output_path],
|
||||
uuid,
|
||||
"MEDIAPIPE",
|
||||
Some(MEDIAPIPE_TIMEOUT),
|
||||
)
|
||||
.await
|
||||
.with_context(|| format!("Failed to run {:?}", script_path))?;
|
||||
|
||||
let json_str =
|
||||
std::fs::read_to_string(output_path).context("Failed to read MEDIAPIPE output")?;
|
||||
|
||||
let result: MediaPipeResult =
|
||||
serde_json::from_str(&json_str).context("Failed to parse MEDIAPIPE output")?;
|
||||
|
||||
tracing::info!("[MEDIAPIPE] Result: {} frames", result.frames.len());
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
@@ -1,203 +0,0 @@
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
use tokio::process::Command;
|
||||
use tokio::time::timeout;
|
||||
|
||||
use super::executor::PythonExecutor;
|
||||
|
||||
const MEDIAPIPE_TIMEOUT: Duration = Duration::from_secs(7200);
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipeResult {
|
||||
pub metadata: MediaPipeMetadata,
|
||||
pub frames: HashMap<String, MediaPipeDictEntry>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipeMetadata {
|
||||
pub fps: f64,
|
||||
pub total_frames: i64,
|
||||
pub processed_frames: i64,
|
||||
pub sample_interval: i64,
|
||||
pub width: i64,
|
||||
pub height: i64,
|
||||
pub processor: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipeDictEntry {
|
||||
pub frame_number: i64,
|
||||
pub timestamp: f64,
|
||||
pub persons: Vec<MediaPipePerson>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipePerson {
|
||||
pub person_id: i64,
|
||||
#[serde(default)]
|
||||
pub bbox: Option<MediaPipeBBox>,
|
||||
pub face_mesh: Option<serde_json::Value>,
|
||||
pub pose: Option<serde_json::Value>,
|
||||
pub hands: MediaPipeHands,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipeBBox {
|
||||
pub x: i64,
|
||||
pub y: i64,
|
||||
pub width: i64,
|
||||
pub height: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct MediaPipeHands {
|
||||
pub left: Option<serde_json::Value>,
|
||||
pub right: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
pub async fn process_mediapipe_v2(
|
||||
video_path: &str,
|
||||
output_path: &str,
|
||||
uuid: Option<&str>,
|
||||
frames: Option<&[i64]>,
|
||||
) -> Result<MediaPipeResult> {
|
||||
let executor = PythonExecutor::new()?;
|
||||
let script_path = executor.script_path("mediapipe_holistic_processor.py");
|
||||
|
||||
tracing::info!("[MEDIAPIPE] Starting MediaPipe Holistic: {}", video_path);
|
||||
|
||||
if !script_path.exists() {
|
||||
anyhow::bail!("mediapipe_holistic_processor.py not found");
|
||||
}
|
||||
|
||||
let mut cmd = Command::new(executor.python_path());
|
||||
cmd.arg(&script_path).arg(video_path).arg(output_path);
|
||||
|
||||
// Use explicit frame list if provided, otherwise calculate sample_interval for ~8Hz
|
||||
if let Some(frames) = frames {
|
||||
let frames_str = frames
|
||||
.iter()
|
||||
.map(|f| f.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
cmd.arg("--frames").arg(&frames_str);
|
||||
tracing::info!("[MEDIAPIPE] 8Hz sampling: {} frames", frames.len());
|
||||
} else {
|
||||
let sample_interval = calculate_sample_interval(video_path).await;
|
||||
cmd.arg("--sample-interval")
|
||||
.arg(sample_interval.to_string());
|
||||
}
|
||||
|
||||
if let Some(u) = uuid {
|
||||
cmd.arg("--uuid").arg(u);
|
||||
}
|
||||
|
||||
cmd.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::piped());
|
||||
|
||||
let child = cmd.spawn().context("Failed to run MEDIAPIPE processor")?;
|
||||
|
||||
let output = match timeout(MEDIAPIPE_TIMEOUT, child.wait_with_output()).await {
|
||||
Ok(Ok(output)) => output,
|
||||
Ok(Err(e)) => return Err(e).context("Failed to run MEDIAPIPE processor"),
|
||||
Err(_) => anyhow::bail!(
|
||||
"MEDIAPIPE processing timed out after {:?}",
|
||||
MEDIAPIPE_TIMEOUT
|
||||
),
|
||||
};
|
||||
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
|
||||
for line in stderr.lines() {
|
||||
let trimmed = line.trim();
|
||||
if trimmed.starts_with("MEDIAPIPE_START") {
|
||||
tracing::info!("[MEDIAPIPE] Loading model...");
|
||||
} else if trimmed.starts_with("MEDIAPIPE_FRAME:") {
|
||||
let count = trimmed.trim_start_matches("MEDIAPIPE_FRAME:");
|
||||
tracing::info!("[MEDIAPIPE] Processed {} frames...", count);
|
||||
} else if trimmed.starts_with("MEDIAPIPE_COMPLETE:") {
|
||||
let count = trimmed.trim_start_matches("MEDIAPIPE_COMPLETE:");
|
||||
tracing::info!("[MEDIAPIPE] Completed! Total: {} frames", count);
|
||||
} else if trimmed.starts_with("MEDIAPIPE_INFO:") {
|
||||
let info = trimmed.trim_start_matches("MEDIAPIPE_INFO:");
|
||||
tracing::info!("[MEDIAPIPE] {}", info);
|
||||
} else if trimmed.starts_with("MEDIAPIPE_ERROR:") {
|
||||
let err = trimmed.trim_start_matches("MEDIAPIPE_ERROR:");
|
||||
tracing::error!("[MEDIAPIPE] {}", err);
|
||||
}
|
||||
}
|
||||
tracing::info!("[MEDIAPIPE] stderr output:\n{}", stderr);
|
||||
|
||||
if !output.status.success() {
|
||||
anyhow::bail!("MEDIAPIPE failed: {}", stderr);
|
||||
}
|
||||
|
||||
let json_str =
|
||||
std::fs::read_to_string(output_path).context("Failed to read MEDIAPIPE output")?;
|
||||
|
||||
let result: MediaPipeResult =
|
||||
serde_json::from_str(&json_str).context("Failed to parse MEDIAPIPE output")?;
|
||||
|
||||
tracing::info!("[MEDIAPIPE] Result: {} frames", result.frames.len());
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn calculate_sample_interval(video_path: &str) -> i64 {
|
||||
// Try ffprobe to get FPS, calculate sample_interval for ~8Hz
|
||||
let probe_cmd = Command::new("ffprobe")
|
||||
.args([
|
||||
"-v",
|
||||
"quiet",
|
||||
"-print_format",
|
||||
"json",
|
||||
"-show_streams",
|
||||
video_path,
|
||||
])
|
||||
.output()
|
||||
.await;
|
||||
|
||||
if let Ok(output) = probe_cmd {
|
||||
if output.status.success() {
|
||||
if let Ok(json_str) = String::from_utf8(output.stdout) {
|
||||
if let Ok(probe_data) = serde_json::from_str::<serde_json::Value>(&json_str) {
|
||||
if let Some(streams) = probe_data["streams"].as_array() {
|
||||
for stream in streams {
|
||||
if stream["codec_type"] == "video" {
|
||||
if let Some(fps_str) = stream["r_frame_rate"].as_str() {
|
||||
// Parse "30000/1001" style fps
|
||||
if let Some(fps) = parse_fractional_fps(fps_str) {
|
||||
let interval = (fps / 8.0).round() as i64;
|
||||
return interval.max(1);
|
||||
}
|
||||
}
|
||||
if let Some(fps_val) = stream["avg_frame_rate"].as_str() {
|
||||
if let Some(fps) = parse_fractional_fps(fps_val) {
|
||||
let interval = (fps / 8.0).round() as i64;
|
||||
return interval.max(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
4 // Default: assume 30fps / 8 = ~4
|
||||
}
|
||||
|
||||
fn parse_fractional_fps(s: &str) -> Option<f64> {
|
||||
let parts: Vec<&str> = s.split('/').collect();
|
||||
if parts.len() == 2 {
|
||||
let num: f64 = parts[0].parse().ok()?;
|
||||
let den: f64 = parts[1].parse().ok()?;
|
||||
if den > 0.0 {
|
||||
return Some(num / den);
|
||||
}
|
||||
}
|
||||
s.parse::<f64>().ok()
|
||||
}
|
||||
@@ -11,11 +11,9 @@ pub mod face_clustering;
|
||||
pub mod face_recognition;
|
||||
pub mod hand;
|
||||
pub mod heuristic_scene;
|
||||
pub mod mediapipe_v2;
|
||||
pub mod ocr;
|
||||
pub mod pose;
|
||||
pub mod scene_classification;
|
||||
pub mod story;
|
||||
pub mod tkg;
|
||||
pub mod yolo;
|
||||
|
||||
@@ -48,17 +46,12 @@ pub use heuristic_scene::{
|
||||
build_heuristic_scene_meta, generate_scene_meta, CrowdSize, HeuristicSceneMeta,
|
||||
SceneSegmentMeta,
|
||||
};
|
||||
pub use mediapipe_v2::{
|
||||
process_mediapipe_v2, MediaPipeBBox, MediaPipeDictEntry, MediaPipeHands, MediaPipeMetadata,
|
||||
MediaPipePerson, MediaPipeResult,
|
||||
};
|
||||
pub use ocr::{process_ocr, OcrFrame, OcrResult, OcrText};
|
||||
pub use pose::{process_pose, Bbox, Keypoint, PersonPose, PoseFrame, PoseResult};
|
||||
pub use scene_classification::{
|
||||
load_scene_from_file, process_scene_classification, SceneClassificationResult, ScenePrediction,
|
||||
SceneSegment,
|
||||
};
|
||||
pub use story::{process_story, StoryChildChunk, StoryParentChunk, StoryResult, StoryStats};
|
||||
pub use tkg::{
|
||||
build_tkg, query_auto_representative_frame, FrameTraceInfo, MainIdentityInfo,
|
||||
RepresentativeFrameResult, TkgResult,
|
||||
|
||||
@@ -1,690 +0,0 @@
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::executor::PythonExecutor;
|
||||
|
||||
const STORY_TIMEOUT: Duration = Duration::from_secs(3600);
|
||||
|
||||
// ── Input data structs (from JSON files) ──────────────────────────
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AsrData {
|
||||
segments: Vec<AsrSegmentInput>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AsrSegmentInput {
|
||||
#[serde(default, alias = "start")]
|
||||
start_time: f64,
|
||||
#[serde(default, alias = "end")]
|
||||
end_time: f64,
|
||||
#[serde(default)]
|
||||
text: String,
|
||||
#[serde(default)]
|
||||
confidence: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CutData {
|
||||
scenes: Vec<CutSceneInput>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CutSceneInput {
|
||||
scene_number: Option<i64>,
|
||||
#[allow(dead_code)]
|
||||
start_frame: Option<i64>,
|
||||
#[allow(dead_code)]
|
||||
end_frame: Option<i64>,
|
||||
start_time: Option<f64>,
|
||||
end_time: Option<f64>,
|
||||
}
|
||||
|
||||
// ── Output data structs ───────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct StoryResult {
|
||||
pub child_chunks: Vec<StoryChildChunk>,
|
||||
pub parent_chunks: Vec<StoryParentChunk>,
|
||||
pub stats: StoryStats,
|
||||
#[serde(default)]
|
||||
pub metadata: serde_json::Value,
|
||||
#[serde(default)]
|
||||
pub parent_chunk_size: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct StoryStats {
|
||||
pub total_child_chunks: usize,
|
||||
pub total_parent_chunks: usize,
|
||||
pub asr_children: usize,
|
||||
pub cut_children: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct StoryChildChunk {
|
||||
pub chunk_id: String,
|
||||
pub chunk_type: String,
|
||||
pub source: String,
|
||||
pub start_time: f64,
|
||||
pub end_time: f64,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub text_content: Option<String>,
|
||||
pub content: serde_json::Value,
|
||||
#[serde(default)]
|
||||
pub child_chunk_ids: Vec<String>,
|
||||
pub parent_chunk_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct StoryParentChunk {
|
||||
pub chunk_id: String,
|
||||
pub chunk_type: String,
|
||||
pub source: String,
|
||||
pub start_time: f64,
|
||||
pub end_time: f64,
|
||||
pub text_content: String,
|
||||
pub content: serde_json::Value,
|
||||
#[serde(default)]
|
||||
pub child_chunk_ids: Vec<String>,
|
||||
pub parent_chunk_id: Option<String>,
|
||||
}
|
||||
|
||||
// ── Public API ────────────────────────────────────────────────────
|
||||
|
||||
pub async fn process_story(
|
||||
video_path: &str,
|
||||
output_path: &str,
|
||||
uuid: Option<&str>,
|
||||
) -> Result<StoryResult> {
|
||||
// Try native Rust implementation first
|
||||
let result = try_native_story(video_path, output_path, uuid);
|
||||
if let Ok(r) = result {
|
||||
return Ok(r);
|
||||
}
|
||||
|
||||
// Fallback: Python script
|
||||
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");
|
||||
|
||||
if !script_path.exists() {
|
||||
return Ok(StoryResult {
|
||||
child_chunks: vec![],
|
||||
parent_chunks: vec![],
|
||||
stats: StoryStats {
|
||||
total_child_chunks: 0,
|
||||
total_parent_chunks: 0,
|
||||
asr_children: 0,
|
||||
cut_children: 0,
|
||||
},
|
||||
metadata: serde_json::json!({}),
|
||||
parent_chunk_size: 5,
|
||||
});
|
||||
}
|
||||
|
||||
executor
|
||||
.run(
|
||||
"story_processor.py",
|
||||
&[video_path, output_path],
|
||||
uuid,
|
||||
"STORY",
|
||||
Some(STORY_TIMEOUT),
|
||||
)
|
||||
.await
|
||||
.with_context(|| format!("Failed to run {:?}", script_path))?;
|
||||
|
||||
let json_str = std::fs::read_to_string(output_path).context("Failed to read STORY output")?;
|
||||
let result: StoryResult =
|
||||
serde_json::from_str(&json_str).context("Failed to parse STORY output")?;
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
// ── Native implementation ─────────────────────────────────────────
|
||||
|
||||
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()
|
||||
.and_then(|s| s.to_str())
|
||||
.and_then(|s| s.split('.').next())
|
||||
.unwrap_or("unknown");
|
||||
|
||||
let asr_path = output_dir.join(format!("{}.asr.json", basename));
|
||||
let cut_path = output_dir.join(format!("{}.cut.json", basename));
|
||||
|
||||
// ASR data is required; CUT is optional
|
||||
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))?
|
||||
} else {
|
||||
AsrData { segments: vec![] }
|
||||
};
|
||||
|
||||
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))?
|
||||
} else {
|
||||
CutData { scenes: vec![] }
|
||||
};
|
||||
|
||||
let parent_chunk_size: usize = 5;
|
||||
|
||||
// ── Build child chunks ────────────────────────────────────────
|
||||
let mut child_chunks: Vec<StoryChildChunk> = Vec::new();
|
||||
|
||||
// ASR child chunks
|
||||
for seg in &asr_data.segments {
|
||||
let chunk_id = format!("asr_{:.1}_{:.1}", seg.start_time, seg.end_time);
|
||||
child_chunks.push(StoryChildChunk {
|
||||
chunk_id,
|
||||
chunk_type: "asr".to_string(),
|
||||
source: "asr".to_string(),
|
||||
start_time: seg.start_time,
|
||||
end_time: seg.end_time,
|
||||
text_content: Some(seg.text.clone()),
|
||||
content: serde_json::json!({
|
||||
"text": seg.text,
|
||||
"confidence": seg.confidence,
|
||||
}),
|
||||
child_chunk_ids: vec![],
|
||||
parent_chunk_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
// CUT child chunks
|
||||
for scene in &cut_data.scenes {
|
||||
let scene_num = scene.scene_number.unwrap_or(0);
|
||||
let start_time = scene.start_time.unwrap_or(0.0);
|
||||
let end_time = scene.end_time.unwrap_or(0.0);
|
||||
let chunk_id = format!("cut_{}", scene_num);
|
||||
child_chunks.push(StoryChildChunk {
|
||||
chunk_id,
|
||||
chunk_type: "cut".to_string(),
|
||||
source: "cut".to_string(),
|
||||
start_time,
|
||||
end_time,
|
||||
text_content: Some(format!("Scene {}", scene_num)),
|
||||
content: serde_json::json!({
|
||||
"scene_number": scene_num,
|
||||
"start_time": start_time,
|
||||
"end_time": end_time,
|
||||
}),
|
||||
child_chunk_ids: vec![],
|
||||
parent_chunk_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
let asr_child_ids: Vec<String> = child_chunks
|
||||
.iter()
|
||||
.filter(|c| c.source == "asr")
|
||||
.map(|c| c.chunk_id.clone())
|
||||
.collect();
|
||||
|
||||
let cut_child_ids: Vec<String> = child_chunks
|
||||
.iter()
|
||||
.filter(|c| c.source == "cut")
|
||||
.map(|c| c.chunk_id.clone())
|
||||
.collect();
|
||||
|
||||
// ── Build parent chunks from ASR ──────────────────────────────
|
||||
let mut parent_chunks: Vec<StoryParentChunk> = Vec::new();
|
||||
|
||||
for (i, batch) in asr_child_ids.chunks(parent_chunk_size).enumerate() {
|
||||
if batch.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let mut texts: Vec<String> = Vec::new();
|
||||
let mut times: Vec<(f64, f64)> = Vec::new();
|
||||
|
||||
for child_id in batch {
|
||||
if let Some(child) = child_chunks.iter().find(|c| &c.chunk_id == child_id) {
|
||||
if let Some(ref t) = child.text_content {
|
||||
texts.push(t.clone());
|
||||
}
|
||||
times.push((child.start_time, child.end_time));
|
||||
}
|
||||
}
|
||||
|
||||
let start_time = times.first().map(|t| t.0).unwrap_or(0.0);
|
||||
let end_time = times.last().map(|t| t.1).unwrap_or(0.0);
|
||||
|
||||
let narrative = generate_narrative(&texts, &[], start_time, end_time);
|
||||
|
||||
let chunk_id = format!("story_asr_{:04}", i);
|
||||
parent_chunks.push(StoryParentChunk {
|
||||
chunk_id: chunk_id.clone(),
|
||||
chunk_type: "story".to_string(),
|
||||
source: "story_asr".to_string(),
|
||||
start_time,
|
||||
end_time,
|
||||
text_content: narrative.clone(),
|
||||
content: serde_json::json!({
|
||||
"description": narrative,
|
||||
"child_count": batch.len(),
|
||||
"speech_preview": texts.iter().take(3).cloned().collect::<Vec<_>>().join(" "),
|
||||
}),
|
||||
child_chunk_ids: batch.to_vec(),
|
||||
parent_chunk_id: None,
|
||||
});
|
||||
|
||||
// Link children to parent
|
||||
for child in &mut child_chunks {
|
||||
if batch.contains(&child.chunk_id) {
|
||||
child.parent_chunk_id = Some(chunk_id.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Build parent chunks from CUT ──────────────────────────────
|
||||
for (i, batch) in cut_child_ids.chunks(parent_chunk_size).enumerate() {
|
||||
if batch.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let mut times: Vec<(f64, f64)> = Vec::new();
|
||||
for child_id in batch {
|
||||
if let Some(child) = child_chunks.iter().find(|c| &c.chunk_id == child_id) {
|
||||
times.push((child.start_time, child.end_time));
|
||||
}
|
||||
}
|
||||
|
||||
let start_time = times.first().map(|t| t.0).unwrap_or(0.0);
|
||||
let end_time = times.last().map(|t| t.1).unwrap_or(0.0);
|
||||
|
||||
let narrative = generate_scene_narrative(&[], start_time, end_time, batch.len());
|
||||
|
||||
let chunk_id = format!("story_cut_{:04}", i);
|
||||
parent_chunks.push(StoryParentChunk {
|
||||
chunk_id: chunk_id.clone(),
|
||||
chunk_type: "story".to_string(),
|
||||
source: "story_cut".to_string(),
|
||||
start_time,
|
||||
end_time,
|
||||
text_content: narrative.clone(),
|
||||
content: serde_json::json!({
|
||||
"description": narrative,
|
||||
"child_count": batch.len(),
|
||||
"scenes": batch,
|
||||
}),
|
||||
child_chunk_ids: batch.to_vec(),
|
||||
parent_chunk_id: None,
|
||||
});
|
||||
|
||||
for child in &mut child_chunks {
|
||||
if batch.contains(&child.chunk_id) {
|
||||
child.parent_chunk_id = Some(chunk_id.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Build result ──────────────────────────────────────────────
|
||||
let total_child = asr_child_ids.len() + cut_child_ids.len();
|
||||
let total_parent = parent_chunks.len();
|
||||
let asr_count = asr_child_ids.len();
|
||||
let cut_count = cut_child_ids.len();
|
||||
|
||||
let result = StoryResult {
|
||||
child_chunks,
|
||||
parent_chunks,
|
||||
stats: StoryStats {
|
||||
total_child_chunks: total_child,
|
||||
total_parent_chunks: total_parent,
|
||||
asr_children: asr_count,
|
||||
cut_children: cut_count,
|
||||
},
|
||||
metadata: serde_json::json!({}),
|
||||
parent_chunk_size,
|
||||
};
|
||||
|
||||
// Write output (for compatibility with Python path)
|
||||
let json_str = serde_json::to_string_pretty(&result)?;
|
||||
std::fs::write(output_path, &json_str)
|
||||
.with_context(|| format!("Failed to write {:?}", output_path))?;
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
// ── Narrative generation (matching Python logic) ──────────────────
|
||||
|
||||
fn generate_narrative(texts: &[String], objects: &[String], start: f64, end: f64) -> String {
|
||||
if texts.is_empty() && objects.is_empty() {
|
||||
return format!("Video segment from {:.1}s to {:.1}s", start, end);
|
||||
}
|
||||
|
||||
let mut parts: Vec<String> = Vec::new();
|
||||
|
||||
if !texts.is_empty() {
|
||||
let combined = texts.join(" ");
|
||||
let truncated = if combined.len() > 150 {
|
||||
format!("{}...", &combined[..150])
|
||||
} else {
|
||||
combined
|
||||
};
|
||||
parts.push(format!("Speech: {}", truncated));
|
||||
}
|
||||
|
||||
if !objects.is_empty() {
|
||||
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(", ");
|
||||
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 {
|
||||
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
|
||||
)
|
||||
} else {
|
||||
format!("[{:.0}s-{:.0}s] {} video scenes.", start, end, scene_count)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_generate_narrative_with_text() {
|
||||
let text = generate_narrative(
|
||||
&["Hello world".to_string()],
|
||||
&["person".to_string()],
|
||||
0.0,
|
||||
5.0,
|
||||
);
|
||||
assert!(text.contains("[0s-5s]"));
|
||||
assert!(text.contains("Speech:"));
|
||||
assert!(text.contains("Visuals:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_generate_narrative_empty() {
|
||||
let text = generate_narrative(&[], &[], 10.0, 20.0);
|
||||
assert!(text.contains("10.0s to 20.0s"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_generate_scene_narrative() {
|
||||
let text = generate_scene_narrative(&["person".to_string()], 0.0, 10.0, 3);
|
||||
assert!(text.contains("3 scenes"));
|
||||
assert!(text.contains("person"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_generate_scene_narrative_empty() {
|
||||
let text = generate_scene_narrative(&[], 0.0, 10.0, 1);
|
||||
assert!(text.contains("1 video scenes"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_narrative_truncation() {
|
||||
let long_text = "a".repeat(200);
|
||||
let text = generate_narrative(&[long_text], &[], 0.0, 5.0);
|
||||
assert!(text.len() < 200 + 50); // truncated with "..."
|
||||
assert!(text.ends_with("..."));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_story_result_serialization() {
|
||||
let result = StoryResult {
|
||||
child_chunks: vec![StoryChildChunk {
|
||||
chunk_id: "asr_0001".to_string(),
|
||||
chunk_type: "sentence".to_string(),
|
||||
source: "asr".to_string(),
|
||||
start_time: 0.0,
|
||||
end_time: 5.0,
|
||||
text_content: Some("Hello world".to_string()),
|
||||
content: serde_json::json!({}),
|
||||
child_chunk_ids: vec![],
|
||||
parent_chunk_id: Some("story_asr_0000".to_string()),
|
||||
}],
|
||||
parent_chunks: vec![StoryParentChunk {
|
||||
chunk_id: "story_asr_0000".to_string(),
|
||||
chunk_type: "story".to_string(),
|
||||
source: "story_asr".to_string(),
|
||||
start_time: 0.0,
|
||||
end_time: 25.0,
|
||||
text_content: "[0s-25s] Hello world...".to_string(),
|
||||
content: serde_json::json!({
|
||||
"description": "[0s-25s] Hello world...",
|
||||
"child_count": 5
|
||||
}),
|
||||
child_chunk_ids: vec!["asr_0001".to_string()],
|
||||
parent_chunk_id: None,
|
||||
}],
|
||||
stats: StoryStats {
|
||||
total_child_chunks: 10,
|
||||
total_parent_chunks: 2,
|
||||
asr_children: 10,
|
||||
cut_children: 0,
|
||||
},
|
||||
metadata: serde_json::json!({}),
|
||||
parent_chunk_size: 5,
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&result).unwrap();
|
||||
assert!(json.contains("asr_0001"));
|
||||
assert!(json.contains("story_asr_0000"));
|
||||
assert!(json.contains("Hello world"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_story_result_deserialization() {
|
||||
let json = r#"{
|
||||
"child_chunks": [{
|
||||
"chunk_id": "asr_0001",
|
||||
"chunk_type": "sentence",
|
||||
"source": "asr",
|
||||
"start_time": 0.0,
|
||||
"end_time": 5.0,
|
||||
"text_content": "Hello",
|
||||
"content": {},
|
||||
"child_chunk_ids": [],
|
||||
"parent_chunk_id": null
|
||||
}],
|
||||
"parent_chunks": [{
|
||||
"chunk_id": "story_asr_0000",
|
||||
"chunk_type": "story",
|
||||
"source": "story_asr",
|
||||
"start_time": 0.0,
|
||||
"end_time": 5.0,
|
||||
"text_content": "Hello segment",
|
||||
"content": {"description": "Hello segment"},
|
||||
"child_chunk_ids": ["asr_0001"],
|
||||
"parent_chunk_id": null
|
||||
}],
|
||||
"stats": {
|
||||
"total_child_chunks": 1,
|
||||
"total_parent_chunks": 1,
|
||||
"asr_children": 1,
|
||||
"cut_children": 0
|
||||
},
|
||||
"metadata": {},
|
||||
"parent_chunk_size": 5
|
||||
}"#;
|
||||
|
||||
let result: StoryResult = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(result.child_chunks.len(), 1);
|
||||
assert_eq!(result.parent_chunks.len(), 1);
|
||||
assert_eq!(result.stats.total_child_chunks, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parent_child_relationship() {
|
||||
let result = StoryResult {
|
||||
child_chunks: vec![
|
||||
StoryChildChunk {
|
||||
chunk_id: "asr_0001".to_string(),
|
||||
chunk_type: "sentence".to_string(),
|
||||
source: "asr".to_string(),
|
||||
start_time: 0.0,
|
||||
end_time: 5.0,
|
||||
text_content: Some("First".to_string()),
|
||||
content: serde_json::json!({}),
|
||||
child_chunk_ids: vec![],
|
||||
parent_chunk_id: Some("story_asr_0000".to_string()),
|
||||
},
|
||||
StoryChildChunk {
|
||||
chunk_id: "asr_0002".to_string(),
|
||||
chunk_type: "sentence".to_string(),
|
||||
source: "asr".to_string(),
|
||||
start_time: 5.0,
|
||||
end_time: 10.0,
|
||||
text_content: Some("Second".to_string()),
|
||||
content: serde_json::json!({}),
|
||||
child_chunk_ids: vec![],
|
||||
parent_chunk_id: Some("story_asr_0000".to_string()),
|
||||
},
|
||||
],
|
||||
parent_chunks: vec![StoryParentChunk {
|
||||
chunk_id: "story_asr_0000".to_string(),
|
||||
chunk_type: "story".to_string(),
|
||||
source: "story_asr".to_string(),
|
||||
start_time: 0.0,
|
||||
end_time: 10.0,
|
||||
text_content: "Combined narrative".to_string(),
|
||||
content: serde_json::json!({}),
|
||||
child_chunk_ids: vec!["asr_0001".to_string(), "asr_0002".to_string()],
|
||||
parent_chunk_id: None,
|
||||
}],
|
||||
stats: StoryStats {
|
||||
total_child_chunks: 2,
|
||||
total_parent_chunks: 1,
|
||||
asr_children: 2,
|
||||
cut_children: 0,
|
||||
},
|
||||
metadata: serde_json::json!({}),
|
||||
parent_chunk_size: 5,
|
||||
};
|
||||
|
||||
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.parent_chunks[0].parent_chunk_id.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_native_story_empty_data() {
|
||||
// Write empty ASR and CUT files, then test try_native_story
|
||||
let dir = std::env::temp_dir().join("story_test_empty");
|
||||
let _ = std::fs::create_dir_all(&dir);
|
||||
|
||||
let basename = "test_video";
|
||||
let asr_path = dir.join(format!("{}.asr.json", basename));
|
||||
let cut_path = dir.join(format!("{}.cut.json", basename));
|
||||
let out_path = dir.join(format!("{}.story.json", basename));
|
||||
|
||||
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();
|
||||
|
||||
assert_eq!(result.stats.total_child_chunks, 0);
|
||||
assert_eq!(result.stats.total_parent_chunks, 0);
|
||||
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_native_story_with_data() {
|
||||
let dir = std::env::temp_dir().join("story_test_data");
|
||||
let _ = std::fs::create_dir_all(&dir);
|
||||
|
||||
let basename = "test_video";
|
||||
let asr_path = dir.join(format!("{}.asr.json", basename));
|
||||
let cut_path = dir.join(format!("{}.cut.json", basename));
|
||||
let out_path = dir.join(format!("{}.story.json", basename));
|
||||
|
||||
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();
|
||||
|
||||
std::fs::write(&cut_path, r#"{
|
||||
"scenes": [
|
||||
{"scene_number": 1, "start_frame": 0, "end_frame": 150, "start_time": 0.0, "end_time": 5.0},
|
||||
{"scene_number": 2, "start_frame": 150, "end_frame": 300, "start_time": 5.0, "end_time": 10.0}
|
||||
]
|
||||
}"#).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);
|
||||
assert_eq!(result.stats.total_child_chunks, 5);
|
||||
|
||||
// 3 ASR segments, parent_chunk_size=5 → 1 parent
|
||||
// 2 CUT scenes, parent_chunk_size=5 → 1 parent
|
||||
assert_eq!(result.stats.total_parent_chunks, 2);
|
||||
|
||||
// Verify child-parent linking
|
||||
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_"));
|
||||
}
|
||||
}
|
||||
|
||||
// Verify output file was written
|
||||
assert!(out_path.exists());
|
||||
let content = std::fs::read_to_string(&out_path).unwrap();
|
||||
assert!(content.contains("Hello"));
|
||||
assert!(content.contains("World"));
|
||||
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
}
|
||||
+5
-1113
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user