Skip to content

Commit

Permalink
feat: add dump-raw option
Browse files Browse the repository at this point in the history
this dump images from subtitle before transforming them for OCR
  • Loading branch information
gwen-lg committed Aug 11, 2024
1 parent 09d6252 commit 7cb20f7
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 3 deletions.
22 changes: 19 additions & 3 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand All @@ -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;
Expand Down Expand Up @@ -134,6 +134,13 @@ pub fn process_pgs(opt: &Opt) -> Result<(Vec<TimeSpan>, Vec<GrayImage>), Error>
.map_err(Error::PgsParsing)?
};

if opt.dump_raw {
let images = rle_images
.iter()
.map(|rle_img| RleToImage::new(rle_img, |pix: LumaA<u8>| 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 = {
Expand Down Expand Up @@ -175,6 +182,15 @@ pub fn process_vobsub(opt: &Opt) -> Result<(Vec<TimeSpan>, Vec<GrayImage>), 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");

Expand Down
4 changes: 4 additions & 0 deletions src/opt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 7cb20f7

Please sign in to comment.