读取ISIC2019数据集

import os
import pandas as pd
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision import transforms
import numpy as np

# 1. 定义 ISIC2019Dataset 类
class ISIC2019Dataset(Dataset):
    def __init__(self, images_dir, labels_file, transform=None):
        """
        初始化 ISIC 2019 数据集。

        参数:
            images_dir (str): 图像文件夹的路径
            labels_file (str): 标签文件的路径 (CSV 格式)
            transform (callable, optional): 对图像进行的变换
        """
        self.images_dir = images_dir
        self.labels_df = pd.read_csv(labels_file)
        
        # 忽略最后一列 "UNK"
        self.labels_df = self.labels_df.iloc[:, :-1]
        self.transform = transform
        self.labels = self.labels_df.iloc[:, 1:].values.astype('float')
        
        # 统计图片文件夹中的总图像数
        self.total_images_in_dir = len(os.listdir(images_dir))
        # 统计标签文件中的记录数
        self.total_images_in_csv = len(self.labels_df)

    def __len__(self):
        """
        返回数据集的大小。
        """
        return len(self.labels_df)

    def __getitem__(self, idx):
        """
        根据索引获取数据。

        参数:
            idx (int): 索引

        返回:
            image (Tensor): 变换后的图像
            label (Tensor): 标签
        """
        # 获取图像路径
        img_name = os.path.join(self.images_dir, self.labels_df.iloc[idx, 0] + '.jpg')
        image = Image.open(img_name).convert('RGB')
        
        # 获取标签
        label = torch.tensor(self.labels[idx], dtype=torch.float32)

        # 应用变换
        if self.transform:
            image = self.transform(image)
        
        return image, label
    
    def get_class_distribution(self):
        """
        计算数据集中每个类别的样本数量
        """
        class_counts = self.labels_df.iloc[:, 1:].sum(axis=0)
        return class_counts

    def print_image_counts(self):
        """
        打印图片文件夹中的总图像数和标签文件中的记录数
        """
        print(f"图片文件夹中的图像总数: {self.total_images_in_dir}")
        print(f"标签文件中的图像总数: {self.total_images_in_csv}")

# 2. 设置图像转换
data_transforms = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小
    transforms.ToTensor(),  # 转换为 Tensor
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 标准化
])

# 3. 定义图像和标签的路径
images_dir = '../../dataset/isic2019classification/ISIC_2019_Training_Input'  # 替换为实际图像路径
labels_file = '../../dataset/isic2019classification/ISIC_2019_Training_GroundTruth.csv'  # 替换为实际标签文件路径

# 4. 创建数据集实例
isic_dataset = ISIC2019Dataset(images_dir=images_dir, labels_file=labels_file, transform=data_transforms)

# 打印图片和标签的数量
isic_dataset.print_image_counts()

# 5. 根据 80%/20% 的比例划分数据集
train_size = int(0.8 * len(isic_dataset))  # 80% 作为训练集
test_size = len(isic_dataset) - train_size  # 20% 作为测试集

train_dataset, test_dataset = random_split(isic_dataset, [train_size, test_size])

# 5. 打印整个数据集的类别分布
class_distribution = isic_dataset.get_class_distribution()
print("数据集中每个类别的图像数量:")
print(class_distribution)

# 6. 统计并打印训练集和测试集的类别分布
def calculate_class_distribution(dataset, dataset_indices, num_classes):
    labels = np.array([dataset.labels[idx] for idx in dataset_indices])
    class_counts = labels.sum(axis=0)
    return class_counts

num_classes = isic_dataset.labels.shape[1]  # 获取类别数量
train_class_distribution = calculate_class_distribution(isic_dataset, train_dataset.indices, num_classes)
test_class_distribution = calculate_class_distribution(isic_dataset, test_dataset.indices, num_classes)

print("训练集中每个类别的图像数量:")
print(train_class_distribution)

print("测试集中每个类别的图像数量:")
print(test_class_distribution)

# 7. 创建 DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

# 8. 使用 DataLoader 进行训练和测试
# 示例代码展示如何从 DataLoader 中提取数据
for images, labels in train_loader:
    print("训练集 - 图像大小:", images.size(), "标签大小:", labels.size())
    # 在这里添加训练代码

for images, labels in test_loader:
    print("测试集 - 图像大小:", images.size(), "标签大小:", labels.size())
    # 在这里添加测试代码