Add experimental multithreading support

Signed-off-by: Ethan Dye <[email protected]>
ecdye · Oct 17, 2024 · 2036216 · 2036216
1 parent 1e621f4
commit 2036216
Show file tree

Hide file tree

Showing 8 changed files with 205 additions and 129 deletions.
diff --git a/Package.swift b/Package.swift
@@ -5,7 +5,7 @@ import PackageDescription
 let package = Package(
     name: "macSubtitleOCR",
     platforms: [
-        .macOS("13.0")
+        .macOS("14.0")
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.5.0")

diff --git a/Sources/macSubtitleOCR/Subtitles/FFmpeg/FFmpeg.swift b/Sources/macSubtitleOCR/Subtitles/FFmpeg/FFmpeg.swift
@@ -90,7 +90,7 @@ struct FFmpeg {
                 var trackSubtitles = subtitleTracks[streamNumber] ?? []
                 for i in 0 ..< Int(subtitle.num_rects) {
                     let rect = subtitle.rects[i]!
-                    let sub = extractImageData(from: rect)
+                    var sub = extractImageData(from: rect)
                     let pts = convertToTimeInterval(packet!.pointee.pts, timeBase: stream.timeBase)
                     sub.startTimestamp = pts + convertToTimeInterval(subtitle.start_display_time, timeBase: timeBase)
                     sub.endTimestamp = pts + convertToTimeInterval(subtitle.end_display_time, timeBase: timeBase)
@@ -103,7 +103,7 @@ struct FFmpeg {
     }
 
     private func extractImageData(from rect: UnsafeMutablePointer<AVSubtitleRect>) -> Subtitle {
-        let subtitle = Subtitle(numberOfColors: Int(rect.pointee.nb_colors))
+        var subtitle = Subtitle(numberOfColors: Int(rect.pointee.nb_colors))
 
         // Check if the subtitle is an image (bitmap)
         if rect.pointee.type == SUBTITLE_BITMAP {

diff --git a/Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift b/Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift
@@ -43,7 +43,7 @@ struct PGS {
     private mutating func parseData() throws {
         var headerData = data.extractBytes(pgsHeaderLength)
         while data.count > 0 {
-            guard let subtitle = try parseNextSubtitle(headerData: &headerData)
+            guard var subtitle = try parseNextSubtitle(headerData: &headerData)
             else {
                 if data.count < pgsHeaderLength { break }
                 headerData = data.extractBytes(pgsHeaderLength)

diff --git a/Sources/macSubtitleOCR/Subtitles/SRT/SRT.swift b/Sources/macSubtitleOCR/Subtitles/SRT/SRT.swift
@@ -7,11 +7,13 @@
 //
 
 import Foundation
+import os
 
 struct SRT {
     // MARK: - Properties
 
     private var subtitles: [Subtitle] = []
+    private let logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category: "SRT")
 
     // MARK: - Getters / Setters
 
@@ -38,8 +40,20 @@ struct SRT {
         var srtContent = ""
 
         for subtitle in subtitles {
+            var endTimestamp = subtitle.endTimestamp ?? 0
+            if subtitle.index! + 1 < subtitles.count {
+                let nextSubtitle = subtitles[subtitle.index! + 1]
+                if nextSubtitle.startTimestamp! <= subtitle.endTimestamp! {
+                    logger.warning("Fixing subtitle index \(subtitle.index!) end timestamp!")
+                    if nextSubtitle.startTimestamp! - subtitle.startTimestamp! > 5 {
+                        endTimestamp = subtitle.startTimestamp! + 5
+                    } else {
+                        endTimestamp = nextSubtitle.startTimestamp! - 0.1
+                    }
+                }
+            }
             let startTime = formatTime(subtitle.startTimestamp!)
-            let endTime = formatTime(subtitle.endTimestamp!)
+            let endTime = formatTime(endTimestamp)
 
             srtContent += "\(subtitle.index!)\n"
             srtContent += "\(startTime) --> \(endTime)\n"

diff --git a/Sources/macSubtitleOCR/Subtitles/Subtitle.swift b/Sources/macSubtitleOCR/Subtitles/Subtitle.swift
@@ -9,7 +9,7 @@
 import CoreGraphics
 import Foundation
 
-class Subtitle {
+struct Subtitle {
     var index: Int?
     var text: String?
     var startTimestamp: TimeInterval?
@@ -48,7 +48,7 @@ class Subtitle {
     // MARK: - Functions
 
     // Converts the RGBA data to a CGImage
-    func createImage(_ invert: Bool) -> CGImage? {
+    mutating func createImage(_ invert: Bool) -> CGImage? {
         // Convert the image data to RGBA format using the palette
         let rgbaData = imageDataToRGBA()
 
@@ -81,7 +81,7 @@ class Subtitle {
     // MARK: - Methods
 
     // Converts the image data to RGBA format using the palette
-    private func imageDataToRGBA() -> Data {
+    private mutating func imageDataToRGBA() -> Data {
         let bytesPerPixel = 4
         imageHeight = imageData!.count / imageWidth!
         var rgbaData = Data(capacity: imageWidth! * imageHeight! * bytesPerPixel)

diff --git a/Sources/macSubtitleOCR/Subtitles/VobSub/VobSubParser.swift b/Sources/macSubtitleOCR/Subtitles/VobSub/VobSubParser.swift
@@ -29,7 +29,7 @@ struct VobSubParser {
 
     // MARK: - Methods
 
-    func readSubFrame(subFile: FileHandle, offset: UInt64, nextOffset: UInt64, idxPalette _: [UInt8]) {
+    mutating func readSubFrame(subFile: FileHandle, offset: UInt64, nextOffset: UInt64, idxPalette _: [UInt8]) {
         var firstPacketFound = false
         var controlOffset: Int?
         var controlSize: Int?
@@ -74,7 +74,6 @@ struct VobSubParser {
                 presentationTimestamp += UInt64(ptsData[1]) << 22
                 presentationTimestamp += UInt64(ptsData[0] & 0x0E) << 29
                 subtitle.startTimestamp = TimeInterval(presentationTimestamp) / 90 / 1000
-                logger.debug("Got \(subtitle.startTimestamp!) as timestamp")
             }
 
             subFile.readData(ofLength: 1) // Stream ID
@@ -115,7 +114,7 @@ struct VobSubParser {
         parseCommandHeader(controlHeader, offset: relativeControlOffset)
     }
 
-    private func parseCommandHeader(_ header: Data, offset: Int) {
+    private mutating func parseCommandHeader(_ header: Data, offset: Int) {
         let relativeEndTimestamp = TimeInterval(Int(header.value(ofType: UInt16.self)!)) * 1024 / 90000 / fps
         let endOfControl = Int(header.value(ofType: UInt16.self)!) - 4 - offset
         subtitle.endTimestamp = subtitle.startTimestamp! + relativeEndTimestamp
@@ -160,7 +159,6 @@ struct VobSubParser {
                 index += 1
                 subtitle.imageAlpha![1] = byte >> 4
                 subtitle.imageAlpha![0] = byte & 0x0F
-                logger.debug("Alpha: \(subtitle.imageAlpha!)")
             case 5:
                 if subtitle.imageXOffset != nil || subtitle.imageYOffset != nil {
                     break // Don't overwrite the offsets if they're already set, only happens in bad files
@@ -173,20 +171,17 @@ struct VobSubParser {
                 subtitle.imageHeight = (Int(header[index + 1] & 0x0F) << 8 | Int(header[index + 2])) - subtitle
                     .imageYOffset! + 1
                 index += 3
-                logger.debug("Image size: \(subtitle.imageWidth!)x\(subtitle.imageHeight!)")
-                logger.debug("X Offset: \(subtitle.imageXOffset!), Y Offset: \(subtitle.imageYOffset!)")
             case 6:
                 subtitle.evenOffset = Int(header.value(ofType: UInt16.self, at: index)! - 4)
                 subtitle.oddOffset = Int(header.value(ofType: UInt16.self, at: index + 2)! - 4)
                 index += 4
-                logger.debug("Even offset: \(subtitle.evenOffset!), Odd offset: \(subtitle.oddOffset!)")
             default:
                 break
             }
         }
     }
 
-    private func decodePalette() {
+    private mutating func decodePalette() {
         var palette = [UInt8](repeating: 0, count: 4 * 4)
 
         for i in 0 ..< 4 {
@@ -200,7 +195,7 @@ struct VobSubParser {
         subtitle.imagePalette = palette
     }
 
-    private func decodeImage() {
+    private mutating func decodeImage() {
         var rleData = RLEData(
             data: subtitle.imageData ?? Data(),
             width: subtitle.imageWidth ?? 0,