如何在Halide中使用Domains计算任意通道数图像的通道最大/最小值

背景 #

我目前正在试用Halide，尝试计算图像所有通道的最大值与最小值，希望该实现可适配通道数仅在运行时才能确定的任意图像。

已验证的可行实现 #

我已经通过显式循环写出了可以正常运行的版本：

#include "Halide.h"
#include "halide_image_io.h"
using namespace Halide::Tools;

int main(int argc, char **argv) {
    Halide::Buffer<uint8_t> input = load_image("rgb.png");

    Halide::Var x, y;
    Halide::Func max_channels, min_channels;
    max_channels(x, y) = input(x, y, 0);
    min_channels(x, y) = input(x, y, 0);
    for (int i = 1; i < input.channels(); ++i)
    {
        max_channels(x, y) = max(max_channels(x, y), input(x, y, i));
        min_channels(x, y) = min(min_channels(x, y), input(x, y, i));
    }

    {
        Halide::Buffer<uint8_t> output =
            max_channels.realize({input.width(), input.height(), 1});
        save_image(output, "maximum.png");
    }
    {
        Halide::Buffer<uint8_t> output =
            min_channels.realize({input.width(), input.height(), 1});
        save_image(output, "minimum.png");
    }
    printf("Success!\n");

    return 0;
}

问题描述 #

我希望去掉代码里的显式for循环，查阅Halide文档后认为Halide::RDom可以实现该需求，于是写出了如下代码：

#include "Halide.h"
#include "halide_image_io.h"
using namespace Halide::Tools;

int main(int argc, char **argv) {
    Halide::Buffer<uint8_t> input = load_image("rgb.png");

    Halide::Var x, y;
    Halide::Func max_channels, min_channels;
    Halide::RDom r(input);
    min_channels(x, y) = uint8_t{255};
    max_channels(x, y) = uint8_t{0};
    min_channels(r.x, r.y) = minimum(input(r.x, r.y, r.z));
    max_channels(r.x, r.y) = maximum(input(r.x, r.y, r.z));

    {
        Halide::Buffer<uint8_t> output =
            max_channels.realize({input.width(), input.height(), 1});
        save_image(output, "maximum.png");
    }
    {
        Halide::Buffer<uint8_t> output =
            min_channels.realize({input.width(), input.height(), 1});
        save_image(output, "minimum.png");
    }
    printf("Success!\n");

    return 0;
}

这段代码可以正常编译，但运行时崩溃，错误信息如下：

terminate called after throwing an instance of 'Halide::CompileError'
what():  Error: In update definition 0 of Func "f1":
Tuple element 0 of update definition has type uint8, but pure definition has type int32

Aborted (core dumped)

我确认代码里所有显式定义的变量都是uint8_t类型，不清楚int32_t类型是哪里来的。想知道用RDom计算逐像素多通道最值的正确写法，还是该需求只能用显式for循环实现？

解答 #

报错原因 #

代码存在两处核心问题：

• RDom定义错误：RDom r(input)会遍历输入图像的x、y、z三个全维度，而你需要的是仅遍历z（通道）维度，对固定(x,y)坐标下的所有通道值求最值。
• 类型隐式转换问题：Halide内置的minimum/maximum规约函数默认返回int32类型，你没有做显式类型转换，就会出现类型不匹配报错。

正确实现 #

只需要将RDom限定在通道维度，再用迭代更新的方式计算最值即可，完全不需要显式宿主侧for循环，也能适配运行时才确定的通道数：

#include "Halide.h"
#include "halide_image_io.h"
using namespace Halide::Tools;

int main(int argc, char **argv) {
    Halide::Buffer<uint8_t> input = load_image("rgb.png");

    Halide::Var x, y;
    Halide::Func max_channels, min_channels;
    // RDom仅遍历通道维度，范围是0到输入的通道数
    Halide::RDom r(0, input.channels());
    // 初始值保持uint8类型
    min_channels(x, y) = uint8_t{255};
    max_channels(x, y) = uint8_t{0};
    // 对每个(x,y)遍历所有通道，迭代更新最值
    min_channels(x, y) = min(min_channels(x, y), input(x, y, r));
    max_channels(x, y) = max(max_channels(x, y), input(x, y, r));

    {
        Halide::Buffer<uint8_t> output =
            max_channels.realize({input.width(), input.height(), 1});
        save_image(output, "maximum.png");
    }
    {
        Halide::Buffer<uint8_t> output =
            min_channels.realize({input.width(), input.height(), 1});
        save_image(output, "minimum.png");
    }
    printf("Success!\n");

    return 0;
}

内容的提问来源于stack exchange，提问作者jan.sende