Refactor and improve structure

Signed-off-by: Ethan Dye <[email protected]>

Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift

@@ -126,7 +126,7 @@ struct PGS {
     }
 
     private func getSegmentTimestamp(from pointer: UnsafeRawBufferPointer, offset: Int) -> TimeInterval {
-        TimeInterval(pointer.loadUnaligned(fromByteOffset: offset + 2, as: UInt32.self).bigEndian) / 90000 // 90 kHz clock
+        TimeInterval(pointer.loadUnaligned(fromByteOffset: offset + 2, as: UInt32.self).bigEndian) / 90000
     }
 
     private func getSegmentLength(from pointer: UnsafeRawBufferPointer, offset: Int) -> Int {

Sources/macSubtitleOCR/Subtitles/Subtitle.swift

@@ -95,7 +95,7 @@ class Subtitle: @unchecked Sendable {
                     continue
                 }
 
-                let paletteOffset = colorIndex * 4
+                let paletteOffset = colorIndex * bytesPerPixel
                 rgbaData.append(contentsOf: [
                     imagePalette![paletteOffset],
                     imagePalette![paletteOffset + 1],
@@ -179,38 +179,35 @@ class Subtitle: @unchecked Sendable {
         return extendedImage.cropping(to: croppedRect)
     }
 
-    func invertColors(of image: CGImage) -> CGImage? {
-        // Get the width, height, and color space of the image
-        let width = image.width
-        let height = image.height
+    private func invertColors(of image: CGImage) -> CGImage? {
         let colorSpace = CGColorSpaceCreateDeviceRGB()
 
         // Create a context with the same dimensions as the image
         guard let context = CGContext(
             data: nil,
-            width: width,
-            height: height,
+            width: image.width,
+            height: image.height,
             bitsPerComponent: 8,
-            bytesPerRow: width * 4,
+            bytesPerRow: image.width * 4,
             space: colorSpace,
             bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) else {
             return nil
         }
 
         // Draw the image into the context
-        context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))
+        context.draw(image, in: CGRect(x: 0, y: 0, width: image.width, height: image.height))
 
         // Get the pixel data from the context
         guard let pixelBuffer = context.data else {
             return nil
         }
 
-        let pixelData = pixelBuffer.bindMemory(to: UInt8.self, capacity: width * height * 4)
+        let pixelData = pixelBuffer.bindMemory(to: UInt8.self, capacity: image.width * image.height * 4)
 
         // Iterate through the pixel data and invert the colors
-        for y in 0 ..< height {
-            for x in 0 ..< width {
-                let pixelIndex = (y * width + x) * 4
+        for y in 0 ..< image.height {
+            for x in 0 ..< image.width {
+                let pixelIndex = (y * image.width + x) * 4
                 pixelData[pixelIndex] = 255 - pixelData[pixelIndex] // Red
                 pixelData[pixelIndex + 1] = 255 - pixelData[pixelIndex + 1] // Green
                 pixelData[pixelIndex + 2] = 255 - pixelData[pixelIndex + 2] // Blue

Sources/macSubtitleOCR/SubtitleProcessor.swift → Sources/macSubtitleOCR/Subtitles/SubtitleProcessor.swift

@@ -12,7 +12,6 @@ import os
 import UniformTypeIdentifiers
 import Vision
 
-private let logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category: "SubtitleProcessor")
 
 actor SubtitleAccumulator {
     var subtitles: [Subtitle] = []
@@ -47,16 +46,31 @@ actor AsyncSemaphore {
 }
 
 struct SubtitleProcessor {
-    let subtitles: [Subtitle]
-    let trackNumber: Int
-    let invert: Bool
-    let saveImages: Bool
-    let language: String
-    let fastMode: Bool
-    let disableLanguageCorrection: Bool
-    let forceOldAPI: Bool
-    let outputDirectory: String
-    let maxConcurrentTasks: Int
+    private let subtitles: [Subtitle]
+    private let trackNumber: Int
+    private let invert: Bool
+    private let saveImages: Bool
+    private let language: String
+    private let fastMode: Bool
+    private let disableLanguageCorrection: Bool
+    private let forceOldAPI: Bool
+    private let outputDirectory: String
+    private let maxConcurrentTasks: Int
+    private let logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category: "SubtitleProcessor")
+
+    init(subtitles: [Subtitle], trackNumber: Int, invert: Bool, saveImages: Bool, language: String, fastMode: Bool, disableLanguageCorrection: Bool,
+         forceOldAPI: Bool, outputDirectory: String, maxConcurrentTasks: Int) {
+        self.subtitles = subtitles
+        self.trackNumber = trackNumber
+        self.invert = invert
+        self.saveImages = saveImages
+        self.language = language
+        self.fastMode = fastMode
+        self.disableLanguageCorrection = disableLanguageCorrection
+        self.forceOldAPI = forceOldAPI
+        self.outputDirectory = outputDirectory
+        self.maxConcurrentTasks = maxConcurrentTasks
+    }
 
     func process() async throws -> macSubtitleOCRResult {
         let accumulator = SubtitleAccumulator()

Sources/macSubtitleOCR/macSubtitleOCR.swift

@@ -69,7 +69,7 @@ struct macSubtitleOCR: AsyncParsableCommand {
     // MARK: - Entrypoint
 
     func run() async throws {
-        let fileHandler = FileHandler(outputDirectory: outputDirectory)
+        let fileHandler = macSubtitleOCRFileHandler(outputDirectory: outputDirectory)
         let results = try await processInput()
         try await saveResults(fileHandler: fileHandler, results: results)
     }
@@ -162,7 +162,7 @@ struct macSubtitleOCR: AsyncParsableCommand {
             maxConcurrentTasks: maxThreads)
     }
 
-    private func saveResults(fileHandler: FileHandler, results: [macSubtitleOCRResult]) async throws {
+    private func saveResults(fileHandler: macSubtitleOCRFileHandler, results: [macSubtitleOCRResult]) async throws {
         for result in results {
             autoreleasepool {
                 try? fileHandler.saveSRTFile(for: result)

Sources/macSubtitleOCR/FileHandler.swift → Sources/macSubtitleOCR/macSubtitleOCRFileHandler.swift

@@ -8,7 +8,7 @@
 
 import Foundation
 
-struct FileHandler {
+struct macSubtitleOCRFileHandler {
     private let outputDirectory: URL
 
     init(outputDirectory: String) {