Skip to content

Commit

Permalink
Refractor SRT into object
Browse files Browse the repository at this point in the history 
Signed-off-by: Ethan Dye <[email protected]>
  • Loading branch information
@ecdye
ecdye committed Sep 19, 2024
1 parent 11f96bc commit 4b9e064
Show file tree
Hide file tree
Showing 6 changed files with 106 additions and 137 deletions.
80 changes: 40 additions & 40 deletions Sources/macSubtitleOCR/PGS/PGS.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand All @@ -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,
Expand All @@ -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
Expand Down
91 changes: 20 additions & 71 deletions Sources/macSubtitleOCR/SRT/SRT.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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,
Expand All @@ -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
}
15 changes: 15 additions & 0 deletions Sources/macSubtitleOCR/SRT/SRTError.swift
View file
Original file line number Diff line number Diff line change
@@ -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
}
4 changes: 2 additions & 2 deletions Sources/macSubtitleOCR/SRT/SrtSubtitle.swift → Sources/macSubtitleOCR/SRT/SRTSubtitle.swift
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
//
// SrtSubtitle.swift
// SRTSubtitle.swift
// macSubtitleOCR
//
// Created by Ethan Dye on 9/16/24.
Expand All @@ -8,7 +8,7 @@

import Foundation

public struct SrtSubtitle {
public struct SRTSubtitle {
public var index: Int
public var startTime: TimeInterval
public var endTime: TimeInterval
Expand Down
51 changes: 27 additions & 24 deletions Sources/macSubtitleOCR/macSubtitleOCR.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -75,35 +74,32 @@ 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
}

// Save subtitle image as PNG if imageDirectory is provided
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
Expand All @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
}
}
}
2 changes: 2 additions & 0 deletions Sources/macSubtitleOCR/macSubtitleOCRError.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
public enum macSubtitleOCRError: Error {
case invalidFormat
case fileReadError
case fileCreationError
case fileWriteError
case unsupportedFormat
case invalidFile
}

0 comments on commit 4b9e064

Please sign in to comment.