%reload_ext autoreload
%autoreload 2
%matplotlib inline

from fastai import *
from fastai.vision import *

tfms = get_transforms()
bs = 64

def show_image_batch_supercharged(dl:DataLoader, classes:Collection[str], rows:int=None, figsize:Tuple[int,int]=(12,15),
                     denorm:Callable=None)->None:
    b_idx = next(iter(dl.batch_sampler))
    if rows is None: rows = int(math.sqrt(len(b_idx)))
    fig, axs = plt.subplots(rows,rows,figsize=figsize)
    for i, ax in zip(b_idx[:rows*rows], axs.flatten()):
        x,y = dl.dataset[i]
        x.show(ax=ax, y=y, classes=classes)
    plt.tight_layout()

def show_batch_supercharged(data:DataBunch, rows:int=3, figsize:Tuple[int,int]=(9,10), is_train:bool=True)->None:
        show_image_batch_supercharged(data.train_dl if is_train else data.valid_dl, data.classes,
            denorm=getattr(data,'denorm',None), figsize=figsize, rows=rows)

path = untar_data(URLs.PETS)

data = (ImageFileList.from_folder(path/'images')
        .label_from_re(r'^(.*)_\d+.jpg$')
        .random_split_by_pct(0.2)
        .datasets(ImageClassificationDataset)
        .transform(tfms, size=224)
        .databunch(bs=bs, path=path))

data = (ImageFileList.from_folder(path/'images')
        .label_from_re(r'^(.*)_\d+.jpg$')
        .random_split_by_pct(0.2))

data.show_batch(rows=3, figsize=(7,7))

show_batch_supercharged(data, rows=3, figsize=(7,7))

path = untar_data(URLs.PLANET_SAMPLE)

data = (InputList.from_folder(path)
        .label_from_csv('labels.csv', sep=' ', suffix='.jpg', folder='train')
        .random_split_by_pct(0.2)
        .datasets(ImageMultiDataset)
        .transform(tfms, size=128)
        .databunch(bs=bs, path=path))

data.show_batch(rows=3, figsize=(10,10))

show_batch_supercharged(data, rows=3, figsize=(10,10))

path = untar_data(URLs.CAMVID)
path_lbl = path/'labels'
path_img = path/'images'

get_y_fn = lambda x: path_lbl/f'{x.stem}_P{x.suffix}'
codes = np.loadtxt(path/'codes.txt', dtype=str); codes

data = (ImageFileList.from_folder(path_img)
        .label_from_func(get_y_fn)
        .random_split_by_pct()
        .datasets(SegmentationDataset, classes=codes)
        .transform(get_transforms(), size=128, tfm_y=True)
        .databunch(bs=64))

data.show_batch(rows=2, figsize=(7,7))

show_batch_supercharged(data, rows=2, figsize=(7,7))

path = untar_data(URLs.COCO_TINY)

images, lbl_bbox = get_annotations(path/'train.json')
img2bbox = {img:bb for img, bb in zip(images, lbl_bbox)}
get_y_func = lambda o:img2bbox[o.name]

data = (ImageFileList.from_folder(path)
        .label_from_func(get_y_func)
        .random_split_by_pct()
        .datasets(ObjectDetectDataset)
        .transform(tfms, size=128, tfm_y=True)
        .databunch(bs=16, collate_fn=bb_pad_collate))

data.show_batch(rows=3)