# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02b_data_birds.ipynb (unless otherwise specified).

__all__ = ['BertTokenizer', 'Datasets', 'DataLoaders']

# Cell
import random
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
from fastcore.all import *
from transformers import AutoTokenizer

# Internal Cell
def preprocess_text(s):
    "remove weird text ��, only keep english words, lowercase"
    s = s.replace("\ufffd\ufffd", " ")
    s = re.sub(r'[^\w\s]','',s)
    return s.lower()

# Internal Cell
def get_caps(caps_dir, img_name):
    "returns list of caps for this img"
    caps = (caps_dir/img_name).with_suffix('.txt').read().splitlines()
    caps = [preprocess_text(cap) for cap in caps]
    caps = [cap for cap in caps if len(cap.split())>1]
    return caps

# Internal Cell
def get_items(data_dir, pct=1):
    " returns: List of (split, img_loc, bbox, caps) "
    # bbox: [x-left, y-top, width, height]
#     split_loc = data_dir/'train_test_split.txt'
#     splits = [o.split()[1] for o in split_loc.read().splitlines()]

    img_names_loc = data_dir/'images.txt'
    img_names = [o.split()[1] for o in img_names_loc.read().splitlines()]
    img_dir = data_dir/'imgs'
    img_locs = [img_dir/img_name for img_name in img_names]

    bbox_loc = data_dir/'bounding_boxes.txt'
    bboxes = [L(o.split()[1:]).map(float) for o in bbox_loc.read().splitlines()]

    caps_dir = data_dir/'caps'
    caps = [get_caps(caps_dir, img_name) for img_name in img_names]

    num = int(len(img_names)*pct)
#     return [(splits[i], img_locs[i], bboxes[i], caps[i]) for i in range(num)]
    return [(img_locs[i], bboxes[i], caps[i]) for i in range(num)], []

# Cell
class BertTokenizer():
    def __init__(self):
        self.max_seq_len = 60
        self.tok = AutoTokenizer.from_pretrained('albert-base-v1')
        self.bos_id = self.tok.cls_token_id
        self.eos_id = self.tok.sep_token_id
        self.pad_id = self.tok.pad_token_id
    def encode(self, o):
        return self.tok.encode(o, max_length=self.max_seq_len, pad_to_max_length=True)
    def decode(self, o):
        return self.tok.decode(o, skip_special_tokens=True)

# Internal Cell
def crop_by_bbox(img, bbox):
    width, height = img.size
    r = int(np.maximum(bbox[2], bbox[3]) * 0.75)
    center_x = int((2 * bbox[0] + bbox[2]) / 2)
    center_y = int((2 * bbox[1] + bbox[3]) / 2)
    y1 = np.maximum(0, center_y - r)
    y2 = np.minimum(height, center_y + r)
    x1 = np.maximum(0, center_x - r)
    x2 = np.minimum(width, center_x + r)
    img = img.crop([x1, y1, x2, y2])
    return img

# Internal Cell
class BirdsDataset(Dataset):
    def __init__(self, items):
        self.items = items
        self.tokenizer = BertTokenizer()
    def __len__(self):
        return len(self.items)
    def __getitem__(self, idx):
        return self.tfm(self.items[idx])
    def tfm(self, item):
        ''' item: (img_loc, bbox, caps), returns: cap(60,), img(256, 256, 3) '''
        img_loc, bbox, caps = item
        cap = self.tokenizer.encode(random.choice(caps))
        img = np.array(crop_by_bbox(Image.open(img_loc), bbox).resize((256, 256)))
        if len(img.shape)==2:
            img = np.repeat(img[...,None], 3, axis=2)
        return torch.tensor(cap), torch.tensor(img)

# Cell
class Datasets():
    def __init__(self, data_dir, pct=1):
        train_items, valid_items = get_items(data_dir, pct=pct)
        self.train = BirdsDataset(train_items)
        self.valid = BirdsDataset(valid_items)

# Cell
class DataLoaders():
    def __init__(self, dsets, bs=64):
        self.dsets = dsets
        self.train = DataLoader(dsets.train, batch_size=bs, shuffle=True, num_workers=2)
        self.valid = DataLoader(dsets.valid, batch_size=bs, shuffle=False, num_workers=2)