278 lines
11 KiB
Rust
278 lines
11 KiB
Rust
use crate::models::ExportDoc;
|
||
use chrono::{DateTime, NaiveDateTime, Utc};
|
||
use sqlx::PgPool;
|
||
use std::collections::HashMap;
|
||
use std::path::{Path, PathBuf};
|
||
use walkdir::WalkDir;
|
||
|
||
const VIDEO_EXTS: &[&str] = &["mp4", "mkv", "avi", "mov", "m4v", "webm"];
|
||
|
||
fn is_video(p: &Path) -> bool {
|
||
p.extension()
|
||
.and_then(|e| e.to_str())
|
||
.map(|e| VIDEO_EXTS.contains(&e.to_lowercase().as_str()))
|
||
.unwrap_or(false)
|
||
}
|
||
|
||
/// Find the sibling export-JSON for a video. Supports `<file>_annotations.json`
|
||
/// (video-annotator's default export name), `<file>.json`, and `<stem>.json`.
|
||
fn sibling_json(video: &Path) -> Option<PathBuf> {
|
||
let dir = video.parent()?;
|
||
let file_name = video.file_name()?.to_str()?;
|
||
let stem = video.file_stem()?.to_str()?;
|
||
let candidates = [
|
||
format!("{file_name}_annotations.json"),
|
||
format!("{file_name}.json"),
|
||
format!("{stem}.json"),
|
||
format!("{stem}_annotations.json"),
|
||
];
|
||
for c in candidates {
|
||
let p = dir.join(c);
|
||
if p.is_file() {
|
||
return Some(p);
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
fn parse_ts(s: &str) -> Option<DateTime<Utc>> {
|
||
if s.is_empty() {
|
||
return None;
|
||
}
|
||
if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
|
||
return Some(dt.with_timezone(&Utc));
|
||
}
|
||
NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S")
|
||
.ok()
|
||
.map(|n| n.and_utc())
|
||
}
|
||
|
||
/// Scan `dir`, upsert every video (and its annotations) into the project.
|
||
/// Idempotent: re-running updates rows; never downgrades a 'verified' video.
|
||
pub async fn ingest_project(db: &PgPool, project_id: i32, dir: &Path) -> anyhow::Result<usize> {
|
||
if !dir.is_dir() {
|
||
anyhow::bail!("ingest dir not found: {}", dir.display());
|
||
}
|
||
let mut count = 0usize;
|
||
|
||
for entry in WalkDir::new(dir).into_iter().filter_map(|e| e.ok()) {
|
||
let path = entry.path();
|
||
if !entry.file_type().is_file() || !is_video(path) {
|
||
continue;
|
||
}
|
||
let rel_path = path
|
||
.strip_prefix(dir)
|
||
.unwrap_or(path)
|
||
.to_string_lossy()
|
||
.to_string();
|
||
let file_name = path
|
||
.file_name()
|
||
.map(|n| n.to_string_lossy().to_string())
|
||
.unwrap_or_default();
|
||
|
||
// Never clobber a verified push (Phase 3): a worker's pushed annotations +
|
||
// completer/verify-time are authoritative, so skip re-ingesting this video.
|
||
// Also never clobber a video that is *actively claimed* (a worker is mid-edit)
|
||
// — a stale NAS JSON must not overwrite in-progress work. Both checks read the
|
||
// current row's status + claim lease.
|
||
let existing: Option<(String, Option<String>, Option<DateTime<Utc>>)> = sqlx::query_as(
|
||
"SELECT status, claimed_by, lease_expires_at FROM videos WHERE project_id=$1 AND rel_path=$2",
|
||
)
|
||
.bind(project_id)
|
||
.bind(&rel_path)
|
||
.fetch_optional(db)
|
||
.await?;
|
||
if let Some((status, claimed_by, lease)) = existing.as_ref() {
|
||
if status == "verified" {
|
||
tracing::debug!("skip re-ingest of verified video: {rel_path}");
|
||
continue;
|
||
}
|
||
let active_claim = claimed_by.is_some()
|
||
&& lease.map(|l| l > Utc::now()).unwrap_or(false);
|
||
if active_claim {
|
||
tracing::debug!("skip re-ingest of actively-claimed video: {rel_path}");
|
||
continue;
|
||
}
|
||
}
|
||
|
||
let json_path = sibling_json(path);
|
||
let has_json = json_path.is_some();
|
||
|
||
// Parse the sibling JSON if present.
|
||
let doc: Option<ExportDoc> = json_path.as_ref().and_then(|jp| {
|
||
std::fs::read_to_string(jp)
|
||
.ok()
|
||
.and_then(|t| serde_json::from_str::<ExportDoc>(&t).ok())
|
||
});
|
||
|
||
let (width, height, fps, frame_count, time_ms) = match doc.as_ref().and_then(|d| d.video.as_ref()) {
|
||
Some(v) => (
|
||
v.width.map(|x| x as i32),
|
||
v.height.map(|x| x as i32),
|
||
v.fps,
|
||
v.frame_count.map(|x| x as i32),
|
||
v.annotation_time_ms,
|
||
),
|
||
None => (None, None, None, None, 0i64),
|
||
};
|
||
|
||
// Aggregate annotation_count, primary annotator, annotated_at.
|
||
let mut ann_count: i32 = 0;
|
||
let mut by_freq: HashMap<String, i32> = HashMap::new();
|
||
let mut latest: Option<DateTime<Utc>> = None;
|
||
if let Some(d) = doc.as_ref() {
|
||
for a in &d.fixed_annotations {
|
||
ann_count += 1;
|
||
if !a.annotated_by.is_empty() {
|
||
*by_freq.entry(a.annotated_by.clone()).or_default() += 1;
|
||
}
|
||
if let Some(t) = parse_ts(&a.created_at) {
|
||
latest = Some(latest.map_or(t, |l| l.max(t)));
|
||
}
|
||
}
|
||
for a in &d.range_annotations {
|
||
ann_count += 1;
|
||
if !a.annotated_by.is_empty() {
|
||
*by_freq.entry(a.annotated_by.clone()).or_default() += 1;
|
||
}
|
||
if let Some(t) = parse_ts(&a.created_at) {
|
||
latest = Some(latest.map_or(t, |l| l.max(t)));
|
||
}
|
||
}
|
||
}
|
||
let primary_annotator = by_freq
|
||
.into_iter()
|
||
.max_by_key(|(_, c)| *c)
|
||
.map(|(u, _)| u)
|
||
.unwrap_or_default();
|
||
let status = if has_json && ann_count > 0 { "annotated" } else { "pending" };
|
||
let raw_json: Option<serde_json::Value> = json_path.as_ref().and_then(|jp| {
|
||
std::fs::read_to_string(jp).ok().and_then(|t| serde_json::from_str(&t).ok())
|
||
});
|
||
|
||
// Upsert the video + replace its annotation rows atomically: either the new
|
||
// metadata AND the new annotation set both land, or neither does. Prevents a
|
||
// failed insert mid-loop from leaving the video with partial/zero annotations.
|
||
let mut tx = db.begin().await?;
|
||
// Upsert the video. Keep 'verified' status sticky (Phase 3 pushes).
|
||
let video_id: i32 = sqlx::query_scalar(
|
||
r#"
|
||
INSERT INTO videos
|
||
(project_id, file_name, rel_path, has_json, width, height, fps, frame_count,
|
||
annotation_count, annotation_time_ms, primary_annotator, status, annotated_at, raw_json,
|
||
imported_count)
|
||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15)
|
||
ON CONFLICT (project_id, rel_path) DO UPDATE SET
|
||
file_name=EXCLUDED.file_name,
|
||
has_json=EXCLUDED.has_json,
|
||
width=EXCLUDED.width, height=EXCLUDED.height, fps=EXCLUDED.fps,
|
||
frame_count=EXCLUDED.frame_count,
|
||
annotation_count=EXCLUDED.annotation_count,
|
||
annotation_time_ms=EXCLUDED.annotation_time_ms,
|
||
primary_annotator=EXCLUDED.primary_annotator,
|
||
status = CASE WHEN videos.status='verified' THEN videos.status ELSE EXCLUDED.status END,
|
||
annotated_at=EXCLUDED.annotated_at,
|
||
raw_json=EXCLUDED.raw_json,
|
||
-- imported baseline = the NAS-ingested count (skipped for verified videos
|
||
-- via the early continue above, so a pushed pass's baseline is preserved).
|
||
imported_count=EXCLUDED.imported_count,
|
||
ingested_at=now()
|
||
RETURNING id
|
||
"#,
|
||
)
|
||
.bind(project_id)
|
||
.bind(&file_name)
|
||
.bind(&rel_path)
|
||
.bind(has_json)
|
||
.bind(width)
|
||
.bind(height)
|
||
.bind(fps)
|
||
.bind(frame_count)
|
||
.bind(ann_count)
|
||
.bind(time_ms)
|
||
.bind(&primary_annotator)
|
||
.bind(status)
|
||
.bind(latest)
|
||
.bind(raw_json)
|
||
.bind(ann_count)
|
||
.fetch_one(&mut *tx)
|
||
.await?;
|
||
|
||
// Replace detailed annotation rows for this video (same transaction).
|
||
if let Some(d) = doc.as_ref() {
|
||
replace_annotations(&mut tx, video_id, d).await?;
|
||
} else {
|
||
sqlx::query("DELETE FROM annotations WHERE video_id=$1")
|
||
.bind(video_id)
|
||
.execute(&mut *tx)
|
||
.await?;
|
||
}
|
||
tx.commit().await?;
|
||
count += 1;
|
||
}
|
||
Ok(count)
|
||
}
|
||
|
||
/// Replace all annotation rows for a video with the contents of an export doc.
|
||
/// Shared by ingest (NAS JSON) and the Phase 3 worker push. Returns the row count
|
||
/// written (each range annotation contributes 1–2 rows: start, and end if present).
|
||
///
|
||
/// Runs on a caller-supplied transaction connection so the delete + all inserts are
|
||
/// **atomic** — either every annotation row lands or none do. Callers (`ingest`, worker
|
||
/// `push`) own the transaction so the surrounding metadata update commits together with
|
||
/// the annotation set. A failed insert mid-loop rolls back the delete, so annotations
|
||
/// can never be partially lost.
|
||
pub async fn replace_annotations(
|
||
conn: &mut sqlx::PgConnection,
|
||
video_id: i32,
|
||
d: &ExportDoc,
|
||
) -> anyhow::Result<i32> {
|
||
sqlx::query("DELETE FROM annotations WHERE video_id=$1")
|
||
.bind(video_id)
|
||
.execute(&mut *conn)
|
||
.await?;
|
||
let mut written = 0i32;
|
||
for a in &d.fixed_annotations {
|
||
insert_ann(&mut *conn, video_id, a.frame_number, &a.label, &a.side, &a.shape_type,
|
||
&a.vertices, &a.annotated_by, &a.review_status, &a.remark, &a.subclass, &a.created_at).await?;
|
||
written += 1;
|
||
}
|
||
for a in &d.range_annotations {
|
||
insert_ann(&mut *conn, video_id, a.start_frame, &a.label, &a.side, &a.shape_type,
|
||
&a.start_vertices, &a.annotated_by, &a.review_status, &a.remark, &a.subclass, &a.created_at).await?;
|
||
written += 1;
|
||
if let (Some(ef), Some(ev)) = (a.end_frame, a.end_vertices.as_ref()) {
|
||
insert_ann(&mut *conn, video_id, ef, &a.label, &a.side, &a.shape_type,
|
||
ev, &a.annotated_by, &a.review_status, &a.remark, &a.subclass, &a.created_at).await?;
|
||
written += 1;
|
||
}
|
||
}
|
||
Ok(written)
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
async fn insert_ann(
|
||
conn: &mut sqlx::PgConnection, video_id: i32, frame: i64, label: &str, side: &str, shape: &str,
|
||
vertices: &serde_json::Value, by: &str, review: &str, remark: &str, subclass: &str, created_at: &str,
|
||
) -> anyhow::Result<()> {
|
||
sqlx::query(
|
||
r#"INSERT INTO annotations
|
||
(video_id, frame_number, label, side, shape_type, vertices, annotated_by, review_status, remark, subclass, created_at)
|
||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)"#,
|
||
)
|
||
.bind(video_id)
|
||
.bind(frame as i32)
|
||
.bind(label)
|
||
.bind(side)
|
||
.bind(if shape.is_empty() { "bbox" } else { shape })
|
||
.bind(vertices)
|
||
.bind(by)
|
||
.bind(if review.is_empty() { "none" } else { review })
|
||
.bind(remark)
|
||
.bind(subclass)
|
||
.bind(created_at)
|
||
.execute(&mut *conn)
|
||
.await?;
|
||
Ok(())
|
||
}
|