前几天手搓了一个照片调色的小 App,还发了个视频出来。虽然这个视频没有人感兴趣,不过这个开发的过程中有一些点我觉得还是值得分享出来的。[视频链接]
这个 App 大致的功能就是:
这个主要的功能在开发的过程中也是遇到了一些坑,最终效果呢距离预期也有那么一点点差距。这个事儿之后找机会再分享吧。总之,做到了现在这个——效果上将将可以接受的程度。
我遇到的问题是什么呢?
说起来挺简单的哈,做的时候其实也不是特别难,没花多少时间就撸出来了。跑起来一看,妈耶~,我都惊呆了!!
跑起来之后这个功能是正常的,但是呢,一张普通的 4K 图片,四千乘三千那种,跑了有6分多钟。对,我没说错,是6分钟,超过360秒。
虽然时间是挺长的,但是至少它是自动化的对不对?比一直用手扒拉(bu lu)那小滑块也还是方便了不少的,是吧?
我当然是不会就这样善罢甘休的,就开始了一些小小的优化。最终,达到了一张 4K 图片可以在 20秒左右处理完成,这基本上就达到了一个可以接受的程度。最终效果可以看我之前发的那个[视频],就是那个样子。
接下来就是比较技术的内容了,稍微分享一下这次优化的思路和技术实现。
我首先做的一件事就是把可以预处理的部分提取出来,我用来传给模型的参数里有一些是通过对原图片的颜色和直方图进行运算得出的数据,所以这些对原图的运算是可以提前到图片载入的时候就提前进行运算的。当然,这个操作只是给运算压力换一个位置,并没有解决问题。然而通过遍历图片所有像素来做这些运算是非常非常耗时的。
然后我就跟 GPT 沟通了一下,怹给了我一些建议,可以试试 Accelerate 这个库。
Accelerate 这个是非常强大的,专门为大规模数学运算和图像运算设计的一个库。而且苹果的文档也是很给力的,有针对具体使用场景提供的文档。
我对预处理部分由该后的代码片段:
分离 RGB 各个通道
vImageConvert_ARGB8888toPlanar8
var sourceBuffer = try! vImage_Buffer(cgImage: cgImage)
var rChannel = [UInt8] (repeating: 0, count: pixelCount)
var gChannel = [UInt8] (repeating: 0, count: pixelCount)
var bChannel = [UInt8] (repeating: 0, count: pixelCount)
var aChannel = [UInt8] (repeating: 0, count: pixelCount)
rChannel.withUnsafeMutableBytes { rChannelPointer in
gChannel.withUnsafeMutableBytes { gChannelPointer in
bChannel.withUnsafeMutableBytes { bChannelPointer in
aChannel.withUnsafeMutableBytes { aChannelPointer in
var redBuffer = vImage_Buffer(data: rChannelPointer.baseAddress, height: vImagePixelCount(height), width: vImagePixelCount(width), rowBytes: Int(width))
var greenBuffer = vImage_Buffer(data: gChannelPointer.baseAddress, height: vImagePixelCount(height), width: vImagePixelCount(width), rowBytes: Int(width))
var blueBuffer = vImage_Buffer(data: bChannelPointer.baseAddress, height: vImagePixelCount(height), width: vImagePixelCount(width), rowBytes: Int(width))
var alphaBuffer = vImage_Buffer(data: aChannelPointer.baseAddress, height: vImagePixelCount(height), width: vImagePixelCount(width), rowBytes: Int(width))
// 分离 RGB 通道
vImageConvert_ARGB8888toPlanar8(&sourceBuffer, &redBuffer, &greenBuffer, &blueBuffer, &alphaBuffer, vImage_Flags(kvImageNoFlags))
}
}
}
}
计算各个通道的直方图
vImageHistogramCalculation_ARGB8888
var buffer = try! vImage_Buffer(cgImage: image)
var histogramRed = [vImagePixelCount](repeating: 0, count: 256)
var histogramGreen = [vImagePixelCount](repeating: 0, count: 256)
var histogramBlue = [vImagePixelCount](repeating: 0, count: 256)
var histogramAlpha = [vImagePixelCount](repeating: 0, count: 256)
histogramRed.withUnsafeMutableBufferPointer { zeroPtr in
histogramGreen.withUnsafeMutableBufferPointer { onePtr in
histogramBlue.withUnsafeMutableBufferPointer { twoPtr in
histogramAlpha.withUnsafeMutableBufferPointer { threePtr in
var histogramBins = [zeroPtr.baseAddress, onePtr.baseAddress,
twoPtr.baseAddress, threePtr.baseAddress]
histogramBins.withUnsafeMutableBufferPointer { histogramBinsPtr in
// `buffer` is a `vImage_Buffer` structure.
_ = vImageHistogramCalculation_ARGB8888(&buffer,
histogramBinsPtr.baseAddress!,
vImage_Flags(kvImageNoFlags))
}
}
}
}
}
数据归一化
vDSP_normalize
let floatChannelData = channel.map { Float($0) }
// 计算均值和标准差
var mean: Float = 0
var stdDev: Float = 0
vDSP_normalize(floatChannelData, 1, nil, 1, &mean, &stdDev, vDSP_Length(floatChannelData.count))
用了这个工具,直接让我的数据预处理部分的耗时从分钟级下降到了毫秒级。
接下来就到了重新组织数据的部分,这部分也是所有环节中最耗时的一个。之前用的是传统的方法,把像素数据放到数据数组里之后,遍历之前预处理的数据,一个一个复制到这个数据数组里。这无疑就是在逐像素遍历的基础上又在每个图块都又遍历了一遍额外的数据,这个预处理数据也有上百个元素呢,就是一个相当沉重的负担了。
这部分的优化主要做了两件事,一个是隐性的,把像素写入数据数组的代码尽量写成对并行友好的形式,这样的话,一方面编译器的优化能力可以发挥一些作用,能在一定程度上提高多核的并行度;另一方面以后可以再用多线程去进一步优化。
nonisolated func copyBlockDataToMLArray(pixelData: [UInt8], mlArray: MLMultiArray, blockX: Int, blockY: Int, blockSize: Int, blockIndex: Int, width: Int, height: Int, bytesPerPixel: Int) {
for ih in 0..<blockSize {
for iw in 0..<blockSize {
let x = blockX * blockSize + iw
let y = blockY * blockSize + ih
let isOut = x >= width || y >= height
let pixelIndex = MyHelper.getIndexFromXY(x: x, y: y, width: width) * bytesPerPixel
let r = isOut ? -1 : Double(pixelData[pixelIndex])
let g = isOut ? -1 : Double(pixelData[pixelIndex + 1])
let b = isOut ? -1 : Double(pixelData[pixelIndex + 2])
let di = MyHelper.getIndexFromXY(x: iw, y: ih, width: blockSize)
let batchIndex = blockIndex * inputWidth + di
mlArray[batchIndex ] = NSNumber(value: r)
mlArray[batchIndex + 1 * blockSize * blockSize] = NSNumber(value: g)
mlArray[batchIndex + 2 * blockSize * blockSize] = NSNumber(value: b)
}
}
}
另外一件事,是用 Swift 的指针操作直接将预处理数据拷贝到数据数组,时间复杂度直接从O(N) 变成 O(1) 。
当然,Swift 的指针没有 C++ 那么自由,写起来也需要适应一下。但是至少它可以用,并且是好用的。
for ibh in 0..<bh {
for ibw in 0..<bw {
copyBlockDataToMLArray(pixelData: pixelData, mlArray: inputArr, blockX: ibw, blockY: ibh, blockSize: blockSize, blockIndex: blockedCount, width: width, height: height, bytesPerPixel: bytesPerPixel)
// 将遍历赋值改为整块内存拷贝,速度大幅提升,4K图片从6分钟缩短为20秒
let extraDataHeadIndex = blockedCount * inputWidth + blockSize * blockSize * 3
inputArr.withUnsafeMutableBufferPointer(ofType: Double.self) { inputArrPtr, _ in
extraDatas.withUnsafeBufferPointer { extraDataPtr in
inputArrPtr.baseAddress!
.advanced(by: extraDataHeadIndex)
.update(from: extraDatas, count: extraDatas.count)
}
}
blockedCount += 1
// .......
// 其他代码
}
}
最后的部分就是将模型输出的数据重新拼接成图片,这里的思路是,每一个图块写到图片的像素缓存里可以没有固定的顺序,完全是可以并行执行的。
所以我就用了一个线程池,(所谓的线程池就是一个 OperationQueue),让每一个图块的写入任务并行起来,这块儿最终也提升了一些运行时间。
let fullOutputQueue = OperationQueue()
fullOutputQueue.maxConcurrentOperationCount = 10
// ......
// 其他代码
fullOutputQueue.addOperation {
// outPixelData 应该保持弱引用,暂时不管警告
self.copyMLArrayToOutputArray(input: result, output: &outPixelData, width: width, height: height, blockSize: self.blockSize, blockWidthCount: bw, blockedCount: _blockedCount, blockStartIndex: blockStartIndex, bytesPerPixel: bytesPerPixel)
}
最终,效果就像我之前发的[视频]里的那样,可以在20秒左右的时间处理完成一张 4K 的照片。达到了一个基本可用的状态。
最后总结一下关键点:
1. 使用 Accelerate 库优化大规模运算
2. 整块拷贝内存
3. 多线程并行处理
Xcode 16.1
Swift 5