|
// |
|
// SEINormalizerConfig.swift |
|
// VTBugApp |
|
// |
|
// Created by jason van cleave on 9/28/25. |
|
// |
|
|
|
import CoreMedia |
|
import Foundation |
|
|
|
/// Configuration for SEINormalizer. |
|
public struct SEINormalizerConfig: Sendable |
|
{ |
|
/// Drop all SEI (type 6) NAL units entirely. |
|
public var dropAllSEI: Bool = false |
|
/// If non-empty, only keep SEI payload types in this set (e.g., [1] for pic_timing). |
|
/// Ignored when `dropAllSEI` is true. |
|
public var keepPayloadTypes: Set<Int> = [] |
|
/// Drop all user_data_unregistered payloads (payloadType = 5). |
|
public var dropAllUserData: Bool = true |
|
/// If not dropping all user data, drop only these UUIDs (exactly 16 bytes). |
|
public var dropUserDataUUIDs: Set<Data> = [] |
|
/// If true, keep only the first occurrence per UUID and drop repeats. |
|
public var dedupeUserDataUUIDs: Bool = true |
|
|
|
public init() {} |
|
} |
|
|
|
/// Normalizes H.264 SEI (type 6) NAL units in AVCC (length-prefixed) form. |
|
/// Call `normalize(seiNAL:pts:log:)` with a single SEI NAL (including its 1-byte header). |
|
public final class SEINormalizer |
|
{ |
|
private let cfg: SEINormalizerConfig |
|
private var seenUUIDs = Set<Data>() // for dedupe of user_data_unregistered |
|
|
|
public init(_ cfg: SEINormalizerConfig) |
|
{ |
|
self.cfg = cfg |
|
} |
|
|
|
/// Normalize a full SEI NAL (including its 1-byte NAL header). |
|
/// - Returns: A normalized SEI NAL (still includes 1-byte header) or `nil` to drop it. |
|
public func normalize(seiNAL fullNal: Data, |
|
pts: CMTime? = nil, |
|
log: ((String) -> Void)? = nil) -> Data? |
|
{ |
|
if cfg.dropAllSEI { return nil } |
|
guard fullNal.count >= 1 else { return nil } |
|
|
|
let nalHeader = fullNal[0] // forbidden_zero_bit | nal_ref_idc | nal_unit_type (= 6) |
|
let ebsp = fullNal.dropFirst() |
|
let rbsp = rbspFromEbsp(Data(ebsp)) |
|
|
|
var kept: [(Int, Data)] = [] |
|
|
|
forEachSEIPayload(rbsp: rbsp) |
|
{ payloadType, payload in |
|
var keep = true |
|
|
|
// Allow-list by payload type if requested |
|
if !cfg.keepPayloadTypes.isEmpty && !cfg.keepPayloadTypes.contains(payloadType) |
|
{ |
|
keep = false |
|
} |
|
|
|
// Special handling for user_data_unregistered (payloadType 5) |
|
if payloadType == 5 |
|
{ |
|
if cfg.dropAllUserData { keep = false } |
|
else if payload.count >= 16 |
|
{ |
|
let uuid = Data(payload.prefix(16)) |
|
if cfg.dropUserDataUUIDs.contains(uuid) { keep = false } |
|
if cfg.dedupeUserDataUUIDs && !seenUUIDs.insert(uuid).inserted { keep = false } |
|
if let log, keep |
|
{ |
|
let ts = pts.map { String(format: "PTS=%.6f ", CMTimeGetSeconds($0)) } ?? "" |
|
log("\(ts)SEI keep user_data_unregistered uuid=\(uuidHex(uuid)) user_len=\(payload.count - 16)") |
|
} |
|
} |
|
else |
|
{ |
|
keep = false |
|
} |
|
} |
|
|
|
if keep { kept.append((payloadType, payload)) } |
|
} |
|
|
|
guard !kept.isEmpty else { return nil } |
|
|
|
// Rebuild RBSP: [payload headers + payload]* + rbsp_trailing_bits |
|
var newRBSP = Data() |
|
for (t, p) in kept |
|
{ |
|
writeSEIPayloadHeader(type: t, size: p.count, into: &newRBSP) |
|
newRBSP.append(p) |
|
} |
|
newRBSP.append(0x80) // rbsp_trailing_bits: '1' followed by padding zeros |
|
|
|
// RBSP → EBSP (insert emulation-prevention bytes), then restore NAL header |
|
let newEBSP = rbspToEbsp(newRBSP) |
|
var out = Data(capacity: 1 + newEBSP.count) |
|
out.append(nalHeader) |
|
out.append(newEBSP) |
|
return out |
|
} |
|
|
|
// MARK: - Private helpers (self-contained) |
|
|
|
/// EBSP → RBSP: remove 0x03 after 0x0000 sequences. |
|
private func rbspFromEbsp(_ ebsp: Data) -> Data |
|
{ |
|
if ebsp.isEmpty { return Data() } |
|
var out = Data(); out.reserveCapacity(ebsp.count) |
|
ebsp.withUnsafeBytes |
|
{ raw in |
|
guard let base = raw.bindMemory(to: UInt8.self).baseAddress else { return } |
|
let n = raw.count |
|
var i = 0, zero = 0 |
|
while i < n |
|
{ |
|
let b = base[i] |
|
if zero >= 2 && b == 0x03 |
|
{ // skip emulation-prevention byte |
|
i += 1 |
|
zero = 0 |
|
continue |
|
} |
|
out.append(b) |
|
zero = (b == 0) ? (zero + 1) : 0 |
|
i += 1 |
|
} |
|
} |
|
return out |
|
} |
|
|
|
/// RBSP → EBSP: insert 0x03 after 0x0000 if next would form 0x000000/01/02/03. |
|
func rbspToEbsp(_ rbsp: Data) -> Data |
|
{ |
|
if rbsp.isEmpty { return Data() } |
|
var out = Data(); out.reserveCapacity(rbsp.count + rbsp.count / 128 + 8) |
|
var zero = 0 |
|
for b in rbsp |
|
{ |
|
if zero >= 2 && (b & 0xFC) == 0x00 { out.append(0x03) } |
|
out.append(b) |
|
zero = (b == 0) ? (zero + 1) : 0 |
|
} |
|
return out |
|
} |
|
|
|
/// Iterate SEI payloads inside an SEI RBSP buffer (handles 0xFF extension bytes). |
|
private func forEachSEIPayload(rbsp: Data, _ body: (_ payloadType: Int, _ payload: Data) -> Void) |
|
{ |
|
let bytes = Array(rbsp) |
|
var i = 0 |
|
while i + 2 <= bytes.count |
|
{ |
|
// payloadType |
|
var t = 0 |
|
var b = bytes[i]; i += 1 |
|
while b == 0xFF, i < bytes.count { t += 255; b = bytes[i]; i += 1 } |
|
t += Int(b) |
|
guard i < bytes.count else { break } |
|
|
|
// payloadSize |
|
var sz = 0 |
|
b = bytes[i]; i += 1 |
|
while b == 0xFF, i < bytes.count { sz += 255; b = bytes[i]; i += 1 } |
|
sz += Int(b) |
|
guard i + sz <= bytes.count else { break } |
|
|
|
body(t, Data(bytes[i ..< i + sz])) |
|
i += sz |
|
// rbsp_trailing_bits (0x80) may follow; safe to ignore here. |
|
} |
|
} |
|
|
|
/// Write SEI payload header (type/size with 0xFF extension bytes). |
|
private func writeSEIPayloadHeader(type: Int, size: Int, into data: inout Data) |
|
{ |
|
var t = type |
|
while t >= 255 { data.append(0xFF); t -= 255 } |
|
data.append(UInt8(t)) |
|
var s = size |
|
while s >= 255 { data.append(0xFF); s -= 255 } |
|
data.append(UInt8(s)) |
|
} |
|
|
|
/// Hex of a 16-byte UUID (no dashes). |
|
private func uuidHex(_ d: Data) -> String |
|
{ |
|
d.map { String(format: "%02X", $0) }.joined() |
|
} |
|
} |