C# 中图像和 OnnxRuntime.Tensor 的互转

因毕设需要，尝试了将 PyTorch 模型转为 ONNX 然后用 C# 做推理，记录一下经验。

总的来说，C# 对深度学习的支持远不如 Python，缺少很多必要的库，不少代码需要自己编写。

思路

毕设做的是 image-to-image 的 low-level 视觉任务，因此需要 3 个主要步骤：

从硬盘读取图片并转成张量（image to tensor）
使用 OnnxRuntime 完成推理（tensor to tensor）
将张量转回图片并保存（tensor to image）

找了一圈居然没发现 C# 的官方图像库。System.Drawing 勉强算一个，但这个玩意是 Windows only 的，.NET Core 不支持。那我们自然不要。

在第三方图像库里，SixLabors.ImageSharp 可能是最流行的那个，在网上能找到的一些 C# 做视觉相关的 ONNX 推理任务的 demo 中，这个库也比较常用。因此选定这个库作为图像库。

但是，ImageSharp 并没有提供到 Microsoft.ML.OnnxRuntime.Tensor 的转换接口，反之 Microsoft.ML.OnnxRuntime.Tensor 也没有提供来自 ImageSharp 的转换接口，因此转换的代码必须自己编写。（吐槽：在 Python 中，就是一句 np.array() 的事）

编写转换代码

在使用以下代码之前，确保你安装了 Microsoft.ML.OnnxRuntime 和 SixLabors.ImageSharp 这两个 NuGet 包。

Image to Tensor

其实就是遍历 Image 的所有像素，然后按顺序拷贝到一个新的 DenseTensor 中。注意使用 DenseTensor，因为它在内存中是连续存储的。此外还要注意除以 255 和数据类型转换（隐式完成），因为 Image 对象的像素值是 \([0, 255]\) 的 8 位整型，而绝大多数深度学习模型接受的张量值范围是 \([0, 1]\) 的 32 位浮点型。

public static DenseTensor<float> RgbImageToTensor(Image<Rgb24> image)
{
    // 大多数深度学习视觉模型接受的输入格式是 BCHW，即 batch-channel-height-width。
    // 在这里，batch 固定为1（在模型训练时它们可能是一个较大的值）
    // channel 因为是 RGB 的关系所以固定为 3
    // height 和 width 则可以变动。请注意，这种可变输入尺寸需要你的 ONNX 模型在导出时显式指定。
    var tensor = new DenseTensor<float>(new[] {1, 3, image.Height, image.Width});
    image.ProcessPixelRows(accessor =>
    {
        for (var y = 0; y < image.Height; y++)
        {
            var pixelSpan = accessor.GetRowSpan(y);
            for (var x = 0; x < image.Width; x++)
            {
                tensor[0, 0, y, x] = pixelSpan[x].R / 255f;
                tensor[0, 1, y, x] = pixelSpan[x].G / 255f;
                tensor[0, 2, y, x] = pixelSpan[x].B / 255f;
            }
        }
    });

    return tensor;
}

上面的代码适用于 RGB 图像。如果是灰度图像，那么只需要拷贝一个通道即可。

Clamp

在图像通过深度学习模型之后，图像张量中的一些值可能会超过 \([0, 1]\) 的范围。在 PyTorch 中，一般都会在模型推理之后紧跟一个 clamp 操作：

pred = model(input_img)

# 计算 PSNR/SSIM 等评价指标：略
# ...

pred = torch.clamp(pred, 0, 1)

# 转成图像格式后保存：略
# ...

torch.clamp 可以将图像的值限制在指定的区间之内。对于上面的示例代码， torch.clamp 将

问了一下 GPT，说是 C# 里没有类似的操作，以 Microsoft.ML.OnnxRuntime 和 SixLabors.ImageSharp 为关键字搜索了一下，也没找到。所以只好自己写了。对性能要求不太高的话很好写：

private static float Clamp(float value, float min, float max)
{
    return (value < min) ? min : (value > max) ? max : value;
}

在绝大多数情况下，value 的值都会介于 min 和 max 之间。对于这种情况，上面的代码会执行两次比较。我尝试了下面这种写法：

private static float Clamp(float value, float min, float max)
{
    if (value < min || value > max)
    {
        return (value < min) ? min : max;
    }
    return value;
}

经过实测，性能没有明显的提升。可能编译器替我们做了某些优化。

Tensor to Image

有了 Clamp 之后，可以编写 Tensor to Image 的代码了：

public static Image<Rgb24> TensorToRgbImage(DenseTensor<float> tensor)
{
    var height = tensor.Dimensions[2];
    var width = tensor.Dimensions[3];
    var image = new Image<Rgb24>(width, height);

    image.ProcessPixelRows(pixelAccessor =>
    {
        for (var y = 0; y < height; y++)
        {
            var rowSpan = pixelAccessor.GetRowSpan(y);

            for (var x = 0; x < width; x++)
            {
                var r = Clamp(tensor[0, 0, y, x] * 255, 0, 255);
                var g = Clamp(tensor[0, 1, y, x] * 255, 0, 255);
                var b = Clamp(tensor[0, 2, y, x] * 255, 0, 255);

                rowSpan[x] = new Rgb24((byte)r, (byte)g, (byte)b);
            }
        }
    });

    return image;
}

