[PYTHON] Konvertieren Sie RGB und HSV mit PyTorch in teilbare Form

Ich konnte es nicht unerwartet finden, daher ist es in HSV Color Space --Wikipedia geschrieben. Basierend auf der Formel implementiert.

Wenn Sie nur RGB und HSV konvertieren möchten, können Sie andere Bibliotheken (PIL, OpenCV usw.) verwenden, aber Sie können nicht davon unterscheiden, sodass Sie es nicht in das neuronale Netz integrieren können. Also habe ich eine Funktion geschrieben, um zwischen RGB und HSV nur mit der PyTorch-Funktion zu konvertieren.

Code

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

Wie benutzt man

Sowohl rgb2hsv als auch hsv2rgb nehmen Bilder in einem Mini-Stapel von Bildern (NCHW-Format) als Eingabe auf. Der Bereich von H (Farbton) reicht von 0 bis 360 (außerhalb der Bereichsschleifen). Der Bereich für RGB und SV liegt zwischen 0 und 1.