aboutsummaryrefslogtreecommitdiff
path: root/src/subtitle_extraction/mod.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/subtitle_extraction/mod.rs')
-rw-r--r--src/subtitle_extraction/mod.rs159
1 files changed, 0 insertions, 159 deletions
diff --git a/src/subtitle_extraction/mod.rs b/src/subtitle_extraction/mod.rs
deleted file mode 100644
index 9e7fff4..0000000
--- a/src/subtitle_extraction/mod.rs
+++ /dev/null
@@ -1,159 +0,0 @@
-/// Extraction of embedded subtitles
-mod embedded;
-/// Synthesis of subtitles from audio using whisper.cpp
-mod whisper;
-
-use std::{collections::BTreeMap, sync::mpsc, thread};
-
-use ffmpeg::Rational;
-use relm4::{ComponentSender, Worker};
-
-use crate::tracks::{SUBTITLE_TRACKS, StreamIndex, SubtitleCue, SubtitleTrack, TrackMetadata};
-
-pub struct SubtitleExtractor {}
-
-#[derive(Debug)]
-pub enum SubtitleExtractorMsg {
- ExtractFromUrl {
- url: String,
- // the index of the audio stream on which to run a whisper transcription
- whisper_stream_index: Option<usize>,
- },
-}
-
-#[derive(Debug)]
-pub enum SubtitleExtractorOutput {
- NewCue(StreamIndex, SubtitleCue),
- ExtractionComplete,
-}
-
-impl Worker for SubtitleExtractor {
- type Init = ();
- type Input = SubtitleExtractorMsg;
- type Output = SubtitleExtractorOutput;
-
- fn init(_init: Self::Init, _sender: ComponentSender<Self>) -> Self {
- Self {}
- }
-
- fn update(&mut self, msg: SubtitleExtractorMsg, sender: ComponentSender<Self>) {
- match msg {
- SubtitleExtractorMsg::ExtractFromUrl {
- url,
- whisper_stream_index: whisper_audio_stream_ix,
- } => {
- self.handle_extract_from_url(url, whisper_audio_stream_ix, sender);
- }
- }
- }
-}
-
-impl SubtitleExtractor {
- fn handle_extract_from_url(
- &mut self,
- url: String,
- whisper_audio_stream_ix: Option<usize>,
- sender: ComponentSender<Self>,
- ) {
- // Clear existing tracks
- SUBTITLE_TRACKS.write().clear();
-
- match self.extract_subtitles(&url, whisper_audio_stream_ix, sender.clone()) {
- Ok(_) => {
- log::info!("Subtitle extraction completed successfully");
- sender
- .output(SubtitleExtractorOutput::ExtractionComplete)
- .unwrap();
- }
- Err(e) => {
- log::error!("Subtitle extraction failed: {}", e);
- }
- }
- }
-
- fn extract_subtitles(
- &self,
- url: &str,
- whisper_audio_stream_ix: Option<usize>,
- sender: ComponentSender<Self>,
- ) -> anyhow::Result<()> {
- let mut input = ffmpeg::format::input(&url)?;
-
- let mut subtitle_extractors = BTreeMap::new();
-
- // create extractor for each subtitle stream
- for stream in input.streams() {
- let stream_ix = stream.index();
-
- if stream.parameters().medium() == ffmpeg::media::Type::Subtitle {
- let metadata = TrackMetadata::from_ffmpeg_stream(&stream);
- let track = SubtitleTrack {
- metadata,
- cues: Vec::new(),
- };
-
- SUBTITLE_TRACKS.write().insert(stream_ix, track);
-
- let context = ffmpeg::codec::Context::from_parameters(stream.parameters())?;
- let (packet_tx, packet_rx) = mpsc::channel();
- let time_base = stream.time_base();
- let sender = sender.clone();
- let join_handle = thread::spawn(move || {
- embedded::extract_embedded_subtitles(
- stream_ix, context, time_base, packet_rx, sender,
- )
- });
-
- subtitle_extractors.insert(stream_ix, (packet_tx, join_handle));
- }
- }
-
- if let Some(stream_ix) = whisper_audio_stream_ix {
- let stream = input.stream(stream_ix).unwrap();
-
- let mut metadata = TrackMetadata::from_ffmpeg_stream(&stream);
- metadata.title = Some(match metadata.title {
- Some(title) => format!("Auto-generated from audio (Whisper): {}", title),
- None => "Auto-generated from audio (Whisper)".to_string(),
- });
-
- let track = SubtitleTrack {
- metadata,
- cues: Vec::new(),
- };
-
- SUBTITLE_TRACKS.write().insert(stream_ix, track);
-
- let context = ffmpeg::codec::Context::from_parameters(stream.parameters())?;
- let (packet_tx, packet_rx) = mpsc::channel();
- let time_base = stream.time_base();
- let sender = sender.clone();
- let join_handle = thread::spawn(move || {
- whisper::generate_whisper_subtitles(
- stream_ix, context, time_base, packet_rx, sender,
- )
- });
-
- subtitle_extractors.insert(stream_ix, (packet_tx, join_handle));
- }
-
- // process packets
- for (stream, packet) in input.packets() {
- let stream_index = stream.index();
-
- if let Some((packet_tx, _)) = subtitle_extractors.get_mut(&stream_index) {
- packet_tx.send(packet).unwrap();
- }
- }
-
- // wait for extraction to complete
- for (_, (_, join_handle)) in subtitle_extractors {
- join_handle
- .join()
- .unwrap()
- .unwrap_or_else(|e| log::error!("error running subtitle extraction: {}", e));
- }
-
- Ok(())
- }
-}