PyTorchでRGBとHSVを微分可能な形で変換する

意外と見つからなかったので、HSV色空間 - Wikipediaに書いてある式をもとに実装しました。

ただRGBとHSVを変換するだけなら他のライブラリ（PIL、OpenCVとか）でもできますが、それでは微分できないので、ニューラルネットの内部に組み込むことができません。
そこで、PyTorchの関数のみを使ってRGBとHSVを相互に変換する関数を書きました。

2020-10-31 追記
よく調べてみたらkorniaという便利なライブラリがあるようです。
korniaはコンピュータービジョン系の処理をPyTorchを使って微分可能な形で実装したライブラリで、そこにRGB-HSVの変換もあります。
多分korniaの実装のほうがいろいろ良いので、korniaを使いましょう。

コード

color_convert.py

import torch


def rgb2hsv(input, epsilon=1e-10):
    assert(input.shape[1] == 3)

    r, g, b = input[:, 0], input[:, 1], input[:, 2]
    max_rgb, argmax_rgb = input.max(1)
    min_rgb, argmin_rgb = input.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)


def hsv2rgb(input):
    assert(input.shape[1] == 3)

    h, s, v = input[:, 0], input[:, 1], input[:, 2]
    h_ = (h - torch.floor(h / 360) * 360) / 60
    c = s * v
    x = c * (1 - torch.abs(torch.fmod(h_, 2) - 1))

    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

使い方

rgb2hsvとhsv2rgbはどちらも画像のミニバッチ（NCHW形式）の画像を入力とします。
H（色相）の範囲は0から360までです（範囲外はループします）。
RGBとSVの範囲は0から1までです。