Add support for VobSub in MKV files

.github/workflows/build.yml

@@ -32,11 +32,11 @@ jobs:
         run: xcrun swift test -Xswiftc -DGITHUB_ACTIONS list
 
       - name: Test FFmpeg Decoder
-        timeout-minutes: 5
+        timeout-minutes: 7
         run: xcrun swift test --skip-build --filter ffmpegDecoder
 
       - name: Test Internal Decoder
-        timeout-minutes: 5
+        timeout-minutes: 7
         run: xcrun swift test --skip-build --filter internalDecoder
 
       - name: Periphery

.gitignore

@@ -5,5 +5,5 @@
 /.VSCodeCounter/
 /.index-build/
 Tests/Resources/*
-!Tests/Resources/sintel.*
+!Tests/Resources/sintel*.*
 !Tests/Resources/README.md

README.md

@@ -22,7 +22,7 @@ For more details on performance, refer to the [Accuracy](#accuracy) section belo
 #### Supported Formats
 
 - PGS (`.mkv`, `.sup`)
-- VobSub (`.sub`, `.idx`)
+- VobSub (`.mkv`, `.sub`, `.idx`)
 
 ### Building the Project
 

Sources/macSubtitleOCR/Extensions/BinaryIntegerExtensions.swift

@@ -0,0 +1,14 @@
+//
+// BinaryIntegerExtensions.swift
+// macSubtitleOCR
+//
+// Created by Ethan Dye on 10/20/24.
+// Copyright © 2024 Ethan Dye. All rights reserved.
+//
+
+extension BinaryInteger {
+    /// Returns a formatted hexadecimal string with `0x` prefix.
+    func hex() -> String {
+        String(format: "0x%0\(MemoryLayout<Self>.size)X", self as! CVarArg)
+    }
+}

Sources/macSubtitleOCR/FileHandler.swift

@@ -2,7 +2,7 @@
 // FileHandler.swift
 // macSubtitleOCR
 //
-// Created by Ethan Dye on 10/16/24.
+// Created by Ethan Dye on 10/17/24.
 // Copyright © 2024 Ethan Dye. All rights reserved.
 //
 
@@ -16,6 +16,9 @@ struct FileHandler {
     }
 
     func saveSRTFile(for result: macSubtitleOCRResult) throws {
+        if result.srt.isEmpty {
+            return
+        }
         let srtFilePath = URL(fileURLWithPath: outputDirectory).appendingPathComponent("track_\(result.trackNumber).srt")
         let srt = SRT(subtitles: result.srt.sorted { $0.index < $1.index })
         srt.write(toFileAt: srtFilePath)

Sources/macSubtitleOCR/MKV/EBML/EBML.swift

@@ -14,6 +14,7 @@ enum EBML {
     static let chapters: UInt32 = 0x1043_A770
     static let cluster: UInt32 = 0x1F43_B675
     static let codecID: UInt32 = 0x86
+    static let codecPrivate: UInt32 = 0x63A2
     static let segmentID: UInt32 = 0x1853_8067
     static let simpleBlock: UInt32 = 0xA3
     static let timestamp: UInt32 = 0xE7

Sources/macSubtitleOCR/MKV/EBML/EBMLParser.swift

@@ -9,39 +9,37 @@
 import Foundation
 import os
 
-private let logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category: "ebml")
+private let logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category: "EBML")
 
 // Helper function to read variable-length integers (VINT) from MKV (up to 8 bytes)
-func readVINT(from fileHandle: FileHandle, unmodified: Bool = false) -> UInt64 {
+func readVINT(from fileHandle: FileHandle, elementSize: Bool = false) -> UInt64 {
     guard let firstByte = fileHandle.readData(ofLength: 1).first else { return 0 }
 
     var length: UInt8 = 1
     var mask: UInt8 = 0x80
 
     // Find how many bytes are needed for the VINT (variable integer)
-    while (firstByte & mask) == 0 {
+    while mask != 0, firstByte & mask == 0 {
         length += 1
         mask >>= 1
     }
 
-    // Adjust mask based on length and unmodified flag
-    mask = (mask == 0x10) ? 0xFF : (length == 1 && !unmodified) ? firstByte : mask - 1
-
-    var value = UInt64(firstByte & mask)
+    var value = UInt64(firstByte)
+    if elementSize {
+        value &= ~UInt64(mask)
+    }
 
-    if length > 1 {
-        for byte in fileHandle.readData(ofLength: Int(length - 1)) {
-            value = (value << 8) | UInt64(byte)
-        }
+    for byte in fileHandle.readData(ofLength: Int(length - 1)) {
+        value = (value << 8) | UInt64(byte)
     }
-    logger.debug("VINT: 0x\(String(format: "%08X", value))")
+    logger.debug("VINT: \(value.hex())")
 
     return value
 }
 
 // Helper function to read an EBML element's ID and size
-func readEBMLElement(from fileHandle: FileHandle, unmodified: Bool = false) -> (elementID: UInt32, elementSize: UInt64) {
-    let elementID = readVINT(from: fileHandle, unmodified: unmodified)
-    let elementSize = readVINT(from: fileHandle, unmodified: true)
+func readEBMLElement(from fileHandle: FileHandle) -> (elementID: UInt32, elementSize: UInt64) {
+    let elementID = readVINT(from: fileHandle)
+    let elementSize = readVINT(from: fileHandle, elementSize: true)
     return (UInt32(elementID), elementSize)
 }

Sources/macSubtitleOCR/MKV/MKVFileHandler.swift

@@ -81,7 +81,7 @@ class MKVFileHandler {
                 return (elementSize, elementID)
             } else {
                 // Skip over the element's data by seeking to its end
-                logger.debug("Found: \(elementID), but not \(targetID), skipping element")
+                logger.debug("\(elementID.hex()) != \(targetID.hex()), skipping element")
                 fileHandle.seek(toFileOffset: fileHandle.offsetInFile + elementSize)
             }
             previousOffset = fileHandle.offsetInFile

Sources/macSubtitleOCR/MKV/MKVHelpers.swift

@@ -9,28 +9,33 @@
 import Foundation
 
 // Function to read a fixed length number of bytes and convert in into a (Un)signed integer
-func readFixedLengthNumber(fileHandle: FileHandle, length: Int, signed: Bool = false) -> Int64 {
+func readFixedLengthNumber(fileHandle: FileHandle, length: Int) -> Int64 {
     let data = fileHandle.readData(ofLength: length)
     var result: Int64 = 0
-
     for byte in data {
         result = result << 8 | Int64(byte)
     }
-
-    if signed, data.first! & 0x80 != 0 {
-        result -= Int64(1) << (8 * length) // Apply two's complement for signed integers
-    }
-
     return result
 }
 
 // Encode the absolute timestamp as 4 bytes in big-endian format for PGS
-func encodePTSForPGS(_ timestamp: Int64) -> [UInt8] {
+func encodePTSForPGS(_ timestamp: UInt64) -> [UInt8] {
     withUnsafeBytes(of: UInt32(timestamp).bigEndian) { Array($0) }
 }
 
-// Calculate the absolute timestamp with 90 kHz accuracy for PGS format
-func calcAbsPTSForPGS(_ clusterTimestamp: Int64, _ blockTimestamp: Int64, _ timestampScale: Double) -> Int64 {
+func encodePTSForVobSub(_ timestamp: UInt64) -> [UInt8] {
+    var buffer = [UInt8](repeating: 0, count: 5) // 5-byte buffer
+
+    buffer[0] = (buffer[0] & 0xF1) | UInt8((timestamp >> 29) & 0x0E)
+    buffer[1] = UInt8((timestamp >> 22) & 0xFF)
+    buffer[2] = UInt8(((timestamp >> 14) & 0xFE) | 1)
+    buffer[3] = UInt8((timestamp >> 7) & 0xFF)
+    buffer[4] = UInt8((timestamp << 1) & 0xFF)
+    return buffer
+}
+
+// Calculate the absolute timestamp with 90 kHz accuracy
+func calcAbsPTS(_ clusterTimestamp: Int64, _ blockTimestamp: Int64) -> UInt64 {
     // The block timestamp is relative, so we add it to the cluster timestamp
-    Int64(((Double(clusterTimestamp) + Double(blockTimestamp)) / timestampScale) * 90000000)
+    UInt64((Double(clusterTimestamp) + Double(blockTimestamp)) * 90)
 }

Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift

@@ -11,12 +11,24 @@ import os
 
 class MKVSubtitleExtractor: MKVTrackParser {
     func saveSubtitleTrackData(trackNumber: Int, outputDirectory: URL) {
-        let trackPath = outputDirectory.appendingPathComponent("\(trackNumber)").appendingPathExtension("sup").path
+        let codecType = tracks[trackNumber].codecId
+        let fileExtension = (codecType == "S_HDMV/PGS") ? "sup" : "sub"
+        let trackPath = outputDirectory.appendingPathComponent("track_\(trackNumber)").appendingPathExtension(fileExtension)
+            .path
 
         if FileManager.default.createFile(atPath: trackPath, contents: tracks[trackNumber].trackData, attributes: nil) {
             logger.debug("Created file at path: \(trackPath)")
         } else {
             logger.error("Failed to create file at path: \(trackPath)!")
         }
+
+        if fileExtension == "sub" {
+            let idxPath = outputDirectory.appendingPathComponent("track_\(trackNumber)").appendingPathExtension("idx")
+            do {
+                try tracks[trackNumber].idxData?.write(to: idxPath, atomically: true, encoding: .utf8)
+            } catch {
+                logger.error("Failed to write idx file at path: \(idxPath)")
+            }
+        }
     }
 }

Sources/macSubtitleOCR/MKV/MKVTrack.swift

@@ -12,4 +12,5 @@ struct MKVTrack {
     var trackNumber: Int
     var codecId: String
     var trackData: Data
+    var idxData: String?
 }