pytorch图像验证码识别数据库构建代码

以下是使用PyTorch构建图像验证码识别数据库的一个基本示例代码：

import os
import random
from PIL import Image
import torch.utils.data as data
import torchvision.transforms as transforms


class CaptchaDataset(data.Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.labels = []
        self.images = []

        for filename in os.listdir(root_dir):
            label = filename.split('.')[0]
            self.labels.append(label)
            img_path = os.path.join(root_dir, filename)
            self.images.append(img_path)

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

    def __getitem__(self, index):
        img_path = self.images[index]
        img = Image.open(img_path)

        if self.transform:
            img = self.transform(img)

        label = self.labels[index]
        return img, label


if __name__ == '__main__':
    root_dir = 'captcha_dataset'
    transform = transforms.Compose([
        transforms.Resize((32, 128)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])

    dataset = CaptchaDataset(root_dir, transform=transform)
    print(f"Number of samples: {len(dataset)}")

    # Visualize a random sample
    idx = random.randint(0, len(dataset))
    sample = dataset[idx]
    img, label = sample
    print(f"Image shape: {img.shape}")
    print(f"Label: {label}")
    img = transforms.functional.to_pil_image(img)
    img.show()

在此示例代码中，我们首先定义了一个名为`CaptchaDataset`的自定义数据集类，该类从指定的文件夹中加载图像并相应地解析标签。我们还定义了一个名为`transform`的变换对象，该对象将图像调整为规定的大小、将其转换为张量并进行标准化。

在`CaptchaDataset`类中，`__len__()`方法返回数据集的大小，`__getitem__()`方法获取具有给定索引的图像及其标签。我们使用`PIL`库打开图像文件，然后应用我们定义的变换对象。最后将标签返回。

在`__main__()`函数中，我们使用示例数据集来创建一个`CaptchaDataset`对象，并打印数据集的大小。我们还随机选择一个样本并可视化它，以确保我们正确地加载了数据集。