Hardware accelerated GIF to MP4 converter in Swift using VideoToolbox

main.swift

import Foundation

let inputPath = "gif.gif"
let outputPath = "mp4.mp4"
let inputURL = NSURL(fileURLWithPath:inputPath)
let outputURL = NSURL(fileURLWithPath:outputPath)

if NSFileManager.defaultManager().fileExistsAtPath(outputPath)
{
	try! NSFileManager.defaultManager().removeItemAtPath(outputPath)
}

print("Can do hardware H264 encoding: \(NYXAVCEncoder.canPerformHardwareH264Compression())")

if let encoder = NYXGIFToMP4Encoder(inputURL:inputURL, outputURL:outputURL)
{
	encoder.outputSize = (encoder.inputSize.width, encoder.inputSize.height) // optional
	let ret = encoder.convert()
}

NYXAVCEncoder.swift

import VideoToolbox
import AVFoundation

private var __canHWAVC: Bool = false
private var __tokenHWAVC: dispatch_once_t = 0

public protocol NYXAVCEncoderDelegate : class
{
	func didEncodeFrame(frame: CMSampleBuffer)
	func didFailToEncodeFrame()
}


public final class NYXAVCEncoder : NSObject
{
	// MARK: - Default encoder values
	#if os(OSX)
	static let defaultAttributes:[NSString: AnyObject] = [
		kCVPixelBufferPixelFormatTypeKey: Int(kCVPixelFormatType_32ARGB),
		kCVPixelBufferIOSurfacePropertiesKey: [:],
	]
	#elseif os(iOS)
	static let defaultAttributes:[NSString: AnyObject] = [
	kCVPixelBufferPixelFormatTypeKey: Int(kCVPixelFormatType_32ARGB),
	kCVPixelBufferIOSurfacePropertiesKey: [:],
	kCVPixelBufferOpenGLESCompatibilityKey: true,
	]
	#endif
	static let defaultFPS: Int = 30
	
	// MARK: - Private properties
	// Input image width
	private var inputWidth: Int
	// Input image height
	private var inputHeight: Int
	// H.264 compression session
	private var compressionSession: VTCompressionSession?
	
	// MARK: - Public properties
	// Delegate
	public weak var delegate: NYXAVCEncoderDelegate?
	// Currently encoding flag
	public private(set) var working: Bool = false
	// Number of frame pushed to the encoder
	public private(set) var pushedFrames: Int64 = 0
	// Number of encoded frames
	public private(set) var encodedFrames: Int64 = 0
	// Output image width
	public var outputWidth: Int
	// Output image height
	public var outputHeight: Int
	// FPS
	public var fps = NYXAVCEncoder.defaultFPS
	// Default interval for keyframes
	public var keyframeInterval = NYXAVCEncoder.defaultFPS * 2
	
	// MARK: - Init
	public init(inputWidth: Int, inputHeight: Int)
	{
		self.inputWidth = inputWidth
		self.inputHeight = inputHeight
		self.outputWidth = inputWidth
		self.outputHeight = inputHeight
		
		super.init()
	}
	
	// MARK: - Public
	public func beginEncode() -> Bool
	{
		if self.working
		{
			return false
		}
		self.pushedFrames = 0
		self.encodedFrames = 0
		
		// input image attributes
		var sourceImageBufferAttributes = NYXAVCEncoder.defaultAttributes
		sourceImageBufferAttributes[kCVPixelBufferWidthKey] = self.inputWidth
		sourceImageBufferAttributes[kCVPixelBufferHeightKey] = self.inputHeight
		
		// Create session, enable hw
		var status = VTCompressionSessionCreate(nil, Int32(self.outputWidth), Int32(self.outputHeight), CMVideoCodecType(kCMVideoCodecType_H264), [kVTVideoEncoderSpecification_EnableHardwareAcceleratedVideoEncoder as String : true], sourceImageBufferAttributes, nil, vt_compression_callback, unsafeBitCast(self, UnsafeMutablePointer<Void>.self), &self.compressionSession)
		if status != noErr
		{
			let error = NSError(domain:NSOSStatusErrorDomain, code:Int(status), userInfo:nil)
			print("VTCompressionSessionCreate: \(error.localizedDescription)")
			return false
		}
		
		// Encoding properties
		let properties:[NSString: NSObject] = [
			kVTCompressionPropertyKey_RealTime: kCFBooleanTrue,
			kVTCompressionPropertyKey_ProfileLevel: kVTProfileLevel_H264_High_AutoLevel,
			/*kVTCompressionPropertyKey_AverageBitRate: 40000,*/
			kVTCompressionPropertyKey_ExpectedFrameRate: self.fps,
			kVTCompressionPropertyKey_MaxKeyFrameInterval: self.keyframeInterval,
			kVTCompressionPropertyKey_AllowFrameReordering: true,
			kVTCompressionPropertyKey_H264EntropyMode: kVTH264EntropyMode_CABAC,
			kVTCompressionPropertyKey_PixelTransferProperties: [
				kVTPixelTransferPropertyKey_ScalingMode as NSString: kVTScalingMode_Trim
			]
		]
		status = VTSessionSetProperties(self.compressionSession!, properties)
		if status != noErr
		{
			let error = NSError(domain:NSOSStatusErrorDomain, code:Int(status), userInfo:nil)
			print("VTSessionSetProperties: \(error.localizedDescription)")
			return false
		}
		
		// Resource allocation (optional)
		VTCompressionSessionPrepareToEncodeFrames(self.compressionSession!)
		
		self.working = true
		
		return true
	}
	
