diff --git a/src/lib_ccx/hardsubx.c b/src/lib_ccx/hardsubx.c index 411377f95..49b160621 100644 --- a/src/lib_ccx/hardsubx.c +++ b/src/lib_ccx/hardsubx.c @@ -10,6 +10,7 @@ #include #include +#ifdef DISABLE_RUST int hardsubx_process_data(struct lib_hardsubx_ctx *ctx, struct lib_ccx_ctx *ctx_normal) { // Get the required media attributes and initialize structures @@ -119,7 +120,7 @@ int hardsubx_process_data(struct lib_hardsubx_ctx *ctx, struct lib_ccx_ctx *ctx_ avcodec_close(ctx->codec_ctx); avformat_close_input(&ctx->format_ctx); } - +#endif void _hardsubx_params_dump(struct ccx_s_options *options, struct lib_hardsubx_ctx *ctx) { // Print the relevant passed parameters onto the screen @@ -317,6 +318,7 @@ void _dinit_hardsubx(struct lib_hardsubx_ctx **ctx) freep(ctx); } +#ifdef DISABLE_RUST void hardsubx(struct ccx_s_options *options, struct lib_ccx_ctx *ctx_normal) { // This is similar to the 'main' function in ccextractor.c, but for hard subs @@ -344,5 +346,5 @@ void hardsubx(struct ccx_s_options *options, struct lib_ccx_ctx *ctx_normal) // Free all allocated memory for the data structures _dinit_hardsubx(&ctx); } - +#endif #endif diff --git a/src/rust/Cargo.lock b/src/rust/Cargo.lock index 34d668338..928606a01 100644 --- a/src/rust/Cargo.lock +++ b/src/rust/Cargo.lock @@ -133,6 +133,9 @@ dependencies = [ "iconv", "leptonica-sys", "log", + "num", + "num-derive", + "num-traits", "palette", "tesseract-sys", ] @@ -208,9 +211,9 @@ dependencies = [ [[package]] name = "ffmpeg-sys-next" -version = "5.0.1" +version = "5.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ba12dea33516e30c160ce557c7e43dd857276368eb1cd0eef4fce6529f2dee5" +checksum = "d780b36e092254367e2f1f21191992735c8e23f31a5a5a8678db3a79f775021f" dependencies = [ "bindgen 0.59.2", "cc", @@ -351,6 +354,84 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43db66d1170d347f9a065114077f7dccb00c1b9478c89384490a3425279a4606" +dependencies = [ + "num-bigint", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + +[[package]] +name = "num-bigint" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f93ab6289c7b344a8a9f60f88d80aa20032336fe78da341afc91c8a2341fc75f" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ae39348c8bc5fbd7f40c727a9925f03517afd2ab27d46702108b6a7e5414c19" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-derive" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "876a53fff98e03a936a674b29568b0e605f06b29372c2489ff4de23f1949743d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "num-integer" +version = "0.1.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" +dependencies = [ + "autocfg", + "num-traits", +] + +[[package]] +name = "num-iter" +version = "0.1.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d03e6c028c5dc5cac6e2dec0efda81fc887605bb3d884578bb6d6bf7514e252" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0638a1c9d0a3c0914158145bc76cff373a75a627e6ecbfb71cbe6f453a5a19b0" +dependencies = [ + "autocfg", + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.15" diff --git a/src/rust/Cargo.toml b/src/rust/Cargo.toml index ea39d2d3c..8f4221a3e 100644 --- a/src/rust/Cargo.toml +++ b/src/rust/Cargo.toml @@ -18,6 +18,10 @@ palette = "0.6.0" ffmpeg-sys-next = { version = "5.0.1", optional = true, default-features = false, features = ["avcodec", "avformat", "swscale", "build"]} tesseract-sys = { version = "0.5.12", optional = true, default-features = false} leptonica-sys = { version = "0.4.1", optional = true, default-features = false} +num = "0.4" +num-derive = "0.3" +num-traits = "0.2" + [build-dependencies] bindgen = "0.58.1" diff --git a/src/rust/build.rs b/src/rust/build.rs index a93ea3058..b43b57acd 100644 --- a/src/rust/build.rs +++ b/src/rust/build.rs @@ -16,6 +16,29 @@ fn main() { "convert_pts_to_.*", "av_rescale_q", "mprint", + "av_packet_unref", + "av_read_frame", + "activity_progress", + "avcodec_decode_video2", + "probe_tessdata_location", + "get_file_extension", + "init_encoder", + "dinit_encoder", + "encode_sub", + "add_cc_sub_text", + "position_sanity_check", + "process_non_multiprogram_general_loop", + "get_fts", + "segment_output_file", + "net_check_conn", + "update_encoder_list", + "general_get_more_data", + "wtv_get_more_data", + "asf_get_more_data", + "ts_get_more_data", + "ps_get_more_data", + "is_decoder_processed_enough", + "flush_cc_decode", ]); let mut allowlist_types = vec![ @@ -24,6 +47,10 @@ fn main() { "lib_cc_decode", "cc_subtitle", "ccx_output_format", + "ccx_s_options", + "lib_ccx_ctx", + "encoder_ctx", + "demuxer_data", ]; #[cfg(feature = "hardsubx_ocr")] diff --git a/src/rust/src/hardsubx/classifier.rs b/src/rust/src/hardsubx/classifier.rs index d8572832a..d08e9e5d1 100644 --- a/src/rust/src/hardsubx/classifier.rs +++ b/src/rust/src/hardsubx/classifier.rs @@ -12,33 +12,31 @@ pub type subdatatype = ::std::os::raw::c_uint; pub type subtype = ::std::os::raw::c_uint; pub type ccx_encoding_type = ::std::os::raw::c_uint; -use crate::hardsubx::lib_hardsubx_ctx; -use crate::utils::string_to_c_char; - use std::os::raw::c_char; use log::warn; +use super::HardsubxContext; + /// # Safety /// The function accepts and dereferences a raw pointer /// The function also makes calls to functions whose safety is not guaranteed -/// The function returns a raw pointer which is a string made in C /// ctx should be not null #[no_mangle] -pub unsafe extern "C" fn get_ocr_text_simple_threshold( - ctx: *mut lib_hardsubx_ctx, +pub unsafe fn get_ocr_text_simple_threshold( + ctx: &mut HardsubxContext, image: *mut Pix, threshold: std::os::raw::c_float, -) -> *mut ::std::os::raw::c_char { +) -> String { let mut text_out: *mut ::std::os::raw::c_char; - TessBaseAPISetImage2((*ctx).tess_handle, image); + TessBaseAPISetImage2(ctx.tess_handle, image); - if TessBaseAPIRecognize((*ctx).tess_handle, null::() as *mut ETEXT_DESC) != 0 { + if TessBaseAPIRecognize(ctx.tess_handle, null::() as *mut ETEXT_DESC) != 0 { warn!("Error in Tesseract recognition, skipping frame\n"); - null::() as *mut c_char + String::new() } else { - text_out = TessBaseAPIGetUTF8Text((*ctx).tess_handle); + text_out = TessBaseAPIGetUTF8Text(ctx.tess_handle); if text_out == null::() as *mut c_char { warn!("Error getting text, skipping frame\n"); @@ -46,31 +44,17 @@ pub unsafe extern "C" fn get_ocr_text_simple_threshold( if threshold > 0.0 { // non-zero conf, only then we'll make the call to check for confidence - let conf = TessBaseAPIMeanTextConf((*ctx).tess_handle); + let conf = TessBaseAPIMeanTextConf(ctx.tess_handle); if (conf as std::os::raw::c_float) < threshold { text_out = null::() as *mut c_char; } else { - (*ctx).cur_conf = conf as std::os::raw::c_float; + ctx.cur_conf = conf as std::os::raw::c_float; } } - text_out - } -} -/// basically the get_oct_text_simple function without threshold -/// This function is being kept only for backwards compatibility reasons -/// # Safety -/// The function accepts and dereferences a raw pointer -/// The function also makes calls to functions whose safety is not guaranteed -/// The function returns a raw pointer which is a string made in C -/// ctx should be not null -#[no_mangle] -pub unsafe extern "C" fn get_ocr_text_simple( - ctx: *mut lib_hardsubx_ctx, - image: *mut Pix, -) -> *mut ::std::os::raw::c_char { - get_ocr_text_simple_threshold(ctx, image, 0.0) + ffi::CStr::from_ptr(text_out).to_string_lossy().into_owned() + } } /// Function extracts string from tess iterator object @@ -100,25 +84,23 @@ unsafe fn _tess_string_helper(it: *mut TessResultIterator, level: TessPageIterat /// # Safety /// The function dereferences a raw pointer /// The function also calls other functions whose safety is not guaranteed -/// The function returns a raw pointer of a String created in Rust -/// This has to be deallocated at some point using from_raw() lest it be a memory leak /// ctx should be not null #[no_mangle] -pub unsafe extern "C" fn get_ocr_text_wordwise_threshold( - ctx: *mut lib_hardsubx_ctx, +pub unsafe fn get_ocr_text_wordwise_threshold( + ctx: &mut HardsubxContext, image: *mut Pix, threshold: std::os::raw::c_float, -) -> *mut ::std::os::raw::c_char { +) -> String { let mut text_out = String::new(); - TessBaseAPISetImage2((*ctx).tess_handle, image); + TessBaseAPISetImage2(ctx.tess_handle, image); - if TessBaseAPIRecognize((*ctx).tess_handle, null::() as *mut ETEXT_DESC) != 0 { + if TessBaseAPIRecognize(ctx.tess_handle, null::() as *mut ETEXT_DESC) != 0 { warn!("Error in Tesseract recognition, skipping word\n"); - return null::() as *mut c_char; + return String::new(); } - let it: *mut TessResultIterator = TessBaseAPIGetIterator((*ctx).tess_handle); + let it: *mut TessResultIterator = TessBaseAPIGetIterator(ctx.tess_handle); let level: TessPageIteratorLevel = TessPageIteratorLevel_RIL_WORD; let mut prev_ital: bool = false; @@ -132,7 +114,7 @@ pub unsafe extern "C" fn get_ocr_text_wordwise_threshold( if first_iter { first_iter = false; } else if TessPageIteratorNext(it as *mut TessPageIterator, level) == 0 { - if (*ctx).detect_italics == 1 && prev_ital { + if ctx.detect_italics && prev_ital { // if there are italics words at the end text_out = format!("{}", text_out); } @@ -155,7 +137,7 @@ pub unsafe extern "C" fn get_ocr_text_wordwise_threshold( num_words += 1; } - if (*ctx).detect_italics != 0 { + if ctx.detect_italics { let mut italic: i32 = 0; let mut dummy: i32 = 0; @@ -185,50 +167,34 @@ pub unsafe extern "C" fn get_ocr_text_wordwise_threshold( if threshold > 0.0 { // any confidence calculation has happened if and only if threshold > 0 - (*ctx).cur_conf = total_conf / (num_words as std::os::raw::c_float); + ctx.cur_conf = total_conf / (num_words as std::os::raw::c_float); } TessResultIteratorDelete(it); - string_to_c_char(&text_out) + text_out } /// # Safety /// The function dereferences a raw pointer /// The function also calls other functions whose safety is not guaranteed -/// The function returns a raw pointer of a String created in Rust -/// This has to be deallocated at some point using from_raw() lest it be a memory leak /// ctx should be not null #[no_mangle] -pub unsafe extern "C" fn get_ocr_text_wordwise( - ctx: *mut lib_hardsubx_ctx, - image: *mut Pix, -) -> *mut ::std::os::raw::c_char { - get_ocr_text_wordwise_threshold(ctx, image, 0.0) -} - -/// # Safety -/// The function dereferences a raw pointer -/// The function also calls other functions whose safety is not guaranteed -/// The function returns a raw pointer of a String created in Rust -/// This has to be deallocated at some point using from_raw() lest it be a memory leak -/// ctx should be not null -#[no_mangle] -pub unsafe extern "C" fn get_ocr_text_letterwise_threshold( - ctx: *mut lib_hardsubx_ctx, +pub unsafe fn get_ocr_text_letterwise_threshold( + ctx: &mut HardsubxContext, image: *mut Pix, threshold: std::os::raw::c_float, -) -> *mut ::std::os::raw::c_char { +) -> String { let mut text_out: String = String::new(); - TessBaseAPISetImage2((*ctx).tess_handle, image); + TessBaseAPISetImage2(ctx.tess_handle, image); - if TessBaseAPIRecognize((*ctx).tess_handle, null::() as *mut ETEXT_DESC) != 0 { + if TessBaseAPIRecognize(ctx.tess_handle, null::() as *mut ETEXT_DESC) != 0 { warn!("Error in Tesseract recognition, skipping symbol\n"); - return null::() as *mut c_char; + return String::new(); } - let it: *mut TessResultIterator = TessBaseAPIGetIterator((*ctx).tess_handle); + let it: *mut TessResultIterator = TessBaseAPIGetIterator(ctx.tess_handle); let level: TessPageIteratorLevel = TessPageIteratorLevel_RIL_SYMBOL; let mut total_conf: std::os::raw::c_float = 0.0; @@ -262,24 +228,10 @@ pub unsafe extern "C" fn get_ocr_text_letterwise_threshold( if threshold > 0.0 { // No confidence calculation has been done if threshold is 0 or less - (*ctx).cur_conf = total_conf / (num_characters as std::os::raw::c_float); + ctx.cur_conf = total_conf / (num_characters as std::os::raw::c_float); } TessResultIteratorDelete(it); - string_to_c_char(&text_out) -} - -/// # Safety -/// The function dereferences a raw pointer -/// The function also calls other functions whose safety is not guaranteed -/// The function returns a raw pointer of a String created in Rust -/// This has to be deallocated at some point using from_raw() lest it be a memory leak -/// ctx should be not null -#[no_mangle] -pub unsafe extern "C" fn get_ocr_text_letterwise( - ctx: *mut lib_hardsubx_ctx, - image: *mut Pix, -) -> *mut ::std::os::raw::c_char { - get_ocr_text_letterwise_threshold(ctx, image, 0.0) + text_out } diff --git a/src/rust/src/hardsubx/decoder.rs b/src/rust/src/hardsubx/decoder.rs index de8c6a307..714bbaf21 100644 --- a/src/rust/src/hardsubx/decoder.rs +++ b/src/rust/src/hardsubx/decoder.rs @@ -1,72 +1,84 @@ #[cfg(feature = "hardsubx_ocr")] +use ffmpeg_sys_next::*; +#[cfg(feature = "hardsubx_ocr")] use leptonica_sys::*; +use num_derive::FromPrimitive; +use num_traits::FromPrimitive; // #[cfg(feature = "hardsubx_ocr")] // use ffmpeg_sys_next::*; use std::convert::TryInto; -use std::eprintln; use std::ffi; use std::format; use std::os::raw::c_char; -use std::process::exit; +use std::os::raw::c_int; use std::ptr::null; +use crate::bindings::{ + activity_progress, add_cc_sub_text, asf_get_more_data, cc_subtitle, demuxer_data, encode_sub, + encoder_ctx, flush_cc_decode, general_get_more_data, get_fts, is_decoder_processed_enough, + lib_cc_decode, lib_ccx_ctx, net_check_conn, position_sanity_check, + process_non_multiprogram_general_loop, ps_get_more_data, segment_output_file, ts_get_more_data, + update_encoder_list, wtv_get_more_data, +}; #[cfg(feature = "hardsubx_ocr")] // use crate::bindings::{hardsubx_ocr_mode_HARDSUBX_OCRMODE_WORD}; -use crate::bindings::AVFrame; use crate::hardsubx::classifier::*; use crate::hardsubx::imgops::{rgb_to_hsv, rgb_to_lab}; -use crate::hardsubx::lib_hardsubx_ctx; use crate::utils::string_to_c_char; -static EXIT_MALFORMED_PARAMETER: i32 = 7; +use super::ccx_options; +use super::hardsubx_color_type; +use super::hardsubx_ocr_mode; +use super::utility::*; +use super::HardsubxContext; +use crate::ccx_encoding_type::*; + +use std::cmp; + +use std::process; -// TODO: turn into an enum definition when the hardsubx context is rewritten -// static HARDSUBX_OCRMODE_FRAME: i32 = 0; -static HARDSUBX_OCRMODE_WORD: i32 = 1; -// static HARDSUBX_OCRMODE_LETTER: i32 = 2; +extern "C" { + pub static mut terminate_asap: c_int; +} -// enum hardsubx_ocr_mode { -// HARDSUBX_OCRMODE_FRAME, -// HARDSUBX_OCRMODE_WORD, -// HARDSUBX_OCRMODE_LETTER -// }; +#[derive(PartialEq, FromPrimitive)] +enum ccx_stream_mode_enum { + CCX_SM_ELEMENTARY_OR_NOT_FOUND = 0, + CCX_SM_TRANSPORT = 1, + CCX_SM_PROGRAM = 2, + CCX_SM_ASF = 3, + CCX_SM_MCPOODLESRAW = 4, + CCX_SM_RCWT = 5, // Raw Captions With Time, not used yet. + CCX_SM_MYTH = 6, // Use the myth loop + CCX_SM_MP4 = 7, // MP4, ISO- + CCX_SM_HEX_DUMP = 8, // Hexadecimal dump generated by wtvccdump + CCX_SM_WTV = 9, + CCX_SM_FFMPEG = 10, + CCX_SM_GXF = 11, + CCX_SM_MKV = 12, + CCX_SM_MXF = 13, + + CCX_SM_AUTODETECT = 16, +} /// # Safety /// dereferences a raw pointer /// calls functions that are not necessarily safe -pub unsafe fn dispatch_classifier_functions(ctx: *mut lib_hardsubx_ctx, im: *mut Pix) -> String { +pub unsafe fn dispatch_classifier_functions(ctx: &mut HardsubxContext, im: *mut Pix) -> String { // function that calls the classifier functions - match (*ctx).ocr_mode { - 0 => { - let ret_char_arr = get_ocr_text_wordwise_threshold(ctx, im, (*ctx).conf_thresh); - ffi::CStr::from_ptr(ret_char_arr) - .to_string_lossy() - .into_owned() - } - 1 => { - let ret_char_arr = get_ocr_text_letterwise_threshold(ctx, im, (*ctx).conf_thresh); - let text_out_result = ffi::CString::from_raw(ret_char_arr).into_string(); - match text_out_result { - Ok(T) => T, - Err(_E) => "".to_string(), - } + match ctx.ocr_mode { + hardsubx_ocr_mode::HARDSUBX_OCRMODE_WORD => { + get_ocr_text_wordwise_threshold(ctx, im, (*ctx).conf_thresh) } - 2 => { - let ret_char_arr = get_ocr_text_simple_threshold(ctx, im, (*ctx).conf_thresh); - let text_out_result = ffi::CString::from_raw(ret_char_arr).into_string(); - match text_out_result { - Ok(T) => T, - Err(_E) => "".to_string(), - } + hardsubx_ocr_mode::HARDSUBX_OCRMODE_LETTER => { + get_ocr_text_letterwise_threshold(ctx, im, (*ctx).conf_thresh) } - _ => { - eprintln!("Invalid OCR Mode"); - exit(EXIT_MALFORMED_PARAMETER); - // "".to_string() + hardsubx_ocr_mode::HARDSUBX_OCRMODE_FRAME => { + get_ocr_text_simple_threshold(ctx, im, (*ctx).conf_thresh) } } } @@ -77,13 +89,13 @@ pub unsafe fn dispatch_classifier_functions(ctx: *mut lib_hardsubx_ctx, im: *mut /// The function returns a raw pointer of a String created in Rust /// This has to be deallocated at some point using from_raw() lest it be a memory leak #[no_mangle] -pub unsafe extern "C" fn _process_frame_white_basic( - ctx: *mut lib_hardsubx_ctx, +pub unsafe fn _process_frame_white_basic( + ctx: &mut HardsubxContext, frame: *mut AVFrame, width: ::std::os::raw::c_int, height: ::std::os::raw::c_int, _index: ::std::os::raw::c_int, -) -> *mut ::std::os::raw::c_char { +) -> String { let mut im: *mut Pix = pixCreate(width, height, 32); let mut lum_im: *mut Pix = pixCreate(width, height, 32); let frame_deref = *frame; @@ -102,7 +114,7 @@ pub unsafe extern "C" fn _process_frame_white_basic( rgb_to_lab(r as f32, g as f32, b as f32, &mut L, &mut A, &mut B); - if L > (*ctx).lum_thresh { + if L > ctx.lum_thresh { pixSetRGBPixel(lum_im, j, i, 255, 255, 255); } else { pixSetRGBPixel(lum_im, j, i, 0, 0, 0); @@ -134,8 +146,8 @@ pub unsafe extern "C" fn _process_frame_white_basic( } } - if (*ctx).detect_italics != 0 { - (*ctx).ocr_mode = HARDSUBX_OCRMODE_WORD; + if ctx.detect_italics { + ctx.ocr_mode = hardsubx_ocr_mode::HARDSUBX_OCRMODE_WORD; } let subtitle_text = dispatch_classifier_functions(ctx, feat_im); @@ -148,7 +160,7 @@ pub unsafe extern "C" fn _process_frame_white_basic( pixDestroy(&mut lum_im as *mut *mut Pix); pixDestroy(&mut feat_im as *mut *mut Pix); - string_to_c_char(&subtitle_text) + subtitle_text } /// # Safety @@ -157,13 +169,13 @@ pub unsafe extern "C" fn _process_frame_white_basic( /// The function returns a raw pointer of a String created in Rust /// This has to be deallocated at some point using from_raw() lest it be a memory leak #[no_mangle] -pub unsafe extern "C" fn _process_frame_color_basic( - ctx: *mut lib_hardsubx_ctx, +pub unsafe fn _process_frame_color_basic( + ctx: &mut HardsubxContext, frame: *mut AVFrame, width: ::std::os::raw::c_int, height: ::std::os::raw::c_int, _index: ::std::os::raw::c_int, -) -> *mut ::std::os::raw::c_char { +) -> String { let mut im: *mut Pix = pixCreate(width, height, 32); let mut hue_im: *mut Pix = pixCreate(width, height, 32); let frame_deref = *frame; @@ -182,7 +194,7 @@ pub unsafe extern "C" fn _process_frame_color_basic( rgb_to_hsv(r as f32, g as f32, b as f32, &mut H, &mut S, &mut V); - if ((H - (*ctx).hue).abs()) < 20.0 { + if ((H - ctx.hue).abs()) < 20.0 { pixSetRGBPixel(hue_im, j, i, r, g, b); } } @@ -227,8 +239,8 @@ pub unsafe extern "C" fn _process_frame_color_basic( } } - if (*ctx).detect_italics != 0 { - (*ctx).ocr_mode = HARDSUBX_OCRMODE_WORD; + if ctx.detect_italics { + ctx.ocr_mode = hardsubx_ocr_mode::HARDSUBX_OCRMODE_WORD; } let subtitle_text = dispatch_classifier_functions(ctx, feat_im); @@ -246,20 +258,20 @@ pub unsafe extern "C" fn _process_frame_color_basic( // This is a memory leak // the returned thing needs to be deallocated by caller - string_to_c_char(&subtitle_text) + subtitle_text } /// # Safety /// The function accepts and dereferences a raw pointer /// The function also makes calls to functions whose safety is not guaranteed /// The function returns a raw pointer which is a string made in C #[no_mangle] -pub unsafe extern "C" fn _process_frame_tickertext( - ctx: *mut lib_hardsubx_ctx, +pub unsafe fn _process_frame_tickertext( + ctx: &mut HardsubxContext, frame: *mut AVFrame, width: ::std::os::raw::c_int, height: ::std::os::raw::c_int, index: ::std::os::raw::c_int, -) -> *mut ::std::os::raw::c_char { +) -> String { let mut im: *mut Pix = pixCreate(width, height, 32); let mut lum_im: *mut Pix = pixCreate(width, height, 32); let frame_deref = *frame; @@ -278,7 +290,7 @@ pub unsafe extern "C" fn _process_frame_tickertext( rgb_to_lab(r as f32, g as f32, b as f32, &mut L, &mut A, &mut B); - if L > (*ctx).lum_thresh { + if L > ctx.lum_thresh { pixSetRGBPixel(lum_im, j, i, 255, 255, 255); } else { pixSetRGBPixel(lum_im, j, i, 0, 0, 0); @@ -332,3 +344,664 @@ pub unsafe extern "C" fn _process_frame_tickertext( subtitle_text } + +/// # Safety +/// dereferences a raw pointer +/// calls potentially unsafe C functions +pub unsafe fn hardsubx_process_frames_linear(ctx: &mut HardsubxContext, enc_ctx: *mut encoder_ctx) { + let mut prev_sub_encoded: bool = true; + let mut got_frame = 0; + let mut dist = 0; + let mut cur_sec = 0; + let mut total_sec; + let mut progress; + + let mut frame_number = 0; + + let mut prev_begin_time: i64 = 0; + let mut prev_end_time: i64 = 0; + + let mut prev_packet_pts: i64 = 0; + + let mut subtitle_text = String::new(); + let mut prev_subtitle_text: String = String::new(); + + while av_read_frame(ctx.format_ctx, &mut ctx.packet as *mut AVPacket) >= 0 { + if ctx.packet.stream_index == ctx.video_stream_id { + frame_number += 1; + + let mut status = avcodec_send_packet(ctx.codec_ctx, &mut ctx.packet as *mut AVPacket); + // status = avcodec_receive_frame(ctx.codec_ctx, ctx.frame); + + if status >= 0 || status == AVERROR(EAGAIN) || status == AVERROR_EOF { + if status >= 0 { + ctx.packet.size = 0; + } + + status = avcodec_receive_frame(ctx.codec_ctx, ctx.frame); + + if status == 0 { + got_frame = 1; + } + } + + if got_frame != 0 && frame_number % 25 == 0 { + let diff = convert_pts_to_ms( + ctx.packet.pts - prev_packet_pts, + (**(*ctx.format_ctx) + .streams + .offset(ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + + if (diff.abs() as f32) < 1000.0 * ctx.min_sub_duration { + continue; + } + + sws_scale( + ctx.sws_ctx, + (*ctx.frame).data.as_ptr() as *const *const u8, + (*ctx.frame).linesize.as_mut_ptr(), + 0, + (*ctx.codec_ctx).height, + (*ctx.rgb_frame).data.as_mut_ptr(), + (*ctx.rgb_frame).linesize.as_mut_ptr(), + ); + + subtitle_text = match ctx.subcolor { + hardsubx_color_type::HARDSUBX_COLOR_WHITE => _process_frame_white_basic( + ctx, + ctx.rgb_frame, + (*ctx.codec_ctx).width, + (*ctx.codec_ctx).height, + frame_number, + ), + _ => _process_frame_color_basic( + ctx, + ctx.rgb_frame, + (*ctx.codec_ctx).width, + (*ctx.codec_ctx).height, + frame_number, + ), + }; + + cur_sec = convert_pts_to_s( + ctx.packet.pts, + (**(*ctx.format_ctx) + .streams + .offset(ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + total_sec = convert_pts_to_s((*ctx.format_ctx).duration, AV_TIME_BASE_Q); + + progress = (cur_sec * 100) / total_sec; + activity_progress( + progress.try_into().unwrap(), + (cur_sec / 60).try_into().unwrap(), + (cur_sec % 60).try_into().unwrap(), + ); + + if subtitle_text.is_empty() && prev_subtitle_text.is_empty() { + prev_end_time = convert_pts_to_ms( + ctx.packet.pts, + (**(*ctx.format_ctx) + .streams + .offset(ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + } + + if !subtitle_text.is_empty() { + let double_enter = subtitle_text.find("\n\n"); + + if let Some(T) = double_enter { + subtitle_text = subtitle_text[0..T].to_string(); + } + } + + if !prev_sub_encoded && !prev_subtitle_text.is_empty() { + if !subtitle_text.is_empty() { + dist = edit_distance_string(&subtitle_text, &prev_subtitle_text); + if (dist as f32) + < 0.2 + * (cmp::min( + subtitle_text.chars().count(), + prev_subtitle_text.chars().count(), + ) as f32) + { + dist = -1; + subtitle_text = String::new(); + prev_end_time = convert_pts_to_ms( + ctx.packet.pts, + (**(*ctx.format_ctx) + .streams + .offset(ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + } + } + + if dist != -1 { + let sub_text_chr = string_to_c_char(&subtitle_text); + let prev_text_chr = string_to_c_char(&prev_subtitle_text); + let empty_chr = string_to_c_char(""); + let mode_chr = string_to_c_char("BURN"); + add_cc_sub_text( + &mut *ctx.dec_sub as *mut cc_subtitle, + prev_text_chr, + prev_begin_time, + prev_end_time, + empty_chr, + mode_chr, + CCX_ENC_UTF_8 as u32, + ); + + encode_sub(enc_ctx, &mut *ctx.dec_sub); + + // Deallocation + subtitle_text = ffi::CString::from_raw(sub_text_chr) + .to_string_lossy() + .into_owned(); + ffi::CString::from_raw(prev_text_chr) + .to_string_lossy() + .into_owned(); + ffi::CString::from_raw(empty_chr); + ffi::CString::from_raw(mode_chr); + + prev_begin_time = prev_end_time + 1; + prev_subtitle_text = String::new(); + prev_sub_encoded = true; + prev_end_time = convert_pts_to_ms( + ctx.packet.pts, + (**(*ctx.format_ctx) + .streams + .offset(ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + + if !subtitle_text.is_empty() { + prev_subtitle_text = subtitle_text.clone(); + prev_sub_encoded = false; + } + } + dist = 0; + } + + if prev_subtitle_text.is_empty() && !subtitle_text.is_empty() { + prev_begin_time = prev_end_time + 1; + prev_end_time = convert_pts_to_ms( + ctx.packet.pts, + (**(*ctx.format_ctx) + .streams + .offset(ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + prev_subtitle_text = subtitle_text.clone(); + prev_sub_encoded = false; + } + prev_packet_pts = ctx.packet.pts; + } + } + av_packet_unref(&mut ctx.packet as *mut AVPacket); + } + + if !prev_sub_encoded { + let sub_text_chr = string_to_c_char(&subtitle_text); + let prev_text_chr = string_to_c_char(&prev_subtitle_text); + let empty_chr = string_to_c_char(""); + let mode_chr = string_to_c_char("BURN"); + add_cc_sub_text( + &mut *ctx.dec_sub as *mut cc_subtitle, + prev_text_chr, + prev_begin_time, + prev_begin_time, + empty_chr, + mode_chr, + CCX_ENC_UTF_8 as u32, + ); + + // Deallocation + ffi::CString::from_raw(sub_text_chr) + .to_string_lossy() + .into_owned(); + ffi::CString::from_raw(prev_text_chr) + .to_string_lossy() + .into_owned(); + ffi::CString::from_raw(empty_chr); + ffi::CString::from_raw(mode_chr); + + encode_sub(enc_ctx, &mut *ctx.dec_sub as *mut cc_subtitle); + } + + activity_progress( + 100, + (cur_sec / 60).try_into().unwrap(), + (cur_sec % 60).try_into().unwrap(), + ); +} + +/// # Safety +/// dereferences a raw pointer +/// calls potentially unsafe C functions +pub unsafe fn hardsubx_process_frames_tickertext( + ctx: &mut HardsubxContext, + _enc_ctx: *mut encoder_ctx, +) { + let mut got_frame: bool = false; + + let mut cur_sec: i64 = 0; + let mut total_sec: i64; + let mut progress: i64; + + let mut frame_number = 0; + + let mut ticker_text; + while av_read_frame(ctx.format_ctx, &mut ctx.packet) >= 0 { + if ctx.packet.stream_index == ctx.video_stream_id { + frame_number += 1; + + let mut status = avcodec_send_packet(ctx.codec_ctx, &mut ctx.packet as *mut AVPacket); + // status = avcodec_receive_frame(ctx.codec_ctx, ctx.frame); + + if status >= 0 || status == AVERROR(EAGAIN) || status == AVERROR_EOF { + if status >= 0 { + ctx.packet.size = 0; + } + + status = avcodec_receive_frame(ctx.codec_ctx, ctx.frame); + + if status == 0 { + got_frame = true; + } + } + + if got_frame && frame_number % 1000 == 0 { + sws_scale( + ctx.sws_ctx, + (*ctx.frame).data.as_ptr() as *const *const u8, + (*ctx.frame).linesize.as_mut_ptr(), + 0, + (*ctx.codec_ctx).height, + (*ctx.rgb_frame).data.as_ptr() as *const *mut u8, + (*ctx.rgb_frame).linesize.as_mut_ptr(), + ); + + ticker_text = _process_frame_tickertext( + ctx, + ctx.rgb_frame, + (*ctx.codec_ctx).width, + (*ctx.codec_ctx).height, + frame_number, + ); + println!("frame_number: {}", frame_number); + + if !ticker_text.is_empty() { + println!("{}", ticker_text); + } + + cur_sec = convert_pts_to_ms( + ctx.packet.pts, + (**(*ctx.format_ctx) + .streams + .offset(ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + + total_sec = convert_pts_to_s((*ctx.format_ctx).duration, AV_TIME_BASE_Q); + progress = (cur_sec * 100) / total_sec; + + activity_progress( + progress.try_into().unwrap(), + (cur_sec / 60).try_into().unwrap(), + (cur_sec % 60).try_into().unwrap(), + ); + } + } + + av_packet_unref(&mut ctx.packet); + } + + activity_progress( + 100, + (cur_sec / 60).try_into().unwrap(), + (cur_sec % 60).try_into().unwrap(), + ); +} + +/// # Safety +/// dereferences a raw pointer +/// calls potentially unsafe C functions +unsafe fn _get_more_data( + ctx_norm: *mut lib_ccx_ctx, + ppdata: *mut *mut demuxer_data, + stream_mode: &ccx_stream_mode_enum, +) -> i32 { + match stream_mode { + ccx_stream_mode_enum::CCX_SM_ELEMENTARY_OR_NOT_FOUND => { + general_get_more_data(ctx_norm, ppdata) + } + + ccx_stream_mode_enum::CCX_SM_TRANSPORT => ts_get_more_data(ctx_norm, ppdata), + + ccx_stream_mode_enum::CCX_SM_PROGRAM => ps_get_more_data(ctx_norm, ppdata), + + ccx_stream_mode_enum::CCX_SM_ASF => asf_get_more_data(ctx_norm, ppdata), + + ccx_stream_mode_enum::CCX_SM_WTV => wtv_get_more_data(ctx_norm, ppdata), + + _ => { + eprint!("In general_loop: Impossible value for stream_mode"); + process::exit(1000); + } + } +} + +/// # Safety +/// dereferences a raw pointer +/// calls potentially unsafe C functions +pub unsafe fn process_hardsubx_linear_frames_and_normal_subs( + ctx: *mut lib_ccx_ctx, + hard_ctx: &mut HardsubxContext, + mut enc_ctx: *mut encoder_ctx, +) { + let mut min_pts: u64 = std::u64::MAX; + let mut prev_sub_encoded_hard = true; + let mut got_frame: i32 = 0; + let mut dist = 0; + let mut cur_sec; + let mut total_sec; + let mut progress; + let mut datalist: *mut demuxer_data = null::() as *mut demuxer_data; + let mut datanode: *mut demuxer_data = null::() as *mut demuxer_data; + let mut caps = 0; + let mut dec_ctx: *mut lib_cc_decode = null::() as *mut lib_cc_decode; + let mut frame_number = 0; + + let mut prev_begin_time_hard: i64 = 0; + let mut prev_end_time_hard: i64 = 0; + let mut prev_packet_pts_hard: i64 = 0; + let mut ret: i32 = 0; + + let mut subtitle_text_hard: String; + let mut prev_subtitle_text_hard: String = String::new(); + + let stream_mode: ccx_stream_mode_enum = match (*(*ctx).demux_ctx).get_stream_mode { + Some(T) => match FromPrimitive::from_i32(T((*ctx).demux_ctx)) { + Some(L) => L, + None => panic!(), + }, + None => ccx_stream_mode_enum::CCX_SM_ELEMENTARY_OR_NOT_FOUND, + }; + + if stream_mode == ccx_stream_mode_enum::CCX_SM_TRANSPORT && ((*ctx).write_format as u32 == 5) { + (*ctx).multiprogram = 1; + } + + let mut end_of_file = false; + let mut status = 0; + let mut last_cc_encoded = false; + + loop { + if status < 0 && end_of_file { + break; + } + + if terminate_asap == 0 + && !end_of_file + && is_decoder_processed_enough(ctx) == 0 + && ret != -102 + { + position_sanity_check((*ctx).demux_ctx); + ret = _get_more_data(ctx, &mut datalist, &stream_mode); + + end_of_file = if ret == -101 { true } else { end_of_file }; + + if datalist != null::() as *mut demuxer_data { + position_sanity_check((*ctx).demux_ctx); + + if (*ctx).multiprogram == 0 + && dec_ctx == null::() as *mut lib_cc_decode + || hard_ctx.dec_sub.start_time >= (*dec_ctx).dec_sub.start_time + || status < 0 + { + process_non_multiprogram_general_loop( + ctx, + &mut datalist, + &mut datanode, + &mut dec_ctx, + &mut enc_ctx, + &mut min_pts, + ret, + &mut caps, + ); + } + + if (*ctx).live_stream != 0 { + let cur_sec = + (get_fts((*dec_ctx).timing, (*dec_ctx).current_field) / 1000) as i32; + let th = cur_sec / 10; + if (*ctx).last_reported_progress != th { + activity_progress(-1, cur_sec / 60, cur_sec % 60); + (*ctx).last_reported_progress = th; + } + } else if (*ctx).total_inputsize > 255 { + let progress = ((((*ctx).total_past + (*(*ctx).demux_ctx).past) >> 8) * 100) + / ((*ctx).total_inputsize >> 8); + + if (*ctx).last_reported_progress as i64 != progress { + let mut t: i64 = get_fts((*dec_ctx).timing, (*dec_ctx).current_field); + if t == 0 && (*(*ctx).demux_ctx).global_timestamp_inited != 0 { + t = (*(*ctx).demux_ctx).global_timestamp + - (*(*ctx).demux_ctx).min_global_timestamp; + } + let cur_sec = (t / 1000) as i32; + activity_progress(progress.try_into().unwrap(), cur_sec / 60, cur_sec % 60); + (*ctx).last_reported_progress = progress as i32; + } + } + segment_output_file(ctx, dec_ctx); + + if ccx_options.send_to_srv != 0 { + net_check_conn(); + } + } + } + + if end_of_file && !last_cc_encoded { + enc_ctx = update_encoder_list(ctx); + flush_cc_decode(dec_ctx, &mut ((*dec_ctx).dec_sub)); + encode_sub(enc_ctx, &mut (*dec_ctx).dec_sub); + last_cc_encoded = true; + } + + if hard_ctx.dec_sub.start_time <= ((*dec_ctx).dec_sub).start_time || end_of_file { + status = av_read_frame(hard_ctx.format_ctx, &mut hard_ctx.packet); + + if status >= 0 && hard_ctx.packet.stream_index == hard_ctx.video_stream_id { + frame_number += 1; + + let mut status = + avcodec_send_packet(hard_ctx.codec_ctx, &mut hard_ctx.packet as *mut AVPacket); + // status = avcodec_receive_frame(ctx.codec_ctx, ctx.frame); + + if status >= 0 || status == AVERROR(EAGAIN) || status == AVERROR_EOF { + if status >= 0 { + hard_ctx.packet.size = 0; + } + + status = avcodec_receive_frame(hard_ctx.codec_ctx, hard_ctx.frame); + + if status == 0 { + got_frame = 1; + } + } + + if got_frame != 0 && frame_number % 25 == 0 { + let diff = convert_pts_to_ms( + hard_ctx.packet.pts - prev_packet_pts_hard, + (**(*hard_ctx.format_ctx) + .streams + .offset(hard_ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + + if diff.abs() as f32 >= 1000.0 * hard_ctx.min_sub_duration { + // sws_scale is used to convert the pixel format to RGB24 from all other cases + + sws_scale( + hard_ctx.sws_ctx, + (*hard_ctx.frame).data.as_ptr() as *const *const u8, + (*hard_ctx.frame).linesize.as_mut_ptr(), + 0, + (*hard_ctx.codec_ctx).height, + (*hard_ctx.rgb_frame).data.as_ptr() as *const *mut u8, + (*hard_ctx.rgb_frame).linesize.as_mut_ptr(), + ); + + subtitle_text_hard = match hard_ctx.subcolor { + hardsubx_color_type::HARDSUBX_COLOR_WHITE => { + _process_frame_white_basic( + hard_ctx, + hard_ctx.rgb_frame, + (*hard_ctx.codec_ctx).width, + (*hard_ctx.codec_ctx).height, + frame_number, + ) + } + _ => _process_frame_color_basic( + hard_ctx, + hard_ctx.rgb_frame, + (*hard_ctx.codec_ctx).width, + (*hard_ctx.codec_ctx).height, + frame_number, + ), + }; + + cur_sec = convert_pts_to_s( + hard_ctx.packet.pts, + (**(*hard_ctx.format_ctx) + .streams + .offset(hard_ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + + total_sec = + convert_pts_to_s((*hard_ctx.format_ctx).duration, AV_TIME_BASE_Q); + + progress = (cur_sec * 100) / total_sec; + activity_progress( + progress.try_into().unwrap(), + (cur_sec / 60).try_into().unwrap(), + (cur_sec % 60).try_into().unwrap(), + ); + + // progress on burnt-in extraction + if subtitle_text_hard.is_empty() && prev_subtitle_text_hard.is_empty() { + prev_end_time_hard = convert_pts_to_ms( + hard_ctx.packet.pts, + (**(*hard_ctx.format_ctx) + .streams + .offset(hard_ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + } + + if !subtitle_text_hard.is_empty() { + let double_enter = subtitle_text_hard.find("\n\n"); + + if let Some(T) = double_enter { + subtitle_text_hard = subtitle_text_hard[0..T].to_string(); + } + } + + if !prev_sub_encoded_hard && !prev_subtitle_text_hard.is_empty() { + if !subtitle_text_hard.is_empty() { + dist = edit_distance_string( + &subtitle_text_hard, + &prev_subtitle_text_hard, + ); + if (dist as f32) + < 0.2 + * (cmp::min( + subtitle_text_hard.chars().count(), + prev_subtitle_text_hard.chars().count(), + ) as f32) + { + dist = -1; + subtitle_text_hard = String::new(); + prev_end_time_hard = convert_pts_to_ms( + hard_ctx.packet.pts, + (**(*hard_ctx.format_ctx) + .streams + .offset(hard_ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + } + } + + if dist != -1 { + let sub_text_chr = string_to_c_char(&subtitle_text_hard); + let prev_text_chr = string_to_c_char(&prev_subtitle_text_hard); + let empty_chr = string_to_c_char(""); + let mode_chr = string_to_c_char("BURN"); + add_cc_sub_text( + &mut *hard_ctx.dec_sub as *mut cc_subtitle, + prev_text_chr, + prev_begin_time_hard, + prev_end_time_hard, + empty_chr, + mode_chr, + CCX_ENC_UTF_8 as u32, + ); + + encode_sub(enc_ctx, &mut *hard_ctx.dec_sub); + + // Deallocation + subtitle_text_hard = ffi::CString::from_raw(sub_text_chr) + .to_string_lossy() + .into_owned(); + ffi::CString::from_raw(prev_text_chr) + .to_string_lossy() + .into_owned(); + ffi::CString::from_raw(empty_chr); + ffi::CString::from_raw(mode_chr); + + prev_begin_time_hard = prev_end_time_hard + 1; + prev_subtitle_text_hard = String::new(); + prev_sub_encoded_hard = true; + prev_end_time_hard = convert_pts_to_ms( + hard_ctx.packet.pts, + (**(*hard_ctx.format_ctx) + .streams + .offset(hard_ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + + if !subtitle_text_hard.is_empty() { + prev_subtitle_text_hard = subtitle_text_hard.clone(); + prev_sub_encoded_hard = false; + } + } + dist = 0; + } + + if prev_subtitle_text_hard.is_empty() && !subtitle_text_hard.is_empty() { + prev_begin_time_hard = prev_end_time_hard + 1; + prev_end_time_hard = convert_pts_to_ms( + hard_ctx.packet.pts, + (**(*hard_ctx.format_ctx) + .streams + .offset(hard_ctx.video_stream_id.try_into().unwrap())) + .time_base, + ); + prev_subtitle_text_hard = subtitle_text_hard.clone(); + prev_sub_encoded_hard = false; + } + prev_packet_pts_hard = hard_ctx.packet.pts; + } + } + } + av_packet_unref(&mut hard_ctx.packet as *mut AVPacket); + } + } +} diff --git a/src/rust/src/hardsubx/main.rs b/src/rust/src/hardsubx/main.rs new file mode 100644 index 000000000..5460fc4f6 --- /dev/null +++ b/src/rust/src/hardsubx/main.rs @@ -0,0 +1,215 @@ +use super::decoder::hardsubx_process_frames_linear; +use super::HardsubxContext; +#[cfg(feature = "hardsubx_ocr")] +use crate::bindings::{ + ccx_s_options, dinit_encoder, encoder_cfg, encoder_ctx, init_encoder, lib_ccx_ctx, +}; +use crate::hardsubx::decoder::hardsubx_process_frames_tickertext; +use crate::hardsubx::decoder::process_hardsubx_linear_frames_and_normal_subs; + +#[cfg(feature = "hardsubx_ocr")] +use crate::hardsubx::{EXIT_NOT_ENOUGH_MEMORY, EXIT_READ_ERROR}; +use crate::utils::string_to_c_char; +#[cfg(feature = "hardsubx_ocr")] +use ffmpeg_sys_next::AVMediaType::AVMEDIA_TYPE_VIDEO; +#[cfg(feature = "hardsubx_ocr")] +use ffmpeg_sys_next::AVPixelFormat::*; +#[cfg(feature = "hardsubx_ocr")] +use ffmpeg_sys_next::*; +use std::convert::TryInto; +use std::ffi; +use std::os::raw::c_void; +use std::process; +use std::ptr::null; +use std::time::{Duration, Instant}; + +#[cfg(feature = "hardsubx_ocr")] +use tesseract_sys::*; +extern "C" { + pub static mut ccx_options: ccx_s_options; +} + +/// # Safety +/// Dereferences a raw pointer (datamember of the object ctx) +/// calls potentially unsafe C functions +pub unsafe fn hardsubx_process_data(ctx: &mut HardsubxContext, ctx_normal: *mut lib_ccx_ctx) { + let mut format_ctx_tmp: *mut AVFormatContext = + null::() as *mut AVFormatContext; + let file_name = string_to_c_char(&ctx.inputfile[0]); + if avformat_open_input( + &mut format_ctx_tmp as *mut *mut AVFormatContext, + file_name, + null::() as *const AVInputFormat, + null::() as *mut *mut AVDictionary, + ) != 0 + { + eprintln!("Error reading input file!\n"); + process::exit(EXIT_READ_ERROR); + } + + ctx.format_ctx = format_ctx_tmp; + + let format_ctx_ref = &mut *format_ctx_tmp; + if avformat_find_stream_info( + format_ctx_ref as *mut AVFormatContext, + null::<*mut AVDictionary>() as *mut *mut AVDictionary, + ) < 0 + { + eprintln!("Error reading input stream!\n"); + + process::exit(EXIT_READ_ERROR); + } + + av_dump_format(format_ctx_tmp as *mut AVFormatContext, 0, file_name, 0); + + ffi::CString::from_raw(file_name); //deallocation + + ctx.video_stream_id = -1; + + for i in 0..format_ctx_ref.nb_streams { + if (*(**format_ctx_ref.streams.offset(i.try_into().unwrap())).codecpar).codec_type + == AVMEDIA_TYPE_VIDEO + { + ctx.video_stream_id = i as i32; + break; + } + } + + if ctx.video_stream_id == -1 { + eprintln!("Video Stream not found!\n"); + + process::exit(EXIT_READ_ERROR); + } + + let mut codec_ctx: *mut AVCodecContext = avcodec_alloc_context3(null::()); + + let codec_par_ptr: *mut AVCodecParameters = + (*(*format_ctx_ref.streams).offset(ctx.video_stream_id.try_into().unwrap())).codecpar; + let codec_par_ref: &mut AVCodecParameters = &mut (*codec_par_ptr); + + let status = + avcodec_parameters_to_context(codec_ctx, codec_par_ref as *const AVCodecParameters); + + if status < 0 { + eprintln!("Creation of AVCodecContext failed."); + process::exit(status); + } else { + ctx.codec_ctx = codec_ctx; + } + + let codec_ptr = avcodec_find_decoder((*codec_ctx).codec_id) as *mut AVCodec; + + if codec_ptr == null::() as *mut AVCodec { + eprintln!("Input codec is not supported!\n"); + + process::exit(EXIT_READ_ERROR); + } else { + ctx.codec = codec_ptr; + } + + if avcodec_open2(ctx.codec_ctx, ctx.codec, &mut ctx.options_dict) < 0 { + eprintln!("Error opening input codec!\n"); + process::exit(EXIT_READ_ERROR); + } + + let mut frame_ptr = av_frame_alloc(); + let mut rgb_frame_ptr = av_frame_alloc(); + + if frame_ptr == null::() as *mut AVFrame + || rgb_frame_ptr == null::() as *mut AVFrame + { + eprintln!("Not enough memory to initialize frame!"); + + process::exit(EXIT_NOT_ENOUGH_MEMORY); + } + + ctx.frame = frame_ptr; + ctx.rgb_frame = rgb_frame_ptr; + + let frame_bytes: i32 = av_image_get_buffer_size( + AVPixelFormat::AV_PIX_FMT_RGB24, + (*codec_ctx).width, + (*codec_ctx).height, + 16, + ); + + let rgb_buffer_ptr = av_malloc((frame_bytes * 8).try_into().unwrap()) as *mut u8; + ctx.rgb_buffer = rgb_buffer_ptr; + + ctx.sws_ctx = sws_getContext( + (*codec_ctx).width, + (*codec_ctx).height, + (*codec_ctx).pix_fmt, + (*codec_ctx).width, + (*codec_ctx).height, + AVPixelFormat::AV_PIX_FMT_RGB24, + SWS_BILINEAR, + null::() as *mut SwsFilter, + null::() as *mut SwsFilter, + null::(), + ); + + av_image_fill_arrays( + (*rgb_frame_ptr).data.as_mut_ptr(), + (*rgb_frame_ptr).linesize.as_mut_ptr(), + rgb_buffer_ptr, + AV_PIX_FMT_RGB24, + (*codec_ctx).width, + (*codec_ctx).height, + 1, + ); + + let mut enc_ctx_ptr = init_encoder((&mut ccx_options.enc_cfg) as *mut encoder_cfg); + + println!("Beginning burned-in subtitle detection...\n"); + + if ctx.tickertext { + hardsubx_process_frames_tickertext(ctx, enc_ctx_ptr); + } else if ctx.hardsubx_and_common { + process_hardsubx_linear_frames_and_normal_subs(ctx_normal, ctx, enc_ctx_ptr) + } else { + hardsubx_process_frames_linear(ctx, enc_ctx_ptr); + } + + dinit_encoder(&mut enc_ctx_ptr as *mut *mut encoder_ctx, 0); + + av_free(rgb_buffer_ptr as *mut u8 as *mut c_void); + + if ctx.frame != null::() as *mut AVFrame { + av_frame_free(&mut frame_ptr as *mut *mut AVFrame); + } + + if ctx.rgb_frame != null::() as *mut AVFrame { + av_frame_free(&mut rgb_frame_ptr as *mut *mut AVFrame); + } + + avcodec_close(codec_ctx); + avformat_close_input( + &mut (format_ctx_ref as *mut AVFormatContext) as *mut *mut AVFormatContext, + ); + avcodec_free_context(&mut codec_ctx as *mut *mut AVCodecContext); +} + +/// # Safety +/// calls potentially unsafe C functions +pub unsafe fn _dinit_hardsubx(ctx: &mut HardsubxContext) { + // Free OCR related stuff + TessBaseAPIEnd(ctx.tess_handle); + TessBaseAPIDelete(ctx.tess_handle); +} + +/// # Safety +/// Dereferences a raw pointer (datamember of the object ctx) +#[no_mangle] +pub unsafe extern "C" fn hardsubx(options: *mut ccx_s_options, ctx_normal: *mut lib_ccx_ctx) { + println!("HardsubX (Hard Subtitle Extractor) - Burned-in subtitle extraction subsystem\n"); + + let mut ctx = HardsubxContext::new(options); + + let start = Instant::now(); + hardsubx_process_data(&mut ctx, ctx_normal); + let duration: Duration = start.elapsed(); + + println!("Processing time = {:?}", duration); + _dinit_hardsubx(&mut ctx); +} diff --git a/src/rust/src/hardsubx/mod.rs b/src/rust/src/hardsubx/mod.rs index f94e14f3c..1f81bd207 100644 --- a/src/rust/src/hardsubx/mod.rs +++ b/src/rust/src/hardsubx/mod.rs @@ -1,6 +1,7 @@ pub mod classifier; pub mod decoder; pub mod imgops; +pub mod main; pub mod utility; #[cfg(feature = "hardsubx_ocr")] @@ -13,7 +14,68 @@ use tesseract_sys::*; use leptonica_sys::*; use crate::bindings; -use crate::bindings::{cc_subtitle, ccx_output_format}; +use crate::bindings::{ + cc_subtitle, ccx_output_format, ccx_s_options, get_file_extension, probe_tessdata_location, +}; +use crate::utils::string_to_c_char; +use std::boxed::Box; +use std::convert::TryInto; +use std::ffi; +use std::matches; +use std::os::raw::c_char; +use std::os::raw::c_void; +use std::process; +use std::ptr::null; +use std::vec::Vec; + +// definitions taken from ccx_common_common.h +static EXIT_NOT_ENOUGH_MEMORY: i32 = 500; +static EXIT_READ_ERROR: i32 = 8; + +static EXIT_MALFORMED_PARAMETER: i32 = 7; + +extern "C" { + pub static mut ccx_options: ccx_s_options; +} + +pub enum hardsubx_color_type { + HARDSUBX_COLOR_WHITE, + HARDSUBX_COLOR_YELLOW, + HARDSUBX_COLOR_GREEN, + HARDSUBX_COLOR_CYAN, + HARDSUBX_COLOR_BLUE, + HARDSUBX_COLOR_MAGENTA, + HARDSUBX_COLOR_RED, + HARDSUBX_COLOR_CUSTOM, +} + +pub enum hardsubx_ocr_mode { + HARDSUBX_OCRMODE_FRAME, + HARDSUBX_OCRMODE_WORD, + HARDSUBX_OCRMODE_LETTER, +} + +impl Default for cc_subtitle { + fn default() -> Self { + Self { + data: null::() as *mut std::os::raw::c_void, + datatype: 0, + nb_data: 0, + type_: 0, + enc_type: 0, + start_time: 0, + end_time: 0, + flags: 0, + lang_index: 0, + got_output: 0, + mode: [0, 0, 0, 0, 0], + info: [0, 0, 0, 0], + time_out: 0, + next: null::() as *mut cc_subtitle, + prev: null::() as *mut cc_subtitle, + } + } +} #[repr(C)] #[derive(Debug, Copy, Clone)] @@ -53,3 +115,280 @@ pub struct lib_hardsubx_ctx { pub hue: f32, pub lum_thresh: f32, } + +pub struct HardsubxContext { + pub cc_to_stdout: i32, + pub subs_delay: i64, + pub last_displayed_subs_ms: i64, + pub basefilename: String, + pub extension: String, + pub current_file: i32, + pub inputfile: Vec, + pub num_input_files: i32, + pub system_start_time: i64, + pub write_format: ccx_output_format, + pub format_ctx: *mut AVFormatContext, + + pub codec_par: *mut AVCodecParameters, + // Codec context will be generated from the code parameters + pub codec_ctx: *mut AVCodecContext, + + pub codec: *mut AVCodec, + pub frame: *mut AVFrame, + pub rgb_frame: *mut AVFrame, + + pub packet: AVPacket, + + pub options_dict: *mut AVDictionary, + pub sws_ctx: *mut SwsContext, + pub rgb_buffer: *mut u8, + + pub video_stream_id: i32, + pub im: *mut Pix, + pub tess_handle: *mut TessBaseAPI, + + pub cur_conf: f32, + pub prev_conf: f32, + + pub tickertext: bool, + pub hardsubx_and_common: bool, + + pub dec_sub: Box, + pub ocr_mode: hardsubx_ocr_mode, + pub subcolor: hardsubx_color_type, + + pub min_sub_duration: f32, + pub detect_italics: bool, + + pub conf_thresh: f32, + pub hue: f32, + pub lum_thresh: f32, +} + +impl Default for HardsubxContext { + fn default() -> Self { + Self { + cc_to_stdout: 0, + subs_delay: 0, + last_displayed_subs_ms: 0, + basefilename: String::new(), + extension: String::new(), + current_file: 0, + inputfile: Vec::::new(), + num_input_files: 0, + system_start_time: -1, + write_format: ccx_output_format::CCX_OF_RAW, + format_ctx: null::() as *mut AVFormatContext, + codec_par: null::() as *mut AVCodecParameters, + codec_ctx: null::() as *mut AVCodecContext, + codec: null::() as *mut AVCodec, + frame: null::() as *mut AVFrame, + rgb_frame: null::() as *mut AVFrame, + packet: AVPacket { + buf: null::() as *mut AVBufferRef, + pts: 0, + dts: 0, + data: null::() as *mut u8, + size: 0, + stream_index: 0, + flags: 0, + side_data: null::() as *mut AVPacketSideData, + side_data_elems: 0, + duration: 0, + pos: 0, + opaque: null::() as *mut c_void, + opaque_ref: null::() as *mut AVBufferRef, + time_base: AV_TIME_BASE_Q, + // convergence_duration: 0 + }, + options_dict: null::() as *mut AVDictionary, + sws_ctx: null::() as *mut SwsContext, + rgb_buffer: null::() as *mut u8, + video_stream_id: 0, + im: null::() as *mut Pix, + tess_handle: null::() as *mut TessBaseAPI, + cur_conf: 0.0, + prev_conf: 0.0, + + tickertext: false, + hardsubx_and_common: false, + + dec_sub: { + let tmp = cc_subtitle::default(); + Box::new(tmp) + }, + ocr_mode: hardsubx_ocr_mode::HARDSUBX_OCRMODE_FRAME, + subcolor: hardsubx_color_type::HARDSUBX_COLOR_WHITE, + + min_sub_duration: 0.0, + detect_italics: false, + + conf_thresh: 0.0, + hue: 0.0, + lum_thresh: 0.0, + } + } +} + +impl HardsubxContext { + /// # Safety + /// dereferences a raw pointer + /// Calls C functions and sends them raw pointers, their safety is not guaranteed + pub unsafe fn new(options: *mut ccx_s_options) -> Self { + let tess_handle = &mut (*TessBaseAPICreate()) as &mut TessBaseAPI; + + let lang: String = { + if (*options).ocrlang != null::() as *mut c_char { + ffi::CStr::from_ptr((*options).ocrlang) + .to_string_lossy() + .into_owned() + } else { + "eng".to_string() + } + }; + + let lang_cstr = string_to_c_char(&lang); + + let tessdata_path = { + let path = probe_tessdata_location(lang_cstr); + + if path == (null::() as *mut c_char) && lang != *"eng" { + let eng_cstr = string_to_c_char("eng"); + let tmp = probe_tessdata_location(eng_cstr); + ffi::CString::from_raw(eng_cstr); // deallocation + if tmp != null::() as *mut c_char { + println!("{}.traineddata not found! Switching to English\n", lang); + tmp + } else { + eprintln!("eng.traineddata not found! No Switching Possible\n"); + process::exit(0); + } + } else if path == null::() as *mut c_char { + eprintln!("eng.traineddata not found! No Switching Possible\n"); + process::exit(0); + } else { + path + } + }; + + let mut tess_data_str: String = ffi::CStr::from_ptr(tessdata_path) + .to_string_lossy() + .into_owned(); + + let mut pars_vec: *mut c_char = string_to_c_char("debug_file"); + let mut pars_values: *mut c_char = string_to_c_char("/dev/null"); + + let ret = { + let tess_version = ffi::CStr::from_ptr(TessVersion()) + .to_string_lossy() + .into_owned(); + + if tess_version[..2].eq(&("4.".to_string())) { + tess_data_str = format!("{}/tessdata", tess_data_str); + if ccx_options.ocr_oem < 0 { + ccx_options.ocr_oem = 1; + } + } else if ccx_options.ocr_oem < 0 { + ccx_options.ocr_oem = 0; + } + + let tessdata_path_cstr = string_to_c_char(&tess_data_str); + let ret = TessBaseAPIInit4( + tess_handle, + tessdata_path_cstr, + lang_cstr, + ccx_options.ocr_oem.try_into().unwrap(), + null::<*mut u8>() as *mut *mut i8, + 0, + &mut pars_vec, + &mut pars_values, + 1, + 0, + ); + + ffi::CString::from_raw(tessdata_path_cstr); // deallocation + + ret + }; + + // deallocation + ffi::CString::from_raw(pars_vec); + ffi::CString::from_raw(pars_values); + + if ret != 0 { + eprintln!("Not enough memory to initialize"); + process::exit(EXIT_NOT_ENOUGH_MEMORY); + } + + ffi::CString::from_raw(lang_cstr); //deallocate + // function to be used for only converting the C struct to rust + Self { + tess_handle, + basefilename: { + if (*options).output_filename == null::() as *mut c_char { + "".to_string() + } else { + ffi::CStr::from_ptr((*options).output_filename) + .to_string_lossy() + .into_owned() + } + }, + current_file: -1, + inputfile: { + let mut vec_inputfile: Vec = Vec::new(); + for i in 0..(*options).num_input_files { + vec_inputfile.push( + ffi::CStr::from_ptr(*(*options).inputfile.offset(i.try_into().unwrap())) + .to_string_lossy() + .into_owned(), + ) + } + vec_inputfile + }, + num_input_files: (*options).num_input_files, + extension: ffi::CStr::from_ptr(get_file_extension((*options).write_format)) + .to_string_lossy() + .into_owned(), + write_format: (*options).write_format, + subs_delay: (*options).subs_delay, + cc_to_stdout: (*options).cc_to_stdout, + tickertext: !matches!((*options).tickertext, 0), + cur_conf: 0.0, + prev_conf: 0.0, + ocr_mode: { + if (*options).hardsubx_ocr_mode == 0 { + hardsubx_ocr_mode::HARDSUBX_OCRMODE_FRAME + } else if (*options).hardsubx_ocr_mode == 1 { + hardsubx_ocr_mode::HARDSUBX_OCRMODE_WORD + } else if (*options).hardsubx_ocr_mode == 2 { + hardsubx_ocr_mode::HARDSUBX_OCRMODE_LETTER + } else { + eprintln!("Invalid OCR Mode"); + process::exit(EXIT_MALFORMED_PARAMETER); + } + }, + + subcolor: match (*options).hardsubx_subcolor { + 0 => hardsubx_color_type::HARDSUBX_COLOR_WHITE, + 1 => hardsubx_color_type::HARDSUBX_COLOR_YELLOW, + 2 => hardsubx_color_type::HARDSUBX_COLOR_GREEN, + 3 => hardsubx_color_type::HARDSUBX_COLOR_CYAN, + 4 => hardsubx_color_type::HARDSUBX_COLOR_BLUE, + 5 => hardsubx_color_type::HARDSUBX_COLOR_MAGENTA, + 6 => hardsubx_color_type::HARDSUBX_COLOR_RED, + 7 => hardsubx_color_type::HARDSUBX_COLOR_CUSTOM, + _ => hardsubx_color_type::HARDSUBX_COLOR_WHITE, // white is default + }, + + min_sub_duration: (*options).hardsubx_min_sub_duration, + detect_italics: !matches!((*options).hardsubx_detect_italics, 0), + + conf_thresh: (*options).hardsubx_conf_thresh, + hue: (*options).hardsubx_hue, + + lum_thresh: (*options).hardsubx_lum_thresh, + hardsubx_and_common: !matches!((*options).hardsubx_and_common, 0), + ..Default::default() + } + } +} diff --git a/src/rust/src/hardsubx/utility.rs b/src/rust/src/hardsubx/utility.rs index d2b8e1509..45aca2002 100644 --- a/src/rust/src/hardsubx/utility.rs +++ b/src/rust/src/hardsubx/utility.rs @@ -86,3 +86,16 @@ pub unsafe extern "C" fn edit_distance( &mut dp_array, ) as c_int } + +pub fn edit_distance_string(word1: &str, word2: &str) -> i32 { + let len1 = word1.chars().count(); + let len2 = word2.chars().count(); + + // kinda assuming ASCII + let word1_chararray = word1.as_bytes(); + let word2_chararray = word2.as_bytes(); + + let mut dp_array = vec![vec![-1; len2 + 1]; len1 + 1]; + + _edit_distance_rec(word1_chararray, word2_chararray, len1, len2, &mut dp_array) +} diff --git a/src/rust/src/lib.rs b/src/rust/src/lib.rs index a76faa419..77105ce46 100644 --- a/src/rust/src/lib.rs +++ b/src/rust/src/lib.rs @@ -28,6 +28,13 @@ use utils::is_true; use env_logger::{builder, Target}; use log::{warn, LevelFilter}; +pub enum ccx_encoding_type { + CCX_ENC_UNICODE = 0, + CCX_ENC_LATIN_1 = 1, + CCX_ENC_UTF_8 = 2, + CCX_ENC_ASCII = 3, +} + extern "C" { static mut cb_708: c_int; static mut cb_field1: c_int; diff --git a/src/rust/wrapper.h b/src/rust/wrapper.h index b429ed1e3..90e552ede 100644 --- a/src/rust/wrapper.h +++ b/src/rust/wrapper.h @@ -7,3 +7,4 @@ #include "../lib_ccx/lib_ccx.h" #include "../lib_ccx/hardsubx.h" #include "../lib_ccx/utility.h" +#include "../lib_ccx/ocr.h"