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

feat: ASRX hybrid pipeline, identity history, worker fixes, checkpoint system

Browse Source
This commit is contained in:
Accusys
2026-06-02 07:13:23 +08:00
parent e3066c3f49
commit e1572907ae
198 changed files with 43705 additions and 8910 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core/processor/asrx.rs
+38 -12
View File
@@ -18,12 +18,22 @@ pub struct AsrxResult {
#[derive(Debug, Serialize, Deserialize)]
pub struct AsrxSegment {
#[serde(alias = "start")]
pub start_time: f64,
#[serde(alias = "end")]
pub end_time: f64,
#[serde(default)]
pub start_frame: u64,
#[serde(default)]
pub end_frame: u64,
pub text: String,
pub speaker_id: Option<String>,
#[serde(default)]
pub language: Option<String>,
#[serde(default)]
pub lang_prob: Option<f64>,
#[serde(default)]
pub quality: Option<f64>,
}
pub async fn process_asrx(
@@ -32,24 +42,16 @@ pub async fn process_asrx(
uuid: Option<&str>,
) -> Result<AsrxResult> {
let executor = PythonExecutor::new()?;
let script_path = executor.script_path("asrx_processor_custom.py");
let script_path = executor.script_path("asrx_processor.py");
tracing::info!(
"[ASRX] Starting speaker diarization (custom): {}",
"[ASRX] Starting hybrid speaker diarization: {}",
video_path
);
if !script_path.exists() {
tracing::warn!("[ASRX] Custom script not found, falling back to original");
let fallback_path = executor.script_path("asrx_processor.py");
if !fallback_path.exists() {
tracing::warn!("[ASRX] No script found, returning empty result");
return Ok(AsrxResult {
language: None,
segments: vec![],
embeddings: None,
});
}
tracing::error!("[ASRX] Script not found: {:?}", script_path);
anyhow::bail!("asrx_processor.py not found");
}
tracing::info!(
@@ -65,6 +67,7 @@ pub async fn process_asrx(
if let Some(u) = uuid {
cmd.arg("--uuid").arg(u);
cmd.arg("--file-uuid").arg(u);
}
cmd.stdout(std::process::Stdio::piped())
@@ -126,6 +129,9 @@ mod tests {
end_frame: 75,
text: "Hello".to_string(),
speaker_id: Some("SPEAKER_00".to_string()),
language: None,
lang_prob: None,
quality: None,
}],
embeddings: None,
};
@@ -173,7 +179,27 @@ mod tests {
end_frame: 150,
text: "Test".to_string(),
speaker_id: None,
language: None,
lang_prob: None,
quality: None,
};
assert!(segment.end_time > segment.start_time);
}
#[test]
fn test_asrx_backward_compat_old_format() {
let json = r#"{
"language": "en",
"segments": [
{"start": 10.0, "end": 12.5, "text": "Hello", "speaker_id": "SPEAKER_00"}
]
}"#;
let result: AsrxResult = serde_json::from_str(json).unwrap();
assert_eq!(result.segments.len(), 1);
assert_eq!(result.segments[0].start_time, 10.0);
assert_eq!(result.segments[0].end_time, 12.5);
assert_eq!(result.segments[0].text, "Hello");
assert_eq!(result.segments[0].start_frame, 0);
assert_eq!(result.segments[0].end_frame, 0);
}
}
src/core/processor/cut.rs
+24 -12
View File
@@ -43,11 +43,15 @@ pub async fn process_cut(
let script_path = executor.script_path("cut_processor.py");
if !script_path.exists() {
return Ok(CutResult {
let empty_result = CutResult {
frame_count: 0,
fps: 0.0,
scenes: vec![],
});
};
let json = serde_json::to_string_pretty(&empty_result)?;
std::fs::write(output_path, &json)
.with_context(|| format!("Failed to write {:?}", output_path))?;
return Ok(empty_result);
}
executor
@@ -127,18 +131,26 @@ fn try_native_cut(video_path: &str) -> Result<CutResult> {
.context("Failed to run ffmpeg scene detection")?;
let stderr_output = String::from_utf8_lossy(&scene_output.stderr);
let stdout_output = String::from_utf8_lossy(&scene_output.stdout);
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);
}
// Parse ffprobe output for scene changes (check both stderr and stdout)
// Format: pts_time=123.456 or pts_time:123.456
for line in stderr_output.lines().chain(stdout_output.lines()) {
// Try pts_time= format (standard ffprobe output)
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);
}
}
// Try pts_time: format (showinfo filter output)
else 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);
}
}
}
src/core/processor/mod.rs
-2
View File
@@ -11,7 +11,6 @@ pub mod pose;
pub mod scene_classification;
pub mod story;
pub mod tkg;
pub mod visual_chunk;
pub mod yolo;
pub use asr::{process_asr, AsrResult, AsrSegment};
@@ -40,5 +39,4 @@ pub use tkg::{
build_tkg, query_auto_representative_frame, FrameTraceInfo, MainIdentityInfo,
RepresentativeFrameResult, TkgResult,
};
pub use visual_chunk::{process_visual_chunk, process_visual_chunk_advanced, VisualChunkResult};
pub use yolo::{process_yolo, YoloFrame, YoloObject, YoloResult};
src/core/processor/tkg.rs
+153 -41
View File
@@ -38,7 +38,10 @@ fn load_face_pose_data(output_dir: &str, file_uuid: &str) -> Result<Vec<FacePose
let mut poses = Vec::new();
if let Some(frames) = json.get("frames").and_then(|v| v.as_array()) {
for frame_entry in frames {
let frame_num = frame_entry.get("frame").and_then(|v| v.as_i64()).unwrap_or(0);
let frame_num = frame_entry
.get("frame")
.and_then(|v| v.as_i64())
.unwrap_or(0);
if let Some(faces) = frame_entry.get("faces").and_then(|v| v.as_array()) {
for face in faces {
let bbox = match face.get("bbox") {
@@ -68,7 +71,14 @@ fn load_face_pose_data(output_dir: &str, file_uuid: &str) -> Result<Vec<FacePose
/// Match a face from face_detections (frame, x, y, w, h) to its pose in face.json
/// Uses bbox center distance to find the best match when multiple faces per frame.
fn get_pose_for_face(frame: i64, x: f64, y: f64, w: f64, h: f64, poses: &[FacePose]) -> Option<(f64, f64, f64)> {
fn get_pose_for_face(
frame: i64,
x: f64,
y: f64,
w: f64,
h: f64,
poses: &[FacePose],
) -> Option<(f64, f64, f64)> {
let cx = x + w / 2.0;
let cy = y + h / 2.0;
let mut best_dist = f64::MAX;
@@ -86,8 +96,12 @@ fn get_pose_for_face(frame: i64, x: f64, y: f64, w: f64, h: f64, poses: &[FacePo
}
fn detect_mutual_gaze(
bbox_a_x: f64, bbox_a_w: f64, yaw_a: f64,
bbox_b_x: f64, bbox_b_w: f64, yaw_b: f64,
bbox_a_x: f64,
bbox_a_w: f64,
yaw_a: f64,
bbox_b_x: f64,
bbox_b_w: f64,
yaw_b: f64,
threshold: f64,
) -> bool {
let cx_a = bbox_a_x + bbox_a_w / 2.0;
@@ -138,12 +152,16 @@ struct AsrxSegmentEntry {
#[serde(default)]
speaker_id: String,
#[serde(default)]
start_time: f64,
start: f64,
#[serde(default)]
end_time: f64,
end: f64,
#[serde(default)]
text: String,
#[allow(dead_code)]
#[serde(default)]
start_frame: i64,
#[allow(dead_code)]
#[serde(default)]
end_frame: i64,
}
@@ -195,7 +213,10 @@ pub struct TkgResult {
pub async fn build_tkg(db: &PostgresDb, file_uuid: &str, output_dir: &str) -> Result<TkgResult> {
let pool = db.pool();
let pose_data = load_face_pose_data(output_dir, file_uuid).unwrap_or_default();
tracing::info!("[TKG] Loaded {} pose entries from face.json", pose_data.len());
tracing::info!(
"[TKG] Loaded {} pose entries from face.json",
pose_data.len()
);
let n_face = build_face_trace_nodes(pool, file_uuid, &pose_data).await?;
let n_objects = build_yolo_object_nodes(pool, file_uuid, output_dir).await?;
@@ -217,7 +238,11 @@ pub async fn build_tkg(db: &PostgresDb, file_uuid: &str, output_dir: &str) -> Re
// ── Node builders ─────────────────────────────────────────────────
async fn build_face_trace_nodes(pool: &PgPool, file_uuid: &str, pose_data: &[FacePose]) -> Result<usize> {
async fn build_face_trace_nodes(
pool: &PgPool,
file_uuid: &str,
pose_data: &[FacePose],
) -> Result<usize> {
let face_table = t("face_detections");
let nodes_table = t("tkg_nodes");
@@ -257,7 +282,10 @@ async fn build_face_trace_nodes(pool: &PgPool, file_uuid: &str, pose_data: &[Fac
// Group by trace_id: trace_id → Vec<(frame, x, y, w, h)>
let mut trace_frames: HashMap<i64, Vec<(i64, f64, f64, f64, f64)>> = HashMap::new();
for (tid, frame, x, y, w, h) in &frame_rows {
trace_frames.entry(*tid).or_default().push((*frame, *x, *y, *w, *h));
trace_frames
.entry(*tid)
.or_default()
.push((*frame, *x, *y, *w, *h));
}
let mut count = 0;
@@ -274,7 +302,9 @@ async fn build_face_trace_nodes(pool: &PgPool, file_uuid: &str, pose_data: &[Fac
if let Some(frames) = trace_frames.get(&tid) {
for (frame, x, y, w, h) in 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;
@@ -284,7 +314,11 @@ async fn build_face_trace_nodes(pool: &PgPool, file_uuid: &str, pose_data: &[Fac
}
let (avg_yaw, avg_pitch, avg_roll) = if pose_count > 0 {
(yaw_sum / pose_count as f64, pitch_sum / pose_count as f64, roll_sum / pose_count as f64)
(
yaw_sum / pose_count as f64,
pitch_sum / pose_count as f64,
roll_sum / pose_count as f64,
)
} else {
(0.0, 0.0, 0.0)
};
@@ -401,8 +435,44 @@ async fn build_speaker_nodes(pool: &PgPool, file_uuid: &str, output_dir: &str) -
let nodes_table = t("tkg_nodes");
let mut count = 0;
// Group segments by speaker_id
let mut speaker_segments: HashMap<String, Vec<&AsrxSegmentEntry>> = HashMap::new();
for seg in &asrx.segments {
speaker_segments
.entry(seg.speaker_id.clone())
.or_default()
.push(seg);
}
for (sid, stat) in &stats {
let props = serde_json::json!({ "segment_count": stat.count });
let segs = speaker_segments.get(sid);
let (full_text, segments_json) = if let Some(seg_list) = segs {
let full: String = seg_list
.iter()
.map(|s| s.text.trim())
.filter(|t| !t.is_empty())
.collect::<Vec<_>>()
.join(" ");
let segments: Vec<serde_json::Value> = seg_list
.iter()
.map(|s| {
serde_json::json!({
"start": s.start,
"end": s.end,
"text": s.text,
})
})
.collect();
(full, serde_json::Value::Array(segments))
} else {
(String::new(), serde_json::Value::Array(vec![]))
};
let props = serde_json::json!({
"segment_count": stat.count,
"segments": segments_json,
"full_text": full_text,
});
sqlx::query(&format!(
r#"
@@ -576,8 +646,8 @@ async fn build_speaker_face_edges(
// 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
let fps = if last.end > 0.0 {
last.end_frame as f64 / last.end
} else {
30.0
};
@@ -604,8 +674,8 @@ async fn build_speaker_face_edges(
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 seg_start = seg.start;
let seg_end = seg.end;
let overlap_start = face_start_sec.max(seg_start);
let overlap_end = face_end_sec.min(seg_end);
@@ -669,7 +739,11 @@ async fn build_speaker_face_edges(
Ok(edge_count)
}
async fn build_face_face_edges(pool: &PgPool, file_uuid: &str, pose_data: &[FacePose]) -> Result<usize> {
async fn build_face_face_edges(
pool: &PgPool,
file_uuid: &str,
pose_data: &[FacePose],
) -> Result<usize> {
let face_table = t("face_detections");
let nodes_table = t("tkg_nodes");
let edges_table = t("tkg_edges");
@@ -722,8 +796,9 @@ async fn build_face_face_edges(pool: &PgPool, file_uuid: &str, pose_data: &[Face
(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)
}
@@ -770,7 +845,11 @@ async fn build_face_face_edges(pool: &PgPool, file_uuid: &str, pose_data: &[Face
};
let frames: Vec<i64> = frame_data.iter().map(|(f, _)| *f).collect();
let gaze_frames: Vec<i64> = frame_data.iter().filter(|(_, g)| *g).map(|(f, _)| *f).collect();
let gaze_frames: Vec<i64> = frame_data
.iter()
.filter(|(_, g)| *g)
.map(|(f, _)| *f)
.collect();
let gaze_count = gaze_frames.len() as i64;
let has_gaze = gaze_count > 0;
@@ -793,8 +872,13 @@ async fn build_face_face_edges(pool: &PgPool, file_uuid: &str, pose_data: &[Face
}
}
let (avg_ya, avg_yb) = if gaze_sample > 0 {
(yaw_a_sum / gaze_sample as f64, yaw_b_sum / gaze_sample as f64)
} else { (0.0, 0.0) };
(
yaw_a_sum / gaze_sample as f64,
yaw_b_sum / gaze_sample as f64,
)
} else {
(0.0, 0.0)
};
serde_json::json!({
"first_frame": frames[0],
@@ -902,9 +986,14 @@ pub async fn query_auto_representative_frame(
.context("Failed to detect main identities")?;
let main_ids: Vec<(i32, String, String, i64)> = mains;
let main_idents: Vec<MainIdentityInfo> = main_ids.iter().map(|(_, u, n, c)|
MainIdentityInfo { identity_uuid: u.clone(), name: n.clone(), face_count: *c }
).collect();
let main_idents: Vec<MainIdentityInfo> = main_ids
.iter()
.map(|(_, u, n, c)| MainIdentityInfo {
identity_uuid: u.clone(),
name: n.clone(),
face_count: *c,
})
.collect();
let frame_number: Option<i64> = if main_ids.len() >= 2 {
let id_a = main_ids[0].0;
@@ -915,16 +1004,20 @@ pub async fn query_auto_representative_frame(
AND trace_id IS NOT NULL GROUP BY trace_id ORDER BY COUNT(*) DESC LIMIT 1",
fd_table
))
.bind(file_uuid).bind(id_a)
.fetch_optional(pool).await?;
.bind(file_uuid)
.bind(id_a)
.fetch_optional(pool)
.await?;
let trace_b: Option<(i32,)> = sqlx::query_as(&format!(
"SELECT trace_id FROM {} WHERE file_uuid = $1 AND identity_id = $2 \
AND trace_id IS NOT NULL GROUP BY trace_id ORDER BY COUNT(*) DESC LIMIT 1",
fd_table
))
.bind(file_uuid).bind(id_b)
.fetch_optional(pool).await?;
.bind(file_uuid)
.bind(id_b)
.fetch_optional(pool)
.await?;
match (trace_a, trace_b) {
(Some((ta,)), Some((tb,))) => {
@@ -940,11 +1033,18 @@ pub async fn query_auto_representative_frame(
LIMIT 1",
edges_table, nodes_table, nodes_table
))
.bind(file_uuid).bind(ta).bind(tb)
.fetch_optional(pool).await?;
.bind(file_uuid)
.bind(ta)
.bind(tb)
.fetch_optional(pool)
.await?;
if let Some((f,)) = tkg_frame {
if f <= half_frame { Some(f) } else { None }
if f <= half_frame {
Some(f)
} else {
None
}
} else {
sqlx::query_scalar::<_, i64>(&format!(
"SELECT MIN(fd_a.frame_number)::bigint \
@@ -954,8 +1054,12 @@ pub async fn query_auto_representative_frame(
AND fd_b.identity_id = $3 AND fd_a.frame_number <= $4",
fd_table, fd_table
))
.bind(file_uuid).bind(id_a).bind(id_b).bind(half_frame)
.fetch_optional(pool).await?
.bind(file_uuid)
.bind(id_a)
.bind(id_b)
.bind(half_frame)
.fetch_optional(pool)
.await?
}
}
_ => None,
@@ -976,8 +1080,11 @@ pub async fn query_auto_representative_frame(
LIMIT 1",
fd_table
))
.bind(file_uuid).bind(first_id).bind(half_frame)
.fetch_optional(pool).await?
.bind(file_uuid)
.bind(first_id)
.bind(half_frame)
.fetch_optional(pool)
.await?
} else {
None
}
@@ -995,20 +1102,25 @@ pub async fn query_auto_representative_frame(
LIMIT 1",
fd_table
))
.bind(file_uuid).bind(half_frame)
.fetch_optional(pool).await?
.bind(file_uuid)
.bind(half_frame)
.fetch_optional(pool)
.await?
}
};
let frame_number = frame_number.ok_or_else(|| anyhow::anyhow!("No faces found in this file"))?;
let frame_number =
frame_number.ok_or_else(|| anyhow::anyhow!("No faces found in this file"))?;
let face_quality: f64 = sqlx::query_scalar::<_, f64>(&format!(
"SELECT COALESCE(MAX((width::float8 * height::float8) * confidence::float8), 0) \
FROM {} WHERE file_uuid = $1 AND frame_number = $2",
fd_table
))
.bind(file_uuid).bind(frame_number)
.fetch_one(pool).await?;
.bind(file_uuid)
.bind(frame_number)
.fetch_one(pool)
.await?;
let traces: Vec<FrameTraceInfo> = sqlx::query_as::<_, (i32, Option<String>, Option<String>, i32, i32, i32, i32, f64)>(&format!(
"SELECT fd.trace_id, i.uuid::text, i.name, fd.x, fd.y, fd.width, fd.height, fd.confidence::float8 \
src/core/processor/visual_chunk.rs
-594
View File
@@ -1,594 +0,0 @@
//! 視覺分片處理器 (Phase 2.2)
//!
//! 從 YOLO 結果生成視覺分片
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::time::Duration;
use super::executor::PythonExecutor;
use super::yolo::{YoloFrame, YoloResult};
const VISUAL_CHUNK_TIMEOUT: Duration = Duration::from_secs(3600);
/// 視覺分片處理結果
#[derive(Debug, Serialize, Deserialize, Clone, Default)]
pub struct VisualChunkResult {
/// 生成的視覺分片數量
pub chunk_count: u32,
/// 處理的總幀數
pub total_frames: u32,
/// 檢測到的總物件數
pub total_objects: u32,
/// 唯一物件類別數
pub unique_classes: u32,
/// 生成的視覺分片
pub chunks: Vec<crate::core::chunk::Chunk>,
}
/// 從 YOLO 結果生成視覺分片
pub async fn process_visual_chunk(
file_id: i32,
uuid: String,
video_path: &str,
yolo_result: &YoloResult,
chunk_index_offset: u32,
fps: f64,
) -> Result<VisualChunkResult> {
tracing::info!(
"[VisualChunk] Starting visual chunk generation for video: {}, {} frames",
video_path,
yolo_result.frames.len()
);
if yolo_result.frames.is_empty() {
tracing::warn!("[VisualChunk] No YOLO frames to process");
return Ok(VisualChunkResult {
chunk_count: 0,
total_frames: 0,
total_objects: 0,
unique_classes: 0,
chunks: vec![],
});
}
// 策略 1: 固定幀數分片（每 N 幀一個分片）
let chunks = create_fixed_frame_chunks(file_id, &uuid, yolo_result, chunk_index_offset, fps);
// 統計信息
let total_objects: u32 = yolo_result
.frames
.iter()
.map(|f| f.objects.len() as u32)
.sum();
let all_classes: Vec<String> = yolo_result
.frames
.iter()
.flat_map(|f| f.objects.iter().map(|o| o.class_name.clone()))
.collect();
let unique_classes: u32 = all_classes
.iter()
.cloned()
.collect::<std::collections::HashSet<_>>()
.len() as u32;
tracing::info!(
"[VisualChunk] Generated {} visual chunks from {} frames, {} total objects, {} unique classes",
chunks.len(),
yolo_result.frames.len(),
total_objects,
unique_classes
);
Ok(VisualChunkResult {
chunk_count: chunks.len() as u32,
total_frames: yolo_result.frames.len() as u32,
total_objects,
unique_classes,
chunks,
})
}
/// 創建固定幀數分片（每 N 幀一個分片）
fn create_fixed_frame_chunks(
file_id: i32,
uuid: &str,
yolo_result: &YoloResult,
chunk_index_offset: u32,
fps: f64,
) -> Vec<crate::core::chunk::Chunk> {
let mut chunks = Vec::new();
// 配置：每 30 幀創建一個分片（約 1 秒，如果 fps=30）
let frames_per_chunk = 30;
let total_frames = yolo_result.frames.len();
if total_frames == 0 {
return chunks;
}
let mut chunk_index = chunk_index_offset;
let mut start_idx = 0;
while start_idx < total_frames {
let end_idx = std::cmp::min(start_idx + frames_per_chunk, total_frames);
// 獲取這個分片的幀
let chunk_frames: Vec<YoloFrame> = yolo_result.frames[start_idx..end_idx]
.iter()
.cloned()
.collect();
if chunk_frames.is_empty() {
break;
}
// 計算幀範圍
let start_frame = chunk_frames.first().unwrap().frame as i64;
let end_frame = chunk_frames.last().unwrap().frame as i64 + 1; // exclusive
// 創建視覺分片
let chunk = crate::core::chunk::Chunk::from_yolo_frames(
file_id,
uuid.to_string(),
format!("vis_{}", chunk_index),
start_frame,
end_frame,
fps,
chunk_frames,
);
chunks.push(chunk);
// 更新索引
start_idx = end_idx;
chunk_index += 1;
}
chunks
}
/// 基於物件相似度創建分片
fn create_similarity_based_chunks(
file_id: i32,
uuid: &str,
yolo_result: &YoloResult,
chunk_index_offset: u32,
fps: f64,
similarity_threshold: f32,
min_frames_per_chunk: usize,
) -> Vec<crate::core::chunk::Chunk> {
let mut chunks = Vec::new();
if yolo_result.frames.is_empty() {
return chunks;
}
let mut current_chunk_frames: Vec<YoloFrame> = Vec::new();
let mut chunk_index = chunk_index_offset;
let mut current_start_frame = 0;
for (i, frame) in yolo_result.frames.iter().enumerate() {
if current_chunk_frames.is_empty() {
current_chunk_frames.push(frame.clone());
current_start_frame = frame.frame as i64;
continue;
}
// 檢查相似度（簡化版本：檢查物件類別是否相同）
let last_frame = current_chunk_frames.last().unwrap();
let similarity = calculate_frame_similarity(last_frame, frame);
if similarity >= similarity_threshold {
// 相似度高，加入當前分片
current_chunk_frames.push(frame.clone());
} else {
// 相似度低，創建新分片
if current_chunk_frames.len() >= min_frames_per_chunk {
let end_frame = current_chunk_frames.last().unwrap().frame as i64 + 1;
let chunk = crate::core::chunk::Chunk::from_yolo_frames(
file_id,
uuid.to_string(),
format!("vis_{}", chunk_index),
current_start_frame,
end_frame,
fps,
current_chunk_frames.clone(),
);
chunks.push(chunk);
chunk_index += 1;
}
// 開始新的分片
current_chunk_frames = vec![frame.clone()];
current_start_frame = frame.frame as i64;
}
}
// 處理最後一個分片
if current_chunk_frames.len() >= min_frames_per_chunk {
let end_frame = current_chunk_frames.last().unwrap().frame as i64 + 1;
let chunk = crate::core::chunk::Chunk::from_yolo_frames(
file_id,
uuid.to_string(),
format!("vis_{}", chunk_index),
current_start_frame,
end_frame,
fps,
current_chunk_frames,
);
chunks.push(chunk);
}
chunks
}
/// 計算兩個幀之間的相似度（基於物件類別）
fn calculate_frame_similarity(frame1: &YoloFrame, frame2: &YoloFrame) -> f32 {
if frame1.objects.is_empty() && frame2.objects.is_empty() {
return 1.0;
}
if frame1.objects.is_empty() || frame2.objects.is_empty() {
return 0.0;
}
let set1: std::collections::HashSet<String> = frame1
.objects
.iter()
.map(|o| o.class_name.clone())
.collect();
let set2: std::collections::HashSet<String> = frame2
.objects
.iter()
.map(|o| o.class_name.clone())
.collect();
let intersection: Vec<_> = set1.intersection(&set2).collect();
let union: Vec<_> = set1.union(&set2).collect();
if union.is_empty() {
0.0
} else {
intersection.len() as f32 / union.len() as f32
}
}
/// 使用 Python 腳本生成視覺分片（進階版本）
pub async fn process_visual_chunk_advanced(
video_path: &str,
output_path: &str,
uuid: Option<&str>,
) -> Result<VisualChunkResult> {
let executor = PythonExecutor::new()?;
let script_path = executor.script_path("visual_chunk_processor.py");
tracing::info!(
"[VisualChunk] Starting advanced visual chunk generation: {}",
video_path
);
if !script_path.exists() {
tracing::warn!("[VisualChunk] Script not found, using basic generation");
// 這裡可以回退到基本生成方法
return Ok(VisualChunkResult {
chunk_count: 0,
total_frames: 0,
total_objects: 0,
unique_classes: 0,
chunks: vec![],
});
}
let yolo_path = uuid.map(|u| {
std::path::PathBuf::from(crate::core::config::OUTPUT_DIR.as_str())
.join(format!("{}.yolo.json", u))
.to_string_lossy()
.to_string()
});
let args: &[&str] = if let Some(ref yp) = yolo_path {
&[video_path, output_path, "--yolo-result", yp]
} else {
&[video_path, output_path]
};
let result = match executor
.run(
"visual_chunk_processor.py",
args,
uuid,
"VisualChunk",
Some(VISUAL_CHUNK_TIMEOUT),
)
.await
{
Ok(_) => match std::fs::read_to_string(output_path) {
Ok(json_str) => match serde_json::from_str::<VisualChunkResult>(&json_str) {
Ok(r) => r,
Err(e) => {
tracing::warn!(
"[VisualChunk] Failed to parse output ({}), returning empty",
e
);
VisualChunkResult::default()
}
},
Err(e) => {
tracing::warn!(
"[VisualChunk] Failed to read output ({}), returning empty",
e
);
VisualChunkResult::default()
}
},
Err(e) => {
tracing::warn!(
"[VisualChunk] Failed to run script ({}), returning empty",
e
);
VisualChunkResult::default()
}
};
tracing::info!(
"[VisualChunk] Advanced generation result: {} chunks, {} frames",
result.chunk_count,
result.total_frames
);
Ok(result)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_calculate_frame_similarity() {
use crate::core::processor::yolo::{YoloFrame, YoloObject};
let frame1 = YoloFrame {
frame: 0,
timestamp: 0.0,
objects: vec![
YoloObject {
class_name: "person".to_string(),
class_id: 0,
x: 100,
y: 200,
width: 50,
height: 100,
confidence: 0.95,
},
YoloObject {
class_name: "car".to_string(),
class_id: 2,
x: 300,
y: 150,
width: 80,
height: 60,
confidence: 0.87,
},
],
};
let frame2 = YoloFrame {
frame: 1,
timestamp: 0.033,
objects: vec![
YoloObject {
class_name: "person".to_string(),
class_id: 0,
x: 110,
y: 210,
width: 52,
height: 102,
confidence: 0.92,
},
YoloObject {
class_name: "car".to_string(),
class_id: 2,
x: 310,
y: 155,
width: 82,
height: 62,
confidence: 0.85,
},
],
};
let frame3 = YoloFrame {
frame: 2,
timestamp: 0.066,
objects: vec![YoloObject {
class_name: "dog".to_string(),
class_id: 16,
x: 150,
y: 250,
width: 40,
height: 60,
confidence: 0.78,
}],
};
// 相同物件的幀應該高度相似
let similarity_same = calculate_frame_similarity(&frame1, &frame2);
assert!((similarity_same - 1.0).abs() < 0.001);
// 不同物件的幀應該不相似
let similarity_diff = calculate_frame_similarity(&frame1, &frame3);
assert!((similarity_diff - 0.0).abs() < 0.001);
// 空幀應該完全相似
let empty_frame = YoloFrame {
frame: 3,
timestamp: 0.1,
objects: vec![],
};
let similarity_empty = calculate_frame_similarity(&empty_frame, &empty_frame);
assert!((similarity_empty - 1.0).abs() < 0.001);
}
#[tokio::test]
async fn test_create_fixed_frame_chunks() {
use crate::core::processor::yolo::{YoloFrame, YoloObject, YoloResult};
// 創建測試 YOLO 結果（60 幀，每幀都有物件）
let mut frames = Vec::new();
for i in 0..60 {
frames.push(YoloFrame {
frame: i as u64,
timestamp: i as f64 / 30.0, // 假設 fps=30
objects: vec![YoloObject {
class_name: "person".to_string(),
class_id: 0,
x: 100,
y: 200,
width: 50,
height: 100,
confidence: 0.9,
}],
});
}
let yolo_result = YoloResult {
frame_count: 60,
fps: 30.0,
frames,
};
let chunks = create_fixed_frame_chunks(1, "test-uuid", &yolo_result, 0, 30.0);
// 60 幀，每 30 幀一個分片，應該有 2 個分片
assert_eq!(chunks.len(), 2);
// 檢查第一個分片
let first_chunk = &chunks[0];
assert_eq!(
first_chunk.chunk_type,
crate::core::chunk::ChunkType::Visual
);
assert_eq!(first_chunk.start_frame, 0);
assert_eq!(first_chunk.end_frame, 30); // exclusive
assert_eq!(first_chunk.frame_count, 30);
// 檢查第二個分片
let second_chunk = &chunks[1];
assert_eq!(
second_chunk.chunk_type,
crate::core::chunk::ChunkType::Visual
);
assert_eq!(second_chunk.start_frame, 30);
assert_eq!(second_chunk.end_frame, 60); // exclusive
assert_eq!(second_chunk.frame_count, 30);
}
#[test]
fn test_create_similarity_based_chunks() {
use crate::core::processor::yolo::{YoloFrame, YoloObject, YoloResult};
// 創建測試 YOLO 結果
let frames = vec![
YoloFrame {
// 幀 0-4: 都有 person 和 car
frame: 0,
timestamp: 0.0,
objects: vec![
YoloObject {
class_name: "person".to_string(),
class_id: 0,
x: 100,
y: 200,
width: 50,
height: 100,
confidence: 0.9,
},
YoloObject {
class_name: "car".to_string(),
class_id: 2,
x: 300,
y: 150,
width: 80,
height: 60,
confidence: 0.8,
},
],
},
YoloFrame {
// 幀 1
frame: 1,
timestamp: 0.033,
objects: vec![
YoloObject {
class_name: "person".to_string(),
class_id: 0,
x: 110,
y: 210,
width: 52,
height: 102,
confidence: 0.88,
},
YoloObject {
class_name: "car".to_string(),
class_id: 2,
x: 310,
y: 155,
width: 82,
height: 62,
confidence: 0.78,
},
],
},
YoloFrame {
// 幀 5-9: 只有 dog
frame: 5,
timestamp: 0.166,
objects: vec![YoloObject {
class_name: "dog".to_string(),
class_id: 16,
x: 150,
y: 250,
width: 40,
height: 60,
confidence: 0.7,
}],
},
YoloFrame {
// 幀 6
frame: 6,
timestamp: 0.2,
objects: vec![YoloObject {
class_name: "dog".to_string(),
class_id: 16,
x: 155,
y: 255,
width: 42,
height: 62,
confidence: 0.68,
}],
},
];
let yolo_result = YoloResult {
frame_count: 7,
fps: 30.0,
frames,
};
let chunks = create_similarity_based_chunks(
1,
"test-uuid",
&yolo_result,
0,
30.0,
0.5, // similarity threshold
2, // min frames per chunk
);
// 應該有 2 個分片：一個是 person+car，一個是 dog
assert_eq!(chunks.len(), 2);
}
}