feat: 新增 Job Worker 系統與 API 文檔全面更新

This commit is contained in:
Warren
2026-03-26 16:16:34 +08:00
parent 80399b1c12
commit 82955504f3
70 changed files with 3460 additions and 376 deletions

View File

@@ -20,7 +20,7 @@ impl ChunkSplitter {
while current_time < duration {
let end_time = (current_time + self.time_based_duration).min(duration);
chunks.push(Chunk::new(
chunks.push(Chunk::from_seconds(
0, // file_id
uuid.to_string(),
index,
@@ -45,7 +45,7 @@ impl ChunkSplitter {
let mut chunks = Vec::new();
for (index, segment) in asr_segments.iter().enumerate() {
chunks.push(Chunk::new(
chunks.push(Chunk::from_seconds(
0, // file_id
uuid.to_string(),
index as u32,

View File

@@ -1,3 +1,4 @@
use crate::core::time::FrameTime;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
@@ -46,10 +47,11 @@ pub struct Chunk {
pub chunk_index: u32,
pub chunk_type: ChunkType,
pub rule: ChunkRule,
pub start_time: f64,
pub end_time: f64,
/// Frames per second (can be fractional, e.g., 29.97, 23.976)
pub fps: f64,
/// Start frame (0-based)
pub start_frame: i64,
/// End frame (exclusive)
pub end_frame: i64,
pub text_content: Option<String>,
pub content: serde_json::Value,
@@ -62,6 +64,13 @@ pub struct Chunk {
}
impl Chunk {
/// Creates a new chunk from frame counts.
///
/// # Arguments
///
/// * `start_frame` - Start frame (0-based)
/// * `end_frame` - End frame (exclusive)
/// * `fps` - Frames per second (can be fractional)
#[allow(clippy::too_many_arguments)]
pub fn new(
file_id: i32,
@@ -69,13 +78,11 @@ impl Chunk {
chunk_index: u32,
chunk_type: ChunkType,
rule: ChunkRule,
start_time: f64,
end_time: f64,
start_frame: i64,
end_frame: i64,
fps: f64,
content: serde_json::Value,
) -> Self {
let start_frame = (start_time * fps) as i64;
let end_frame = (end_time * fps) as i64;
let chunk_id = format!("{}_{:04}", chunk_type.as_str(), chunk_index);
Self {
file_id,
@@ -84,8 +91,6 @@ impl Chunk {
chunk_index,
chunk_type,
rule,
start_time,
end_time,
fps,
start_frame,
end_frame,
@@ -100,6 +105,95 @@ impl Chunk {
}
}
/// Creates a new chunk from seconds (legacy conversion).
///
/// This is useful for migrating from older systems that store time as seconds.
/// The frame counts are calculated by rounding `seconds * fps`.
#[allow(clippy::too_many_arguments)]
pub fn from_seconds(
file_id: i32,
uuid: String,
chunk_index: u32,
chunk_type: ChunkType,
rule: ChunkRule,
start_time: f64,
end_time: f64,
fps: f64,
content: serde_json::Value,
) -> Self {
let start_frame = (start_time * fps).round() as i64;
let end_frame = (end_time * fps).round() as i64;
Self::new(
file_id,
uuid,
chunk_index,
chunk_type,
rule,
start_frame,
end_frame,
fps,
content,
)
}
/// Returns the start time as a `FrameTime`.
pub fn start_time(&self) -> FrameTime {
FrameTime::from_frames(self.start_frame, self.fps)
}
/// Returns the end time as a `FrameTime`.
pub fn end_time(&self) -> FrameTime {
FrameTime::from_frames(self.end_frame, self.fps)
}
/// Returns the duration in frames.
pub fn duration_frames(&self) -> i64 {
self.end_frame - self.start_frame
}
/// Returns the duration in seconds.
pub fn duration_seconds(&self) -> f64 {
self.duration_frames() as f64 / self.fps
}
/// Formats the start time as "seconds.frame" (e.g., "123.04").
pub fn format_start_sec_frame(&self) -> String {
self.start_time().format_sec_frame()
}
/// Formats the end time as "seconds.frame" (e.g., "456.15").
pub fn format_end_sec_frame(&self) -> String {
self.end_time().format_sec_frame()
}
/// Formats the start time as "HH:MM:SS".
pub fn format_start_hms(&self) -> String {
self.start_time().format_hms()
}
/// Formats the end time as "HH:MM:SS".
pub fn format_end_hms(&self) -> String {
self.end_time().format_hms()
}
/// Formats the start time as "HH:MM:SS.FF".
pub fn format_start_hms_frame(&self) -> String {
self.start_time().format_hms_frame()
}
/// Formats the end time as "HH:MM:SS.FF".
pub fn format_end_hms_frame(&self) -> String {
self.end_time().format_hms_frame()
}
/// Returns a tuple of (start_seconds, end_seconds) for compatibility.
///
/// This is provided for backward compatibility during migration.
/// Prefer using `start_time()` and `end_time()` methods.
pub fn time_range_seconds(&self) -> (f64, f64) {
(self.start_time().seconds(), self.end_time().seconds())
}
pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
self.metadata = Some(metadata);
self

View File

@@ -28,13 +28,15 @@ pub struct ChunkDocument {
impl From<Chunk> for ChunkDocument {
fn from(chunk: Chunk) -> Self {
let start_time = chunk.start_time().seconds();
let end_time = chunk.end_time().seconds();
Self {
uuid: chunk.uuid,
chunk_id: chunk.chunk_id,
chunk_index: chunk.chunk_index,
chunk_type: chunk.chunk_type.as_str().to_string(),
start_time: chunk.start_time,
end_time: chunk.end_time,
start_time,
end_time,
fps: chunk.fps,
start_frame: chunk.start_frame,
end_frame: chunk.end_frame,
@@ -118,8 +120,6 @@ impl MongoDb {
chunk_index: doc.chunk_index,
chunk_type,
rule: ChunkRule::Rule1,
start_time: doc.start_time,
end_time: doc.end_time,
fps: doc.fps,
start_frame: doc.start_frame,
end_frame: doc.end_frame,
@@ -178,8 +178,6 @@ impl MongoDb {
chunk_index: doc.chunk_index,
chunk_type,
rule: ChunkRule::Rule1,
start_time: doc.start_time,
end_time: doc.end_time,
fps: doc.fps,
start_frame: doc.start_frame,
end_frame: doc.end_frame,
@@ -235,8 +233,6 @@ impl MongoDb {
chunk_index: doc.chunk_index,
chunk_type,
rule: ChunkRule::Rule1,
start_time: doc.start_time,
end_time: doc.end_time,
fps: doc.fps,
start_frame: doc.start_frame,
end_frame: doc.end_frame,

View File

@@ -126,8 +126,6 @@ pub struct PreChunk {
pub source_type: String,
pub source_file: Option<String>,
pub chunk_type: String,
pub start_time: f64,
pub end_time: f64,
pub start_frame: i64,
pub end_frame: i64,
pub fps: f64,
@@ -209,7 +207,7 @@ pub struct MonitorJobStats {
pub failed: i32,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
#[serde(rename_all = "snake_case")]
pub enum ProcessorType {
Asr,
@@ -449,6 +447,12 @@ impl PostgresDb {
.parse::<u64>()
.unwrap_or(60);
tracing::info!(
"DB pool config: max_connections={}, acquire_timeout={}s",
max_connections,
acquire_timeout_secs
);
let pool_options = PgPoolOptions::new()
.max_connections(max_connections)
.acquire_timeout(std::time::Duration::from_secs(acquire_timeout_secs));
@@ -1770,8 +1774,8 @@ impl PostgresDb {
.bind(&chunk.chunk_id)
.bind(chunk.chunk_index as i32)
.bind(chunk.chunk_type.as_str())
.bind(chunk.start_time)
.bind(chunk.end_time)
.bind(chunk.start_time().seconds())
.bind(chunk.end_time().seconds())
.bind(chunk.fps)
.bind(chunk.start_frame)
.bind(chunk.end_frame)
@@ -1791,7 +1795,7 @@ impl PostgresDb {
pub async fn get_chunks_by_uuid(&self, uuid: &str) -> Result<Vec<Chunk>> {
let rows = sqlx::query(
"SELECT COALESCE(file_id, 0) as file_id, uuid, chunk_id, chunk_index, chunk_type, start_time, end_time, COALESCE(fps, 24.0) as fps, COALESCE(start_frame, 0) as start_frame, COALESCE(end_frame, 0) as end_frame, text_content, content, metadata, vector_id, COALESCE(frame_count, 0) as frame_count, pre_chunk_ids, parent_chunk_id, child_chunk_ids FROM chunks WHERE uuid = $1 ORDER BY chunk_index"
"SELECT COALESCE(file_id, 0) as file_id, uuid, chunk_id, chunk_index, chunk_type, COALESCE(fps, 24.0) as fps, COALESCE(start_frame, 0) as start_frame, COALESCE(end_frame, 0) as end_frame, text_content, content, metadata, vector_id, COALESCE(frame_count, 0) as frame_count, pre_chunk_ids, parent_chunk_id, child_chunk_ids FROM chunks WHERE uuid = $1 ORDER BY chunk_index"
)
.bind(uuid)
.fetch_all(&self.pool)
@@ -1811,12 +1815,12 @@ impl PostgresDb {
_ => ChunkType::TimeBased,
};
let content: serde_json::Value = r.get(11);
let metadata: Option<serde_json::Value> = r.get(12);
let content: serde_json::Value = r.get(9);
let metadata: Option<serde_json::Value> = r.get(10);
let pre_chunk_ids: Vec<i32> = r.try_get(15).unwrap_or_default();
let parent_chunk_id: Option<String> = r.try_get(16).ok().flatten();
let child_chunk_ids: Vec<String> = r.try_get(17).unwrap_or_default();
let pre_chunk_ids: Vec<i32> = r.try_get(13).unwrap_or_default();
let parent_chunk_id: Option<String> = r.try_get(14).ok().flatten();
let child_chunk_ids: Vec<String> = r.try_get(15).unwrap_or_default();
let (rule, content_data) = if content.get("rule").is_some() {
let rule_str = content
@@ -1844,8 +1848,7 @@ impl PostgresDb {
chunk_index: chunk_index as u32,
chunk_type,
rule,
start_time: r.get("start_time"),
end_time: r.get("end_time"),
fps: r.get("fps"),
start_frame: r.get("start_frame"),
end_frame: r.get("end_frame"),
@@ -1866,7 +1869,7 @@ impl PostgresDb {
pub async fn get_chunk_by_chunk_id(&self, chunk_id: &str) -> Result<Option<Chunk>> {
let row = sqlx::query(
"SELECT COALESCE(file_id, 0) as file_id, uuid, chunk_id, chunk_index, chunk_type, start_time, end_time, COALESCE(fps, 24.0) as fps, COALESCE(start_frame, 0) as start_frame, COALESCE(end_frame, 0) as end_frame, text_content, content, metadata, vector_id, COALESCE(frame_count, 0) as frame_count, pre_chunk_ids, parent_chunk_id, child_chunk_ids FROM chunks WHERE chunk_id = $1"
"SELECT COALESCE(file_id, 0) as file_id, uuid, chunk_id, chunk_index, chunk_type, COALESCE(fps, 24.0) as fps, COALESCE(start_frame, 0) as start_frame, COALESCE(end_frame, 0) as end_frame, text_content, content, metadata, vector_id, COALESCE(frame_count, 0) as frame_count, pre_chunk_ids, parent_chunk_id, child_chunk_ids FROM chunks WHERE chunk_id = $1"
)
.bind(chunk_id)
.fetch_optional(&self.pool)
@@ -1884,12 +1887,12 @@ impl PostgresDb {
_ => ChunkType::TimeBased,
};
let content: serde_json::Value = r.get(11);
let metadata: Option<serde_json::Value> = r.get(12);
let content: serde_json::Value = r.get(9);
let metadata: Option<serde_json::Value> = r.get(10);
let pre_chunk_ids: Vec<i32> = r.try_get(15).unwrap_or_default();
let parent_chunk_id: Option<String> = r.try_get(16).ok().flatten();
let child_chunk_ids: Vec<String> = r.try_get(17).unwrap_or_default();
let pre_chunk_ids: Vec<i32> = r.try_get(13).unwrap_or_default();
let parent_chunk_id: Option<String> = r.try_get(14).ok().flatten();
let child_chunk_ids: Vec<String> = r.try_get(15).unwrap_or_default();
let (rule, content_data) = if content.get("rule").is_some() {
let rule_str = content
@@ -1917,8 +1920,6 @@ impl PostgresDb {
chunk_index: chunk_index as u32,
chunk_type,
rule,
start_time: r.get("start_time"),
end_time: r.get("end_time"),
fps: r.get("fps"),
start_frame: r.get("start_frame"),
end_frame: r.get("end_frame"),
@@ -1942,6 +1943,9 @@ impl PostgresDb {
INSERT INTO pre_chunks (file_id, source_type, source_file, chunk_type, start_time, end_time, start_frame, end_frame, fps, raw_json, text_content, processed, chunk_id)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
ON CONFLICT (file_id, source_type, start_frame, end_frame) DO UPDATE SET
start_time = EXCLUDED.start_time,
end_time = EXCLUDED.end_time,
fps = EXCLUDED.fps,
raw_json = EXCLUDED.raw_json,
text_content = EXCLUDED.text_content,
processed = EXCLUDED.processed,
@@ -1953,8 +1957,8 @@ impl PostgresDb {
.bind(&pre_chunk.source_type)
.bind(&pre_chunk.source_file)
.bind(&pre_chunk.chunk_type)
.bind(pre_chunk.start_time)
.bind(pre_chunk.end_time)
.bind(pre_chunk.start_frame as f64 / pre_chunk.fps)
.bind(pre_chunk.end_frame as f64 / pre_chunk.fps)
.bind(pre_chunk.start_frame)
.bind(pre_chunk.end_frame)
.bind(pre_chunk.fps)
@@ -2108,8 +2112,7 @@ impl PostgresDb {
chunk_index: chunk_index as u32,
chunk_type,
rule,
start_time: r.get("start_time"),
end_time: r.get("end_time"),
fps: r.get("fps"),
start_frame: r.get("start_frame"),
end_frame: r.get("end_frame"),
@@ -2134,7 +2137,7 @@ impl PostgresDb {
}
let rows = sqlx::query(
"SELECT file_id, uuid, chunk_id, chunk_index, chunk_type, start_time, end_time, fps, start_frame, end_frame, text_content, content, metadata, vector_id, frame_count, pre_chunk_ids, parent_chunk_id, child_chunk_ids FROM chunks WHERE chunk_id = ANY($1) ORDER BY chunk_index",
"SELECT file_id, uuid, chunk_id, chunk_index, chunk_type, fps, start_frame, end_frame, text_content, content, metadata, vector_id, frame_count, pre_chunk_ids, parent_chunk_id, child_chunk_ids FROM chunks WHERE chunk_id = ANY($1) ORDER BY chunk_index",
)
.bind(chunk_ids)
.fetch_all(&self.pool)
@@ -2154,12 +2157,12 @@ impl PostgresDb {
_ => ChunkType::TimeBased,
};
let content: serde_json::Value = r.get(11);
let metadata: Option<serde_json::Value> = r.get(12);
let content: serde_json::Value = r.get(9);
let metadata: Option<serde_json::Value> = r.get(10);
let pre_chunk_ids: Vec<i32> = r.try_get(15).unwrap_or_default();
let parent_chunk_id: Option<String> = r.try_get(16).ok().flatten();
let child_chunk_ids: Vec<String> = r.try_get(17).unwrap_or_default();
let pre_chunk_ids: Vec<i32> = r.try_get(13).unwrap_or_default();
let parent_chunk_id: Option<String> = r.try_get(14).ok().flatten();
let child_chunk_ids: Vec<String> = r.try_get(15).unwrap_or_default();
let (rule, content_data) = if content.get("rule").is_some() {
let rule_str = content
@@ -2187,8 +2190,7 @@ impl PostgresDb {
chunk_index: chunk_index as u32,
chunk_type,
rule,
start_time: r.get("start_time"),
end_time: r.get("end_time"),
fps: r.get("fps"),
start_frame: r.get("start_frame"),
end_frame: r.get("end_frame"),
@@ -2337,8 +2339,6 @@ impl PostgresDb {
chunk_index: r.2 as u32,
chunk_type,
rule: ChunkRule::Rule1,
start_time: r.4,
end_time: r.5,
fps: r.6,
start_frame: r.7,
end_frame: r.8,
@@ -2497,8 +2497,8 @@ impl PostgresDb {
chunk_type: chunk_data
.map(|c| c.chunk_type.as_str().to_string())
.unwrap_or_default(),
start_time: chunk_data.map(|c| c.start_time).unwrap_or(0.0),
end_time: chunk_data.map(|c| c.end_time).unwrap_or(0.0),
start_time: chunk_data.map(|c| c.start_time().seconds()).unwrap_or(0.0),
end_time: chunk_data.map(|c| c.end_time().seconds()).unwrap_or(0.0),
text: chunk_data
.and_then(|c| c.text_content.clone())
.unwrap_or_default(),
@@ -2584,6 +2584,7 @@ impl PostgresDb {
error_count, last_error, started_at, updated_at, created_at
FROM monitor_jobs
WHERE status = 'pending'
OR (status = 'running' AND EXISTS (SELECT 1 FROM processor_results WHERE job_id = monitor_jobs.id AND status = 'pending'))
ORDER BY created_at ASC
LIMIT $1
FOR UPDATE SKIP LOCKED
@@ -2619,6 +2620,77 @@ impl PostgresDb {
Ok(jobs)
}
pub async fn get_running_jobs_with_all_processors_done(
&self,
limit: i32,
) -> Result<Vec<MonitorJob>> {
let rows = sqlx::query(
r#"
SELECT id, uuid, video_path, status, current_processor, progress_total, progress_current,
error_count, last_error, started_at, updated_at, created_at
FROM monitor_jobs
WHERE status = 'running'
AND NOT EXISTS (
SELECT 1 FROM processor_results pr
WHERE pr.job_id = monitor_jobs.id
AND pr.status IN ('pending', 'running')
)
ORDER BY updated_at ASC
LIMIT $1
FOR UPDATE SKIP LOCKED
"#
)
.bind(limit)
.fetch_all(&self.pool)
.await?;
let jobs: Vec<MonitorJob> = rows
.into_iter()
.map(|r| {
let status_str: String = r.get(3);
let status =
MonitorJobStatus::from_db_str(&status_str).unwrap_or(MonitorJobStatus::Pending);
MonitorJob {
id: r.get(0),
uuid: r.get(1),
video_path: r.get(2),
status,
current_processor: r.get(4),
progress_total: r.get(5),
progress_current: r.get(6),
error_count: r.get(7),
last_error: r.get(8),
started_at: r.get(9),
updated_at: r.get(10),
created_at: r.get(11),
}
})
.collect();
Ok(jobs)
}
pub async fn update_job_processors_arrays(
&self,
job_id: i32,
completed_processors: Vec<String>,
failed_processors: Vec<String>,
) -> Result<()> {
sqlx::query(
"UPDATE monitor_jobs
SET completed_processors = $1,
failed_processors = $2,
updated_at = CURRENT_TIMESTAMP
WHERE id = $3",
)
.bind(completed_processors)
.bind(failed_processors)
.bind(job_id)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn update_job_status(&self, job_id: i32, status: MonitorJobStatus) -> Result<()> {
sqlx::query(
"UPDATE monitor_jobs SET status = $1, updated_at = CURRENT_TIMESTAMP WHERE id = $2",
@@ -2660,6 +2732,7 @@ impl PostgresDb {
r#"
INSERT INTO processor_results (job_id, processor, status)
VALUES ($1, $2, 'pending')
ON CONFLICT (job_id, processor) DO UPDATE SET job_id = EXCLUDED.job_id
RETURNING id
"#,
)
@@ -2685,8 +2758,8 @@ impl PostgresDb {
SET status = $1,
error_message = $2,
output_data = $3,
started_at = CASE WHEN $1 = 'running' AND started_at IS NULL THEN CURRENT_TIMESTAMP ELSE started_at END,
completed_at = CASE WHEN $1 IN ('completed', 'failed', 'skipped') THEN CURRENT_TIMESTAMP ELSE completed_at END,
duration_secs = CASE WHEN $1 IN ('completed', 'failed', 'skipped') THEN EXTRACT(EPOCH FROM (CURRENT_TIMESTAMP - started_at)) ELSE duration_secs END,
updated_at = CURRENT_TIMESTAMP
WHERE id = $4
"#,
@@ -2705,7 +2778,7 @@ impl PostgresDb {
r#"
SELECT id, job_id, processor, status, output_path, started_at, completed_at,
error_message, progress_total, progress_current, last_checkpoint,
created_at, updated_at, duration_secs
created_at, updated_at, duration_secs
FROM processor_results
WHERE job_id = $1
ORDER BY created_at ASC

View File

@@ -26,8 +26,8 @@ impl SyncDb {
let uuid = chunk.uuid.clone();
let chunk_id = chunk.chunk_id.clone();
let chunk_type = chunk.chunk_type.as_str().to_string();
let start_time = chunk.start_time;
let end_time = chunk.end_time;
let start_time = chunk.start_time().seconds();
let end_time = chunk.end_time().seconds();
let vector = self.embed_text(text).await?;
@@ -117,7 +117,7 @@ impl SyncDb {
"language_probability": asr_result.language_probability,
});
let chunk = Chunk::new(
let chunk = Chunk::from_seconds(
0, // file_id - will be set later
uuid.to_string(),
i as u32,

View File

@@ -9,3 +9,4 @@ pub mod probe;
pub mod processor;
pub mod storage;
pub mod thumbnail;
pub mod time;

View File

@@ -4,7 +4,7 @@ use std::time::Duration;
use super::executor::PythonExecutor;
const ASR_TIMEOUT: Duration = Duration::from_secs(3600);
const ASR_TIMEOUT: Duration = Duration::from_secs(1800); // 30 minutes
#[derive(Debug, Serialize, Deserialize)]
pub struct AsrResult {

View File

@@ -1,4 +1,5 @@
use anyhow::{Context, Result};
use libc;
use std::path::PathBuf;
use std::process::Stdio;
use std::time::Duration;
@@ -159,12 +160,16 @@ impl PythonExecutor {
cmd.stdout(Stdio::piped());
cmd.stderr(Stdio::piped());
cmd.kill_on_drop(true);
// Create new process group for clean termination
cmd.process_group(0);
tracing::info!("[{}] Starting: {:?}", log_prefix, script_name);
let mut child = cmd
.spawn()
.with_context(|| format!("Failed to run {}", script_name))?;
let child_pid = child.id();
let stdout = child.stdout.take().context("Failed to capture stdout")?;
let stderr = child.stderr.take().context("Failed to capture stderr")?;
@@ -220,6 +225,13 @@ impl PythonExecutor {
Ok(Ok(())) => {}
Ok(Err(e)) => return Err(e),
Err(_) => {
// Try to kill the entire process group
if let Some(pid) = child_pid {
let pgid = pid as i32;
unsafe {
libc::killpg(pgid, libc::SIGKILL);
}
}
child.kill().await.context("Failed to kill process")?;
anyhow::bail!("{} timed out after {:?}", script_name, duration);
}

View File

@@ -1,5 +1,6 @@
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Duration;
use super::executor::PythonExecutor;
@@ -31,6 +32,90 @@ pub struct YoloObject {
pub confidence: f32,
}
// New structs for parsing Python YOLO output
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct YoloPythonMetadata {
video_path: String,
fps: f64,
width: i32,
height: i32,
total_frames: i64,
total_duration: f64,
processed_at: String,
auto_save_interval: i32,
auto_save_frames: i32,
status: String,
last_saved_at: String,
last_saved_frame: i64,
completed_at: Option<String>,
processing_time: Option<f64>,
total_detections: Option<i64>,
auto_save_count: Option<i32>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct YoloPythonDetection {
class_name: String,
confidence: f32,
x1: f32,
y1: f32,
x2: f32,
y2: f32,
width: i32,
height: i32,
class_id: Option<u32>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct YoloPythonFrame {
frame_number: u64,
time_seconds: f64,
time_formatted: String,
detections: Vec<YoloPythonDetection>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct YoloPythonResult {
metadata: YoloPythonMetadata,
frames: HashMap<String, YoloPythonFrame>,
}
impl YoloPythonResult {
pub fn to_yolo_result(&self) -> YoloResult {
let mut frames = Vec::new();
// Sort frames by frame number (key is string, but we parse as u64)
let mut frame_entries: Vec<_> = self.frames.iter().collect();
frame_entries.sort_by_key(|(key, _)| key.parse::<u64>().unwrap_or(0));
for (_, frame) in frame_entries {
let mut objects = Vec::new();
for detection in &frame.detections {
objects.push(YoloObject {
class_name: detection.class_name.clone(),
class_id: detection.class_id.unwrap_or(0),
x: detection.x1 as i32,
y: detection.y1 as i32,
width: detection.width,
height: detection.height,
confidence: detection.confidence,
});
}
frames.push(YoloFrame {
frame: frame.frame_number,
timestamp: frame.time_seconds,
objects,
});
}
YoloResult {
frame_count: frames.len() as u64,
fps: self.metadata.fps,
frames,
}
}
}
pub async fn process_yolo(
video_path: &str,
output_path: &str,
@@ -63,9 +148,11 @@ pub async fn process_yolo(
let json_str = std::fs::read_to_string(output_path).context("Failed to read YOLO output")?;
let result: YoloResult =
let python_result: YoloPythonResult =
serde_json::from_str(&json_str).context("Failed to parse YOLO output")?;
let result = python_result.to_yolo_result();
tracing::info!(
"[YOLO] Result: {} frames, {:.2} fps",
result.frame_count,
@@ -150,4 +237,75 @@ mod tests {
};
assert!(result.frames.is_empty());
}
#[test]
fn test_yolo_python_result_parsing() {
// Sample JSON matching Python script output
let json = r#"{
"metadata": {
"video_path": "/test/video.mp4",
"fps": 22.0,
"width": 640,
"height": 360,
"total_frames": 3512,
"total_duration": 159.63636363636363,
"processed_at": "2026-03-26T05:20:48.230143",
"auto_save_interval": 30,
"auto_save_frames": 300,
"status": "completed",
"last_saved_at": "2026-03-26T05:23:22.791673",
"last_saved_frame": 0,
"completed_at": "2026-03-26T05:23:22.791666",
"processing_time": 154.5577518939972,
"total_detections": 12786,
"auto_save_count": 11
},
"frames": {
"13": {
"frame_number": 13,
"time_seconds": 0.545,
"time_formatted": "00:00:00",
"detections": [
{
"class_id": 0,
"class_name": "person",
"confidence": 0.8424218893051147,
"x1": 473.4156494140625,
"y1": 79.5609359741211,
"x2": 639.77783203125,
"y2": 303.8714294433594,
"width": 166,
"height": 224
}
]
}
}
}"#;
let python_result: YoloPythonResult = serde_json::from_str(json).unwrap();
assert_eq!(python_result.metadata.fps, 22.0);
assert_eq!(python_result.frames.len(), 1);
let frame = python_result.frames.get("13").unwrap();
assert_eq!(frame.frame_number, 13);
assert_eq!(frame.detections.len(), 1);
let detection = &frame.detections[0];
assert_eq!(detection.class_id, Some(0));
assert_eq!(detection.class_name, "person");
assert!((detection.confidence - 0.8424218893051147).abs() < 0.0001);
assert!((detection.x1 - 473.4156494140625).abs() < 0.0001);
// Convert to YoloResult
let yolo_result = python_result.to_yolo_result();
assert_eq!(yolo_result.frames.len(), 1);
assert_eq!(yolo_result.frames[0].frame, 13);
assert_eq!(yolo_result.frames[0].objects.len(), 1);
let obj = &yolo_result.frames[0].objects[0];
assert_eq!(obj.class_name, "person");
assert_eq!(obj.class_id, 0);
assert_eq!(obj.x, 473);
assert_eq!(obj.y, 79);
assert_eq!(obj.width, 166);
assert_eq!(obj.height, 224);
assert!((obj.confidence - 0.842421889).abs() < 0.0001);
}
}

383
src/core/time.rs Normal file
View File

@@ -0,0 +1,383 @@
//! Frame-based time representation for video processing.
//!
//! This module provides a `FrameTime` struct that stores time as frame count
//! with a given FPS (frames per second). This avoids floating-point precision
//! issues when converting between seconds and frames.
//!
//! # Examples
//!
//! ```
//! use momentry_core::time::FrameTime;
//!
//! // Create a FrameTime from frames
//! let time = FrameTime::from_frames(1234, 30.0);
//! assert_eq!(time.seconds(), 41.13333333333333);
//! assert_eq!(time.format_sec_frame(), "41.04");
//!
//! // Create from seconds (useful for migration)
//! let time = FrameTime::from_seconds(41.133333, 30.0);
//! assert_eq!(time.frames(), 1234);
//! ```
use serde::{Deserialize, Serialize};
use std::fmt;
/// Frame-based time representation.
///
/// Stores time as an integer frame count with a floating-point FPS.
/// All calculations are performed using integer frame counts to avoid
/// floating-point precision issues.
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct FrameTime {
/// Frame count (0-based)
frames: i64,
/// Frames per second (can be fractional, e.g., 29.97, 23.976)
fps: f64,
}
impl FrameTime {
/// Creates a new `FrameTime` from frame count and FPS.
///
/// If `fps <= 0.0` or `fps.is_nan()`, defaults to 30.0 FPS.
pub fn from_frames(frames: i64, fps: f64) -> Self {
let fps = if fps <= 0.0 || !fps.is_finite() {
30.0
} else {
fps
};
Self { frames, fps }
}
/// Creates a new `FrameTime` from seconds and FPS.
///
/// This is useful for migrating from existing time representations.
/// The frame count is calculated as `(seconds * fps).round() as i64`
/// to minimize precision loss.
///
/// If `fps <= 0.0` or `fps.is_nan()`, defaults to 30.0 FPS.
pub fn from_seconds(seconds: f64, fps: f64) -> Self {
let fps = if fps <= 0.0 || !fps.is_finite() {
30.0
} else {
fps
};
let frames = (seconds * fps).round() as i64;
Self { frames, fps }
}
/// Returns the frame count.
pub fn frames(&self) -> i64 {
self.frames
}
/// Returns the FPS (frames per second).
pub fn fps(&self) -> f64 {
self.fps
}
/// Returns the time in seconds as a floating-point value.
///
/// Note: This may have precision limitations for fractional FPS values.
/// For display purposes, use `format_sec_frame()` or `format_hms()` instead.
pub fn seconds(&self) -> f64 {
self.frames as f64 / self.fps
}
/// Formats the time as "seconds.frame" with fixed two-digit frame number.
///
/// The frame number is displayed as a zero-padded two-digit number
/// representing the frame within the current second.
///
/// # Examples
///
/// - `123.04` = 123 seconds, frame 4 (at 30 FPS, frame 4 = 0.133 seconds)
/// - `5.29` = 5 seconds, frame 29 (at 30 FPS, last frame of that second)
pub fn format_sec_frame(&self) -> String {
let total_seconds = self.frames as f64 / self.fps;
let seconds = total_seconds.floor() as i64;
// For fractional FPS, use ceil of fps for modulo operation
let fps_ceil = self.fps.ceil() as i64;
// Ensure fps_ceil > 0
let frames_in_second = if fps_ceil == 0 {
0
} else {
self.frames % fps_ceil
};
// Handle negative frames
let frames_in_second = if frames_in_second < 0 {
// This shouldn't happen in practice
0
} else {
frames_in_second
};
format!("{}.{:02}", seconds, frames_in_second)
}
/// Formats the time as "HH:MM:SS" (hours, minutes, seconds).
///
/// This displays whole seconds only, without frame information.
/// Useful for human-readable time displays.
pub fn format_hms(&self) -> String {
let total_seconds = self.seconds();
let hours = (total_seconds / 3600.0) as i64;
let minutes = ((total_seconds % 3600.0) / 60.0) as i64;
let seconds = (total_seconds % 60.0) as i64;
format!("{:02}:{:02}:{:02}", hours, minutes, seconds)
}
/// Formats the time as "HH:MM:SS.FF" (hours, minutes, seconds, frames).
///
/// Displays full time with frame information. Frames are shown as
/// zero-padded two-digit numbers.
pub fn format_hms_frame(&self) -> String {
let total_seconds = self.seconds();
let hours = (total_seconds / 3600.0) as i64;
let minutes = ((total_seconds % 3600.0) / 60.0) as i64;
let seconds = (total_seconds % 60.0) as i64;
// For fractional FPS, use ceil of fps for modulo operation
let fps_ceil = self.fps.ceil() as i64;
let frames_in_second = if fps_ceil == 0 {
0
} else {
self.frames % fps_ceil
};
let frames_in_second = if frames_in_second < 0 {
0
} else {
frames_in_second
};
format!(
"{:02}:{:02}:{:02}.{:02}",
hours, minutes, seconds, frames_in_second
)
}
/// Adds frames to this time, returning a new `FrameTime`.
///
/// # Panics
///
/// Panics if the FPS doesn't match.
pub fn add_frames(&self, frames: i64) -> Self {
Self {
frames: self.frames + frames,
fps: self.fps,
}
}
/// Subtracts frames from this time, returning a new `FrameTime`.
///
/// # Panics
///
/// Panics if the FPS doesn't match or if the result would be negative.
pub fn sub_frames(&self, frames: i64) -> Self {
assert!(
self.frames >= frames,
"Cannot subtract more frames than available"
);
Self {
frames: self.frames - frames,
fps: self.fps,
}
}
/// Adds seconds to this time, returning a new `FrameTime`.
///
/// # Panics
///
/// Panics if the FPS doesn't match.
pub fn add_seconds(&self, seconds: f64) -> Self {
let frames_to_add = (seconds * self.fps).round() as i64;
self.add_frames(frames_to_add)
}
/// Subtracts seconds from this time, returning a new `FrameTime`.
///
/// # Panics
///
/// Panics if the FPS doesn't match or if the result would be negative.
pub fn sub_seconds(&self, seconds: f64) -> Self {
let frames_to_sub = (seconds * self.fps).round() as i64;
self.sub_frames(frames_to_sub)
}
/// Returns the duration between two `FrameTime` instances.
///
/// # Panics
///
/// Panics if the FPS values don't match.
pub fn duration(&self, other: &FrameTime) -> FrameDuration {
assert!(
(self.fps - other.fps).abs() < f64::EPSILON * 2.0,
"FPS mismatch: {} != {}",
self.fps,
other.fps
);
let frame_diff = (self.frames - other.frames).abs();
FrameDuration::from_frames(frame_diff, self.fps)
}
}
impl fmt::Display for FrameTime {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.format_sec_frame())
}
}
/// Duration between two frame times.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct FrameDuration {
frames: i64,
fps: f64,
}
impl FrameDuration {
/// Creates a duration from frame count and FPS.
/// If `fps <= 0.0` or `fps.is_nan()`, defaults to 30.0 FPS.
pub fn from_frames(frames: i64, fps: f64) -> Self {
let fps = if fps <= 0.0 || !fps.is_finite() {
30.0
} else {
fps
};
Self { frames, fps }
}
/// Creates a duration from seconds and FPS.
/// If `fps <= 0.0` or `fps.is_nan()`, defaults to 30.0 FPS.
pub fn from_seconds(seconds: f64, fps: f64) -> Self {
let fps = if fps <= 0.0 || !fps.is_finite() {
30.0
} else {
fps
};
let frames = (seconds * fps).round() as i64;
Self { frames, fps }
}
/// Returns the duration in frames.
pub fn frames(&self) -> i64 {
self.frames
}
/// Returns the duration in seconds.
pub fn seconds(&self) -> f64 {
self.frames as f64 / self.fps
}
/// Formats the duration as "seconds.frame" (same as `FrameTime`).
pub fn format_sec_frame(&self) -> String {
let temp_time = FrameTime::from_frames(self.frames, self.fps);
temp_time.format_sec_frame()
}
/// Formats the duration as "HH:MM:SS".
pub fn format_hms(&self) -> String {
let temp_time = FrameTime::from_frames(self.frames, self.fps);
temp_time.format_hms()
}
}
impl fmt::Display for FrameDuration {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.format_sec_frame())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_from_frames() {
let time = FrameTime::from_frames(150, 30.0);
assert_eq!(time.frames(), 150);
assert_eq!(time.fps(), 30.0);
assert_eq!(time.seconds(), 5.0);
}
#[test]
fn test_from_seconds() {
let time = FrameTime::from_seconds(5.0, 30.0);
assert_eq!(time.frames(), 150);
assert_eq!(time.seconds(), 5.0);
}
#[test]
fn test_format_sec_frame() {
let time = FrameTime::from_frames(123, 30.0);
assert_eq!(time.format_sec_frame(), "4.03");
let time = FrameTime::from_frames(29, 30.0);
assert_eq!(time.format_sec_frame(), "0.29");
let time = FrameTime::from_frames(60, 30.0);
assert_eq!(time.format_sec_frame(), "2.00");
}
#[test]
fn test_format_sec_frame_fractional_fps() {
// 29.97 fps (NTSC)
let time = FrameTime::from_frames(30, 29.97);
// 30 frames at 29.97 fps = 1.001 seconds = 1 second, frame 0
assert_eq!(time.format_sec_frame(), "1.00");
let time = FrameTime::from_frames(60, 29.97);
// 60 frames at 29.97 fps = 2.002 seconds = 2 seconds, frame 0
assert_eq!(time.format_sec_frame(), "2.00");
}
#[test]
fn test_format_hms() {
let time = FrameTime::from_frames(3661, 30.0); // 122.033 seconds = 2 minutes 2 seconds
assert_eq!(time.format_hms(), "00:02:02");
let time = FrameTime::from_frames(4500, 30.0); // 150 seconds = 2 minutes 30 seconds
assert_eq!(time.format_hms(), "00:02:30");
}
#[test]
fn test_format_hms_frame() {
let time = FrameTime::from_frames(123, 30.0); // 4 seconds, 3 frames
assert_eq!(time.format_hms_frame(), "00:00:04.03");
}
#[test]
fn test_add_sub_frames() {
let time = FrameTime::from_frames(100, 30.0);
let new_time = time.add_frames(50);
assert_eq!(new_time.frames(), 150);
let new_time = time.sub_frames(30);
assert_eq!(new_time.frames(), 70);
}
#[test]
fn test_add_sub_seconds() {
let time = FrameTime::from_frames(100, 30.0);
let new_time = time.add_seconds(2.0);
assert_eq!(new_time.frames(), 160); // 100 + 60
let new_time = time.sub_seconds(1.0);
assert_eq!(new_time.frames(), 70); // 100 - 30
}
#[test]
fn test_duration() {
let time1 = FrameTime::from_frames(200, 30.0);
let time2 = FrameTime::from_frames(150, 30.0);
let duration = time1.duration(&time2);
assert_eq!(duration.frames(), 50);
assert_eq!(duration.seconds(), 50.0 / 30.0);
}
#[test]
fn test_frame_duration() {
let duration = FrameDuration::from_frames(90, 30.0);
assert_eq!(duration.seconds(), 3.0);
assert_eq!(duration.format_sec_frame(), "3.00");
assert_eq!(duration.format_hms(), "00:00:03");
}
}