UIImageの特定の色だけ変更したい

Swift で画像をいじっている時、特定の色だけ変更したいことがあって詰まってしまったので実装を共有します。
順序は以下の通りです。

• 1. red, green, blue, alpha を持つ Pixel 構造体を定義する。
• 2. UIImage を [Pixel] に変換する
• 3. [Pixel] のうち、対象の色を for で回して色の変更を行う
• 4. [Pixel] を [UInt8] にする
• 5. UIImage にする

かなり大変なのでもし簡単な方法があったら教えていただけると嬉しいです。

1. red, green, blue, alpha を持つ Pixel 構造体を定義する。

これは簡単ですね。

struct Pixel {
    var red: UInt8
    var green: UInt8
    var blue: UInt8
    var alpha: UInt8
    
    init(red: UInt8, green: UInt8, blue: UInt8, alpha: UInt8) {
        self.red = red
        self.green = green
        self.blue = blue
        self.alpha = alpha
    }

    func isBlack() -> Bool {
        return self.red == 0 && self.green == 0 && self.blue == 0 && self.alpha == 255
    }
}

特定の色の変更を行いたいのでその実装を行います。今回は特定の色=黒で行きましょう。

2. UIImage を [Pixel] に変換する

RGBA形式を想定しているので、並びが違う場合は適宜変えてください。

extension UIImage {
    func pixelArrayFromImage() -> [Pixel] {
        let data = self.cgImage?.dataProvider?.data
        let length = CFDataGetLength(data)
        var rawData = [UInt8](repeating: 0, count: length)
        CFDataGetBytes(data, CFRange(location: 0, length: length), &rawData)
        
        var pixelArray: [Pixel] = []
        // どんどん Pixel にして詰めていく
        for i in stride(from: 0, to: rawData.count, by: 4) {
            pixelArray.append(Pixel.init(red: rawData[i], green: rawData[i+1], blue: rawData[i+2], alpha: rawData[i+3]))
        }
        
        return pixelArray
    }
}

3. [Pixel] のうち、対象の色を for で回して色の変更を行う

今回は黒を透明にしたいので、こんな感じの実装にします。
Pixel 構造体に以下を付け足してください。
(新たにPixelをインスタンス化しているのなんとかしたい...)

    func blackToClear() -> Pixel {
        if self.isBlack() {
            return Pixel(red: self.red, green: self.green, blue: self.blue, alpha: 0)
        }
        return self
    }

4. [Pixel] を [UInt8] にする

Pixel 構造体に以下を足しましょう。

    func toBuffer() -> [UInt8] {
        return [self.red, self.green, self.blue, self.alpha]
    }

5. UIImage にする

こんな感じですね。
これは UIImage への拡張に足しましょう。
また、vImage_hoge や @State を使うために Accelrate と SwiftUI を import しておきましょう。

    static func imageFromPixelBuffer(@State pixelBuffer: [UInt8], width: Int = 500, height: Int = 490) -> UIImage? {
        let argbImage: CGImage? = pixelBuffer.withUnsafeMutableBufferPointer { pixelPointer in
            let buffer = vImage_Buffer(data: pixelPointer.baseAddress!,
                                       height: vImagePixelCount(height),
                                       width: vImagePixelCount(width),
                                       rowBytes: width * 4)
            guard let format = vImage_CGImageFormat(bitsPerComponent: 8,
                                                    bitsPerPixel: 8 * 4,
                                                    colorSpace: CGColorSpaceCreateDeviceRGB(),
                                                    bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.last.rawValue | CGImageByteOrderInfo.orderDefault.rawValue)) else {
                return nil
            }
            return try? buffer.createCGImage(format: format)
        }
        
        return self.init(cgImage: argbImage!)
    }

以上のコードをエントリポイントから呼び出します。

let image = UIImage(named: "hoge.jpg")

let pixelArray = image?.pixelArrayFromImage()

var pixelBuffer: [UInt8] = []
for pixel in pixelArray! {
    var pixel = pixel.blackToClear()
    pixelBuffer.append(contentsOf: pixel.toBuffer())
}
let changedImage = UIImage.imageFromPixelBuffer(pixelBuffer: pixelBuffer)

以上になります。読んでいただきありがとうございました。