From c6e0975b556362a64201b44b75dccfbaf00f66fb Mon Sep 17 00:00:00 2001 From: Ethan Dye Date: Wed, 23 Oct 2024 01:11:07 -0600 Subject: [PATCH] Refactor and improve structure Signed-off-by: Ethan Dye --- .../macSubtitleOCR/Subtitles/PGS/PGS.swift | 2 +- .../macSubtitleOCR/Subtitles/Subtitle.swift | 23 ++++++------ .../{ => Subtitles}/SubtitleProcessor.swift | 36 +++++++++++++------ Sources/macSubtitleOCR/macSubtitleOCR.swift | 4 +-- ....swift => macSubtitleOCRFileHandler.swift} | 2 +- 5 files changed, 39 insertions(+), 28 deletions(-) rename Sources/macSubtitleOCR/{ => Subtitles}/SubtitleProcessor.swift (87%) rename Sources/macSubtitleOCR/{FileHandler.swift => macSubtitleOCRFileHandler.swift} (98%) diff --git a/Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift b/Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift index eca58d2..8329319 100644 --- a/Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift +++ b/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 { diff --git a/Sources/macSubtitleOCR/Subtitles/Subtitle.swift b/Sources/macSubtitleOCR/Subtitles/Subtitle.swift index db579d5..3ad95fd 100644 --- a/Sources/macSubtitleOCR/Subtitles/Subtitle.swift +++ b/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 diff --git a/Sources/macSubtitleOCR/SubtitleProcessor.swift b/Sources/macSubtitleOCR/Subtitles/SubtitleProcessor.swift similarity index 87% rename from Sources/macSubtitleOCR/SubtitleProcessor.swift rename to Sources/macSubtitleOCR/Subtitles/SubtitleProcessor.swift index 44b0989..5ca6905 100644 --- a/Sources/macSubtitleOCR/SubtitleProcessor.swift +++ b/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() diff --git a/Sources/macSubtitleOCR/macSubtitleOCR.swift b/Sources/macSubtitleOCR/macSubtitleOCR.swift index d16f44e..3873519 100644 --- a/Sources/macSubtitleOCR/macSubtitleOCR.swift +++ b/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) diff --git a/Sources/macSubtitleOCR/FileHandler.swift b/Sources/macSubtitleOCR/macSubtitleOCRFileHandler.swift similarity index 98% rename from Sources/macSubtitleOCR/FileHandler.swift rename to Sources/macSubtitleOCR/macSubtitleOCRFileHandler.swift index 7727e92..d303fdc 100644 --- a/Sources/macSubtitleOCR/FileHandler.swift +++ b/Sources/macSubtitleOCR/macSubtitleOCRFileHandler.swift @@ -8,7 +8,7 @@ import Foundation -struct FileHandler { +struct macSubtitleOCRFileHandler { private let outputDirectory: URL init(outputDirectory: String) {