Add boilerplate support for VobSub tracks in MKV

Signed-off-by: Ethan Dye <mrtops03@gmail.com>

Sources/macSubtitleOCR/MKV/MKVTrack.swift

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

Sources/macSubtitleOCR/MKV/MKVTrackParser.swift

@@ -16,7 +16,7 @@ class MKVTrackParser: MKVFileHandler {
 
     // MARK: - Functions
 
-    func parseTracks(codec: String) throws {
+    func parseTracks(codec: [String]) throws {
         guard findElement(withID: EBML.segmentID) as? (UInt64, UInt32) != nil else {
             fatalError("Segment element not found in file: \(filePath)")
         }
@@ -27,13 +27,18 @@ class MKVTrackParser: MKVFileHandler {
 
         let endOfTracksOffset = fileHandle.offsetInFile + tracksSize
 
-        var trackNumbers = [Int]()
+        var trackNumbersPGS = [Int]()
+        var trackNumbersVobSub = [Int]()
         while fileHandle.offsetInFile < endOfTracksOffset {
             if let (elementID, elementSize) = tryParseElement() {
                 if elementID == EBML.trackEntryID {
                     logger.debug("Found TrackEntry element")
                     if let track = parseTrackEntry(codec: codec) {
-                        trackNumbers.append(track)
+                        if track.1 == "S_HDMV/PGS" {
+                            trackNumbersPGS.append(track.0)
+                        } else {
+                            trackNumbersVobSub.append(track.0)
+                        }
                     }
                 } else if elementID == EBML.chapters {
                     break
@@ -43,13 +48,19 @@ class MKVTrackParser: MKVFileHandler {
             }
         }
 
-        let trackData = extractTrackData(trackNumber: trackNumbers)
+        let trackData = extractTrackDataPGS(trackNumber: trackNumbersPGS)
         trackData?.enumerated().forEach { index, data in
-            tracks.append(MKVTrack(trackNumber: index, codecId: codec, trackData: data))
+            tracks.append(MKVTrack(trackNumber: index, codecId: codec[0], trackData: data))
         }
+        /*
+        let trackDataVobSub = extractTrackDataVobSub(trackNumber: trackNumbersVobSub)
+        trackDataVobSub?.enumerated().forEach { index, data in
+            tracks.append(MKVTrack(trackNumber: index, codecId: codec[1], trackData: data))
+        }
+        */
     }
 
-    func extractTrackData(trackNumber: [Int]) -> [Data]? {
+    func extractTrackDataPGS(trackNumber: [Int]) -> [Data]? {
         fileHandle.seek(toFileOffset: 0)
 
         // Step 1: Locate the Segment element
@@ -72,7 +83,7 @@ class MKVTrackParser: MKVFileHandler {
             }
 
             // Step 4: Parse Blocks (SimpleBlock or Block) within each Cluster
-            parseBlocks(
+            parseBlocksPGS(
                 within: clusterEndOffset,
                 trackNumber: trackNumber,
                 clusterTimestamp: clusterTimestamp,
@@ -84,7 +95,7 @@ class MKVTrackParser: MKVFileHandler {
 
     // MARK: - Methods
 
-    private func parseTrackEntry(codec: String) -> Int? {
+    private func parseTrackEntry(codec: [String]) -> (Int, String)? {
         var trackNumber: Int?
         var trackType: UInt8?
         var codecId: String?
@@ -109,8 +120,8 @@ class MKVTrackParser: MKVFileHandler {
         }
 
         if let trackNumber, let codecId {
-            if codecId == codec {
-                return trackNumber
+            if codec.contains(codecId) {
+                return (trackNumber, codecId)
             }
         }
         return nil
@@ -123,7 +134,7 @@ class MKVTrackParser: MKVFileHandler {
         return nil
     }
 
-    private func parseBlocks(within clusterEndOffset: UInt64, trackNumber: [Int], clusterTimestamp: Int64,
+    private func parseBlocksPGS(within clusterEndOffset: UInt64, trackNumber: [Int], clusterTimestamp: Int64,
                              trackData: inout [Data]) {
         while fileHandle.offsetInFile < clusterEndOffset {
             // swiftformat:disable:next redundantSelf

Sources/macSubtitleOCR/Subtitles/VobSub/VobSub.swift

@@ -22,11 +22,16 @@ struct VobSub {
         let subData = try subFile.readToEnd()!
         subFile.closeFile()
         let idx = VobSubIDX(URL(filePath: idx))
-        subData.withUnsafeBytes { (pointer: UnsafeRawBufferPointer) in
-            extractSubtitleImages(buffer: pointer, idx: idx)
+        subData.withUnsafeBytes { (buffer: UnsafeRawBufferPointer) in
+            extractSubtitleImages(buffer: buffer, idx: idx)
         }
     }
 
+    init(_ buffer: UnsafeRawBufferPointer, _ idxData: String) throws {
+        let idx = VobSubIDX(idxData)
+        extractSubtitleImages(buffer: buffer, idx: idx)
+    }
+
     // MARK: - Methods
 
     private mutating func extractSubtitleImages(buffer: UnsafeRawBufferPointer, idx: VobSubIDX) {

Sources/macSubtitleOCR/Subtitles/VobSub/VobSubIDX.swift

@@ -26,6 +26,14 @@ struct VobSubIDX {
         }
     }
 
+    init(_ idxData: String) {
+        do {
+            try parseIdxFile(idxData: idxData)
+        } catch {
+            fatalError("Failed to parse IDX file: \(error)")
+        }
+    }
+
     // MARK: - Methods
 
     private mutating func parseIdxFile(idxData: String) throws {

Sources/macSubtitleOCR/macSubtitleOCR.swift

@@ -95,7 +95,7 @@ struct macSubtitleOCR: AsyncParsableCommand {
             results.append(result)
         } else if input.hasSuffix(".mkv") {
             let mkvStream = MKVSubtitleExtractor(filePath: input)
-            try mkvStream.parseTracks(codec: "S_HDMV/PGS")
+            try mkvStream.parseTracks(codec: ["S_HDMV/PGS", "S_VOBSUB"])
             for track in mkvStream.tracks {
                 logger.debug("Found subtitle track: \(track.trackNumber), Codec: \(track.codecId)")
                 if experimentalOptions.saveSubtitleFile {
@@ -104,13 +104,22 @@ struct macSubtitleOCR: AsyncParsableCommand {
                         outputDirectory: URL(fileURLWithPath: outputDirectory))
                 }
 
-                // Open the PGS data stream
-                let pgs: PGS = try mkvStream.tracks[track.trackNumber].trackData
-                    .withUnsafeBytes { (buffer: UnsafeRawBufferPointer) in
-                        try PGS(buffer)
-                    }
-                let result = try await processSubtitle(pgs.subtitles, trackNumber: track.trackNumber)
-                results.append(result)
+
+                if track.codecId == "S_HDMV/PGS" {
+                    let pgs: PGS = try track.trackData
+                        .withUnsafeBytes { (buffer: UnsafeRawBufferPointer) in
+                            try PGS(buffer)
+                        }
+                    let result = try await processSubtitle(pgs.subtitles, trackNumber: track.trackNumber)
+                    results.append(result)
+                } else if track.codecId == "S_VOBSUB" {
+                    let vobSub: VobSub = try track.trackData
+                        .withUnsafeBytes { (buffer: UnsafeRawBufferPointer) in
+                            try VobSub(buffer, track.idxData ?? "")
+                        }
+                    let result = try await processSubtitle(vobSub.subtitles, trackNumber: track.trackNumber)
+                    results.append(result)
+                }
             }
         } else if input.hasSuffix(".sup") {
             // Open the PGS data stream