subtile_ocr/
lib.rs

1#![doc = include_str!("../README.md")]
2
3mod ocr;
4mod opt;
5
6pub use crate::{ocr::process, ocr::Error as OcrError, ocr::OcrOpt, opt::Opt};
7
8use image::GrayImage;
9use log::warn;
10use rayon::{
11    iter::{IntoParallelRefIterator as _, ParallelIterator as _},
12    ThreadPoolBuildError,
13};
14use std::{
15    fs::File,
16    io::{self, BufReader, BufWriter},
17    path::{Path, PathBuf},
18};
19use subtile::{
20    image::{dump_images, luma_a_to_luma, ToImage as _, ToOcrImage as _, ToOcrImageOpt},
21    pgs::{self, DecodeTimeImage, RleToImage},
22    srt,
23    time::TimeSpan,
24    vobsub::{
25        self, conv_to_rgba, palette_rgb_to_luminance, VobSubError, VobSubIndexedImage,
26        VobSubOcrImage, VobSubToImage,
27    },
28    SubtileError,
29};
30use thiserror::Error;
31
32/// Gather different `Error`s in a dedicated enum.
33#[expect(missing_docs)]
34#[derive(Error, Debug)]
35pub enum Error {
36    #[error("failed to create a rayon ThreadPool")]
37    RayonThreadPool(#[from] ThreadPoolBuildError),
38
39    #[error("the file extension '{extension}' is not managed")]
40    InvalidFileExtension { extension: String },
41
42    #[error("the file doesn't have an extension, can't choose a parser")]
43    NoFileExtension,
44
45    #[error("the file doesn't have an utf8 extension, can't choose a parser")]
46    NotUtf8Extension,
47
48    #[error("failed to open `Index` file")]
49    IndexOpen(#[source] VobSubError),
50
51    #[error("failed to open `Sub` file")]
52    SubOpen(#[source] VobSubError),
53
54    #[error("failed to create PgsParser from file")]
55    PgsParserFromFile(#[source] pgs::PgsError),
56
57    #[error("failed to parse Pgs")]
58    PgsParsing(#[source] pgs::PgsError),
59
60    #[error("failed to dump subtitles images")]
61    DumpImage(#[source] SubtileError),
62
63    #[error("could not perform OCR on subtitles.")]
64    Ocr(#[from] ocr::Error),
65
66    #[error("error happen during OCR on {0} subtitles images")]
67    OcrFails(u32),
68
69    #[error("could not generate SRT file: {message}")]
70    GenerateSrt { message: String },
71
72    #[error("could not write SRT file {}", path.display())]
73    WriteSrtFile { path: PathBuf, source: io::Error },
74
75    #[error("could not write SRT on stdout.")]
76    WriteSrtStdout { source: io::Error },
77}
78
79/// Run OCR for `opt`.
80///
81/// # Errors
82///
83/// Will return [`Error::RayonThreadPool`] if `build_global` of the `ThreadPool` rayon failed.
84/// Will return [`Error::InvalidFileExtension`] if the file extension is not managed.
85/// Will return [`Error::NoFileExtension`] if the file have no extension.
86/// Will return [`Error::NotUtf8Extension`] if the file have an extension which is not utf8.
87/// Will return [`Error::WriteSrtFile`] of [`Error::WriteSrtStdout`] if failed to write subtitles as `srt`.
88/// Will forward error from `ocr` processing and [`check_subtitles`] if any.
89#[profiling::function]
90pub fn run(opt: &Opt) -> Result<(), Error> {
91    rayon::ThreadPoolBuilder::new()
92        .thread_name(|idx| format!("Rayon_{idx}"))
93        .build_global()
94        .map_err(Error::RayonThreadPool)?;
95
96    let (times, images) = match extract_extension(&opt.input)? {
97        "sup" => process_pgs(opt),
98        "sub" | "idx" => process_vobsub(opt),
99        ext => Err(Error::InvalidFileExtension {
100            extension: ext.into(),
101        }),
102    }?;
103
104    // Dump images if requested.
105    if opt.dump {
106        dump_images("dumps", &images).map_err(Error::DumpImage)?;
107    }
108
109    let ocr_opt = OcrOpt::new(&opt.tessdata_dir, opt.lang.as_str(), &opt.config, opt.dpi);
110    let texts = ocr::process(images, &ocr_opt)?;
111    let subtitles = check_subtitles(times.into_iter().zip(texts))?;
112
113    // Create subtitle file.
114    write_srt(opt.output.as_deref(), &subtitles)?;
115
116    Ok(())
117}
118
119/// Extract extension of a path
120///
121/// # Errors
122///
123/// Will return [`Error::NoFileExtension`] if the file have no extension.
124/// Will return [`Error::NotUtf8Extension`] if the file have an extension which is not utf8.
125pub fn extract_extension(path: &Path) -> Result<&str, Error> {
126    path.extension()
127        .ok_or(Error::NoFileExtension)?
128        .to_str()
129        .ok_or(Error::NotUtf8Extension)
130}
131
132/// Process `PGS` subtitle file
133///
134/// # Errors
135///
136/// Will return [`Error::PgsParserFromFile`] if `SupParser` failed to be init from file.
137/// Will return [`Error::PgsParsing`] if the parsing of subtitles failed.
138/// Will return [`Error::DumpImage`] if the dump of raw image failed.
139#[profiling::function]
140pub fn process_pgs(opt: &Opt) -> Result<(Vec<TimeSpan>, Vec<GrayImage>), Error> {
141    let parser = {
142        profiling::scope!("Create PGS parser");
143        subtile::pgs::SupParser::<BufReader<File>, DecodeTimeImage>::from_file(&opt.input)
144            .map_err(Error::PgsParserFromFile)?
145    };
146
147    let (times, rle_images) = {
148        profiling::scope!("Parse PGS file");
149        parser
150            .collect::<Result<(Vec<_>, Vec<_>), _>>()
151            .map_err(Error::PgsParsing)?
152    };
153
154    if opt.dump_raw {
155        let images = rle_images
156            .iter()
157            .map(|rle_img| RleToImage::new(rle_img, |pix| pix).to_image());
158        dump_images("dumps_raw", images).map_err(Error::DumpImage)?;
159    }
160
161    let conv_fn = luma_a_to_luma::<_, _, 100, 100>; // Hardcoded value for alpha and luma threshold than work not bad.
162
163    let images = {
164        profiling::scope!("Convert images for OCR");
165        let ocr_opt = ocr_opt(opt);
166        rle_images
167            .par_iter()
168            .map(|rle_img| RleToImage::new(rle_img, &conv_fn).image(&ocr_opt))
169            .collect::<Vec<_>>()
170    };
171
172    Ok((times, images))
173}
174
175/// Process `VobSub` subtitle file
176///
177/// # Errors
178///
179/// Will return [`Error::IndexOpen`] if the subtitle files can't be opened.
180/// Will return [`Error::DumpImage`] if the dump of raw image failed.
181#[profiling::function]
182pub fn process_vobsub(opt: &Opt) -> Result<(Vec<TimeSpan>, Vec<GrayImage>), Error> {
183    let mut input_path = opt.input.clone();
184    let sub = {
185        profiling::scope!("Open sub");
186        input_path.set_extension("sub");
187        vobsub::Sub::open(&input_path).map_err(Error::SubOpen)?
188    };
189    let idx = {
190        profiling::scope!("Open idx");
191        input_path.set_extension("idx");
192        vobsub::Index::open(&input_path).map_err(Error::IndexOpen)?
193    };
194    let (times, images): (Vec<_>, Vec<_>) = {
195        profiling::scope!("Parse subtitles");
196        sub.subtitles::<(TimeSpan, VobSubIndexedImage)>()
197            .filter_map(|sub| match sub {
198                Ok(sub) => Some(sub),
199                Err(e) => {
200                    warn!(
201        "warning: unable to read subtitle: {e}. (This can usually be safely ignored.)"
202    );
203                    None
204                }
205            })
206            .unzip()
207    };
208
209    if opt.dump_raw {
210        let images = images.iter().map(|rle_img| {
211            let image: image::RgbaImage =
212                VobSubToImage::new(rle_img, idx.palette(), conv_to_rgba).to_image();
213            image
214        });
215        dump_images("dumps_raw", images).map_err(Error::DumpImage)?;
216    }
217
218    let images_for_ocr = {
219        profiling::scope!("Convert images for OCR");
220
221        let ocr_opt = ocr_opt(opt);
222        let palette = palette_rgb_to_luminance(idx.palette());
223        images
224            .par_iter()
225            .map(|vobsub_img| {
226                let converter = VobSubOcrImage::new(vobsub_img, &palette);
227                converter.image(&ocr_opt)
228            })
229            .collect::<Vec<_>>()
230    };
231
232    Ok((times, images_for_ocr))
233}
234
235/// Create [`ToOcrImageOpt`] from [`Opt`]
236fn ocr_opt(opt: &Opt) -> ToOcrImageOpt {
237    ToOcrImageOpt {
238        border: opt.border,
239        ..Default::default()
240    }
241}
242
243/// Log errors and remove bad results.
244///
245/// # Errors
246///  Will return [`Error::OcrFails`] if the ocr return an error for at least one image.
247#[profiling::function]
248pub fn check_subtitles<In>(subtitles: In) -> Result<Vec<(TimeSpan, String)>, Error>
249where
250    In: IntoIterator<Item = (TimeSpan, Result<String, ocr::Error>)>,
251{
252    let mut ocr_error_count = 0;
253    let subtitles = subtitles
254        .into_iter()
255        .enumerate()
256        .filter_map(|(idx, (time, maybe_text))| match maybe_text {
257            Ok(text) => Some((time, text)),
258            Err(e) => {
259                let err = anyhow::Error::new(e); // warp in anyhow::Error to display the error stack with :#
260                warn!(
261                    "Error while running OCR on subtitle image ({} - {time:?}):\n\t {err:#}",
262                    idx + 1,
263                );
264                ocr_error_count += 1;
265                None
266            }
267        })
268        .collect::<Vec<_>>();
269
270    if ocr_error_count > 0 {
271        Err(Error::OcrFails(ocr_error_count))
272    } else {
273        Ok(subtitles)
274    }
275}
276
277#[profiling::function]
278fn write_srt(path: Option<&Path>, subtitles: &[(TimeSpan, String)]) -> Result<(), Error> {
279    if let Some(path) = &path {
280        write_to_file(path, subtitles).map_err(|source| Error::WriteSrtFile {
281            path: path.to_path_buf(),
282            source,
283        })?;
284    } else {
285        // Write to stdout.
286        let mut stdout = io::stdout();
287        srt::write_srt(&mut stdout, subtitles)
288            .map_err(|source| Error::WriteSrtStdout { source })?;
289    }
290    Ok(())
291}
292
293// Write to file.
294fn write_to_file(path: &Path, subtitles: &[(TimeSpan, String)]) -> Result<(), io::Error> {
295    let subtitle_file = File::create(path)?;
296    let mut stream = BufWriter::new(subtitle_file);
297    srt::write_srt(&mut stream, subtitles)
298}
subtile_ocr/lib.rs

subtile_ocr/
lib.rs
