text: # 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 (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.tok = AutoTokenizer.from_pretrained('albert-base-v2')
  #         self.bos_id = self.tok.cls_token_id
  #         self.eos_id = self.tok.sep_token_id
          self.max_seq_len = 25
          self.pad_id = self.tok.pad_token_id
          self.vocab_sz = self.tok.vocab_size
      def encode(self, o):
          toks = self.tok.encode(o, max_length=self.max_seq_len, pad_to_max_length=True, add_special_tokens=False)
          tok_len = sum([1 if tok!=self.pad_id else 0 for tok in toks])
          return toks, tok_len
      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: (25,), cap_len: (), img: (256, 256, 3) '''
          img_loc, bbox, caps = item
          cap, cap_len = 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(cap_len), 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, drop_last=True)
          self.valid = DataLoader(dsets.valid, batch_size=bs, shuffle=False, num_workers=2)