sjs/SJSAssetExportSession
1
0
Fork 0
mirror of https://github.com/samsonjs/SJSAssetExportSession.git synced 2026-04-15 11:55:46 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add a test and make it actually work

Browse source
This commit is contained in:
Sami Samhuri 2024-07-06 22:08:29 -07:00
parent 36f055d36f
commit 750b77210c
No known key found for this signature in database
5 changed files with 140 additions and 117 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
SJSAssetExportSession.xcodeproj/project.pbxproj
View file

@ -13,6 +13,7 @@
7B9BC0192C305D2C00C160C2 /* SJSAssetExportSessionTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7B9BC0182C305D2C00C160C2 /* SJSAssetExportSessionTests.swift */; };
7B9BC01A2C305D2C00C160C2 /* SJSAssetExportSession.h in Headers */ = {isa = PBXBuildFile; fileRef = 7B9BC00C2C305D2C00C160C2 /* SJSAssetExportSession.h */; settings = {ATTRIBUTES = (Public, ); }; };
7B9BC0282C30612C00C160C2 /* ExportSession.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7B9BC0272C30612C00C160C2 /* ExportSession.swift */; };
7BC5FC712C3A52B50090B757 /* test.mov in Resources */ = {isa = PBXBuildFile; fileRef = 7BC5FC702C3A52B50090B757 /* test.mov */; };
/* End PBXBuildFile section */
/* Begin PBXContainerItemProxy section */
@ -33,6 +34,7 @@
7B9BC0132C305D2C00C160C2 /* SJSAssetExportSessionTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = SJSAssetExportSessionTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
7B9BC0182C305D2C00C160C2 /* SJSAssetExportSessionTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SJSAssetExportSessionTests.swift; sourceTree = "<group>"; };
7B9BC0272C30612C00C160C2 /* ExportSession.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ExportSession.swift; sourceTree = "<group>"; };
7BC5FC702C3A52B50090B757 /* test.mov */ = {isa = PBXFileReference; lastKnownFileType = video.quicktime; path = test.mov; sourceTree = "<group>"; };
/* End PBXFileReference section */
/* Begin PBXFrameworksBuildPhase section */
@ -86,11 +88,20 @@
7B9BC0172C305D2C00C160C2 /* SJSAssetExportSessionTests */ = {
isa = PBXGroup;
children = (
7BC5FC6D2C3A525A0090B757 /* Resources */,
7B9BC0182C305D2C00C160C2 /* SJSAssetExportSessionTests.swift */,
);
path = SJSAssetExportSessionTests;
sourceTree = "<group>";
};
7BC5FC6D2C3A525A0090B757 /* Resources */ = {
isa = PBXGroup;
children = (
7BC5FC702C3A52B50090B757 /* test.mov */,
);
path = Resources;
sourceTree = "<group>";
};
/* End PBXGroup section */
/* Begin PBXHeadersBuildPhase section */
@ -190,6 +201,7 @@
isa = PBXResourcesBuildPhase;
buildActionMask = 2147483647;
files = (
7BC5FC712C3A52B50090B757 /* test.mov in Resources */,
);
runOnlyForDeploymentPostprocessing = 0;
};

4
SJSAssetExportSession/ExportSession.swift
View file