	public func endEncode()
	{
		if self.working
		{
			VTCompressionSessionCompleteFrames(self.compressionSession!, kCMTimeInvalid)
			VTCompressionSessionInvalidate(self.compressionSession!)
			self.compressionSession = nil
		}
		self.working = false
	}
	
	public func encodeFrame(frame: CVPixelBufferRef, frameNumber: Int64) -> Bool
	{
		if !self.working
		{
			return false
		}
		
		let status = VTCompressionSessionEncodeFrame(
			self.compressionSession! /*compression session*/,
			frame /*video frame*/,
			CMTimeMake(frameNumber, Int32(NYXAVCEncoder.defaultFPS)) /*presentation timestamp*/,
			kCMTimeInvalid /*duration*/,
			nil /*frameProperties*/,
			nil /*sourceFrameRefCon*/,
			nil /*infoFlagsOut*/)
		self.pushedFrames += 1
		return status == noErr
	}
	
	// MARK: - Compression callback
	private var vt_compression_callback:VTCompressionOutputCallback = {(outputCallbackRefCon: UnsafeMutablePointer<Void>, sourceFrameRefCon: UnsafeMutablePointer<Void>, status: OSStatus, infoFlags: VTEncodeInfoFlags, sampleBuffer: CMSampleBuffer?) in
		
		let encoder: NYXAVCEncoder = unsafeBitCast(outputCallbackRefCon, NYXAVCEncoder.self)
		encoder.encodedFrames += 1
		if status != noErr
		{
			encoder.delegate?.didFailToEncodeFrame()
			print("VTCompressionOutputCallback: \(status)")
			return
		}
		
		if let frame = sampleBuffer
		{
			encoder.delegate?.didEncodeFrame(frame)
		}
		else
		{
			encoder.delegate?.didFailToEncodeFrame()
		}
	}
	
	// MARK: - Static
	public class func canPerformHardwareH264Compression() -> Bool
	{
		dispatch_once(&__tokenHWAVC)
			{
				// Hardware encode also depends on the image dimensions
				let encoderSpecs = [kVTVideoEncoderSpecification_RequireHardwareAcceleratedVideoEncoder as String : true]
				var sessionRef: VTCompressionSession? = nil
				let status = VTCompressionSessionCreate(nil, 512, 512, CMVideoCodecType(kCMVideoCodecType_H264), encoderSpecs, nil, nil, nil, nil, &sessionRef)
				__canHWAVC = (status == noErr)
				if let s = sessionRef
				{
					VTCompressionSessionInvalidate(s)
				}
		}
		return __canHWAVC
	}
}

NYXGIFToH264Encoder.swift

import VideoToolbox
import AVFoundation


private let kNYXNumberOfComponentsPerARBGPixel: Int = 4


public final class NYXGIFToMP4Encoder : NSObject
{
	// MARK: - Properties
	// File URL to encode
	public private(set) var inputURL: NSURL!
	// URL of encoded file
	public private(set) var outputURL: NSURL!
	// Input image size
	public private(set) var inputSize: (width: Int, height: Int) = (0, 0)
	// Output image size
	public var outputSize: (width: Int, height: Int) = (0, 0)
	// Output file writer
	private var fileWriter: AVAssetWriter!
	// Video track writer
	private var videoWriterInput: AVAssetWriterInput!

	// MARK: - Init
	public init?(inputURL: NSURL, outputURL: NSURL)
	{
		super.init()

		if !NSFileManager().fileExistsAtPath(inputURL.path!)
		{
			return nil
		}

		self.inputURL = inputURL
		self.outputURL = outputURL

		self._findInputImageDimensions()
		self.outputSize = self.inputSize
	}

