diff --git a/Sources/macSubtitleOCR/PGS/PGS.swift b/Sources/macSubtitleOCR/PGS/PGS.swift index 1acfc5d..4c5bf17 100644 --- a/Sources/macSubtitleOCR/PGS/PGS.swift +++ b/Sources/macSubtitleOCR/PGS/PGS.swift @@ -9,44 +9,30 @@ import CoreGraphics import Foundation import ImageIO -import UniformTypeIdentifiers public class PGS { - // MARK: - Functions - - // Parses a `.sup` file and returns an array of `PGSSubtitle` objects - public func parseSupFile(fromFileAt url: URL) throws -> [PGSSubtitle] { - let fileHandle = try FileHandle(forReadingFrom: url) - defer { fileHandle.closeFile() } + // MARK: - Properties - var subtitles = [PGSSubtitle]() - let fileLength = try fileHandle.seekToEnd() - fileHandle.seek(toFileOffset: 0) // Ensure the file handle is at the start - var headerData = fileHandle.readData(ofLength: 13) + private var subtitles = [PGSSubtitle]() - while fileHandle.offsetInFile < fileLength { - guard var subtitle = try parseNextSubtitle(fileHandle: fileHandle, headerData: &headerData) - else { - headerData = fileHandle.readData(ofLength: 13) - continue - } + // MARK: - Lifecycle - // Find the next timestamp to use as our end timestamp - while subtitle.endTimestamp <= subtitle.timestamp { - headerData = fileHandle.readData(ofLength: 13) - subtitle.endTimestamp = parseTimestamp(headerData) - } + init(_ url: URL) throws { + try parseSupFile(fromFileAt: url) + } - subtitles.append(subtitle) - } + // MARK: - Getters - return subtitles + public func getSubtitles() -> [PGSSubtitle] { + subtitles } + // MARK: - Functions + // Converts the RGBA data to a CGImage - public func createImage(from subtitle: inout PGSSubtitle) -> CGImage? { + public func createImage(index: Int) -> CGImage? { // Convert the image data to RGBA format using the palette - let rgbaData = imageDataToRGBA(&subtitle) + let rgbaData = imageDataToRGBA(&subtitles[index]) let bitmapInfo = CGBitmapInfo.byteOrder32Big .union(CGBitmapInfo(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue)) @@ -56,11 +42,11 @@ public class PGS { return nil } - return CGImage(width: subtitle.imageWidth, - height: subtitle.imageHeight, + return CGImage(width: subtitles[index].imageWidth, + height: subtitles[index].imageHeight, bitsPerComponent: 8, bitsPerPixel: 32, - bytesPerRow: subtitle.imageWidth * 4, // 4 bytes per pixel (RGBA) + bytesPerRow: subtitles[index].imageWidth * 4, // 4 bytes per pixel (RGBA) space: colorSpace, bitmapInfo: bitmapInfo, provider: provider, @@ -69,20 +55,34 @@ public class PGS { intent: .defaultIntent) } - public func saveImageAsPNG(image: CGImage, outputPath: URL) throws { - guard let destination = CGImageDestinationCreateWithURL(outputPath as CFURL, UTType.png.identifier as CFString, 1, nil) - else { - throw macSubtitleOCRError.fileReadError - } - CGImageDestinationAddImage(destination, image, nil) + // MARK: - Methods + + // Parses a `.sup` file and populates the subtitles array + private func parseSupFile(fromFileAt url: URL) throws { + let fileHandle = try FileHandle(forReadingFrom: url) + defer { fileHandle.closeFile() } + + let fileLength = try fileHandle.seekToEnd() + fileHandle.seek(toFileOffset: 0) // Ensure the file handle is at the start + var headerData = fileHandle.readData(ofLength: 13) + + while fileHandle.offsetInFile < fileLength { + guard var subtitle = try parseNextSubtitle(fileHandle: fileHandle, headerData: &headerData) + else { + headerData = fileHandle.readData(ofLength: 13) + continue + } + + // Find the next timestamp to use as our end timestamp + while subtitle.endTimestamp <= subtitle.timestamp { + headerData = fileHandle.readData(ofLength: 13) + subtitle.endTimestamp = parseTimestamp(headerData) + } - if !CGImageDestinationFinalize(destination) { - throw macSubtitleOCRError.fileReadError + subtitles.append(subtitle) } } - // MARK: - Methods - private func parseTimestamp(_ data: Data) -> TimeInterval { let pts = (Int(data[2]) << 24 | Int(data[3]) << 16 | Int(data[4]) << 8 | Int(data[5])) return TimeInterval(pts) / 90000.0 // 90 kHz clock diff --git a/Sources/macSubtitleOCR/SRT/SRT.swift b/Sources/macSubtitleOCR/SRT/SRT.swift index 97b322d..2a93373 100644 --- a/Sources/macSubtitleOCR/SRT/SRT.swift +++ b/Sources/macSubtitleOCR/SRT/SRT.swift @@ -9,66 +9,36 @@ import Foundation public class SRT { - // MARK: - Functions + // MARK: - Properties - // Decodes subtitles from a string containing the SRT content - public func decode(from content: String) throws -> [SrtSubtitle] { - var subtitles = [SrtSubtitle]() - - // Split the content by subtitle blocks - let blocks = content.components(separatedBy: "\n\n") - - for block in blocks { - let lines = block.components(separatedBy: .newlines).filter { !$0.isEmpty } - - guard lines.count >= 2 else { - continue - } - - // Parse index - guard let index = Int(lines[0]) else { - throw SRTError.invalidFormat - } - - // Parse times - let timeComponents = lines[1].components(separatedBy: " --> ") - guard timeComponents.count == 2, - let startTime = parseTime(timeComponents[0]), - let endTime = parseTime(timeComponents[1]) - else { - throw SRTError.invalidTimeFormat - } - - // Combine remaining lines as the subtitle text - var text = "" - if lines.count <= 3 { - text = lines[2...].joined(separator: "\n") - } - - // Create and append the subtitle - let subtitle = SrtSubtitle(index: index, startTime: startTime, endTime: endTime, text: text) - subtitles.append(subtitle) - } + private var subtitles: [SRTSubtitle] = [] + + // MARK: - Getters / Setters + + public func getSubtitles() -> [SRTSubtitle] { + subtitles + } - return subtitles + public func appendSubtitle(_ subtitle: SRTSubtitle) { + subtitles.append(subtitle) } - // Decodes subtitles from an SRT file at the given URL - public func decode(fromFileAt url: URL) throws -> [SrtSubtitle] { + // MARK: - Functions + + // Writes the SRT object to the file at the given URL + public func write(toFileAt url: URL) throws { + let srtContent = encode() do { - // Read the file content into a string - let content = try String(contentsOf: url, encoding: .utf8) - // Decode the content into subtitles - return try decode(from: content) + try srtContent.write(to: url, atomically: true, encoding: .utf8) } catch { - throw SRTError.fileReadError + throw SRTError.fileWriteError } } - // MARK: - Re-Encoding + // MARK: - Methods - // Re-encodes an array of `Subtitle` objects into SRT format and returns it as a string - public func encode(subtitles: [SrtSubtitle]) -> String { + // Encodes the SRT object into SRT format and returns it as a string + private func encode() -> String { var srtContent = "" for subtitle in subtitles { @@ -83,19 +53,6 @@ public class SRT { return srtContent } - // Re-encodes an array of `Subtitle` objects into SRT format and writes it to a file at the given URL - public func encode(subtitles: [SrtSubtitle], toFileAt url: URL) throws { - let srtContent = encode(subtitles: subtitles) - - do { - try srtContent.write(to: url, atomically: true, encoding: .utf8) - } catch { - throw SRTError.fileWriteError - } - } - - // MARK: - Helper Methods - private func parseTime(_ timeString: String) -> TimeInterval? { let components = timeString.components(separatedBy: [":", ","]) guard components.count == 4, @@ -119,11 +76,3 @@ public class SRT { return String(format: "%02d:%02d:%02d,%03d", hours, minutes, seconds, milliseconds) } } - -public enum SRTError: Error { - case invalidFormat - case invalidTimeFormat - case fileNotFound - case fileReadError - case fileWriteError -} diff --git a/Sources/macSubtitleOCR/SRT/SRTError.swift b/Sources/macSubtitleOCR/SRT/SRTError.swift new file mode 100644 index 0000000..213173a --- /dev/null +++ b/Sources/macSubtitleOCR/SRT/SRTError.swift @@ -0,0 +1,15 @@ +// +// SRTError.swift +// macSubtitleOCR +// +// Created by Ethan Dye on 9/19/24. +// Copyright © 2024 Ethan Dye. All rights reserved. +// + +public enum SRTError: Error { + case invalidFormat + case invalidTimeFormat + case fileNotFound + case fileReadError + case fileWriteError +} diff --git a/Sources/macSubtitleOCR/SRT/SrtSubtitle.swift b/Sources/macSubtitleOCR/SRT/SRTSubtitle.swift similarity index 84% rename from Sources/macSubtitleOCR/SRT/SrtSubtitle.swift rename to Sources/macSubtitleOCR/SRT/SRTSubtitle.swift index 5773507..dc14253 100644 --- a/Sources/macSubtitleOCR/SRT/SrtSubtitle.swift +++ b/Sources/macSubtitleOCR/SRT/SRTSubtitle.swift @@ -1,5 +1,5 @@ // -// SrtSubtitle.swift +// SRTSubtitle.swift // macSubtitleOCR // // Created by Ethan Dye on 9/16/24. @@ -8,7 +8,7 @@ import Foundation -public struct SrtSubtitle { +public struct SRTSubtitle { public var index: Int public var startTime: TimeInterval public var endTime: TimeInterval diff --git a/Sources/macSubtitleOCR/macSubtitleOCR.swift b/Sources/macSubtitleOCR/macSubtitleOCR.swift index f2eecbd..83a244f 100644 --- a/Sources/macSubtitleOCR/macSubtitleOCR.swift +++ b/Sources/macSubtitleOCR/macSubtitleOCR.swift @@ -57,8 +57,7 @@ struct macSubtitleOCR: ParsableCommand { // Setup data variables var subIndex = 1 var jsonStream: [Any] = [] - var inSubStream: [PGSSubtitle] - var outSubStream: [SrtSubtitle] = [] + let srtStream = SRT() if sup.hasSuffix(".mkv") { let mkvParser = try MKVParser(filePath: sup) @@ -75,19 +74,18 @@ struct macSubtitleOCR: ParsableCommand { mkvParser.closeFile() } - // Initialize the decoder - let PGS = PGS() - inSubStream = try PGS.parseSupFile(fromFileAt: URL(fileURLWithPath: sup)) + // Open the PGS data stream + let PGS = try PGS(URL(fileURLWithPath: sup)) - for var subtitle in inSubStream { + for subtitle in PGS.getSubtitles() { if subtitle.imageWidth == 0, subtitle.imageHeight == 0 { logger.debug("Skipping subtitle index \(subIndex) with empty image data!") continue } - guard let subImage = PGS.createImage(from: &subtitle) + guard let subImage = PGS.createImage(index: subIndex - 1) else { - logger.info("Could not create image from decoded data for index \(subIndex)! Skipping...") + logger.info("Could not create image for index \(subIndex)! Skipping...") continue } @@ -95,15 +93,13 @@ struct macSubtitleOCR: ParsableCommand { if let imageDirectory { let outputDirectory = URL(fileURLWithPath: imageDirectory) do { - try manager.createDirectory(at: outputDirectory, - withIntermediateDirectories: false, - attributes: nil) + try manager.createDirectory(at: outputDirectory, withIntermediateDirectories: false, attributes: nil) } catch CocoaError.fileWriteFileExists { // Folder already existed } let pngPath = outputDirectory.appendingPathComponent("subtitle_\(subIndex).png") - try PGS.saveImageAsPNG(image: subImage, outputPath: pngPath) + try saveImageAsPNG(image: subImage, outputPath: pngPath) } // Perform text recognition @@ -120,9 +116,7 @@ struct macSubtitleOCR: ParsableCommand { let stringRange = string.startIndex ..< string.endIndex let boxObservation = try? candidate?.boundingBox(for: stringRange) let boundingBox = boxObservation?.boundingBox ?? .zero - let rect = VNImageRectForNormalizedRect(boundingBox, - subtitle.imageWidth, - subtitle.imageHeight) + let rect = VNImageRectForNormalizedRect(boundingBox, subtitle.imageWidth, subtitle.imageHeight) let line: [String: Any] = [ "text": string, @@ -149,12 +143,11 @@ struct macSubtitleOCR: ParsableCommand { jsonStream.append(subtitleData) - let newSubtitle = SrtSubtitle(index: subIndex, - startTime: subtitle.timestamp, - endTime: subtitle.endTimestamp, - text: subtitleText) - - outSubStream.append(newSubtitle) + // Append subtitle to SRT stream + srtStream.appendSubtitle(SRTSubtitle(index: subIndex, + startTime: subtitle.timestamp, + endTime: subtitle.endTimestamp, + text: subtitleText)) } request.recognitionLevel = recognitionLevel @@ -184,9 +177,7 @@ struct macSubtitleOCR: ParsableCommand { to: URL(fileURLWithPath: inFile).deletingPathExtension().appendingPathExtension("sup")) } - // Encode subtitles to SRT file - try SRT().encode(subtitles: outSubStream, - toFileAt: URL(fileURLWithPath: srt)) + try srtStream.write(toFileAt: URL(fileURLWithPath: srt)) } // MARK: - Methods @@ -206,4 +197,16 @@ struct macSubtitleOCR: ParsableCommand { VNRecognizeTextRequestRevision2 } } + + private func saveImageAsPNG(image: CGImage, outputPath: URL) throws { + guard let destination = CGImageDestinationCreateWithURL(outputPath as CFURL, UTType.png.identifier as CFString, 1, nil) + else { + throw macSubtitleOCRError.fileCreationError + } + CGImageDestinationAddImage(destination, image, nil) + + if !CGImageDestinationFinalize(destination) { + throw macSubtitleOCRError.fileWriteError + } + } } diff --git a/Sources/macSubtitleOCR/macSubtitleOCRError.swift b/Sources/macSubtitleOCR/macSubtitleOCRError.swift index 3298072..bf5ebca 100644 --- a/Sources/macSubtitleOCR/macSubtitleOCRError.swift +++ b/Sources/macSubtitleOCR/macSubtitleOCRError.swift @@ -9,6 +9,8 @@ public enum macSubtitleOCRError: Error { case invalidFormat case fileReadError + case fileCreationError + case fileWriteError case unsupportedFormat case invalidFile }