MKV Subtitle Extraction Demo

main.rs

//! This is a proof-of-concept for extracting vobsub subtitles from an MKV file.
//! It makes use of some private functions from the vobsub crate, and requires a
//! modified copy to export them.
//!
//! This is primarily created as a testing ground for integrating subtitle extraction
//! into mediacorral. The current version really only works for vobsub, and converts
//! the vobsub images into sixel images, printing them to the terminal.

use matroska_demuxer::*;
use std::{fs::File, time::Duration};

struct IdxData {
    palette: vobsub::Palette,
}

fn parse_idx(data: &[u8]) -> IdxData {
    for line in String::from_utf8_lossy(data).split("\n") {
        if line.trim_start().starts_with("#") {
            continue;
        }
        let (key, value) = line.split_once(": ").unwrap();
        if key == "palette" {
            return IdxData {
                palette: vobsub::palette(value.as_bytes()).unwrap().1,
            };
        }
    }
    panic!();
}

fn main() {
    let file = File::open("test.mkv").unwrap();
    let mut mkv = MatroskaFile::open(file).unwrap();
    let video_track = mkv
        .tracks()
        .iter()
        .find(|t| t.track_type() == TrackType::Subtitle)
        .inspect(|t| {
            dbg!(t.codec_id());
            dbg!(t.codec_name());
        })
        .unwrap()
        .clone();
    let track_num = video_track.track_number().get();
    let idx = parse_idx(video_track.codec_private().unwrap());

    let mut frame = Frame::default();
    let mut last_frame: Option<_> = None;
    while mkv.next_frame(&mut frame).unwrap() {
        if frame.track == track_num {
            match video_track.codec_id() {
                "S_SUBRIP" => {
                    println!(
                        "video frame found: {}",
                        String::from_utf8(frame.data.clone()).unwrap()
                    );
                }
                "S_VOBSUB" => {
                    if let Some(timestamp) = last_frame {
                        std::thread::sleep(Duration::from_millis(frame.timestamp - timestamp));
                    }
                    last_frame = Some(frame.timestamp);
                    let subtitle = vobsub::subtitle(&frame.data, frame.timestamp as _).unwrap();
                    let result = subtitle.to_image(&idx.palette);
                    let mut pix_buf: Vec<u8> = Vec::new();
                    for pixel in result.pixels() {
                        pix_buf.extend_from_slice(&pixel.data);
                    }
                    let encoder = sixel::encoder::Encoder::new().unwrap();
                    encoder
                        .encode_bytes(
                            sixel::encoder::QuickFrameBuilder::new()
                                .width(result.width() as _)
                                .height(result.height() as _)
                                .format(sixel_sys::PixelFormat::RGBA8888)
                                .pixels(pix_buf),
                        )
                        .unwrap();
                    // println!("{:?}", result);
                }
                _ => {
                    break;
                }
            }
        }
    }
}