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

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

# Cell
import pandas as pd
import numpy as np
from PIL import Image
from sklearn.model_selection import train_test_split
from kornia.augmentation.functional import random_hflip
from kornia.geometry.transform import resize
import torch
from torch.utils.data import Dataset, DataLoader
from fastcore.all import *

# Internal Cell
def get_items(data_dir, pct=1, valid_pct=0.2):
    df = pd.read_csv(data_dir/'tags.csv', header=None, names=['id', 'cap'])
    df = df[:int(len(df)*pct)]
    if valid_pct==0:
        return df, pd.DataFrame(data=None, columns=['id', 'cap'])
    train_items, valid_items = train_test_split(df, test_size=valid_pct, random_state=42, shuffle=True, stratify=df.cap)
    return train_items, valid_items

# Cell
class Tokenizer():
    def __init__(self):
        self.vocab = [
            'orange hair', 'white hair', 'aqua hair', 'gray hair','green hair', 'red hair',
            'purple hair', 'pink hair','blue hair', 'black hair', 'brown hair', 'blonde hair', 'black eyes', 'orange eyes',
            'purple eyes', 'pink eyes', 'yellow eyes', 'aqua eyes', 'green eyes', 'brown eyes', 'red eyes', 'blue eyes',
        ]
        self.o2i = {v:k for k,v in enumerate(self.vocab)}
    def encode(self, cap):
        "cap: 'aqua hair aqua eyes', returns: [2, 17]"
        cap = cap.split()
        tags = [' '.join(cap[:2]), ' '.join(cap[2:])]
        return [self.o2i[tags[0]], self.o2i[tags[1]]]
    def decode(self, o):
        "o: [2, 17], returns: 'aqua hair aqua eyes'"
        tags = [self.vocab[o[0]], self.vocab[o[1]]]
        return ' '.join(tags)


# Internal Cell
class AnimeHeadsDataset(Dataset):
    def __init__(self, items, data_dir):
        "items: df of id and cap"
        self.data_dir = data_dir
        self.items = list(items.itertuples(index=False, name=None))
        self.tokenizer = Tokenizer()
    def __len__(self):
        return len(self.items)
    def __getitem__(self, idx):
        return self.tfm(self.items[idx])
    def tfm(self, item):
        ''' item: (0, aqua hair aqua eyes), returns: tag(2,), img64(64, 64, 3) '''
        img_id, cap = item
        tag = self.tokenizer.encode(cap)

        img_path = self.data_dir/f'imgs/{img_id}.jpg'
        img64 = np.array(Image.open(img_path))
        if len(img64.shape)==2:
            img64 = np.repeat(img64[...,None], 3, axis=2)
        return torch.tensor(tag), torch.tensor(img64)

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

# 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)