From ddfd96c985eb7bda595a2703d7be8b23b7357587 Mon Sep 17 00:00:00 2001 From: Ethan Dye Date: Thu, 31 Oct 2024 23:44:07 -0600 Subject: [PATCH] Fix MKV decoding for multiple tracks Fixes #33 Signed-off-by: Ethan Dye --- .../MKV/MKVSubtitleExtractor.swift | 4 +- .../macSubtitleOCR/MKV/MKVTrackParser.swift | 39 +++++++++++-------- .../macSubtitleOCR/Subtitles/PGS/PGS.swift | 4 +- 3 files changed, 27 insertions(+), 20 deletions(-) diff --git a/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift b/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift index add7de5..5146835 100644 --- a/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift +++ b/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift @@ -19,7 +19,7 @@ class MKVSubtitleExtractor: MKVTrackParser { 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)!") + print("Failed to create file at path: \(trackPath)!", to: &stderr) } if fileExtension == "sub" { @@ -27,7 +27,7 @@ class MKVSubtitleExtractor: MKVTrackParser { do { try tracks[trackNumber].idxData?.write(to: idxPath, atomically: true, encoding: .utf8) } catch { - logger.error("Failed to write idx file at path: \(idxPath)") + print("Failed to write idx file at path: \(idxPath)", to: &stderr) } } } diff --git a/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift b/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift index 5a80acb..9c4287e 100644 --- a/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift +++ b/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift @@ -46,30 +46,29 @@ class MKVTrackParser: MKVFileHandler { } let trackData = extractTrackData(for: subtitleTracks) - trackData?.enumerated().forEach { index, data in + trackData?.forEach { trackNumber, data in if data.isEmpty { - print("Found empty track data for track \(index + 1), skipping track!") + print("Found empty track data for track \(trackNumber), skipping track!") return } tracks.append(MKVTrack( - trackNumber: index, - codecID: subtitleTracks[index + 1]!, + trackNumber: trackNumber, + codecID: subtitleTracks[trackNumber]!, trackData: data, - idxData: codecPrivate[index + 1], - language: languages[index + 1])) + idxData: codecPrivate[trackNumber], + language: languages[trackNumber])) } } - func extractTrackData(for tracks: [Int: String]) -> [Data]? { + func extractTrackData(for tracks: [Int: String]) -> [Int: Data]? { fileHandle.seek(toFileOffset: 0) // Step 1: Locate the Segment element guard let segmentSize = locateSegment() else { return nil } let segmentEndOffset = fileHandle.offsetInFile + segmentSize - // swiftformat:disable:next redundantSelf logger.debug("Found Segment, Size: \(segmentSize), End Offset: \(segmentEndOffset)") - var trackData = [Data](repeating: Data(), count: tracks.count) + var trackData = [Int: Data]() // Step 2: Parse Clusters within the Segment while fileHandle.offsetInFile < segmentEndOffset { @@ -156,7 +155,7 @@ class MKVTrackParser: MKVFileHandler { } private func parseBlocks(until clusterEndOffset: UInt64, for tracks: [Int: String], with clusterTimestamp: Int64, - into trackData: inout [Data]) { + into trackData: inout [Int: Data]) { while fileHandle.offsetInFile < clusterEndOffset { // swiftformat:disable:next redundantSelf logger.debug("Looking for Block at Offset: \(self.fileHandle.offsetInFile)/\(clusterEndOffset)") @@ -180,7 +179,7 @@ class MKVTrackParser: MKVFileHandler { } private func handlePGSBlock(for blockTrackNumber: UInt64, with blockTimestamp: Int64, _ blockSize: UInt64, - _ clusterTimestamp: Int64, _ blockStartOffset: UInt64, _ trackData: inout [Data]) { + _ clusterTimestamp: Int64, _ blockStartOffset: UInt64, _ trackData: inout [Int: Data]) { let absPTS = calculateAbsolutePTS(clusterTimestamp, blockTimestamp) let pgsPTS = encodePTSForPGS(absPTS) let pgsHeader = Data([0x50, 0x47] + pgsPTS + [0x00, 0x00, 0x00, 0x00]) @@ -199,11 +198,15 @@ class MKVTrackParser: MKVFileHandler { offset += segmentSize } - trackData[Int(blockTrackNumber - 1)].append(blockData) + let trackNumber = Int(blockTrackNumber) + if trackData[trackNumber] == nil { + trackData[trackNumber] = Data() + } + trackData[trackNumber]?.append(blockData) } private func handleVobSubBlock(for blockTrackNumber: UInt64, with blockTimestamp: Int64, _ blockSize: UInt64, - _ clusterTimestamp: Int64, _ blockStartOffset: UInt64, _ trackData: inout [Data]) { + _ clusterTimestamp: Int64, _ blockStartOffset: UInt64, _ trackData: inout [Int: Data]) { let absolutePTS = calculateAbsolutePTS(clusterTimestamp, blockTimestamp) let vobSubPTS = encodePTSForVobSub(from: absolutePTS) @@ -215,10 +218,14 @@ class MKVTrackParser: MKVFileHandler { appendVobSubSegments(segmentSize: segmentSize, header: &vobSubHeader) - trackData[Int(blockTrackNumber - 1)].append(vobSubHeader) + let trackNumber = Int(blockTrackNumber) + if trackData[trackNumber] == nil { + trackData[trackNumber] = Data() + } + trackData[trackNumber]?.append(vobSubHeader) - let offset = String(format: "%09X", trackData[Int(blockTrackNumber - 1)].count - vobSubHeader.count) - codecPrivate[Int(blockTrackNumber)]?.append("\ntimestamp: \(formatTime(absolutePTS)), filepos: \(offset)") + let offset = String(format: "%09X", trackData[trackNumber]!.count - vobSubHeader.count) + codecPrivate[trackNumber]?.append("\ntimestamp: \(formatTime(absolutePTS)), filepos: \(offset)") } // swiftformat:disable all diff --git a/Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift b/Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift index a67e34e..a8a1567 100644 --- a/Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift +++ b/Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift @@ -77,7 +77,7 @@ struct PGS { pds = try PDS(buffer, offset, segmentLength) offset += segmentLength } catch let macSubtitleOCRError.invalidPDSDataLength(length) { - logger.error("Invalid PDS length: \(length), abandoning remaining segments!") + print("Invalid PDS length: \(length), abandoning remaining segments!", to: &stderr) offset = buffer.count return nil } @@ -97,7 +97,7 @@ struct PGS { offset += segmentLength } } catch let macSubtitleOCRError.invalidODSDataLength(length) { - logger.error("Invalid ODS length: \(length), abandoning remaining segments!") + print("Invalid ODS length: \(length), abandoning remaining segments!", to: &stderr) offset = buffer.count continue }