@ -25,11 +25,11 @@ public final class ExportSession {
}
}
public func export(
public static func export(
asset: sending AVAsset,
audioMix: sending AVAudioMix?,
audioOutputSettings: [String: (any Sendable)],
videoComposition: sending AVVideoComposition?,
videoComposition: sending AVVideoComposition,
videoOutputSettings: [String: (any Sendable)],
timeRange: CMTimeRange? = nil,
optimizeForNetworkUse: Bool = false,

186
SJSAssetExportSession/SampleWriter.swift
View file

@ -6,6 +6,9 @@
//
import AVFoundation.AVAsset
import OSLog
private let log = Logger(subsystem: "SJSAssetExportSession", category: "SampleWriter")
private extension AVAsset {
func sendTracks(withMediaType mediaType: AVMediaType) async throws -> sending [AVAssetTrack] {
@ -24,25 +27,21 @@ actor SampleWriter {
queue.asUnownedSerialExecutor()
}
let audioTracks: [AVAssetTrack]
private let audioMix: AVAudioMix?
let audioMix: AVAudioMix?
private let audioOutputSettings: [String: (any Sendable)]
let audioOutputSettings: [String: (any Sendable)]
private let videoComposition: AVVideoComposition?
let videoTracks: [AVAssetTrack]
private let videoOutputSettings: [String: (any Sendable)]
let videoComposition: AVVideoComposition?
private let reader: AVAssetReader
let videoOutputSettings: [String: (any Sendable)]
private let writer: AVAssetWriter
let reader: AVAssetReader
private let duration: CMTime
let writer: AVAssetWriter
let duration: CMTime
let timeRange: CMTimeRange
private let timeRange: CMTimeRange
private var audioOutput: AVAssetReaderAudioMixOutput?
@ -52,14 +51,18 @@ actor SampleWriter {
private var videoInput: AVAssetWriterInput?
private var pixelBufferAdaptor: AVAssetWriterInputPixelBufferAdaptor?
private lazy var progressStream: AsyncStream<Float> = AsyncStream { continuation in
progressContinuation = continuation
}
private var progressContinuation: AsyncStream<Float>.Continuation?
init(
asset: sending AVAsset,
timeRange: CMTimeRange,
audioMix: AVAudioMix?,
audioOutputSettings: sending [String: (any Sendable)],
videoComposition: AVVideoComposition?,
videoComposition: AVVideoComposition,
videoOutputSettings: sending [String: (any Sendable)],
optimizeForNetworkUse: Bool,
outputURL: URL,
@ -81,26 +84,31 @@ actor SampleWriter {
throw ExportSession.Error.setupFailure(reason: "Cannot apply video output settings")
}
self.audioTracks = try await asset.sendTracks(withMediaType: .audio)
let audioTracks = try await asset.sendTracks(withMediaType: .audio)
let videoTracks = try await asset.sendTracks(withMediaType: .video)
self.audioMix = audioMix
self.audioOutputSettings = audioOutputSettings
self.videoTracks = try await asset.sendTracks(withMediaType: .video)
self.videoComposition = videoComposition
self.videoOutputSettings = videoOutputSettings
self.reader = reader
self.writer = writer
self.duration = duration
self.timeRange = timeRange
try setUpAudio(audioTracks: audioTracks)
try setUpVideo(videoTracks: videoTracks)
}
func writeSamples() async throws {
progressContinuation?.yield(0)
writer.startWriting()
reader.startReading()
writer.startSession(atSourceTime: timeRange.start)
async let audioResult = try encodeAudioTracks(audioTracks)
async let videoResult = try encodeVideoTracks(videoTracks)
_ = try await (audioResult, videoResult)
await encodeAudioTracks()
await encodeVideoTracks()
if reader.status == .cancelled || writer.status == .cancelled {
throw CancellationError()
@ -119,10 +127,12 @@ actor SampleWriter {
}
}
private func encodeAudioTracks(_ audioTracks: [AVAssetTrack]) async throws -> Bool {
guard !audioTracks.isEmpty else { return false }
private func setUpAudio(audioTracks: [AVAssetTrack]) throws {
guard !audioTracks.isEmpty else { return }
let audioOutput = AVAssetReaderAudioMixOutput(audioTracks: audioTracks, audioSettings: nil)
audioOutput.alwaysCopiesSampleData = false
audioOutput.audioMix = audioMix
guard reader.canAdd(audioOutput) else {
throw ExportSession.Error.setupFailure(reason: "Can't add audio output to reader")
}
@ -130,55 +140,23 @@ actor SampleWriter {
self.audioOutput = audioOutput
let audioInput = AVAssetWriterInput(mediaType: .audio, outputSettings: audioOutputSettings)
audioInput.expectsMediaDataInRealTime = false
guard writer.canAdd(audioInput) else {
throw ExportSession.Error.setupFailure(reason: "Can't add audio input to writer")
}
writer.add(audioInput)
self.audioInput = audioInput
return await withCheckedContinuation { continuation in
self.audioInput?.requestMediaDataWhenReady(on: queue) {
let hasMoreSamples = self.assumeIsolated { $0.writeReadyAudioSamples() }
if !hasMoreSamples {
continuation.resume(returning: true)
}
}
}
}
private func writeReadyAudioSamples() -> Bool {
guard let audioOutput, let audioInput else { return true }
while audioInput.isReadyForMoreMediaData {
guard reader.status == .reading && writer.status == .writing,
let sampleBuffer = audioOutput.copyNextSampleBuffer() else {
audioInput.markAsFinished()
NSLog("Finished encoding ready audio samples from \(audioOutput)")
return false
}
guard audioInput.append(sampleBuffer) else {
NSLog("Failed to append audio sample buffer \(sampleBuffer) to input \(audioInput)")
return false
}
private func setUpVideo(videoTracks: [AVAssetTrack]) throws {
guard !videoTracks.isEmpty else {
throw ExportSession.Error.setupFailure(reason: "No video tracks")
}
// Everything was appended successfully, return true indicating there's more to do.
NSLog("Completed encoding ready audio samples, more to come...")
return true
}
private func encodeVideoTracks(_ videoTracks: [AVAssetTrack]) async throws -> Bool {
guard !videoTracks.isEmpty else { return false }
guard let width = videoComposition.map({ Int($0.renderSize.width) })
?? (videoOutputSettings[AVVideoWidthKey] as? NSNumber)?.intValue,
let height = videoComposition.map({ Int($0.renderSize.height) })
?? (videoOutputSettings[AVVideoHeightKey] as? NSNumber)?.intValue else {
throw ExportSession.Error.setupFailure(reason: "Export dimensions must be provided in a video composition or video output settings")
}
let videoOutput = AVAssetReaderVideoCompositionOutput(videoTracks: videoTracks, videoSettings: nil)
let videoOutput = AVAssetReaderVideoCompositionOutput(
videoTracks: videoTracks,
videoSettings: nil
)
videoOutput.alwaysCopiesSampleData = false
videoOutput.videoComposition = videoComposition
guard reader.canAdd(videoOutput) else {
@ -188,72 +166,60 @@ actor SampleWriter {
self.videoOutput = videoOutput
let videoInput = AVAssetWriterInput(mediaType: .video, outputSettings: videoOutputSettings)
videoInput.expectsMediaDataInRealTime = false
guard writer.canAdd(videoInput) else {
throw ExportSession.Error.setupFailure(reason: "Can't add video input to writer")
}
writer.add(videoInput)
self.videoInput = videoInput
let pixelBufferAttributes: [String: Any] = [
kCVPixelBufferPixelFormatTypeKey as String: NSNumber(integerLiteral: Int(kCVPixelFormatType_32RGBA)),
kCVPixelBufferWidthKey as String: NSNumber(integerLiteral: width),
kCVPixelBufferHeightKey as String: NSNumber(integerLiteral: height),
"IOSurfaceOpenGLESTextureCompatibility": NSNumber(booleanLiteral: true),
"IOSurfaceOpenGLESFBOCompatibility": NSNumber(booleanLiteral: true),
]
pixelBufferAdaptor = AVAssetWriterInputPixelBufferAdaptor(
assetWriterInput: videoInput,
sourcePixelBufferAttributes: pixelBufferAttributes
)
return await withCheckedContinuation { continuation in
self.videoInput?.requestMediaDataWhenReady(on: queue) {
let hasMoreSamples = self.assumeIsolated { $0.writeReadyVideoSamples() }
if !hasMoreSamples {
continuation.resume(returning: true)
}
}
}
}
private func writeReadyVideoSamples() -> Bool {
guard let videoOutput, let videoInput, let pixelBufferAdaptor else { return true }
while videoInput.isReadyForMoreMediaData {
private func writeReadySamples(output: AVAssetReaderOutput, input: AVAssetWriterInput) -> Bool {
while input.isReadyForMoreMediaData {
guard reader.status == .reading && writer.status == .writing,
let sampleBuffer = videoOutput.copyNextSampleBuffer() else {
videoInput.markAsFinished()
NSLog("Finished encoding ready video samples from \(videoOutput)")
let sampleBuffer = output.copyNextSampleBuffer() else {
input.markAsFinished()
log.debug("Finished encoding ready audio samples from \(output)")
return false
}
let samplePresentationTime = CMSampleBufferGetPresentationTimeStamp(sampleBuffer) - timeRange.start
let progress = Float(samplePresentationTime.seconds / duration.seconds)
#warning("TODO: publish progress to an AsyncStream")
guard let pixelBufferPool = pixelBufferAdaptor.pixelBufferPool else {
NSLog("No pixel buffer pool available on adaptor \(pixelBufferAdaptor)")
guard input.append(sampleBuffer) else {
log.error("""
Failed to append audio sample buffer \(String(describing: sampleBuffer)) to
input \(input.debugDescription)
""")
return false
}
var toRenderBuffer: CVPixelBuffer?
let result = CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, pixelBufferPool, &toRenderBuffer)
var handled = false
if result == kCVReturnSuccess, let toBuffer = toRenderBuffer {
handled = pixelBufferAdaptor.append(toBuffer, withPresentationTime: samplePresentationTime)
if !handled { return false }
}
if !handled {
#warning("is this really necessary?! seems like a failure scenario...")
guard videoInput.append(sampleBuffer) else {
NSLog("Failed to append video sample buffer \(sampleBuffer) to input \(videoInput)")
return false
}
}
}
// Everything was appended successfully, return true indicating there's more to do.
NSLog("Completed encoding ready video samples, more to come...")
log.debug("Completed encoding ready audio samples, more to come...")
return true
}
private func encodeAudioTracks() async {
return await withCheckedContinuation { continuation in
self.audioInput!.requestMediaDataWhenReady(on: queue) {
let hasMoreSamples = self.assumeIsolated { _self in
_self.writeReadySamples(output: _self.audioOutput!, input: _self.audioInput!)
}
if !hasMoreSamples {
continuation.resume()
}
}
}
}
private func encodeVideoTracks() async {
return await withCheckedContinuation { continuation in
self.videoInput!.requestMediaDataWhenReady(on: queue) {
let hasMoreSamples = self.assumeIsolated { _self in
_self.writeReadySamples(output: _self.videoOutput!, input: _self.videoInput!)
}
if !hasMoreSamples {
continuation.resume()
}
}
}
}
}

BIN
SJSAssetExportSessionTests/Resources/test.mov Normal file
View file

Binary file not shown.

55
SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift
View file

@ -5,13 +5,58 @@
// Created by Sami Samhuri on 2024-06-29.
//
import Testing
import AVFoundation
@testable import SJSAssetExportSession
import Testing
struct SJSAssetExportSessionTests {
final class ExportSessionTests {
@Test func test_encode_h264_720p_30fps() async throws {
let sourceURL = Bundle(for: Self.self).url(forResource: "test", withExtension: "mov")!
let sourceAsset = AVURLAsset(url: sourceURL)
let timestamp = Int(Date.now.timeIntervalSince1970)
let filename = "ExportSessionTests_testEncode_\(timestamp).mp4"
let destinationURL = URL.temporaryDirectory.appending(component: filename)
defer { _ = try? FileManager.default.removeItem(at: destinationURL) }
@Test func testExample() async throws {
// Write your test here and use APIs like `#expect(...)` to check expected conditions.
let size = CGSize(width: 1280, height: 720)
let duration = CMTime(seconds: 1, preferredTimescale: 600)
let videoComposition = try await AVMutableVideoComposition.videoComposition(withPropertiesOf: sourceAsset)
videoComposition.renderSize = size
videoComposition.renderScale = 1
videoComposition.frameDuration = CMTime(value: 1, timescale: 30)
try await ExportSession.export(
asset: sourceAsset,
audioMix: nil,
audioOutputSettings: [
AVFormatIDKey: kAudioFormatMPEG4AAC,
AVNumberOfChannelsKey: NSNumber(value: 2),
AVSampleRateKey: NSNumber(value: 44_100.0),
],
videoComposition: videoComposition,
videoOutputSettings: [
AVVideoCodecKey: AVVideoCodecType.h264.rawValue,
AVVideoWidthKey: NSNumber(value: Int(size.width)),
AVVideoHeightKey: NSNumber(value: Int(size.height)),
AVVideoCompressionPropertiesKey: [
AVVideoAverageBitRateKey: NSNumber(value: 1_000_000),
AVVideoProfileLevelKey: AVVideoProfileLevelH264HighAutoLevel as String,
] as [String: any Sendable],
AVVideoColorPropertiesKey: [
AVVideoColorPrimariesKey: AVVideoColorPrimaries_ITU_R_709_2,
AVVideoTransferFunctionKey: AVVideoTransferFunction_ITU_R_709_2,
AVVideoYCbCrMatrixKey: AVVideoYCbCrMatrix_ITU_R_709_2,
],
],
timeRange: CMTimeRange(start: .zero, duration: duration),
to: destinationURL,
as: .mp4
)
let asset = AVURLAsset(url: destinationURL)
#expect(try await asset.load(.duration) == duration)
try #require(await asset.loadTracks(withMediaType: .video).count == 1)
let videoTrack = try #require(try await asset.loadTracks(withMediaType: .video).first)
#expect(try await videoTrack.load(.naturalSize) == size)
try #require(try await asset.loadTracks(withMediaType: .audio).count == 1)
}
}