From 7cb20f73b078cead545a534b83859076d387220c Mon Sep 17 00:00:00 2001
From: Gwen Lg <me@gwenlg.fr>
Date: Fri, 12 Jul 2024 11:42:35 +0200
Subject: [PATCH] feat: add dump-raw option

this dump images from subtitle before transforming them for OCR
---
 src/lib.rs | 22 +++++++++++++++++++---
 src/opt.rs |  4 ++++
 2 files changed, 23 insertions(+), 3 deletions(-)

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<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 = {
@@ -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");
 
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