# AUTOGENERATED! DO NOT EDIT! File to edit: box.ipynb (unless otherwise specified).

__all__ = ['image_box', 'get_array', 'cat_box', 'mcat_box', 'bbox_box']

# Cell
from fastai.vision.all import *

from .core import *
from .utils import *

# Cell
def image_box(size = 224, folder = '', aug_kwargs = None, **kwargs):
    """Make a DataBox from an ImageBlock."""
    if aug_kwargs is not None:
        kwargs = merge_tfms(
            {'batch_tfms': [*aug_transforms(size=size, **aug_kwargs)]}, kwargs
        )

    return DataBox(
        ImageBlock,
        lambda x: Path(x.parent) / folder / x.name,
        **merge_tfms(kwargs, {'item_tfms': Resize(size)})
    )

# Cell
def get_array(suffix):
    """Return a getter that loads a numpy array with given suffix."""
    return lambda x: np.load(x.with_suffix(f'{suffix}.npy'), allow_pickle=True)

# Cell
def cat_box(suffix=''):
    """Make a Category Box from the categorial array of the dataset."""
    return DataBox(CategoryBlock, pipe(get_array(suffix), str))

# Cell
from ast import literal_eval

def mcat_box(suffix=''):
    """Make a Category Box from the categorial array of the dataset."""
    return DataBox(MultiCategoryBlock, lambda x: literal_eval(get_array(suffix)(x).tolist()))

# Cell
def bbox_box():
    return DataBox(
    (BBoxBlock, BBoxLblBlock),
    [get_array('.box'), get_array('.label')]
)