import numpy as np
import torch
import torch.nn as nn
from torch.optim import Adam
from PIL import Image
from torchvision import transforms
from pycocotools.coco import COCO
import os

# 设置COCO数据集路径
data_dir = 'C:/Users/Yoyo/Desktop/Project_3'
data_type = 'train2014'
ann_file = os.path.join(data_dir, 'annotations', f'captions_{data_type}.json')
img_dir = os.path.join(data_dir, data_type)

# 加载COCO数据集
coco = COCO(ann_file)
img_ids = coco.getImgIds()
images = coco.loadImgs(img_ids)

# 分词函数（简单中文分词）
def tokenize(text):
    return list(text)

# 创建词汇表
vocab = set()
for img in images:
    ann_ids = coco.getAnnIds(imgIds=img['id'])
    anns = coco.loadAnns(ann_ids)
    for ann in anns:
        tokens = tokenize(ann['caption'])
        vocab.update(tokens)
vocab = list(vocab)
word_to_idx = {word: idx for idx, word in enumerate(vocab)}

# 单热编码
def one_hot_encode(tokens, vocab_size):
    encoding = np.zeros((len(tokens), vocab_size), dtype=np.float32)
    for i, token in enumerate(tokens):
        encoding[i, word_to_idx[token]] = 1.0
    return encoding

# 图片处理
def load_image(image_path):
    image = Image.open(image_path).convert("RGB")
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    return transform(image)

# 模拟扩散模型
class SimpleDiffusionModel(nn.Module):
    def __init__(self, vocab_size):
        super(SimpleDiffusionModel, self).__init__()
        self.conv1 = nn.Conv2d(vocab_size, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 3, kernel_size=3, padding=1)
    
    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.sigmoid(self.conv2(x))
        return x

# 计算损失函数
def loss_function(original, generated):
    return nn.MSELoss()(original, generated)

# 主函数
def main():
    vocab_size = len(vocab)
    model = SimpleDiffusionModel(vocab_size)
    optimizer = Adam(model.parameters(), lr=0.001)

    for img in images:
        image_path = os.path.join(img_dir, img['file_name'])
        image = load_image(image_path)
        ann_ids = coco.getAnnIds(imgIds=img['id'])
        anns = coco.loadAnns(ann_ids)

        for ann in anns:
            tokens = tokenize(ann['caption'])
            encoded_text = one_hot_encode(tokens, vocab_size)

            # 调整编码文本形状以匹配图片
            encoded_text = torch.tensor(encoded_text).permute(2, 0, 1)
            encoded_text = encoded_text.unsqueeze(0).repeat(1, 1, image.size(1), image.size(2))

            # 模型训练
            optimizer.zero_grad()
            generated_image = model(encoded_text)
            loss = loss_function(image, generated_image)
            loss.backward()
            optimizer.step()
            
            print(f"Loss: {loss.item()}")

if __name__ == "__main__":
    main()
报错：Traceback (most recent call last):
  File "c:/Users/Yoyo/Desktop/Project_3/andy diffusion.py", line 97, in <module>
    main()
  File "c:/Users/Yoyo/Desktop/Project_3/andy diffusion.py", line 84, in main
    encoded_text = torch.tensor(encoded_text).permute(2, 0, 1)
RuntimeError: permute(sparse_coo): number of dimensions in the tensor input does not match the length of the desired ordering of dimensions i.e. input.dim() = 2 is not equal to len(dims) = 3