Skip to content

Commit

Permalink
Improve VobSub performance and prepare for MKV
Browse files Browse the repository at this point in the history 
Signed-off-by: Ethan Dye <mrtops03@gmail.com>
  • Loading branch information
@ecdye
ecdye committed Oct 17, 2024
1 parent 24c8c63 commit 416dceb
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 43 deletions.
15 changes: 10 additions & 5 deletions Sources/macSubtitleOCR/Subtitles/VobSub/VobSub.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,26 +19,31 @@ struct VobSub {

init(_ sub: String, _ idx: String) throws {
let subFile = try FileHandle(forReadingFrom: URL(filePath: sub))
defer { subFile.closeFile() }
var subData = try subFile.readToEnd()!
subFile.closeFile()
let idx = VobSubIDX(URL(filePath: idx))
extractSubtitleImages(subFile: subFile, idx: idx)
subData.withUnsafeMutableBytes { (pointer: UnsafeMutableRawBufferPointer) in
var pointer = pointer
extractSubtitleImages(subData: &pointer, idx: idx)
}
// extractSubtitleImages(subData: subData, idx: idx)
}

// MARK: - Methods

private mutating func extractSubtitleImages(subFile: FileHandle, idx: VobSubIDX) {
private mutating func extractSubtitleImages(subData: inout UnsafeMutableRawBufferPointer, idx: VobSubIDX) {
for index in idx.offsets.indices {
logger.debug("Index \(index), offset: \(idx.offsets[index]), timestamp: \(idx.timestamps[index])")
let offset = idx.offsets[index]
let timestamp = idx.timestamps[index]
let nextOffset: UInt64 = if index + 1 < idx.offsets.count {
idx.offsets[index + 1]
} else {
subFile.seekToEndOfFile()
UInt64(subData.count)
}
let subtitle = VobSubParser(
index: index + 1,
subFile: subFile,
subData: &subData,
timestamp: timestamp,
offset: offset,
nextOffset: nextOffset,
Expand Down
87 changes: 49 additions & 38 deletions Sources/macSubtitleOCR/Subtitles/VobSub/VobSubParser.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,17 +19,19 @@ struct VobSubParser {

// MARK: - Lifecycle

init(index: Int, subFile: FileHandle, timestamp: TimeInterval, offset: UInt64, nextOffset: UInt64, idxPalette: [UInt8]) {
init(index: Int, subData: inout UnsafeMutableRawBufferPointer, timestamp: TimeInterval, offset: UInt64,
nextOffset: UInt64, idxPalette: [UInt8]) {
subtitle = Subtitle(index: index, startTimestamp: timestamp, imageData: .init(), numberOfColors: 16)
masterPalette = idxPalette
readSubFrame(subFile: subFile, offset: offset, nextOffset: nextOffset, idxPalette: idxPalette)
readSubFrame(subData: &subData, offset: offset, nextOffset: nextOffset, idxPalette: idxPalette)
decodeImage()
decodePalette()
}

// MARK: - Methods

func readSubFrame(subFile: FileHandle, offset: UInt64, nextOffset: UInt64, idxPalette _: [UInt8]) {
func readSubFrame(subData: inout UnsafeMutableRawBufferPointer, offset: UInt64, nextOffset: UInt64,
idxPalette _: [UInt8]) {
var firstPacketFound = false
var controlOffset: Int?
var controlSize: Int?
Expand All @@ -38,71 +40,80 @@ struct VobSubParser {
var relativeControlOffset = 0
var rleLengthFound = 0

subFile.seek(toFileOffset: offset)
var offset = Int(offset)
repeat {
let startOffset = subFile.offsetInFile
guard subFile.readData(ofLength: 4).value(ofType: UInt32.self) == MPEG2PacketType.psPacket else {
fatalError("Failed to find PS packet at offset \(subFile.offsetInFile)")
let startOffset = offset
guard subData.loadUnaligned(fromByteOffset: offset, as: UInt32.self).bigEndian == MPEG2PacketType.psPacket else {
fatalError("Failed to find PS packet at offset \(offset)")
}
offset += 4

subFile.readData(ofLength: 6) // System clock reference
subFile.readData(ofLength: 3) // Multiplexer rate
let stuffingLength = Int(subFile.readData(ofLength: 1)[0] & 7)
subFile.readData(ofLength: stuffingLength) // Stuffing bytes
offset += 6 // System clock reference
offset += 3 // Multiplexer rate
let stuffingLength = subData.loadUnaligned(fromByteOffset: offset, as: UInt8.self) & 7
offset += 1 + Int(stuffingLength)

guard subFile.readData(ofLength: 4).value(ofType: UInt32.self) == MPEG2PacketType.pesPacket else {
fatalError("Failed to find PES packet at offset \(subFile.offsetInFile)")
guard subData.loadUnaligned(fromByteOffset: offset, as: UInt32.self).bigEndian == MPEG2PacketType.pesPacket
else {
fatalError("Failed to find PES packet at offset \(offset)")
}
offset += 4

let pesLength = Int(subFile.readData(ofLength: 2).value(ofType: UInt16.self) ?? 0)
let pesLength = Int(subData.loadUnaligned(fromByteOffset: offset, as: UInt16.self).bigEndian)
if pesLength == 0 {
fatalError("PES packet length is 0 at offset \(subFile.offsetInFile)")
fatalError("PES packet length is 0 at offset \(offset)")
}
let nextPSOffset = subFile.offsetInFile + UInt64(pesLength)
offset += 2
let nextPSOffset = offset + pesLength

subFile.readData(ofLength: 1) // Skip PES miscellaneous data
let extByteOne = subFile.readData(ofLength: 1)[0]
offset += 1 // Skip PES miscellaneous data
let extByteOne = subData.loadUnaligned(fromByteOffset: offset, as: UInt8.self)
offset += 1
let firstPacket = (extByteOne & 0x80) == 0x80 || (extByteOne & 0xC0) == 0xC0

let ptsDataLength = Int(subFile.readData(ofLength: 1)[0])
let ptsData = subFile.readData(ofLength: ptsDataLength)
let ptsDataLength = Int(subData.loadUnaligned(fromByteOffset: offset, as: UInt8.self))
offset += 1
if ptsDataLength == 5 {
var presentationTimestamp: UInt64 = 0
presentationTimestamp = UInt64(ptsData[4]) >> 1
presentationTimestamp += UInt64(ptsData[3]) << 7
presentationTimestamp += UInt64(ptsData[2] & 0xFE) << 14
presentationTimestamp += UInt64(ptsData[1]) << 22
presentationTimestamp += UInt64(ptsData[0] & 0x0E) << 29
presentationTimestamp = UInt64(subData[offset + ptsDataLength - 1]) >> 1
presentationTimestamp += UInt64(subData[offset + ptsDataLength - 2]) << 7
presentationTimestamp += UInt64(subData[offset + ptsDataLength - 3] & 0xFE) << 14
presentationTimestamp += UInt64(subData[offset + ptsDataLength - 4]) << 22
presentationTimestamp += UInt64(subData[offset] & 0x0E) << 29
subtitle.startTimestamp = TimeInterval(presentationTimestamp) / 90 / 1000
}
offset += ptsDataLength

subFile.readData(ofLength: 1) // Stream ID
offset += 1 // Stream ID

var trueHeaderSize = Int(subFile.offsetInFile - startOffset)
var trueHeaderSize = offset - startOffset
if firstPacket, ptsDataLength >= 5 {
let size = Int(subFile.readData(ofLength: 2).value(ofType: UInt16.self) ?? 0)
relativeControlOffset = Int(subFile.readData(ofLength: 2).value(ofType: UInt16.self) ?? 0)
let size = Int(subData.loadUnaligned(fromByteOffset: offset, as: UInt16.self).bigEndian)
offset += 2
relativeControlOffset = Int(subData.loadUnaligned(fromByteOffset: offset, as: UInt16.self).bigEndian)
offset += 2
let rleSize = relativeControlOffset - 2
controlSize = size - rleSize - 4 // 4 bytes for the size and control offset
controlOffset = Int(subFile.offsetInFile) + rleSize
trueHeaderSize = Int(subFile.offsetInFile - startOffset)
controlOffset = offset + rleSize
trueHeaderSize = startOffset - offset
firstPacketFound = true
} else if firstPacketFound {
controlOffset! += trueHeaderSize
}

let savedOffset = subFile.offsetInFile
let difference = max(0, Int(nextPSOffset) - controlOffset! - controlHeaderCopied)
let rleFragmentSize = Int(nextPSOffset - savedOffset) - difference
subtitle.imageData!.append(subFile.readData(ofLength: rleFragmentSize))
let difference = max(0, nextPSOffset - controlOffset! - controlHeaderCopied)
let rleFragmentSize = nextPSOffset - offset - difference
subtitle.imageData!.append(contentsOf: subData[offset ..< offset + rleFragmentSize])
rleLengthFound += rleFragmentSize
offset += rleFragmentSize

let bytesToCopy = max(0, min(difference, controlSize! - controlHeaderCopied))
controlHeader.append(subFile.readData(ofLength: bytesToCopy))
controlHeader.append(contentsOf: subData[offset ..< offset + bytesToCopy])
controlHeaderCopied += bytesToCopy
offset += bytesToCopy

subFile.seek(toFileOffset: nextPSOffset)
} while subFile.offsetInFile < nextOffset && controlHeaderCopied < controlSize!
offset = nextPSOffset
} while offset < nextOffset && controlHeaderCopied < controlSize!

if controlHeaderCopied < controlSize! {
logger.warning("Failed to read control header completely")
Expand Down

0 comments on commit 416dceb

Please sign in to comment.