[PYTHON] Convertissez RVB et HSV sous une forme divisible avec PyTorch

Je n'ai pas pu le trouver de manière inattendue, il est donc écrit dans HSV Color Space --Wikipedia. Mis en œuvre sur la base de la formule.

Si vous souhaitez simplement convertir RVB et HSV, vous pouvez utiliser d'autres bibliothèques (PIL, OpenCV, etc.), mais vous ne pouvez pas faire la différence avec cela, vous ne pouvez donc pas l'incorporer dans le réseau neuronal. J'ai donc écrit une fonction pour convertir entre RVB et HSV en utilisant uniquement la fonction de PyTorch.

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

Comment utiliser

«Rgb2hsv» et «hsv2rgb» prennent des images dans un mini-lot d'images (format NCHW) en entrée. La plage de H (teinte) est comprise entre 0 et 360 (l'extérieur des boucles de plage). La plage pour RVB et SV est de 0 à 1.