107 lines
3.2 KiB
Rust
107 lines
3.2 KiB
Rust
use anyhow::{Context, Result};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::path::{Path, PathBuf};
|
|
use std::process::Command;
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ThumbnailResult {
|
|
pub uuid: String,
|
|
pub count: usize,
|
|
pub files: Vec<String>,
|
|
}
|
|
|
|
pub struct ThumbnailExtractor {
|
|
output_dir: PathBuf,
|
|
count: u32,
|
|
}
|
|
|
|
impl ThumbnailExtractor {
|
|
pub fn new(output_dir: PathBuf, count: u32) -> Self {
|
|
Self { output_dir, count }
|
|
}
|
|
|
|
pub fn extract(&self, video_path: &str, uuid: &str) -> Result<ThumbnailResult> {
|
|
let script_path = Path::new(env!("CARGO_MANIFEST_DIR"))
|
|
.join("scripts")
|
|
.join("thumbnail_extractor.py");
|
|
|
|
let python_path = std::env::var("MOMENTRY_PYTHON_PATH")
|
|
.unwrap_or_else(|_| "/opt/homebrew/bin/python3.11".to_string());
|
|
let python_bin = Path::new(&python_path);
|
|
|
|
let output = Command::new(python_bin)
|
|
.arg(script_path)
|
|
.arg(video_path)
|
|
.arg(uuid)
|
|
.arg("-o")
|
|
.arg(&self.output_dir)
|
|
.arg("-c")
|
|
.arg(self.count.to_string())
|
|
.output()
|
|
.context("Failed to run thumbnail extractor")?;
|
|
|
|
if !output.status.success() {
|
|
let stderr = String::from_utf8_lossy(&output.stderr);
|
|
anyhow::bail!("Thumbnail extraction failed: {}", stderr);
|
|
}
|
|
|
|
let json_str = String::from_utf8_lossy(&output.stdout);
|
|
let result: ThumbnailResult =
|
|
serde_json::from_str(&json_str).context("Failed to parse thumbnail result")?;
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
pub fn get_or_create(&self, video_path: &str, uuid: &str) -> Result<ThumbnailResult> {
|
|
let thumb_dir = self.output_dir.join(uuid);
|
|
|
|
// Check if thumbnails already exist
|
|
if thumb_dir.exists() {
|
|
let files: Vec<String> = (0..self.count)
|
|
.map(|i| thumb_dir.join(format!("thumb_{:03}.jpg", i)))
|
|
.filter(|p| p.exists())
|
|
.map(|p| p.to_string_lossy().to_string())
|
|
.collect();
|
|
|
|
if files.len() as u32 == self.count {
|
|
return Ok(ThumbnailResult {
|
|
uuid: uuid.to_string(),
|
|
count: files.len(),
|
|
files,
|
|
});
|
|
}
|
|
}
|
|
|
|
// Extract new thumbnails
|
|
self.extract(video_path, uuid)
|
|
}
|
|
|
|
pub fn get_thumbnails(&self, uuid: &str) -> Option<Vec<String>> {
|
|
let thumb_dir = self.output_dir.join(uuid);
|
|
|
|
if !thumb_dir.exists() {
|
|
return None;
|
|
}
|
|
|
|
let files: Vec<String> = (0..10)
|
|
.map(|i| thumb_dir.join(format!("thumb_{:03}.jpg", i)))
|
|
.filter(|p| p.exists())
|
|
.map(|p| p.to_string_lossy().to_string())
|
|
.collect();
|
|
|
|
if files.is_empty() {
|
|
None
|
|
} else {
|
|
Some(files)
|
|
}
|
|
}
|
|
|
|
pub fn cleanup(&self, uuid: &str) -> Result<()> {
|
|
let thumb_dir = self.output_dir.join(uuid);
|
|
if thumb_dir.exists() {
|
|
std::fs::remove_dir_all(&thumb_dir).context("Failed to remove thumbnail directory")?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|