Skip to content

Commit

Permalink
Improve readability
Browse files Browse the repository at this point in the history 
Signed-off-by: Ethan Dye <[email protected]>
  • Loading branch information
@ecdye
ecdye committed Oct 21, 2024
1 parent 3028ced commit 822e94a
Show file tree
Hide file tree
Showing 5 changed files with 132 additions and 112 deletions.
6 changes: 3 additions & 3 deletions Sources/macSubtitleOCR/MKV/MKVHelpers.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
}
2 changes: 1 addition & 1 deletion Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion Sources/macSubtitleOCR/MKV/MKVTrack.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import Foundation

struct MKVTrack {
var trackNumber: Int
var codecId: String
var codecID: String
var trackData: Data
var idxData: String?
}
226 changes: 123 additions & 103 deletions Sources/macSubtitleOCR/MKV/MKVTrackParser.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)")
}
Expand All @@ -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
Expand All @@ -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)

Expand All @@ -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 {
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -141,100 +137,124 @@ 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)
}
}
}

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
Expand Down
Loading

0 comments on commit 822e94a

Please sign in to comment.