Pytorch1.3的自定义MPII数据集python3类

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

class MPIIDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.images = os.listdir(root_dir)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, self.images[idx])
        image = Image.open(img_name).convert('RGB')
        
        if self.transform:
            image = self.transform(image)
        
        # 这里假设每个图像文件名中包含了对应的标签信息
        label = int(self.images[idx].split('.')[0])  # 示例标签提取方式
        
        return image, label

# 定义一些数据预处理
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 创建数据集实例
dataset = MPIIDataset(root_dir='path_to_mpii_dataset', transform=transform)

# 创建数据加载器
dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True, num_workers=2)

# 使用数据加载器进行迭代
for images, labels in dataloader:
    # 在这里进行模型的训练或其他处理
    pass