Proof-of-concept: takes a random screencap, doesn't post it yet
This commit is contained in:
commit
6873e14f37
9 changed files with 2184 additions and 0 deletions
17
src/config.rs
Normal file
17
src/config.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use std::{collections::HashMap, env, fs, path::PathBuf};
|
||||
|
||||
pub type Shows = HashMap<String, PathBuf>;
|
||||
|
||||
pub struct Config {
|
||||
pub shows: Shows,
|
||||
}
|
||||
|
||||
pub fn load() -> Config {
|
||||
let shows_file = env::var("SCREENCAP_BOT_SHOWS_FILE").unwrap_or(String::from("./shows.yaml"));
|
||||
Config {
|
||||
shows: serde_yaml::from_reader(
|
||||
fs::File::open(shows_file).expect("Failed to open shows file"),
|
||||
)
|
||||
.expect("Failed to parse YAML from shows file"),
|
||||
}
|
||||
}
|
||||
122
src/enumeration.rs
Normal file
122
src/enumeration.rs
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
use std::{
|
||||
collections::{BTreeSet, HashMap},
|
||||
fs,
|
||||
io::ErrorKind,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
||||
use lazy_static::lazy_static;
|
||||
use log::{debug, error, trace, warn};
|
||||
use regex::Regex;
|
||||
|
||||
use crate::config::Shows;
|
||||
|
||||
lazy_static! {
|
||||
static ref NUMBER_REGEX: Regex = Regex::new("[0-9]+").unwrap();
|
||||
}
|
||||
|
||||
/// [episode number, or None for a standalone movie] -> [video file path]
|
||||
type Enumeration = HashMap<Option<u32>, PathBuf>;
|
||||
|
||||
/// [show name] -> [enumeration]
|
||||
type EnumeratedShows = HashMap<String, Enumeration>;
|
||||
|
||||
pub fn enumerate_shows(shows: Shows) -> EnumeratedShows {
|
||||
shows
|
||||
.into_iter()
|
||||
.filter_map(|(name, path)| {
|
||||
debug!("Enumerating {} from {}", name, path.display());
|
||||
enumerate_show(path)
|
||||
.map_err(|e| {
|
||||
error!("Error processing {}: {}", name, e);
|
||||
})
|
||||
.ok()
|
||||
.map(|en| (name, en))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn enumerate_show(path: PathBuf) -> std::io::Result<Enumeration> {
|
||||
let metadata = fs::metadata(&path)?;
|
||||
if metadata.is_file() {
|
||||
debug!("{} is a file, standalone", path.display());
|
||||
Ok(HashMap::from([(None, path)]))
|
||||
} else if metadata.is_dir() {
|
||||
debug!("{} is a directory, enumerating episodes", path.display());
|
||||
enumerate_dir(path)
|
||||
} else {
|
||||
Err(std::io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
format!("Invalid file type for {}", path.display()),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn enumerate_dir(path: PathBuf) -> std::io::Result<Enumeration> {
|
||||
let files: Vec<(PathBuf, String)> = fs::read_dir(&path)?
|
||||
.map(|entry| {
|
||||
let entry = entry?;
|
||||
if !entry.file_type()?.is_file() {
|
||||
debug!("Skipping {}, not a file", entry.path().display());
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some((
|
||||
entry.path(),
|
||||
entry.file_name().into_string().map_err(|o| {
|
||||
std::io::Error::new(ErrorKind::InvalidData, format!("Invalid filename {:?}", o))
|
||||
})?,
|
||||
)))
|
||||
})
|
||||
.filter_map(|r| r.transpose())
|
||||
.collect::<std::io::Result<Vec<(PathBuf, String)>>>()?;
|
||||
|
||||
// For each string prefix which is followed by a number in one or
|
||||
// more filenames, record what the number is for each filename
|
||||
// with that prefix, sorted by number.
|
||||
//
|
||||
// We use a BTreeSet of pairs rather than a BTreeMap so that
|
||||
// duplicates are preserved.
|
||||
let mut prefix_enumerations: HashMap<&str, BTreeSet<(u32, &Path)>> = HashMap::new();
|
||||
for (path, filename) in files.iter() {
|
||||
for m in NUMBER_REGEX.find_iter(filename) {
|
||||
let num = m.as_str().parse();
|
||||
if let Err(e) = num {
|
||||
warn!("Failed to parse candidate number: {}", e);
|
||||
continue;
|
||||
}
|
||||
let num = num.unwrap();
|
||||
let prefix = filename.get(..m.start()).unwrap();
|
||||
trace!("{}: candidate prefix {}, number {}", filename, prefix, num);
|
||||
prefix_enumerations
|
||||
.entry(prefix)
|
||||
.or_default()
|
||||
.insert((num, path));
|
||||
}
|
||||
}
|
||||
|
||||
// Use the prefix enumeration that produces the most consecutive numbers
|
||||
let best_enumeration = prefix_enumerations
|
||||
.into_values()
|
||||
.max_by_key(|set| {
|
||||
set.iter()
|
||||
.zip(set.iter().skip(1))
|
||||
.filter(|((cur, _), (next, _))| next - cur == 1)
|
||||
.count()
|
||||
})
|
||||
.ok_or(std::io::Error::new(
|
||||
ErrorKind::InvalidData,
|
||||
"No valid prefixes found",
|
||||
))?;
|
||||
|
||||
let mut result: Enumeration = HashMap::new();
|
||||
for (num, path) in best_enumeration.into_iter() {
|
||||
debug!("Episode {}: {}", num, path.display());
|
||||
if let Some(dup) = result.insert(Some(num), path.to_path_buf()) {
|
||||
warn!(
|
||||
"Duplicate episode number, discarding file {}",
|
||||
dup.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
80
src/main.rs
Normal file
80
src/main.rs
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
use std::path::Path;
|
||||
|
||||
use enumeration::enumerate_shows;
|
||||
use log::{debug, error, info};
|
||||
use rand::{distributions::Standard, seq::IteratorRandom, Rng};
|
||||
|
||||
mod config;
|
||||
mod enumeration;
|
||||
mod video;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> std::io::Result<()> {
|
||||
env_logger::init();
|
||||
ffmpeg_next::init().unwrap();
|
||||
dotenvy::dotenv().ok();
|
||||
|
||||
let mut rng = rand::thread_rng();
|
||||
|
||||
let conf = config::load();
|
||||
let enumerated_shows = enumerate_shows(conf.shows);
|
||||
|
||||
loop {
|
||||
let (show, episodes) = enumerated_shows
|
||||
.iter()
|
||||
.choose(&mut rng)
|
||||
.expect("No shows found!");
|
||||
let (num, file) = episodes.iter().choose(&mut rng).unwrap();
|
||||
|
||||
let descriptor = format!(
|
||||
"{}{}",
|
||||
show,
|
||||
if let Some(n) = num {
|
||||
format!(" episode {}", n)
|
||||
} else {
|
||||
String::new()
|
||||
}
|
||||
);
|
||||
|
||||
info!("Selected: {} - {}", descriptor, file.display());
|
||||
|
||||
let video_info = video::get_video_info(file, Some("eng")).unwrap();
|
||||
debug!(
|
||||
"Video duration: {}",
|
||||
format_timestamp(video_info.duration_secs, None)
|
||||
);
|
||||
debug!(
|
||||
"Subtitle stream index: {:?}",
|
||||
video_info.subtitle_stream_index
|
||||
);
|
||||
|
||||
let timestamp = video_info.duration_secs * rng.sample::<f64, _>(Standard);
|
||||
let formatted_timestamp = format_timestamp(timestamp, Some(video_info.duration_secs));
|
||||
info!("Taking screencap at {}", formatted_timestamp);
|
||||
if let Err(e) = video::take_screencap(
|
||||
file,
|
||||
timestamp,
|
||||
video_info.subtitle_stream_index,
|
||||
&Path::new("./out.png"),
|
||||
) {
|
||||
error!("Failed to take screencap: {}", e);
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn format_timestamp(timestamp: f64, total_duration: Option<f64>) -> String {
|
||||
let total_duration = total_duration.unwrap_or(timestamp);
|
||||
format!(
|
||||
"{}{:02}:{:05.2}",
|
||||
if total_duration >= 3600.0 {
|
||||
format!("{}:", (timestamp / 3600.0) as u32)
|
||||
} else {
|
||||
String::new()
|
||||
},
|
||||
((timestamp % 3600.0) / 60.0) as u32,
|
||||
timestamp % 60.0
|
||||
)
|
||||
}
|
||||
107
src/video.rs
Normal file
107
src/video.rs
Normal file
|
|
@ -0,0 +1,107 @@
|
|||
use ffmpeg_next::{
|
||||
format::{input, stream::Disposition},
|
||||
media::Type,
|
||||
};
|
||||
use lazy_static::lazy_static;
|
||||
use regex::Regex;
|
||||
use tempfile::tempdir;
|
||||
|
||||
use std::{ffi::OsStr, io::ErrorKind, os::unix::fs, path::Path, process::Command};
|
||||
|
||||
lazy_static! {
|
||||
static ref SUBTITLE_FORBID_REGEX: Regex = Regex::new("(?i)sign|song").unwrap();
|
||||
}
|
||||
|
||||
pub struct VideoInfo {
|
||||
pub duration_secs: f64,
|
||||
// The index among the subtitle streams, not among the streams in general
|
||||
pub subtitle_stream_index: Option<usize>,
|
||||
}
|
||||
|
||||
pub fn get_video_info<P: AsRef<Path>>(
|
||||
source: &P,
|
||||
subtitle_lang: Option<&str>,
|
||||
) -> Result<VideoInfo, ffmpeg_next::Error> {
|
||||
let ctx = input(source)?;
|
||||
|
||||
let duration_secs = ctx.duration() as f64 / f64::from(ffmpeg_next::ffi::AV_TIME_BASE);
|
||||
|
||||
let subtitle_stream_index = subtitle_lang
|
||||
.map(|lang| {
|
||||
ctx.streams()
|
||||
.filter(|stream| {
|
||||
ffmpeg_next::codec::context::Context::from_parameters(stream.parameters())
|
||||
.map(|c| c.medium())
|
||||
== Ok(Type::Subtitle)
|
||||
})
|
||||
.enumerate()
|
||||
.filter(|(_, stream)| {
|
||||
let metadata = stream.metadata();
|
||||
if metadata.get("language") != Some(lang) {
|
||||
return false;
|
||||
}
|
||||
if metadata
|
||||
.get("title")
|
||||
.map(|t| SUBTITLE_FORBID_REGEX.is_match(t))
|
||||
== Some(true)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
true
|
||||
})
|
||||
.min_by_key(|(_, stream)| stream.disposition().contains(Disposition::FORCED))
|
||||
.map(|(idx, _)| idx)
|
||||
})
|
||||
.flatten();
|
||||
|
||||
Ok(VideoInfo {
|
||||
duration_secs,
|
||||
subtitle_stream_index,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn take_screencap<P: AsRef<Path>, Q: AsRef<OsStr>>(
|
||||
source: &P,
|
||||
timestamp_secs: f64,
|
||||
subtitle_stream_index: Option<usize>,
|
||||
dest: &Q,
|
||||
) -> std::io::Result<()> {
|
||||
let ext = source.as_ref().extension().map(|s| s.to_str()).flatten();
|
||||
if ext != Some("mkv") && ext != Some("mp4") {
|
||||
return Err(std::io::Error::new(
|
||||
ErrorKind::Other,
|
||||
"unexpected file extension",
|
||||
));
|
||||
}
|
||||
let tmp_dir = tempdir()?;
|
||||
let link_path = tmp_dir.path().join(format!("input.{}", ext.unwrap()));
|
||||
fs::symlink(source, &link_path)?;
|
||||
let mut cmd = Command::new("ffmpeg");
|
||||
|
||||
cmd.arg("-ss")
|
||||
.arg(format!("{:.2}", timestamp_secs))
|
||||
.arg("-copyts")
|
||||
.arg("-i")
|
||||
.arg(&link_path);
|
||||
|
||||
if let Some(idx) = subtitle_stream_index {
|
||||
cmd.arg("-filter_complex").arg(format!(
|
||||
"[0:v]subtitles={}:si={}",
|
||||
link_path.to_string_lossy(),
|
||||
idx
|
||||
));
|
||||
}
|
||||
|
||||
cmd.args(["-vframes", "1"])
|
||||
.args(["-loglevel", "quiet"])
|
||||
.arg("-y")
|
||||
.arg(&dest);
|
||||
|
||||
if !cmd.status()?.success() {
|
||||
return Err(std::io::Error::new(
|
||||
ErrorKind::Other,
|
||||
"ffmpeg command failed",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue