如何在PyTorch中将RGB图像编码为n_class One热张量

import torch
import torchvision.transforms as transforms
import torch.nn.functional as F

def encode_one_hot(image, n_class):
    # 定义图像转换
    transform = transforms.Compose([
        transforms.ToTensor(),  # 将图像转换为张量
    ])
    
    # 对图像进行转换
    image = transform(image)
    
    # 将图像扩展为(batch_size, channels, height, width)的形状
    image = image.unsqueeze(0)
    
    # 使用torch.argmax函数将图像编码为类别索引
    class_index = torch.argmax(image, dim=1)
    
    # 使用torch.nn.functional.one_hot函数将类别索引编码为One热张量
    one_hot = F.one_hot(class_index, num_classes=n_class)
    
    return one_hot

import PIL.Image as Image

# 读取RGB图像
image = Image.open('image.jpg').convert('RGB')

# 将RGB图像编码为5类的One热张量
one_hot = encode_one_hot(image, n_class=5)