Metal图像处理——直方图均衡化
原创
前言
正文
核心思路
首先，我们用直方图来表示一张图像：横坐标代表的是颜色值，纵坐标代表的是该颜色值在图像中出现次数。
如图，对于某些图像，可能出现颜色值集中分布在某个区间的情况。
直方图均衡化(Histogram Equalization) ，指的是对图像的颜色值进行重新分配，使得颜色值的分布更加均匀。
本文用compute shader对图像的颜色值进行统计，然后计算得出映射关系，由fragment shader进行颜色映射处理。
效果展示
具体步骤
1、Metal的render管道、compute管道配置；
同前文，不再赘述，详见Metal入门教程总结。
2、CPU进行直方图均衡化处理；
- 2.1 把UIImage转成Bytes；
- 2.2 颜色统计；
// CPU进行统计
Byte *color = (Byte *)spriteData;
for (int i = 0; i < width * height; ++i) {
for (int j = 0; j < LY_CHANNEL_NUM; ++j) {
uint c = color[i * 4 + j];
++cpuColorBuffer.channel[j][c];
}
}- 2.3 映射关系；
int rgb[3][LY_CHANNEL_SIZE], sum = (int)(width * height);
int val[3] = {0};
// 颜色映射
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < LY_CHANNEL_SIZE; ++j) {
val[i] += cpuColorBuffer.channel[i][j];
rgb[i][j] = val[i] * 1.0 * (LY_CHANNEL_SIZE - 1) / sum;
}
}- 2.4 颜色值修改；
// 值修改
for (int i = 0; i < width * height; ++i) {
for (int j = 0; j < LY_CHANNEL_NUM; ++j) {
uint c = color[i * 4 + j];
color[i * 4 + j] = rgb[j][c];
}
}最后用处理之后的Bytes生成新图片。
3 GPU进行直方图均衡化处理；
- 3.1 compute shader进行颜色统计；
kernel void
grayKernel(texture2d<float, access::read> sourceTexture [[textureLYKernelTextureIndexSource]], // 纹理输入，
device LYColorBuffer &out [[buffer(LYKernelBufferIndexOutput)]], // 输出的buffer
uint2 grid [[thread_position_in_grid]]) // 格子索引
{
// 边界保护
if(grid.x < sourceTexture.get_width() && grid.y < sourceTexture.get_height())
{
float4 color = sourceTexture.read(grid); // 初始颜色
int3 rgb = int3(color.rgb * SIZE); // 乘以SIZE，得到[0, 255]的颜色值
// 颜色统计，每个像素点计一次
atomic_fetch_add_explicit(&out.channel[0][rgb.r], 1, memory_order_relaxed);
atomic_fetch_add_explicit(&out.channel[1][rgb.g], 1, memory_order_relaxed);
atomic_fetch_add_explicit(&out.channel[2][rgb.b], 1, memory_order_relaxed);
}
}atomic_fetch_add_explicit是用于在多线程进行数据操作，具体的函数解释见这里。
- 3.2 映射关系处理； compute shader回调后，根据GPU统计的颜色分布结果，求出映射关系；
LYLocalBuffer *buffer = (LYLocalBuffer *)strongSelf.colorBuffer.contents; // GPU统计的结果
LYLocalBuffer *convertBuffer = self.convertBuffer.contents; // 颜色转换的buffer
int sum = (int)(self.sourceTexture.width * self.sourceTexture.height); // 总的像素点
int val[3] = {0}; // 累计和
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < LY_CHANNEL_SIZE; ++j) {
val[i] += buffer->channel[i][j]; // 当前[0, j]累计出现的总次数
convertBuffer->channel[i][j] = val[i] * 1.0 * (LY_CHANNEL_SIZE - 1) / sum;
// 对比CPU和GPU处理的结果
if (buffer->channel[i][j] != strongSelf->cpuColorBuffer.channel[i][j]) {
// 如果不相同，则把对应的结果输出
printf("%d, %d, gpuBuffer:%u cpuBuffer:%u \n", i, j, buffer->channel[i][j], strongSelf->cpuColorBuffer.channel[i][j]);
}
}
}
memset(buffer, 0, strongSelf.colorBuffer.length);3.3 根据映射关系处理原图片，并渲染到屏幕上；
fragment float4
samplingShader(RasterizerData input [[stage_in]], // stage_in表示这个数据来自光栅化。（光栅化是顶点处理之后的步骤，业务层无法修改）
texture2d<float> colorTexture [[ texture(LYFragmentTextureIndexSource) ]], // texture表明是纹理数据，LYFragmentTextureIndexSource是索引
device LYLocalBuffer &convertBuffer [[buffer(LYFragmentBufferIndexConvert)]]) // 转换的buffer
{
constexpr sampler textureSampler (mag_filter::linear, min_filter::linear); // sampler是采样器
float4 colorSample = colorTexture.sample(textureSampler, input.textureCoordinate); // 得到纹理对应位置的颜色
int3 rgb = int3(colorSample.rgb * SIZE); // 记得先乘以SIZE
colorSample.rgb = float3(convertBuffer.channel[0][rgb.r], convertBuffer.channel[1][rgb.g], convertBuffer.channel[2][rgb.b]) / SIZE; // 返回的值也要经过归一化处理
return colorSample;
}遇到的问题
1、统计结果集中在头部
问题表现：
统计结果异常，集中在前面两个值。
如下，green通道的颜色集中在r0和r1上：
// 0~255颜色值的分布
28269 4492 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0问题分析：
实际上，gpu里面存着的是0.0~1.0的值；（归一化）
统计的值全部是在前面，是因为没有乘以255！
先用CPU实现了直方图均衡化，在实现shader的时候，参考CPU的代码实现，犯了这个错误。
2、cpu和gpu统计结果相差较多
问题表现：
如下代码，buffer是gpu统计的颜色值分布结果，cpuColorBuffer是cpu统计的颜色值分布结果。
理论上结果应该接近，但实际上printf出来的差异非常多。
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < LY_CHANNEL_SIZE; ++j) {
val[i] += buffer->channel[i][j];
convertBuffer->channel[i][j] = val[i] * 1.0 * (LY_CHANNEL_SIZE - 1) / sum;
// 对比CPU和GPU处理的结果
if (buffer->channel[i][j] != strongSelf->cpuColorBuffer.channel[i][j]) {
// 如果不相同，则把对应的结果输出
printf("%d, %d, gpuBuffer:%u cpuBuffer:%u \n", i, j, buffer->channel[i][j], strongSelf->cpuColorBuffer.channel[i][j]);
}
}
}问题分析：
通过检查代码，先判定cpu统计的结果是正常。（cpu的处理过程就是正常的for循环，不易出错）
仔细观察log的不同：
0, 1, gpuBuffer:763 cpuBuffer:762
结果很接近，但是有细微的差距。
我们知道gpu是浮点数的处理，而cpu是整数型处理，浮点数到整数中间有精度的问题。
此时再看我们的shader，我们是以half来进行计算，这样统计出来的结果会有点误差。
grayKernel(texture2d<half, access::read> sourceTexture [[texture(LYFragmentTextureIndexTextureSource)]],
device LYColorBuffer &out [[buffer(LYKernelBufferIndexOutput)]],
uint2 grid [[thread_position_in_grid]])通过把精度从half改成float，cpu和gpu的统计差异就只有3个：
0, 248, gpuBuffer:23215 cpuBuffer:22854
1, 74, gpuBuffer:23201 cpuBuffer:22840
2, 64, gpuBuffer:23336 cpuBuffer:229753、gpu渲染的图片为白色
问题表现：
在gpu统计的结果与cpu接近的情况下，把映射buffer传给fragment shader，最后进行一次颜色处理。
但是结果是白色的图片，shader的代码如下：
fragment float4
samplingShader(RasterizerData input [[stage_in]], // stage_in表示这个数据来自光栅化。（光栅化是顶点处理之后的步骤，业务层无法修改）
texture2d<float> colorTexture [[ texture(LYFragmentTextureIndexTextureSource) ]], // texture表明是纹理数据，LYFragmentTextureIndexTextureSource是索引
device LYLocalBuffer &localBuffer [[buffer(LYFragmentBufferIndexConvert)]])
{
constexpr sampler textureSampler (mag_filter::linear,
min_filter::linear); // sampler是采样器
float4 colorSample = colorTexture.sample(textureSampler, input.textureCoordinate); // 得到纹理对应位置的颜色
int3 rgb = int3(colorSample.rgb);
colorSample.rgb = float3(localBuffer.channel[0][rgb.r], localBuffer.channel[1][rgb.g], localBuffer.channel[2][rgb.b]);
return colorSample;
}问题分析：
我们先把colorSample.rgb = ...的这行代码屏蔽，发现渲染结果是正常的，那么问题就出现在映射处理上面。
再通过Xcode的Capture GPU Frame工具，查看传入的映射buffer数据，也是正常的数据。
那么问题可能出现int3 rgb的初始化，或者从映射buffer读取数据。
观察到int3 rgb = int3(colorSample.rgb)，是有一个float->int的操作，联想到前面提到的归一化处理，马上明白：在这里的初始化时应该乘以SIZE。
那么问题是否就此解决？不是的。
我们在进行颜色转换的时候，float->int 需要乘以SIZE；
在获取到映射buffer里面对应颜色的值后，仍需要做一次int->float的处理，除以SIZE；
如果下：
float4 colorSample = colorTexture.sample(textureSampler, input.textureCoordinate); // 得到纹理对应位置的颜色
int3 rgb = int3(colorSample.rgb * size);
colorSample.rgb = float3(localBuffer.channel[0][rgb.r], localBuffer.channel[1][rgb.g], localBuffer.channel[2][rgb.b]) / size;4、映射结果异常
问题表现：
问题如下，映射结果应该是0~255的值，但是通过Xcode看到最终的映射值远超过255，甚至接近255*2的数字。
问题分析：
下面是映射的算法
int rgb[3][LY_CHANNEL_SIZE], sum = (int)(width * height);
int val[3] = {0};
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < LY_CHANNEL_SIZE; ++j) {
val[i] += cpuColorBuffer.channel[i][j];
rgb[i][j] = val[i] * 1.0 * (LY_CHANNEL_SIZE - 1) / sum;
}
}sum是固定值，LY_CHANNEL_SIZE是常量值256，那么映射结果超过255的原因就是vali的统计结果太大！
通过Xcode调试，确实如此：
那么，会是什么原因导致？
在看到结果接近255的两倍时，大概猜测可能是重复运算导致。
我们的均衡化处理是在MTKView的回调进行，如下：
- (void)drawInMTKView:(MTKView *)view {
[self customDraw];
}这里会回调多次，从而导致多次执行compute shader的颜色统计，这里可以引入isDrawing的临时变量解决：
- (void)drawInMTKView:(MTKView *)view {
if (!self.isDrawing) {
self.isDrawing = YES;
[self customDraw];
}
}但是，问题并没有彻底解决：首次统计正常，但是第二次处理的时候就会累积上一次的统计值。
如何对值进行清理？（这里实际上只处理一次也行，但是debug过程中需要通过Xcode的GPU Capture Frame工具进行查看，而这个工具需要多次渲染）
我们知道MTLBuffer是cpu、gpu都可以操作的buffer，那么在cpu直接清除这个数据即可。
在 commandBuffer addCompletedHandler:^(){}的结束回调中，使用memset(buffer, 0, strongSelf.colorBuffer.length)清理统计结果。
5、映射结果最大值为256
问题表现：
在踩过上面的各种坑之后，直方图均衡化的效果也已经展现，但是仍有一点小问题：
映射结果buffer的数字范围是0~256，而不是255。
问题分析：
根据直方图均衡化的算法，我们知道是因为像素颜色值的统计，结果稍微偏大。
回顾Compute Shader的代码：
kernel void
grayKernel(texture2d<float, access::read> sourceTexture [[texture(LYFragmentTextureIndexTextureSource)]],
device LYColorBuffer &out [[buffer(LYKernelBufferIndexOutput)]],
uint2 grid [[thread_position_in_grid]]) {
// 边界保护
if(grid.x <= sourceTexture.get_width() && grid.y <= sourceTexture.get_height())
{
float4 color = sourceTexture.read(grid); // 初始颜色
int3 rgb = int3(color.rgb * size); //
atomic_fetch_add_explicit(&out.channel[0][rgb.r], 1, memory_order_relaxed);
atomic_fetch_add_explicit(&out.channel[1][rgb.g], 1, memory_order_relaxed);
atomic_fetch_add_explicit(&out.channel[2][rgb.b], 1, memory_order_relaxed);
}
}Metal的Compute Shader是按每组网格进行处理，那么可能会出现边界超过图像的情况，所以添加了边界保护。
但是，这里存在误判的情况：边界判断不应该是<=，而是<。
因为这个原因，会导致统计的结果偏大，最终出现256的情况。
在修复完这个问题后，Demo再无其他问题，GPU的处理结果也和CPU的处理结果完全一致！
总结
本文是在前文的Metal入门教程基础上进行更复杂的尝试，过程中也遇到较多问题，最终demo也顺利完成，地址在Github。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。