captcha-images.py

import os
from pathlib import Path

import datasets

_DESCRIPTION = """\
Captcha images dataset.
"""
_LICENSE = "wtfpl"
_DATA_DIR = "data"

class Captcha(datasets.GeneratorBasedBuilder):
    def _info(self):
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=datasets.Features(
                {
                        "image": datasets.Image(),
                        "solution": datasets.Value("string"),
                }
            ),
            license=_LICENSE,
        )

    def _split_generators(self, dl_manager):
        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
                gen_kwargs={
                    "images": dl_manager.iter_archive(dl_manager.download(os.path.join(_DATA_DIR, "train.tar.gz"))),
                },
            ),
            datasets.SplitGenerator(
                name=datasets.Split.VALIDATION,
                gen_kwargs={
                    "images": dl_manager.iter_archive(dl_manager.download(os.path.join(_DATA_DIR, "validation.tar.gz"))),
                },
            ),
            datasets.SplitGenerator(
                name=datasets.Split.TEST,
                gen_kwargs={
                    "images": dl_manager.iter_archive(dl_manager.download(os.path.join(_DATA_DIR, "test.tar.gz"))),
                },
            ),
        ]

    def _generate_examples(self, images):
        for file_path, file_obj in images:
            yield file_path, {
                    "image": {"path": file_path, "bytes": file_obj.read()},
                    "solution": Path(file_path).name.split('.')[0],
                }