swift vImage_Buffer包装器

Posted 2021-05-19

//  Copyright © 2015 Indragie Karunaratne. All rights reserved.

import Foundation
import Accelerate

public class Image {
    private var buffer: vImage_Buffer
    
    public var size: CGSize {
        return vImageBuffer_GetSize(&buffer)
    }
    
    public var integralSize: (width: Int, height: Int) {
        let size = self.size
        return (Int(size.width), Int(size.height))
    }
    
    // MARK: Initialization
    
    private init(buffer: vImage_Buffer) {
        self.buffer = buffer
    }
    
    public convenience init(CGImage: CGImageRef) throws {
        var format = ARGBFFFFFormat.vImageFormat
        var inBuf = vImage_Buffer()
        try handleError(vImageBuffer_InitWithCGImage(&inBuf, &format, nil, CGImage, UInt32(kvImageNoFlags)))
        self.init(buffer: inBuf)
    }
    
    private convenience init(imageSource: CGImageSourceRef) throws {
        if let image = CGImageSourceCreateImageAtIndex(imageSource, 0, nil) {
            try self.init(CGImage: image)
        } else {
            throw ImageProcessingError.InvalidImageObject
        }
    }
    
    public convenience init(data: NSData) throws {
        if let source = CGImageSourceCreateWithData(data, nil) {
            try self.init(imageSource: source)
        } else {
            throw ImageProcessingError.InvalidImageObject
        }
    }
    
    public convenience init(URL: NSURL) throws {
        if let source = CGImageSourceCreateWithURL(URL, nil) {
            try self.init(imageSource: source)
        } else {
            throw ImageProcessingError.InvalidImageObject
        }
    }
    
    deinit {
        free(buffer.data)
    }
    
    // MARK: Transformations
    
    internal func transform(width width: Int, height: Int, transform: (inBuf: inout vImage_Buffer, outBuf: inout vImage_Buffer) throws -> ()) rethrows -> Image {
        let bytesPerRow = ARGBFFFFFormat.BytesPerPixel * width
        let data = calloc(width * height * ARGBFFFFFormat.BytesPerPixel, sizeof(Float))
        var outBuffer = vImage_Buffer(data: data, height: UInt(height), width: UInt(width), rowBytes: bytesPerRow)
        try transform(inBuf: &buffer, outBuf: &outBuffer)
        return Image(buffer: outBuffer)
    }
    
    public func toCGImage() throws -> CGImageRef {
        var format = ARGBFFFFFormat.vImageFormat
        var rawError: vImage_Error = kvImageNoError
        let image = vImageCreateCGImageFromBuffer(&buffer, &format, nil, nil, UInt32(kvImageNoFlags), &rawError)
        try handleError(rawError)
        return image.takeRetainedValue()
    }
}

private struct ARGBFFFFFormat {
    static let BitsPerComponent = 32
    static let BytesPerPixel = 16
    static let BitmapInfo = CGBitmapInfo(rawValue:
          CGBitmapInfo.ByteOrder32Little.rawValue
        | CGBitmapInfo.FloatComponents.rawValue
        | CGImageAlphaInfo.PremultipliedFirst.rawValue
    )
    
    static let vImageFormat: vImage_CGImageFormat = {
        var format = vImage_CGImageFormat()
        format.bitsPerComponent = UInt32(ARGBFFFFFormat.BitsPerComponent)
        format.bitsPerPixel = UInt32(ARGBFFFFFormat.BytesPerPixel) * 8
        format.bitmapInfo = ARGBFFFFFormat.BitmapInfo
        return format
    }()
}