	// MARK: - Public
	public func convert() -> Bool
	{
		let encoder = NYXAVCEncoder(inputWidth:self.inputSize.width, inputHeight:self.inputSize.height)
		encoder.delegate = self
		encoder.outputWidth = self.outputSize.width
		encoder.outputHeight = self.outputSize.height

		encoder.beginEncode()

		do {
			self.fileWriter = try AVAssetWriter(URL:self.outputURL, fileType:AVFileTypeMPEG4)
		}
		catch {
			print("AVAssetWriter")
			return false
		}
		var formatHint: CMFormatDescriptionRef? = nil
		let status = CMVideoFormatDescriptionCreate(nil, kCMVideoCodecType_H264, Int32(self.outputSize.width), Int32(self.outputSize.height), nil, &formatHint)
		if status != noErr
		{
			let error = NSError(domain:NSOSStatusErrorDomain, code:Int(status), userInfo:nil)
			print("CMVideoFormatDescriptionCreate: \(error.localizedDescription)")
			return false
		}
		self.videoWriterInput = AVAssetWriterInput(mediaType:AVMediaTypeVideo, outputSettings:nil, sourceFormatHint:formatHint)
		self.fileWriter.addInput(self.videoWriterInput)
		self.fileWriter.startWriting()
		self.fileWriter.startSessionAtSourceTime(kCMTimeZero)

		guard let src = CGImageSourceCreateWithURL(self.inputURL, nil) else {return false}
		let count = CGImageSourceGetCount(src)
		let imgContext = NYXGIFToMP4Encoder.ARGBBitmapContext(width:self.inputSize.width, height:self.inputSize.height, bytesPerRow:kNYXNumberOfComponentsPerARBGPixel * self.inputSize.width, withAlpha:true)
		var curFrame: Int64 = 1
		print("\(count) frames to encode")
		for var i = 0; i < count; ++i
		{
			// Create CGImageRef
			guard let image = CGImageSourceCreateImageAtIndex(src, i, nil) else {return false}
			// get ARGB pixels
			CGContextDrawImage(imgContext, CGRect(x:0.0, y:0.0, width:CGFloat(self.inputSize.width), height:CGFloat(self.inputSize.height)), image)
			let imgBuf = CGBitmapContextGetData(imgContext)

			// Create pixels buffer
			var pixel_buffer: CVPixelBufferRef? = nil
			let ret = CVPixelBufferCreateWithBytes(kCFAllocatorSystemDefault, self.inputSize.width, self.inputSize.height, kCVPixelFormatType_32ARGB, imgBuf, kNYXNumberOfComponentsPerARBGPixel * self.inputSize.width, nil, nil, nil, &pixel_buffer)
			if ret != kCVReturnSuccess
			{
				print("CVPixelBufferCreateWithBytes: \(ret)")
				return false
			}

			// Encode frame
			encoder.encodeFrame(pixel_buffer!, frameNumber:curFrame)
			curFrame++
		}

		// Indicate that there are no more frames to process
		encoder.endEncode()

		// lol ugly, don't do this
		while encoder.encodedFrames != encoder.pushedFrames
		{
			print("not good yet. \(encoder.encodedFrames)/\(encoder.pushedFrames)")
			usleep(1000) // 1ms
		}

		self.videoWriterInput.markAsFinished()
		self.fileWriter.finishWritingWithCompletionHandler { () -> Void in
			if self.fileWriter.status == .Failed
			{
				guard let error = self.fileWriter.error else {return}
				print("finishWritingWithCompletionHandler: \(error)")
			}
			else
			{
				print("Done, \(encoder.encodedFrames) frames")
			}
		}

		return true
	}

	// MARK: - Private
	private func _findInputImageDimensions()
	{
		guard let src = CGImageSourceCreateWithURL(self.inputURL, nil) else {return}
		let count = CGImageSourceGetCount(src)
		if count > 0
		{
			guard let imgRef = CGImageSourceCreateImageAtIndex(src, 0, nil) else {return}
			self.inputSize = (CGImageGetWidth(imgRef), CGImageGetHeight(imgRef))
		}
		else
		{
			self.inputSize = (0, 0)
		}
	}

	private class func ARGBBitmapContext(width width: Int, height: Int, bytesPerRow: Int, withAlpha: Bool) -> CGContextRef?
	{
		let alphaInfo = CGBitmapInfo(rawValue: withAlpha ? CGImageAlphaInfo.PremultipliedFirst.rawValue : CGImageAlphaInfo.NoneSkipFirst.rawValue)
		let bmContext = CGBitmapContextCreate(nil, width, height, 8/*Bits per component*/, bytesPerRow, CGColorSpaceCreateWithName(kCGColorSpaceGenericRGB), alphaInfo.rawValue)
		return bmContext
	}
}

// MARK: - NYXAVCEncoderDelegate
extension NYXGIFToMP4Encoder : NYXAVCEncoderDelegate
{
	public func didFailToEncodeFrame()
	{
		print("didFailToEncodeFrame")
	}

	public func didEncodeFrame(frame: CMSampleBuffer)
	{
		self.videoWriterInput.appendSampleBuffer(frame)
	}
}