1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
|
/// Extraction of embedded subtitles
mod embedded;
/// Synthesis of subtitles from audio using whisper.cpp
mod whisper;
use std::{collections::BTreeMap, fmt::Display, sync::mpsc, thread};
use ffmpeg::Rational;
use relm4::{ComponentSender, Worker};
use crate::subtitles::{StreamIndex, SubtitleCue};
pub struct SubtitleExtractor {}
#[derive(Debug, Clone)]
pub struct ExtractionArgs {
pub url: String,
pub whisper_stream_index: Option<usize>,
}
impl Display for ExtractionArgs {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{} {:?}", self.url, self.whisper_stream_index)
}
}
#[derive(Debug)]
pub enum SubtitleExtractorMsg {
Extract(ExtractionArgs),
}
#[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::Extract(ExtractionArgs {
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>,
) {
match self.extract_subtitles(&url, whisper_audio_stream_ix, sender.clone()) {
Ok(_) => {
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 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 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 (packet_tx, join_handle) in subtitle_extractors.into_values() {
drop(packet_tx);
join_handle
.join()
.unwrap()
.unwrap_or_else(|e| log::error!("error running subtitle extraction: {}", e));
}
Ok(())
}
}
|