Fix VobSub decoding

Signed-off-by: Ethan Dye <mrtops03@gmail.com>
ecdye · Oct 4, 2024 · bb48698 · bb48698
1 parent 327f57f
commit bb48698
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 10 deletions.
diff --git a/Sources/macSubtitleOCR/Subtitles/RLE/RLEData.swift b/Sources/macSubtitleOCR/Subtitles/RLE/RLEData.swift
@@ -100,7 +100,7 @@ struct RLEData {
             let color = UInt8(nibble & 0x03)
             var run = Int(nibble >> 2)
 
-            if image.count % width == 0, color != 0, run <= 15 {
+            if image.count % width == 0, color != 0, run == 15 {
                 i -= 5
                 currentNibbles = [nibbles[i], nibbles[i + 1]]
                 i += 2
@@ -112,6 +112,9 @@ struct RLEData {
                 run += width - x
                 x = 0
                 y += 1
+                if i % 2 != 0 {
+                    _ = getNibble(currentNibbles: &currentNibbles, nibbles: nibbles, i: &i)
+                }
             }
 
             image.append(contentsOf: repeatElement(color, count: run))

diff --git a/Sources/macSubtitleOCR/Subtitles/VobSub/VobSub.swift b/Sources/macSubtitleOCR/Subtitles/VobSub/VobSub.swift
@@ -44,7 +44,6 @@ struct VobSub {
                 nextOffset: nextOffset,
                 idxPalette: idx.palette).subtitle
             logger.debug("Found image at offset \(offset) with timestamp \(timestamp)")
-            logger.debug("Image size: \(subtitle.imageWidth!) x \(subtitle.imageHeight!)")
             subtitles.append(subtitle)
         }
     }

diff --git a/Sources/macSubtitleOCR/Subtitles/VobSub/VobSubParser.swift b/Sources/macSubtitleOCR/Subtitles/VobSub/VobSubParser.swift
@@ -43,7 +43,7 @@ struct VobSubParser {
             guard subFile.readData(ofLength: 4).value(ofType: UInt32.self, at: 0) == MPEG2PacketType.psPacket else {
                 fatalError("Error: Failed to find PS packet at offset \(subFile.offsetInFile)")
             }
-            logger.debug("Found PS packet at offset \(subFile.offsetInFile)")
+            logger.debug("Found PS packet at offset \(startOffset)")
 
             subFile.readData(ofLength: 6) // System clock reference
             subFile.readData(ofLength: 3) // Multiplexer rate
@@ -56,7 +56,7 @@ struct VobSubParser {
             guard subFile.readData(ofLength: 4).value(ofType: UInt32.self, at: 0) == MPEG2PacketType.pesPacket else {
                 fatalError("Error: Failed to find PES packet at offset \(subFile.offsetInFile)")
             }
-            logger.debug("Found PES packet at offset \(subFile.offsetInFile)")
+            logger.debug("Found PES packet at offset \(subFile.offsetInFile - 4)")
 
             let pesLength = Int(subFile.readData(ofLength: 2).value(ofType: UInt16.self, at: 0) ?? 0)
             if pesLength == 0 {
@@ -124,15 +124,14 @@ struct VobSubParser {
         var endOfControl = Int(controlHeader.value(ofType: UInt16.self, at: index)!) - relativeControlOffset - 4
         if endOfControl < 0 || endOfControl > controlSize! {
             logger.warning("Invalid control header size \(endOfControl). Setting to \(controlSize!)")
-            endOfControl = Int(controlSize!)
+            endOfControl = Int(controlSize! - 1)
         }
         index += 2
 
-        var alphaSum = 0
-
         // This is a hacky way to get the end timestamp, but it works somewhat accurately
         let relativeEndTimestamp = controlHeader.value(ofType: UInt16.self, at: endOfControl - 1)! << 10
         subtitle.endTimestamp = subtitle.startTimestamp! + TimeInterval(relativeEndTimestamp) / 90.0 - 9
+        logger.debug("relativeEndTimestamp: \(relativeEndTimestamp), endTimestamp: \(subtitle.endTimestamp!)")
 
         while index < endOfControl {
             let command = controlHeader[index]
@@ -163,16 +162,18 @@ struct VobSubParser {
                 index += 1
                 if subtitle.imageAlpha == nil {
                     subtitle.imageAlpha = [UInt8](repeating: 0, count: 4)
+                } else {
+                    // If the alpha is already set, don't overwrite it. This typically happens when fade in/out is used.
+                    index += 1
+                    break
                 }
                 subtitle.imageAlpha![3] = byte >> 4
                 subtitle.imageAlpha![2] = byte & 0x0F
                 byte = controlHeader[index]
                 index += 1
                 subtitle.imageAlpha![1] = byte >> 4
                 subtitle.imageAlpha![0] = byte & 0x0F
-                for i in 0 ..< 4 {
-                    alphaSum += Int(subtitle.imageAlpha![i])
-                }
+                logger.debug("Alpha: \(subtitle.imageAlpha!)")
             case 5:
                 subtitle.imageXOffset = Int(controlHeader[index]) << 4 | Int(controlHeader[index + 1] >> 4)
                 subtitle.imageWidth = (Int(controlHeader[index + 1] & 0x0F) << 8 | Int(controlHeader[index + 2])) - subtitle
@@ -182,6 +183,7 @@ struct VobSubParser {
                 subtitle.imageHeight = (Int(controlHeader[index + 1] & 0x0F) << 8 | Int(controlHeader[index + 2])) - subtitle
                     .imageYOffset! + 1
                 index += 3
+                logger.debug("Image size: \(subtitle.imageWidth!)x\(subtitle.imageHeight!)")
             default:
                 break
             }