Use language tags in MKV as a hint for OCR

Fixes #28

Signed-off-by: Ethan Dye <[email protected]>

Sources/macSubtitleOCR/MKV/EBML/EBML.swift

@@ -15,6 +15,8 @@ enum EBML {
     static let cluster: UInt32 = 0x1F43_B675
     static let codecID: UInt32 = 0x86
     static let codecPrivate: UInt32 = 0x63A2
+    static let language: UInt32 = 0x22B59C
+    static let languageBCP47: UInt32 = 0x22B59D
     static let segmentID: UInt32 = 0x1853_8067
     static let simpleBlock: UInt32 = 0xA3
     static let timestamp: UInt32 = 0xE7

Sources/macSubtitleOCR/MKV/MKVTrack.swift

@@ -13,4 +13,5 @@ struct MKVTrack {
     var codecID: String
     var trackData: Data
     var idxData: String?
+    var language: String?
 }

Sources/macSubtitleOCR/MKV/MKVTrackParser.swift

@@ -13,7 +13,8 @@ class MKVTrackParser: MKVFileHandler {
     // MARK: - Properties
 
     private(set) var tracks: [MKVTrack] = []
-    private(set) var codecPrivate = [Int: String]()
+    private var codecPrivate = [Int: String]()
+    private var languages = [Int: String]()
 
     // MARK: - Functions
 
@@ -50,7 +51,8 @@ class MKVTrackParser: MKVFileHandler {
                 trackNumber: index,
                 codecID: subtitleTracks[index + 1]!,
                 trackData: data,
-                idxData: codecPrivate[index + 1]))
+                idxData: codecPrivate[index + 1],
+                language: languages[index + 1]))
         }
     }
 
@@ -89,6 +91,7 @@ class MKVTrackParser: MKVFileHandler {
         var trackNumber: Int?
         var trackType: UInt8?
         var codecID: String?
+        var language: String?
 
         while let (elementID, elementSize) = tryParseElement() {
             switch elementID {
@@ -103,6 +106,12 @@ class MKVTrackParser: MKVFileHandler {
                 data.removeNullBytes()
                 codecID = String(data: data, encoding: .ascii)
                 logger.debug("Found codec ID: \(codecID!)")
+            case EBML.language, EBML.languageBCP47:
+                var data = fileHandle.readData(ofLength: Int(elementSize))
+                data.removeNullBytes()
+                language = String(data: data, encoding: .ascii)
+                logger.debug("Found language: \(language!)")
+                languages[trackNumber!] = language
             default:
                 fileHandle.seek(toFileOffset: fileHandle.offsetInFile + elementSize)
             }

Sources/macSubtitleOCR/macSubtitleOCR.swift

@@ -68,23 +68,23 @@ struct macSubtitleOCR: AsyncParsableCommand {
 
     // MARK: - Entrypoint
 
-    func run() async throws {
+    mutating func run() async throws {
         let fileHandler = macSubtitleOCRFileHandler(outputDirectory: outputDirectory)
         let results = try await processInput()
         try await saveResults(fileHandler: fileHandler, results: results)
     }
 
     // MARK: - Methods
 
-    private func processInput() async throws -> [macSubtitleOCRResult] {
+    private mutating func processInput() async throws -> [macSubtitleOCRResult] {
         if ffmpegDecoder {
             try await processFFmpegDecoder()
         } else {
             try await processInternalDecoder()
         }
     }
 
-    private func processInternalDecoder() async throws -> [macSubtitleOCRResult] {
+    private mutating func processInternalDecoder() async throws -> [macSubtitleOCRResult] {
         var results: [macSubtitleOCRResult] = []
 
         if input.hasSuffix(".sub") || input.hasSuffix(".idx") {
@@ -104,6 +104,10 @@ struct macSubtitleOCR: AsyncParsableCommand {
                         outputDirectory: URL(fileURLWithPath: outputDirectory))
                 }
 
+                if track.language != nil {
+                    languages += ",\(track.language!)"
+                }
+
                 if track.codecID == "S_HDMV/PGS" {
                     let pgs: PGS = try track.trackData
                         .withUnsafeBytes { (buffer: UnsafeRawBufferPointer) in

Tests/macSubtitleOCRTests.swift

@@ -36,7 +36,7 @@ private func runTest(with options: [String]) async throws {
     let outputPath = options[1]
 
     // Run tests
-    let runner = try macSubtitleOCR.parse(options)
+    var runner = try macSubtitleOCR.parse(options)
     try await runner.run()
 
     try compareOutputs(with: outputPath, track: 0)