帮我写一个coco_dataset.py

好的，以下是一个简单的coco_dataset.py：

import os
import json
import random

import torch
from torch.utils.data import Dataset
from PIL import Image
import torchvision.transforms as transforms

class CocoDataset(Dataset):
    def __init__(self, root_dir, ann_file, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        
        with open(ann_file, 'r') as f:
            self.annotations = json.load(f)
            
    def __len__(self):
        return len(self.annotations)
    
    def __getitem__(self, index):
        annotation = self.annotations[index]
        img_id = annotation['image_id']
        img_path = os.path.join(self.root_dir, 'train2017', f'{str(img_id).zfill(12)}.jpg')
        image = Image.open(img_path).convert('RGB')
        
        if self.transform:
            image = self.transform(image)
            
        return image

这是一个用于加载 Coco 数据集的 PyTorch 数据集类。加载之前需要将 Coco 数据集下载到本地，并将主目录设置为 root_dir，注释文件设置为 ann_file。

你可以在实例化类时指定 transform 参数，用于定义要对图像进行的任何预处理。在这里，我们只是加载图像并将其转换为 PyTorch 的张量。

希望这能帮到你！