From 822e94acee5fb16cf5c5da12f7a7878295492654 Mon Sep 17 00:00:00 2001 From: Ethan Dye Date: Sun, 20 Oct 2024 21:00:09 -0600 Subject: [PATCH] Improve readability Signed-off-by: Ethan Dye --- Sources/macSubtitleOCR/MKV/MKVHelpers.swift | 6 +- .../MKV/MKVSubtitleExtractor.swift | 2 +- Sources/macSubtitleOCR/MKV/MKVTrack.swift | 2 +- .../macSubtitleOCR/MKV/MKVTrackParser.swift | 226 ++++++++++-------- Sources/macSubtitleOCR/macSubtitleOCR.swift | 8 +- 5 files changed, 132 insertions(+), 112 deletions(-) diff --git a/Sources/macSubtitleOCR/MKV/MKVHelpers.swift b/Sources/macSubtitleOCR/MKV/MKVHelpers.swift index d597437..52f8dbc 100644 --- a/Sources/macSubtitleOCR/MKV/MKVHelpers.swift +++ b/Sources/macSubtitleOCR/MKV/MKVHelpers.swift @@ -23,7 +23,7 @@ func encodePTSForPGS(_ timestamp: UInt64) -> [UInt8] { withUnsafeBytes(of: UInt32(timestamp).bigEndian) { Array($0) } } -func encodePTSForVobSub(_ timestamp: UInt64) -> [UInt8] { +func encodePTSForVobSub(from timestamp: UInt64) -> [UInt8] { var buffer = [UInt8](repeating: 0, count: 5) // 5-byte buffer buffer[0] = (buffer[0] & 0xF1) | UInt8((timestamp >> 29) & 0x0E) @@ -35,7 +35,7 @@ func encodePTSForVobSub(_ timestamp: UInt64) -> [UInt8] { } // Calculate the absolute timestamp with 90 kHz accuracy -func calcAbsPTS(_ clusterTimestamp: Int64, _ blockTimestamp: Int64) -> UInt64 { +func calculateAbsolutePTS(_ clusterTimestamp: Int64, _ blockTimestamp: Int64) -> UInt64 { // The block timestamp is relative, so we add it to the cluster timestamp - UInt64((Double(clusterTimestamp) + Double(blockTimestamp)) * 90) + UInt64(clusterTimestamp + blockTimestamp) * 90 } diff --git a/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift b/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift index 40b93a5..add7de5 100644 --- a/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift +++ b/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift @@ -11,7 +11,7 @@ import os class MKVSubtitleExtractor: MKVTrackParser { func saveSubtitleTrackData(trackNumber: Int, outputDirectory: URL) { - let codecType = tracks[trackNumber].codecId + let codecType = tracks[trackNumber].codecID let fileExtension = (codecType == "S_HDMV/PGS") ? "sup" : "sub" let trackPath = outputDirectory.appendingPathComponent("track_\(trackNumber)").appendingPathExtension(fileExtension) .path diff --git a/Sources/macSubtitleOCR/MKV/MKVTrack.swift b/Sources/macSubtitleOCR/MKV/MKVTrack.swift index 75a4cd8..4a530bc 100644 --- a/Sources/macSubtitleOCR/MKV/MKVTrack.swift +++ b/Sources/macSubtitleOCR/MKV/MKVTrack.swift @@ -10,7 +10,7 @@ import Foundation struct MKVTrack { var trackNumber: Int - var codecId: String + var codecID: String var trackData: Data var idxData: String? } diff --git a/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift b/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift index 439a376..d25aa2f 100644 --- a/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift +++ b/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift @@ -17,7 +17,7 @@ class MKVTrackParser: MKVFileHandler { // MARK: - Functions - func parseTracks(codec: [String]) throws { + func parseTracks(for codecs: [String]) throws { guard findElement(withID: EBML.segmentID) as? (UInt64, UInt32) != nil else { fatalError("Segment element not found in file: \(filePath)") } @@ -28,13 +28,13 @@ class MKVTrackParser: MKVFileHandler { let endOfTracksOffset = fileHandle.offsetInFile + tracksSize - var tracks = [Int: String]() + var subtitleTracks = [Int: String]() while fileHandle.offsetInFile < endOfTracksOffset { if let (elementID, elementSize) = tryParseElement() { if elementID == EBML.trackEntryID { logger.debug("Found TrackEntry element") - if let track = parseTrackEntry(codec: codec) { - tracks[track.0] = track.1 + if let track = parseTrackEntry(for: codecs) { + subtitleTracks[track.0] = track.1 } } else if elementID == EBML.chapters { break @@ -44,24 +44,24 @@ class MKVTrackParser: MKVFileHandler { } } - let trackData = extractTrackData(from: tracks) + let trackData = extractTrackData(for: subtitleTracks) trackData?.enumerated().forEach { index, data in - self.tracks.append(MKVTrack( + tracks.append(MKVTrack( trackNumber: index, - codecId: tracks[index + 1]!, + codecID: subtitleTracks[index + 1]!, trackData: data, idxData: codecPrivate[index + 1])) } } - func extractTrackData(from tracks: [Int: String]) -> [Data]? { + func extractTrackData(for tracks: [Int: String]) -> [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), EOF: \(self.endOfFile)") + logger.debug("Found Segment, Size: \(segmentSize), End Offset: \(segmentEndOffset)") var trackData = [Data](repeating: Data(), count: tracks.count) @@ -77,22 +77,18 @@ class MKVTrackParser: MKVFileHandler { } // Step 4: Parse Blocks (SimpleBlock or Block) within each Cluster - parseBlocks( - within: clusterEndOffset, - trackNumber: tracks, - clusterTimestamp: clusterTimestamp, - trackData: &trackData) + parseBlocks(until: clusterEndOffset, for: tracks, with: clusterTimestamp, into: &trackData) } - return trackData.isEmpty ? nil : trackData + return trackData } // MARK: - Methods - private func parseTrackEntry(codec: [String]) -> (Int, String)? { + private func parseTrackEntry(for codecs: [String]) -> (Int, String)? { var trackNumber: Int? var trackType: UInt8? - var codecId: String? + var codecID: String? while let (elementID, elementSize) = tryParseElement() { switch elementID { @@ -105,16 +101,16 @@ class MKVTrackParser: MKVFileHandler { case EBML.codecID: var data = fileHandle.readData(ofLength: Int(elementSize)) data.removeNullBytes() - codecId = String(data: data, encoding: .ascii) - logger.debug("Found codec ID: \(codecId ?? "nil")") + codecID = String(data: data, encoding: .ascii) + logger.debug("Found codec ID: \(codecID!)") default: fileHandle.seek(toFileOffset: fileHandle.offsetInFile + elementSize) } - if trackNumber != nil, trackType != nil, codecId != nil { break } + if trackNumber != nil, trackType != nil, codecID != nil { break } } - if let trackNumber, let codecId { - if codecId == "S_VOBSUB" { + if let trackNumber, let codecID { + if codecID == "S_VOBSUB" { while let (elementID, elementSize) = tryParseElement() { switch elementID { case EBML.codecPrivate: @@ -127,8 +123,8 @@ class MKVTrackParser: MKVFileHandler { if codecPrivate[trackNumber] != nil { break } } } - if codec.contains(codecId) { - return (trackNumber, codecId) + if codecs.contains(codecID) { + return (trackNumber, codecID) } } return nil @@ -141,93 +137,23 @@ class MKVTrackParser: MKVFileHandler { return nil } - private func parseBlocks(within clusterEndOffset: UInt64, trackNumber: [Int: String], clusterTimestamp: Int64, - trackData: inout [Data]) { + private func parseBlocks(until clusterEndOffset: UInt64, for tracks: [Int: String], with clusterTimestamp: Int64, + into trackData: inout [Data]) { while fileHandle.offsetInFile < clusterEndOffset { // swiftformat:disable:next redundantSelf logger.debug("Looking for Block at Offset: \(self.fileHandle.offsetInFile)/\(clusterEndOffset)") - guard case (var blockSize?, let blockType?) = findElement(withID: EBML.simpleBlock, EBML.blockGroup) - else { break } - var blockStartOffset = fileHandle.offsetInFile + guard let (blockSize, blockStartOffset) = findBlockTypeAndSize() else { break } - if blockType == EBML.blockGroup { - guard let (ns, _) = findElement(withID: EBML.block) as? (UInt64, UInt32) else { return } - blockSize = ns - blockStartOffset = fileHandle.offsetInFile - } - - // Step 5: Read the track number in the block and compare it guard let (blockTrackNumber, blockTimestamp) = readTrackNumber(from: fileHandle) as? (UInt64, Int64) else { continue } - if trackNumber[Int(blockTrackNumber)] == "S_HDMV/PGS" { - // Step 6: Calculate and encode the timestamp as 4 bytes in big-endian (PGS format) - let absPTS = calcAbsPTS(clusterTimestamp, blockTimestamp) - let pgsPTS = encodePTSForPGS(absPTS) - - // Step 7: Read the block data and add needed PGS headers and timestamps - let pgsHeader = Data([0x50, 0x47] + pgsPTS + [0x00, 0x00, 0x00, 0x00]) - var blockData = Data() - let raw = fileHandle.readData(ofLength: Int(blockSize - (fileHandle.offsetInFile - blockStartOffset))) - var offset = 0 - while (offset + 3) <= raw.count { - let segmentSize = min(Int(raw.getUInt16BE(at: offset + 1)! + 3), raw.count - offset) - logger.debug("Segment size \(segmentSize) at \(offset) type \(raw[offset].hex())") - - blockData.append(pgsHeader) - blockData.append(raw.subdata(in: offset ..< segmentSize + offset)) - offset += segmentSize - } - - trackData[Int(blockTrackNumber - 1)].append(blockData) - } else if trackNumber[Int(blockTrackNumber)] == "S_VOBSUB" { - // swiftformat:disable all - // Step 6: Calculate and encode the timestamp as 5 bytes in big-endian (VobSub format) - let absPTS = calcAbsPTS(clusterTimestamp, blockTimestamp) - let vobSubPTS = encodePTSForVobSub(absPTS) - var segmentSize = Int(blockSize - (fileHandle.offsetInFile - blockStartOffset)) - let pesLength = withUnsafeBytes(of: UInt16(min(segmentSize, 2028)).bigEndian) { Array($0) } - // 2028 is the maximum size of a VobSub segment, so we need to split the data into multiple segments - // The first segment will contain the PTS data, while the rest will not, so it only gets 2019 bytes of data - // The rest of the segments will get 2024 bytes of data - - // Step 7: Read the block data and add needed VobSub headers and timestamps - var vobSubHeader = Data([0x00, 0x00, 0x01, 0xBA, // PS packet start code - 0x00, 0x00, 0x00, 0x00, 0x00, 0x0, // Null system clock reference - 0x00, 0x00, 0x00, // Null multiplexer rate - 0x00, // Stuffing length - 0x00, 0x00, 0x01, 0xBD]) // PES packet start code - vobSubHeader.append(contentsOf: pesLength) // PES packet length - vobSubHeader.append(contentsOf: [0x00, // PES miscellaneous data - 0x80, // PTS DTS flag - UInt8(vobSubPTS.count)]) // PTS data length - vobSubHeader.append(contentsOf: vobSubPTS) // PTS data - vobSubHeader.append(contentsOf: [0x00]) // Null stream ID - vobSubHeader.append(fileHandle.readData(ofLength: min(segmentSize, 2019))) - - segmentSize -= min(segmentSize, 2019) - - while segmentSize > 0 { - let nextSegmentSize = min(segmentSize, 2028) - let pesLength = withUnsafeBytes(of: UInt16(nextSegmentSize).bigEndian) { Array($0) } - vobSubHeader.append(contentsOf: [0x00, 0x00, 0x01, 0xBA, // PS packet start code - 0x00, 0x00, 0x00, 0x00, 0x00, 0x0, // Null system clock reference - 0x00, 0x00, 0x00, // Null multiplexer rate - 0x00, // Stuffing length - 0x00, 0x00, 0x01, 0xBD]) // PES packet start code - vobSubHeader.append(contentsOf: pesLength) // PES packet length - vobSubHeader.append(contentsOf: [0x00, // PES miscellaneous data - 0x00, // PTS DTS flag - 0x00]) // PTS data length - vobSubHeader.append(contentsOf: [0x00]) // Null stream ID - vobSubHeader.append(fileHandle.readData(ofLength: min(segmentSize, 2024))) - segmentSize -= min(segmentSize, 2024) - } - trackData[Int(blockTrackNumber - 1)].append(vobSubHeader) - let offset = String(format: "%09X", trackData[Int(blockTrackNumber - 1)].count - vobSubHeader.count) - codecPrivate[Int(blockTrackNumber)]?.append("\ntimestamp: \(formatTime(absPTS)), filepos: \(offset)") - // swiftformat:enable all + if tracks[Int(blockTrackNumber)] == "S_HDMV/PGS" { + handlePGSBlock(for: blockTrackNumber, + with: blockTimestamp, blockSize, clusterTimestamp, blockStartOffset, &trackData) + } else if tracks[Int(blockTrackNumber)] == "S_VOBSUB" { + handleVobSubBlock(for: blockTrackNumber, + with: blockTimestamp, blockSize, clusterTimestamp, blockStartOffset, &trackData) } else { // Skip this block because it's for a different track fileHandle.seek(toFileOffset: blockStartOffset + blockSize) @@ -235,6 +161,100 @@ class MKVTrackParser: MKVFileHandler { } } + private func handlePGSBlock(for blockTrackNumber: UInt64, with blockTimestamp: Int64, _ blockSize: UInt64, + _ clusterTimestamp: Int64, _ blockStartOffset: UInt64, _ trackData: inout [Data]) { + let absPTS = calculateAbsolutePTS(clusterTimestamp, blockTimestamp) + let pgsPTS = encodePTSForPGS(absPTS) + let pgsHeader = Data([0x50, 0x47] + pgsPTS + [0x00, 0x00, 0x00, 0x00]) + + var blockData = Data() + blockData.reserveCapacity(Int(blockSize)) + let raw = fileHandle.readData(ofLength: Int(blockSize - (fileHandle.offsetInFile - blockStartOffset))) + + var offset = 0 + while (offset + 3) <= raw.count { + let segmentSize = min(Int(raw.getUInt16BE(at: offset + 1)! + 3), raw.count - offset) + logger.debug("Segment size \(segmentSize) at \(offset) type \(raw[offset].hex())") + + blockData.append(pgsHeader) + blockData.append(raw.subdata(in: offset ..< segmentSize + offset)) + offset += segmentSize + } + + trackData[Int(blockTrackNumber - 1)].append(blockData) + } + + private func handleVobSubBlock(for blockTrackNumber: UInt64, with blockTimestamp: Int64, _ blockSize: UInt64, + _ clusterTimestamp: Int64, _ blockStartOffset: UInt64, _ trackData: inout [Data]) { + let absolutePTS = calculateAbsolutePTS(clusterTimestamp, blockTimestamp) + let vobSubPTS = encodePTSForVobSub(from: absolutePTS) + + var segmentSize = Int(blockSize - (fileHandle.offsetInFile - blockStartOffset)) + var vobSubHeader = buildVobSubHeader(pts: vobSubPTS, segmentSize: segmentSize) + + vobSubHeader.append(fileHandle.readData(ofLength: min(segmentSize, 2019))) + segmentSize -= min(segmentSize, 2019) + + appendVobSubSegments(segmentSize: segmentSize, header: &vobSubHeader) + + trackData[Int(blockTrackNumber - 1)].append(vobSubHeader) + + let offset = String(format: "%09X", trackData[Int(blockTrackNumber - 1)].count - vobSubHeader.count) + codecPrivate[Int(blockTrackNumber)]?.append("\ntimestamp: \(formatTime(absolutePTS)), filepos: \(offset)") + } + + // swiftformat:disable all + private func buildVobSubHeader(pts: [UInt8], segmentSize: Int) -> Data { + let pesLength = withUnsafeBytes(of: UInt16(min(segmentSize, 2028)).bigEndian) { Array($0) } + var vobSubHeader = Data([0x00, 0x00, 0x01, 0xBA, // PS packet start code + 0x00, 0x00, 0x00, 0x00, 0x00, 0x0, // Null system clock reference + 0x00, 0x00, 0x00, // Null multiplexer rate + 0x00, // Stuffing length + 0x00, 0x00, 0x01, 0xBD]) // PES packet start code + vobSubHeader.append(contentsOf: pesLength) // PES packet length + vobSubHeader.append(contentsOf: [0x00, // PES miscellaneous data + 0x80, // PTS DTS flag + UInt8(pts.count)]) // PTS data length + vobSubHeader.append(contentsOf: pts) // PTS data + vobSubHeader.append(contentsOf: [0x00]) // Null stream ID + return vobSubHeader + } + + private func appendVobSubSegments(segmentSize: Int, header: inout Data) { + var remainingSize = segmentSize + while remainingSize > 0 { + let nextSegmentSize = min(remainingSize, 2028) + let pesLength = withUnsafeBytes(of: UInt16(nextSegmentSize).bigEndian) { Array($0) } + header.append(contentsOf: [0x00, 0x00, 0x01, 0xBA, // PS packet start code + 0x00, 0x00, 0x00, 0x00, 0x00, 0x0, // Null system clock reference + 0x00, 0x00, 0x00, // Null multiplexer rate + 0x00, // Stuffing length + 0x00, 0x00, 0x01, 0xBD]) // PES packet start code + header.append(contentsOf: pesLength) // PES packet length + header.append(contentsOf: [0x00, // PES miscellaneous data + 0x00, // PTS DTS flag + 0x00]) // PTS data length + header.append(contentsOf: [0x00]) // Null stream ID + header.append(fileHandle.readData(ofLength: min(remainingSize, 2024))) + remainingSize -= min(remainingSize, 2024) + } + } + // swiftformat:enable all + + private func findBlockTypeAndSize() -> (blockSize: UInt64, blockStartOffset: UInt64)? { + guard case (var blockSize?, let blockType?) = findElement(withID: EBML.simpleBlock, EBML.blockGroup) else { + return nil + } + + var blockStartOffset = fileHandle.offsetInFile + if blockType == EBML.blockGroup { + guard let (newSize, _) = findElement(withID: EBML.block) as? (UInt64, UInt32) else { return nil } + blockSize = newSize + blockStartOffset = fileHandle.offsetInFile + } + return (blockSize, blockStartOffset) + } + private func formatTime(_ time: UInt64) -> String { let time = TimeInterval(time) / 90000 let hours = Int(time) / 3600 diff --git a/Sources/macSubtitleOCR/macSubtitleOCR.swift b/Sources/macSubtitleOCR/macSubtitleOCR.swift index cc3b049..d16f44e 100644 --- a/Sources/macSubtitleOCR/macSubtitleOCR.swift +++ b/Sources/macSubtitleOCR/macSubtitleOCR.swift @@ -95,23 +95,23 @@ struct macSubtitleOCR: AsyncParsableCommand { results.append(result) } else if input.hasSuffix(".mkv") || input.hasSuffix(".mks") { let mkvStream = MKVSubtitleExtractor(filePath: input) - try mkvStream.parseTracks(codec: ["S_HDMV/PGS", "S_VOBSUB"]) + try mkvStream.parseTracks(for: ["S_HDMV/PGS", "S_VOBSUB"]) for track in mkvStream.tracks { - logger.debug("Found subtitle track: \(track.trackNumber), Codec: \(track.codecId)") + logger.debug("Found subtitle track: \(track.trackNumber), Codec: \(track.codecID)") if experimentalOptions.saveSubtitleFile { mkvStream.saveSubtitleTrackData( trackNumber: track.trackNumber, outputDirectory: URL(fileURLWithPath: outputDirectory)) } - if track.codecId == "S_HDMV/PGS" { + 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" { + } else if track.codecID == "S_VOBSUB" { let vobSub: VobSub = try track.trackData .withUnsafeBytes { (buffer: UnsafeRawBufferPointer) in try VobSub(buffer, track.idxData ?? "")