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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
|
/// 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(())
}
}
|