合起来

把以上三部分代码合起来：

using Microsoft.ML.OnnxRuntime.Tensors;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;

public static class Convert
{
    // 把上面的三个函数塞里面
}

可能的优化

很容易想到，上面的代码的性能是很差的。

可以参考此处。

实际上，上面的代码就是根据这里的代码改来的，但是这份代码我看不太懂。

测试

自己准备一个 image-to-image 的 ONNX 模型，这一类模型里最好找的应该是做超分辨率任务的。你可以在此处找到许多用于超分辨率的预训练模型。你也可以使用自己的模型。

然后，编写执行 ONNX 推理的代码：

using Microsoft.ML.OnnxRuntime;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using System.Diagnostics;

public class OnnxInference
{
    private readonly InferenceSession _session;

    public OnnxInference(string modelPath)
    {
        // check the path
        if (!File.Exists(modelPath))
        {
            throw new FileNotFoundException("Model file not found", modelPath);
        }
        // create inference session
        _session = new InferenceSession(modelPath);
    }

    public void Inference(string inputImagePath, string outputImagePath)
    {
        long startingTime;
        long endingTime;
        
        // check the path
        if (!File.Exists(inputImagePath))
        {
            throw new FileNotFoundException("Image file not found", inputImagePath);
        }
        // read the image
        using var image = Image.Load<Rgb24>(inputImagePath);
        var origHeight = image.Height;
        var origWidth = image.Width;
        var height = origHeight;
        var width = origWidth;
        // 在这里，你需要根据自己模型的要求将输入图像的尺寸调整到合适大小。
        // 例如，对于我的模型，长和宽需要是16的整数倍。
        // 你应该记录原始的图像尺寸，这样在得到结果图像之后可以 resize 回去。
        if (origHeight % 16 != 0 || origWidth % 16 != 0)
        {
            height = origHeight + 16 - origHeight % 16;
            width = origWidth + 16 - origWidth % 16;
            // Resize image
            image.Mutate(x =>
            {
                x.Resize(new ResizeOptions
                {
                    Size = new Size(width, height),
                    Mode = ResizeMode.Stretch
                });
            });
        }
        // 我在测试时记录了一些关键操作的耗时，你可以根据需要删去这些代码。
        startingTime = Stopwatch.GetTimestamp();
        // convert the image to tensor
        var tensor = Convert.RgbImageToTensor(image);
        endingTime = Stopwatch.GetTimestamp();
        Console.WriteLine($"Convert image to tensor: {(endingTime - startingTime) / (double)Stopwatch.Frequency}");
        
        // 这里的输入，需要根据模型导出时指定的输入来写。
        // create input container
        var inputs = new List<NamedOnnxValue>
        {
            NamedOnnxValue.CreateFromTensor("data_with_est", tensor),
            NamedOnnxValue.CreateFromTensor("data", tensor),
        };
        startingTime = Stopwatch.GetTimestamp();
        // run the inference
        using IDisposableReadOnlyCollection<DisposableNamedOnnxValue> results = _session.Run(inputs);
        endingTime = Stopwatch.GetTimestamp();
        Console.WriteLine($"Run inference: {(endingTime - startingTime) / (double)Stopwatch.Frequency}");
        // get the output tensor
        var output = results[0].AsTensor<float>();
        // print the shape of the output tensor
        Console.WriteLine("Output tensor shape:");
        foreach (var dimension in output.Dimensions)
        {
            Console.Write($"{dimension} ");
        }
        
        startingTime = Stopwatch.GetTimestamp();
        // convert the tensor to image
        var resultImage = Convert.TensorToRgbImage(output.ToDenseTensor());
        endingTime = Stopwatch.GetTimestamp();
        Console.WriteLine($"Convert tensor to image: {(endingTime - startingTime) / (double)Stopwatch.Frequency}");
        // 如果图像的尺寸修改过（意味着在推理之前执行了 resize），
        // 那么就要 resize 回去
        if (height != origHeight || width != origWidth)
        {
            // Resize image
            resultImage.Mutate(x =>
            {
                x.Resize(new ResizeOptions
                {
                    Size = new Size(origWidth, origHeight),
                    Mode = ResizeMode.Stretch
                });
            });
        }
        // save the image
        resultImage.Save(outputImagePath);
    }
}

然后使用以下的代码来使用上面的工具代码：

const string modelPath = @"D:\Path\to\your\onnx\model.onnx";
const string inputImagePath = @"D:\path\to\your\image.png";

var model = new OnnxInference(modelPath);
model.Inference(inputImagePath, "output.png");

上面的代码在我的电脑上针对一张 256 * 464 的图像的执行时间如下：

image to tensor: 47ms
inference: 616ms（取决于模型）
tensor to image: 75ms

绝对算不上快，但对我来说已经够了。因为我只需要推理单张图像。

tensor to image 过程稍慢，可能是因为 clamp 操作。

晚些日子我会试着用 Avalonia、ImageSharp 和 OnnxRuntime 写一个手写数字识别的程序，然后公开出来。

posted @ 2024-04-18 21:48 Eslzzyl 阅读(894) 评论(0) 编辑收藏举报

刷新页面返回顶部