Make decoder functional (somewhat)

Signed-off-by: Ethan Dye <mrtops03@gmail.com>

Sources/macSubtitleOCR/Extensions/CollectionExtensions.swift

@@ -0,0 +1,18 @@
+//
+// CollectionExtensions.swift
+// macSubtitleOCR
+//
+// Created by Ethan Dye on 9/24/24.
+// Copyright © 2024 Ethan Dye. All rights reserved.
+//
+
+extension Collection {
+    func unfoldSubSequences(limitedTo maxLength: Int) -> UnfoldSequence<SubSequence, Index> {
+        sequence(state: startIndex) { start in
+            guard start < self.endIndex else { return nil }
+            let end = self.index(start, offsetBy: maxLength, limitedBy: self.endIndex) ?? self.endIndex
+            defer { start = end }
+            return self[start ..< end]
+        }
+    }
+}

Sources/macSubtitleOCR/Extensions/StringProtocolExtensions.swift

@@ -0,0 +1,15 @@
+//
+// StringProtocolExtensions.swift
+// macSubtitleOCR
+//
+// Created by Ethan Dye on 9/24/24.
+// Copyright © 2024 Ethan Dye. All rights reserved.
+//
+
+import Foundation
+
+extension StringProtocol {
+    var byte: UInt8? { UInt8(self, radix: 16) }
+    var hexToBytes: [UInt8] { unfoldSubSequences(limitedTo: 2).compactMap(\.byte) }
+    var hexToData: Data { .init(hexToBytes) }
+}

Sources/macSubtitleOCR/SUB/SUB.swift

