color.rgb_to_hsv can't be converted to simplified onnx model

Describe the bug

I’ve created a simple NN that simply convert RGB to HSV using the color.rgb_to_hsv() function. When I try to convert the model to Onnx and then simplify it (or convert it to blob format using this converter: http://blobconverter.luxonis.com) I get this error message: This is an invalid model. Type Error: Type ‘tensor(bool)’ of input parameter (3) of operator (CumSum) in node (CumSum_4) is invalid.

Reproduction steps

Use this code to reproduced the error:


import torch
from torch import nn

import kornia

import onnx
from onnxsim import simplify

class rgb2hsv(nn.Module):
    def forward(self, image):

        hsv_tensor = kornia.color.rgb_to_hsv(image)

        return hsv_tensor

X = torch.ones((1, 3, 300, 300), dtype=torch.float32)
torch.onnx.export(
    rgb2hsv(),
    (X), # Dummy input for shape
    "model.onnx",
    opset_version=12,
    do_constant_folding=True,
)


onnx_model = onnx.load("model.onnx")
model_simpified, check = simplify(onnx_model)
onnx.save(model_simpified, "model_simplified.onnx

### Expected behavior

Expecting the model to convert correctly

### Environment

```shell
- PyTorch Version (e.g., 1.0): 1.9.1
- OS (e.g., Linux): MacOS
- How you installed PyTorch (`conda`, `pip`, source): Pip
- Build command you used (if compiling from source): NA
- Python version: 3.8
- CUDA/cuDNN version: NA
- GPU models and configuration: NA
- Any other relevant information: NA

Additional context

No response

[PINTO0309]

My workaround for now. There are causes for errors in both kornia and onnx-simplifier. Thus, it removes the error sources for both kornia and onnx-simplifier.

• kornia

  • myriad_compilier fails due to the fact that CumSum is set to a negative axis. Exporting to ONNX and converting to OpenVINO IR can be done without problems. Therefore, I have created an approximation logic that does not make heavy use of boolean masks. OpenVINO does not fail if the axis of CumSum is set to a negative value, but it only fails when myriad_compilier is executed. This is obviously a problem specific to myriad_compilier. It’s not a kornia problem. This test should always complete successfully.

• onnx-simplifier

  • An error occurs when performing an internal optimization operation (Gather). The last two lines of logic are essentially unnecessary, but I left them in to illustrate that there is an option to disable the optimization logic of onnx-simplifier. model.onnx and model_non-opt.onnx will generate the exact same file.

• Wikipedia https://en.wikipedia.org/wiki/HSL_and_HSV

import torch
from torch import nn
import kornia

import onnx
from onnxsim import simplify

class rgb2hsv(nn.Module):
    def forward(self, image):
        hsv_tensor = kornia.color.rgb_to_hsv(image)
        return hsv_tensor

class rgb2hsv_custom(nn.Module):
    def forward(self, image, epsilon=1e-10):
        assert(image.shape[1] == 3)
        r, g, b = image[:, 0], image[:, 1], image[:, 2]
        max_rgb, argmax_rgb = image.max(1)
        min_rgb, argmin_rgb = image.min(1)
        max_min = max_rgb - min_rgb + epsilon
        h1 = 60.0 * (g - r) / max_min + 60.0
        h2 = 60.0 * (b - g) / max_min + 180.0
        h3 = 60.0 * (r - b) / max_min + 300.0
        h = torch.stack((h2, h3, h1), dim=0).gather(dim=0, index=argmin_rgb.unsqueeze(0)).squeeze(0)
        s = max_min / (max_rgb + epsilon)
        v = max_rgb
        return torch.stack((h, s, v), dim=1)

class hsv2rgb_custom(nn.Module):
    def forward(self, image):
        assert(image.shape[1] == 3)
        h, s, v = image[:, 0], image[:, 1], image[:, 2]
        h_ = (h - torch.floor(h / 360) * 360) / 60
        c = s * v
        x = c * (1 - torch.abs(torch.fmod(h_, 2) - 1)) #The abs function causes an error in myriad_compilier.
        zero = torch.zeros_like(c)
        y = torch.stack((
            torch.stack((c, x, zero), dim=1),
            torch.stack((x, c, zero), dim=1),
            torch.stack((zero, c, x), dim=1),
            torch.stack((zero, x, c), dim=1),
            torch.stack((x, zero, c), dim=1),
            torch.stack((c, zero, x), dim=1),
        ), dim=0)
        index = torch.repeat_interleave(torch.floor(h_).unsqueeze(1), 3, dim=1).unsqueeze(0).to(torch.long)
        rgb = (y.gather(dim=0, index=index) + (v - c)).squeeze(0)
        return rgb

X = torch.ones((1, 3, 300, 300), dtype=torch.float32)

torch.onnx.export(
    rgb2hsv_custom(),
    (X),
    "model.onnx",
    opset_version=12,
    do_constant_folding=True,
)
onnx_model = onnx.load("model.onnx")
model_simpified, check = simplify(onnx_model, perform_optimization=False)
onnx.save(model_simpified, "model_non-opt.onnx")

Note that I have not firmly verified that the RGB to HSV conversion custom logic works correctly. This procedure does not dare to optimize ONNX and will produce a very redundant model. Therefore, if you do not like it, please ignore it.

• ONNX

• OpenVINO

• Myriad Blob model.blob.tar.gz

By the way, I don’t know what kind of tool kornia is. One last thing. If you need to convert models between frameworks, it is recommended that you do not use boolean masks. This is because the implicit type conversion between True/False and 0/-1 will fail in frameworks that require strict type specification. For example, Myriad or EdgeTPU, TensorFlow Lite.

[edgarriba]

@simetin let’s keep it open until we merge #1329. We need to create a proper mechanism to test from kornia side that after we export to onnx, later is will be usable. Because in this case, the exporter was was working and later you found that error.