# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/mnist.ipynb.

# %% auto 0
__all__ = ['MNIST', 'TinyMNIST']

# %% ../nbs/mnist.ipynb 3
from functools import partial

from fastai.vision.all import *

import fastai_datasets.patches

# %% ../nbs/mnist.ipynb 4
def MNIST(sample=False):
    return DataBlock(
        blocks=(ImageBlock(PILImageBW), CategoryBlock),
        get_items=get_image_files,
        get_y=parent_label,
        splitter=GrandparentSplitter('training', 'testing') if not sample else GrandparentSplitter()
    ).datasets(untar_data(URLs.MNIST if not sample else URLs.MNIST_SAMPLE))

def TinyMNIST():
    return DataBlock(
        blocks=(ImageBlock(PILImageBW), CategoryBlock),
        get_items=partial(get_image_files, folders=['train', 'valid']),
        get_y=parent_label,
        splitter=GrandparentSplitter(),
    ).datasets(untar_data(URLs.MNIST_TINY))
