diff --git a/src/lib.rs b/src/lib.rs index bb04cd0..354d5c2 100755 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,7 +6,7 @@ mod preprocessor; pub use crate::{ocr::OcrOpt, opt::Opt}; -use image::GrayImage; +use image::{GrayImage, LumaA}; use log::warn; use preprocessor::rgb_palette_to_luminance; use rayon::{ @@ -20,11 +20,11 @@ use std::{ path::PathBuf, }; use subtile::{ - image::{dump_images, luma_a_to_luma, ToOcrImage, ToOcrImageOpt}, + image::{dump_images, luma_a_to_luma, ToImage, ToOcrImage, ToOcrImageOpt}, pgs::{self, DecodeTimeImage, RleToImage}, srt, time::TimeSpan, - vobsub::{self, VobSubError, VobSubIndexedImage, VobSubOcrImage}, + vobsub::{self, conv_to_rgba, VobSubError, VobSubIndexedImage, VobSubOcrImage, VobSubToImage}, SubtileError, }; use thiserror::Error; @@ -134,6 +134,13 @@ pub fn process_pgs(opt: &Opt) -> Result<(Vec, Vec), Error> .map_err(Error::PgsParsing)? }; + if opt.dump_raw { + let images = rle_images + .iter() + .map(|rle_img| RleToImage::new(rle_img, |pix: LumaA| pix).to_image()); + dump_images("dumps_raw", images).map_err(Error::DumpImage)?; + } + let conv_fn = luma_a_to_luma::<_, _, 100, 100>; // Hardcoded value for alpha and luma threshold than work not bad. let images = { @@ -175,6 +182,15 @@ pub fn process_vobsub(opt: &Opt) -> Result<(Vec, Vec), Erro .unzip() }; + if opt.dump_raw { + let images = images.iter().map(|rle_img| { + let image: image::RgbaImage = + VobSubToImage::new(rle_img, idx.palette(), conv_to_rgba).to_image(); + image + }); + dump_images("dumps_raw", images).map_err(Error::DumpImage)?; + } + let images_for_ocr = { profiling::scope!("Convert images for OCR"); diff --git a/src/opt.rs b/src/opt.rs index 6eef30e..758b2cd 100644 --- a/src/opt.rs +++ b/src/opt.rs @@ -66,6 +66,10 @@ pub struct Opt { /// Dump processed subtitle images into the working directory as PNG files. #[clap(long)] pub dump: bool, + + /// Dump raw subtitle images into the working directory as PNG files. + #[clap(long)] + pub dump_raw: bool, } // https://github.com/clap-rs/clap_derive/blob/master/examples/keyvalue.rs