pytorch rgb转hsv

以下是使用PyTorch实现RGB到HSV色彩空间转换的代码：

import torch

def rgb_to_hsv(rgb):
    # 将RGB值归一化到[0, 1]
    r, g, b = rgb[:, 0], rgb[:, 1], rgb[:, 2]
    max_val, max_idx = torch.max(rgb, dim=1)
    min_val, min_idx = torch.min(rgb, dim=1)
    delta = max_val - min_val
    delta_zero = delta == 0

    # 计算Hue
    hue = torch.zeros_like(max_val)
    hue[max_idx == 0] = (g - b) / delta[max_idx == 0]
    hue[max_idx == 1] = 2.0 + (b - r) / delta[max_idx == 1]
    hue[max_idx == 2] = 4.0 + (r - g) / delta[max_idx == 2]
    hue[delta_zero] = 0.0
    hue = (hue / 6.0) % 1.0

    # 计算Saturation
    saturation = delta / max_val
    saturation[delta_zero] = 0.0

    # 计算Value
    value = max_val

    # 将HSV值还原到[0, 255]
    hue = hue * 255.0
    saturation = saturation * 255.0
    value = value * 255.0

    return torch.stack([hue, saturation, value], dim=1)

其中，输入的RGB值应该是一个形状为(N, 3)的张量，每一行代表一张图片的RGB值。输出的HSV值也是一个形状为(N, 3)的张量，每一行代表一张图片的HSV值。