Sacha
342 lines
12 KiB
Swift
342 lines
12 KiB
Swift
import CoreMedia
|
|
import CoreVideo
|
|
import XCTest
|
|
@testable import HaptikaFrameMetrics
|
|
|
|
final class PPGFrameMetricsExtractorTests: XCTestCase {
|
|
func testFullRangeYUVExtractionUsesCentralROIAndUnlocksBuffer() throws {
|
|
let pixelBuffer = try makePixelBuffer(
|
|
width: 64,
|
|
height: 64,
|
|
pixelFormat: kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
|
|
)
|
|
try fillBiPlanarYUV(pixelBuffer, y: 100, cb: 128, cr: 128)
|
|
let metrics = try PPGFrameMetricsExtractor(gridSize: 16, roiFraction: 0.5)
|
|
.extract(from: makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 4.25))
|
|
|
|
XCTAssertEqual(metrics.timestampSeconds, 4.25, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.redMean, 100, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.greenMean, 100, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.blueMean, 100, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.luminanceMean, 100, accuracy: 0.000_001)
|
|
try assertCanRelock(pixelBuffer)
|
|
}
|
|
|
|
func testVideoRangeYUVExtractionAppliesVideoRangeConversion() throws {
|
|
let pixelBuffer = try makePixelBuffer(
|
|
width: 48,
|
|
height: 48,
|
|
pixelFormat: kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
|
|
)
|
|
try fillBiPlanarYUV(pixelBuffer, y: 126, cb: 128, cr: 128)
|
|
let metrics = try PPGFrameMetricsExtractor(gridSize: 12, roiFraction: 0.5)
|
|
.extract(from: makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 2))
|
|
|
|
let expected = 1.164_383 * 110
|
|
XCTAssertEqual(metrics.redMean, expected, accuracy: 0.001)
|
|
XCTAssertEqual(metrics.greenMean, expected, accuracy: 0.001)
|
|
XCTAssertEqual(metrics.blueMean, expected, accuracy: 0.001)
|
|
try assertCanRelock(pixelBuffer)
|
|
}
|
|
|
|
func testBGRAExtractionSamplesOnlyCentralROI() throws {
|
|
let pixelBuffer = try makePixelBuffer(
|
|
width: 64,
|
|
height: 64,
|
|
pixelFormat: kCVPixelFormatType_32BGRA
|
|
)
|
|
try fillBGRA(pixelBuffer) { x, y in
|
|
if (16..<48).contains(x), (16..<48).contains(y) {
|
|
return (blue: 20, green: 80, red: 200, alpha: 255)
|
|
}
|
|
return (blue: 240, green: 10, red: 5, alpha: 255)
|
|
}
|
|
let sampleBuffer = try makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 12.5)
|
|
|
|
let metrics = try PPGFrameMetricsExtractor(gridSize: 16, roiFraction: 0.5)
|
|
.extract(from: sampleBuffer)
|
|
|
|
XCTAssertEqual(metrics.timestampSeconds, 12.5, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.redMean, 200, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.greenMean, 80, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.blueMean, 20, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.redVariance, 0, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.greenVariance, 0, accuracy: 0.000_001)
|
|
XCTAssertEqual(
|
|
metrics.luminanceMean,
|
|
0.2126 * 200 + 0.7152 * 80 + 0.0722 * 20,
|
|
accuracy: 0.000_001
|
|
)
|
|
XCTAssertEqual(metrics.saturationRatio, 0, accuracy: 0.000_001)
|
|
try assertCanRelock(pixelBuffer)
|
|
}
|
|
|
|
func testBGRARedOnlyClippingKeepsUsableGreenChannel() throws {
|
|
let pixelBuffer = try makePixelBuffer(
|
|
width: 48,
|
|
height: 48,
|
|
pixelFormat: kCVPixelFormatType_32BGRA
|
|
)
|
|
try fillBGRA(pixelBuffer) { x, y in
|
|
let isCentral = (12..<36).contains(x) && (12..<36).contains(y)
|
|
return isCentral
|
|
? (blue: 40, green: 100, red: 255, alpha: 255)
|
|
: (blue: 0, green: 0, red: 0, alpha: 255)
|
|
}
|
|
let sampleBuffer = try makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 1)
|
|
|
|
let metrics = try PPGFrameMetricsExtractor(gridSize: 12, roiFraction: 0.5)
|
|
.extract(from: sampleBuffer)
|
|
|
|
XCTAssertEqual(metrics.redMean, 255, accuracy: 0.000_001)
|
|
XCTAssertEqual(metrics.saturationRatio, 0, accuracy: 0.000_001)
|
|
try assertCanRelock(pixelBuffer)
|
|
}
|
|
|
|
func testBGRASaturationRequiresBothUsefulChannelsToClip() throws {
|
|
let pixelBuffer = try makePixelBuffer(
|
|
width: 48,
|
|
height: 48,
|
|
pixelFormat: kCVPixelFormatType_32BGRA
|
|
)
|
|
try fillBGRA(pixelBuffer) { _, _ in
|
|
(blue: 80, green: 255, red: 255, alpha: 255)
|
|
}
|
|
let sampleBuffer = try makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 1)
|
|
|
|
let metrics = try PPGFrameMetricsExtractor(gridSize: 12, roiFraction: 0.5)
|
|
.extract(from: sampleBuffer)
|
|
|
|
XCTAssertEqual(metrics.saturationRatio, 1, accuracy: 0.000_001)
|
|
try assertCanRelock(pixelBuffer)
|
|
}
|
|
|
|
func testUnsupportedPixelFormatUnlocksBufferOnThrow() throws {
|
|
let pixelBuffer = try makePixelBuffer(
|
|
width: 32,
|
|
height: 32,
|
|
pixelFormat: kCVPixelFormatType_OneComponent8
|
|
)
|
|
let sampleBuffer = try makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 2)
|
|
|
|
XCTAssertThrowsError(try PPGFrameMetricsExtractor().extract(from: sampleBuffer)) { error in
|
|
guard case PPGFrameMetricsExtractorError.unsupportedPixelFormat(let format) = error else {
|
|
return XCTFail("Expected unsupportedPixelFormat, got \(error)")
|
|
}
|
|
XCTAssertEqual(format, kCVPixelFormatType_OneComponent8)
|
|
}
|
|
try assertCanRelock(pixelBuffer)
|
|
}
|
|
|
|
func testRepeatedExtractionDoesNotLeavePixelBuffersLocked() throws {
|
|
let extractor = PPGFrameMetricsExtractor(gridSize: 8, roiFraction: 0.5)
|
|
for index in 0..<100 {
|
|
autoreleasepool {
|
|
do {
|
|
let pixelBuffer = try makePixelBuffer(
|
|
width: 32,
|
|
height: 32,
|
|
pixelFormat: kCVPixelFormatType_32BGRA
|
|
)
|
|
try fillBGRA(pixelBuffer) { _, _ in (20, 90, 180, 255) }
|
|
_ = try extractor.extract(
|
|
from: makeSampleBuffer(
|
|
pixelBuffer: pixelBuffer,
|
|
timestampSeconds: Double(index) / 30
|
|
)
|
|
)
|
|
try assertCanRelock(pixelBuffer)
|
|
} catch {
|
|
XCTFail("Repeated extraction failed at \(index): \(error)")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private func makePixelBuffer(
|
|
width: Int,
|
|
height: Int,
|
|
pixelFormat: OSType
|
|
) throws -> CVPixelBuffer {
|
|
if pixelFormat == kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
|
|
|| pixelFormat == kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange {
|
|
return try makeBiPlanarPixelBuffer(width: width, height: height, pixelFormat: pixelFormat)
|
|
}
|
|
var pixelBuffer: CVPixelBuffer?
|
|
let status = CVPixelBufferCreate(
|
|
kCFAllocatorDefault,
|
|
width,
|
|
height,
|
|
pixelFormat,
|
|
nil,
|
|
&pixelBuffer
|
|
)
|
|
guard status == kCVReturnSuccess, let pixelBuffer else {
|
|
throw TestBufferError.pixelBufferCreation(status)
|
|
}
|
|
return pixelBuffer
|
|
}
|
|
|
|
private func makeBiPlanarPixelBuffer(
|
|
width: Int,
|
|
height: Int,
|
|
pixelFormat: OSType
|
|
) throws -> CVPixelBuffer {
|
|
let yBytesPerRow = width
|
|
let uvBytesPerRow = width
|
|
let ySize = yBytesPerRow * height
|
|
let uvHeight = height / 2
|
|
let totalSize = ySize + uvBytesPerRow * uvHeight
|
|
guard let storage = malloc(totalSize) else {
|
|
throw TestBufferError.missingBaseAddress
|
|
}
|
|
var planeAddresses: [UnsafeMutableRawPointer?] = [
|
|
storage,
|
|
storage.advanced(by: ySize),
|
|
]
|
|
var planeWidths = [width, width / 2]
|
|
var planeHeights = [height, uvHeight]
|
|
var planeBytesPerRow = [yBytesPerRow, uvBytesPerRow]
|
|
var pixelBuffer: CVPixelBuffer?
|
|
let status = CVPixelBufferCreateWithPlanarBytes(
|
|
kCFAllocatorDefault,
|
|
width,
|
|
height,
|
|
pixelFormat,
|
|
storage,
|
|
totalSize,
|
|
2,
|
|
&planeAddresses,
|
|
&planeWidths,
|
|
&planeHeights,
|
|
&planeBytesPerRow,
|
|
{ _, dataPointer, _, _, _ in
|
|
free(UnsafeMutableRawPointer(mutating: dataPointer))
|
|
},
|
|
nil,
|
|
nil,
|
|
&pixelBuffer
|
|
)
|
|
guard status == kCVReturnSuccess, let pixelBuffer else {
|
|
free(storage)
|
|
throw TestBufferError.pixelBufferCreation(status)
|
|
}
|
|
return pixelBuffer
|
|
}
|
|
|
|
private func fillBiPlanarYUV(
|
|
_ pixelBuffer: CVPixelBuffer,
|
|
y: UInt8,
|
|
cb: UInt8,
|
|
cr: UInt8
|
|
) throws {
|
|
let status = CVPixelBufferLockBaseAddress(pixelBuffer, [])
|
|
guard status == kCVReturnSuccess else {
|
|
throw TestBufferError.pixelBufferLock(status)
|
|
}
|
|
defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, []) }
|
|
guard CVPixelBufferGetPlaneCount(pixelBuffer) >= 2,
|
|
let yBase = CVPixelBufferGetBaseAddressOfPlane(pixelBuffer, 0),
|
|
let uvBase = CVPixelBufferGetBaseAddressOfPlane(pixelBuffer, 1)
|
|
else {
|
|
throw TestBufferError.missingBaseAddress
|
|
}
|
|
let yHeight = CVPixelBufferGetHeightOfPlane(pixelBuffer, 0)
|
|
let yStride = CVPixelBufferGetBytesPerRowOfPlane(pixelBuffer, 0)
|
|
let yBytes = yBase.assumingMemoryBound(to: UInt8.self)
|
|
for row in 0..<yHeight {
|
|
yBytes.advanced(by: row * yStride).initialize(repeating: y, count: yStride)
|
|
}
|
|
let uvWidth = CVPixelBufferGetWidthOfPlane(pixelBuffer, 1)
|
|
let uvHeight = CVPixelBufferGetHeightOfPlane(pixelBuffer, 1)
|
|
let uvStride = CVPixelBufferGetBytesPerRowOfPlane(pixelBuffer, 1)
|
|
let uvBytes = uvBase.assumingMemoryBound(to: UInt8.self)
|
|
for row in 0..<uvHeight {
|
|
for column in 0..<uvWidth {
|
|
let offset = row * uvStride + column * 2
|
|
uvBytes[offset] = cb
|
|
uvBytes[offset + 1] = cr
|
|
}
|
|
}
|
|
}
|
|
|
|
private func fillBGRA(
|
|
_ pixelBuffer: CVPixelBuffer,
|
|
pixel: (_ x: Int, _ y: Int) -> (blue: UInt8, green: UInt8, red: UInt8, alpha: UInt8)
|
|
) throws {
|
|
let lockStatus = CVPixelBufferLockBaseAddress(pixelBuffer, [])
|
|
guard lockStatus == kCVReturnSuccess else {
|
|
throw TestBufferError.pixelBufferLock(lockStatus)
|
|
}
|
|
defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, []) }
|
|
|
|
guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else {
|
|
throw TestBufferError.missingBaseAddress
|
|
}
|
|
let width = CVPixelBufferGetWidth(pixelBuffer)
|
|
let height = CVPixelBufferGetHeight(pixelBuffer)
|
|
let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
|
|
let bytes = baseAddress.assumingMemoryBound(to: UInt8.self)
|
|
|
|
for y in 0..<height {
|
|
for x in 0..<width {
|
|
let color = pixel(x, y)
|
|
let offset = y * bytesPerRow + x * 4
|
|
bytes[offset] = color.blue
|
|
bytes[offset + 1] = color.green
|
|
bytes[offset + 2] = color.red
|
|
bytes[offset + 3] = color.alpha
|
|
}
|
|
}
|
|
}
|
|
|
|
private func makeSampleBuffer(
|
|
pixelBuffer: CVPixelBuffer,
|
|
timestampSeconds: Double
|
|
) throws -> CMSampleBuffer {
|
|
var formatDescription: CMVideoFormatDescription?
|
|
let formatStatus = CMVideoFormatDescriptionCreateForImageBuffer(
|
|
allocator: kCFAllocatorDefault,
|
|
imageBuffer: pixelBuffer,
|
|
formatDescriptionOut: &formatDescription
|
|
)
|
|
guard formatStatus == noErr, let formatDescription else {
|
|
throw TestBufferError.formatDescriptionCreation(formatStatus)
|
|
}
|
|
|
|
var timing = CMSampleTimingInfo(
|
|
duration: .invalid,
|
|
presentationTimeStamp: CMTime(seconds: timestampSeconds, preferredTimescale: 60_000),
|
|
decodeTimeStamp: .invalid
|
|
)
|
|
var sampleBuffer: CMSampleBuffer?
|
|
let sampleStatus = CMSampleBufferCreateReadyWithImageBuffer(
|
|
allocator: kCFAllocatorDefault,
|
|
imageBuffer: pixelBuffer,
|
|
formatDescription: formatDescription,
|
|
sampleTiming: &timing,
|
|
sampleBufferOut: &sampleBuffer
|
|
)
|
|
guard sampleStatus == noErr, let sampleBuffer else {
|
|
throw TestBufferError.sampleBufferCreation(sampleStatus)
|
|
}
|
|
return sampleBuffer
|
|
}
|
|
|
|
private func assertCanRelock(_ pixelBuffer: CVPixelBuffer) throws {
|
|
let status = CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
|
|
guard status == kCVReturnSuccess else {
|
|
throw TestBufferError.pixelBufferLock(status)
|
|
}
|
|
CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly)
|
|
}
|
|
}
|
|
|
|
private enum TestBufferError: Error {
|
|
case pixelBufferCreation(CVReturn)
|
|
case pixelBufferLock(CVReturn)
|
|
case missingBaseAddress
|
|
case formatDescriptionCreation(OSStatus)
|
|
case sampleBufferCreation(OSStatus)
|
|
}
|