@@ -279,64 +279,81 @@ func readSubFrame(pic: inout SubPictureDVD, subFile: FileHandle, offset: UInt64,
 
 import CoreGraphics
 
-func decodeRLEToCGImage(subPicture: SubPictureDVD, fileData: Data) -> CGImage? {
-    let width = subPicture.imageWidth
-    let height = subPicture.imageHeight
-    let evenOffset = subPicture.evenOffset
-    let oddOffset = subPicture.oddOffset
-    let rleSize = subPicture.rleSize
-
-    guard width > 0, height > 0, rleSize > 0 else {
-        print("Invalid image dimensions or RLE size")
-        return nil
+func decodePalette(subPicture: SubPictureDVD, masterPalette: [UInt8]) -> [UInt8] {
+    var palette = [UInt8](repeating: 0, count: 4 * 4)
+
+    for i in 0 ..< 4 {
+        let index = subPicture.palette[i]
+        palette[4 * i] = masterPalette[3 * index]
+        palette[4 * i + 1] = masterPalette[3 * index + 1]
+        palette[4 * i + 2] = masterPalette[3 * index + 2]
+        // palette[4 * i + 3] = UInt8(subPicture.alpha[i])
+        if subPicture.alpha[i] >= 0 {
+            palette[4 * i + 3] = 255
+        }
     }
 
-    var decodedPixels = [UInt8](repeating: 0, count: width * height)
-
-    func decodeLine(src: Data, offset: Int, length: Int, target: inout [UInt8], targetOffset: Int, stride _: Int) {
-        let end = offset + length
-        var x = 0
-        var pixelIndex = targetOffset
-
-        var srcOffset = offset
+    return palette
+}
 
-        while srcOffset < end, x < width {
-            let byte = src[srcOffset]
-            srcOffset += 1
+// Converts the image data to RGBA format using the palette
+private func imageData(_ image: Data, width: Int, height: Int, palette: [UInt8]) -> Data {
+    let bytesPerPixel = 4
+    let numColors = 16 // There are only 256 possible palette entries in a PGS Subtitle
+    var rgbaData = Data(capacity: width * height * bytesPerPixel)
 
-            let len = Int(byte >> 2)
-            let color = byte & 0x03
+    for y in 0 ..< height {
+        for x in 0 ..< width {
+            let index = Int(y) * width + Int(x)
+            let colorIndex = Int(image[index])
 
-            for _ in 0 ..< len {
-                if x >= width { break }
-                target[pixelIndex + x] = color * 85 // 0, 85, 170, 255 (grayscale)
-                x += 1
+            guard colorIndex < numColors else {
+                continue
             }
+
+            let paletteOffset = colorIndex * 4
+            rgbaData.append(contentsOf: [
+                palette[paletteOffset],
+                palette[paletteOffset + 1],
+                palette[paletteOffset + 2],
+                palette[paletteOffset + 3],
+            ])
         }
     }
 
-    // Decode even lines
-    decodeLine(src: fileData, offset: evenOffset, length: rleSize / 2, target: &decodedPixels, targetOffset: 0, stride: width)
+    return rgbaData
+}
+
+func decodeRLEToCGImage(subPicture: SubPictureDVD, fileData: FileHandle, palette: [UInt8]) throws -> CGImage? {
+    let width = subPicture.imageWidth
+    let height = subPicture.imageHeight
+    let rleSize = subPicture.rleSize
+
+    guard width > 0, height > 0, rleSize > 0 else {
+        print("Invalid image dimensions or RLE size")
+        return nil
+    }
+
+    let decodedPixels = try decodeImage(pic: subPicture, fBuf: fileData)
 
-    // Decode odd lines
-    decodeLine(src: fileData, offset: oddOffset, length: rleSize / 2, target: &decodedPixels, targetOffset: width, stride: width)
+    let imageData = imageData(decodedPixels, width: width, height: height, palette: decodePalette(subPicture: subPicture, masterPalette: palette))
 
     // Create a grayscale CGImage
-    let colorSpace = CGColorSpaceCreateDeviceGray()
-    let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.none.rawValue)
+    let bitmapInfo = CGBitmapInfo.byteOrder32Big
+        .union(CGBitmapInfo(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue))
+    let colorSpace = CGColorSpaceCreateDeviceRGB()
     let bitsPerComponent = 8
-    let bytesPerRow = width
+    let bytesPerRow = width * 4
 
-    guard let provider = CGDataProvider(data: NSData(bytes: &decodedPixels, length: decodedPixels.count * MemoryLayout<UInt8>.size)) else {
-        print("Failed to create CGDataProvider")
+    guard let provider = CGDataProvider(data: imageData as CFData) else {
         return nil
     }
 
     let cgImage = CGImage(
         width: width,
         height: height,
         bitsPerComponent: bitsPerComponent,
-        bitsPerPixel: bitsPerComponent,
+        bitsPerPixel: bitsPerComponent * 4,
         bytesPerRow: bytesPerRow,
         space: colorSpace,
         bitmapInfo: bitmapInfo,

Sources/macSubtitleOCR/SUB/SUBFileHandler.swift

@@ -27,6 +27,7 @@ class VobSubDecoder {
         var yOfs: Int = 0
     }
 
+    var idxPalette = [UInt8]()
     var decodedSubtitles = [DecodedSubtitle]()
 
     func readVobSubFiles(subFileURL: URL, idxFilePath: String) -> (FileHandle, String)? {
@@ -57,6 +58,13 @@ class VobSubDecoder {
 
         for line in lines {
             print("Parsing line: \(line)")
+            if line.starts(with: "palette:") {
+                let entries = line.split(separator: ", ").map { String($0) }
+                for entry in entries {
+                    print(entry.hexToBytes)
+                    idxPalette.append(contentsOf: entry.hexToBytes)
+                }
+            }
             if line.starts(with: "timestamp:") {
                 let timestampMatch = timestampRegex.firstMatch(in: String(line), options: [], range: NSRange(location: 0, length: line.count))
                 let offsetMatch = offsetRegex.firstMatch(in: String(line), options: [], range: NSRange(location: 0, length: line.count))
@@ -92,7 +100,7 @@ class VobSubDecoder {
             let imageData = subFile.readData(ofLength: pic.rleSize)
             print("Found image at offset \(subtitle.offset) with timestamp \(subtitle.timestamp)")
             print("Image size: \(pic.imageWidth) x \(pic.imageHeight)")
-            let cgImage = decodeRLEToCGImage(subPicture: pic, fileData: imageData)
+            let cgImage = try decodeRLEToCGImage(subPicture: pic, fileData: subFile, palette: idxPalette)
             decodedSubtitles.append(DecodedSubtitle(image: cgImage, width: pic.imageWidth, height: pic.imageWidth))
             idx += 1
         }