import os import sys import re import six import math import lmdb import torch from natsort import natsorted from PIL import Image import numpy as np from torch.utils.data import Dataset, ConcatDataset, Subset from torch._utils import _accumulate import torchvision.transforms as transforms class Batch_Balanced_Dataset(object): def __init__(self, opt): """ Modulate the data ratio in the batch. For example, when select_data is "MJ-ST" and batch_ratio is "0.5-0.5", the 50% of the batch is filled with MJ and the other 50% of the batch is filled with ST. """ print('-' * 80) print(f'dataset_root: {opt.train_data}\nopt.select_data: {opt.select_data}\nopt.batch_ratio: {opt.batch_ratio}') assert len(opt.select_data) == len(opt.batch_ratio) _AlignCollate = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD) self.data_loader_list = [] self.dataloader_iter_list = [] batch_size_list = [] Total_batch_size = 0 for selected_d, batch_ratio_d in zip(opt.select_data, opt.batch_ratio): _batch_size = max(round(opt.batch_size * float(batch_ratio_d)), 1) print('-' * 80) _dataset = hierarchical_dataset(root=opt.train_data, opt=opt, select_data=[selected_d]) total_number_dataset = len(_dataset) """ The total number of data can be modified with opt.total_data_usage_ratio. ex) opt.total_data_usage_ratio = 1 indicates 100% usage, and 0.2 indicates 20% usage. See 4.2 section in our paper. """ number_dataset = int(total_number_dataset * float(opt.total_data_usage_ratio)) dataset_split = [number_dataset, total_number_dataset - number_dataset] indices = range(total_number_dataset) _dataset, _ = [Subset(_dataset, indices[offset - length:offset]) for offset, length in zip(_accumulate(dataset_split), dataset_split)] print(f'num total samples of {selected_d}: {total_number_dataset} x {opt.total_data_usage_ratio} (total_data_usage_ratio) = {len(_dataset)}') print(f'num samples of {selected_d} per batch: {opt.batch_size} x {float(batch_ratio_d)} (batch_ratio) = {_batch_size}') batch_size_list.append(str(_batch_size)) Total_batch_size += _batch_size _data_loader = torch.utils.data.DataLoader( _dataset, batch_size=_batch_size, shuffle=True, num_workers=int(opt.workers), collate_fn=_AlignCollate, pin_memory=True) self.data_loader_list.append(_data_loader) self.dataloader_iter_list.append(iter(_data_loader)) print('-' * 80) print('Total_batch_size: ', '+'.join(batch_size_list), '=', str(Total_batch_size)) opt.batch_size = Total_batch_size print('-' * 80) def get_batch(self): balanced_batch_images = [] balanced_batch_texts = [] for i, data_loader_iter in enumerate(self.dataloader_iter_list): try: image, text = data_loader_iter.next() balanced_batch_images.append(image) balanced_batch_texts += text except StopIteration: self.dataloader_iter_list[i] = iter(self.data_loader_list[i]) image, text = self.dataloader_iter_list[i].next() balanced_batch_images.append(image) balanced_batch_texts += text except ValueError: pass balanced_batch_images = torch.cat(balanced_batch_images, 0) return balanced_batch_images, balanced_batch_texts def hierarchical_dataset(root, opt, select_data='/'): """ select_data='/' contains all sub-directory of root directory """ dataset_list = [] print(f'dataset_root: {root}\t dataset: {select_data[0]}') for dirpath, dirnames, filenames in os.walk(root+'/'): if not dirnames: select_flag = False for selected_d in select_data: if selected_d in dirpath: select_flag = True break if select_flag: dataset = LmdbDataset(dirpath, opt) print(f'sub-directory:\t/{os.path.relpath(dirpath, root)}\t num samples: {len(dataset)}') dataset_list.append(dataset) concatenated_dataset = ConcatDataset(dataset_list) return concatenated_dataset class LmdbDataset(Dataset): def __init__(self, root, opt): self.root = root self.opt = opt self.env = lmdb.open(root, max_readers=32, readonly=True, lock=False, readahead=False, meminit=False) if not self.env: print('cannot create lmdb from %s' % (root)) sys.exit(0) with self.env.begin(write=False) as txn: nSamples = int(txn.get('num-samples'.encode())) self.nSamples = nSamples if self.opt.data_filtering_off: # for fast check or benchmark evaluation with no filtering self.filtered_index_list = [index + 1 for index in range(self.nSamples)] else: """ Filtering part If you want to evaluate IC152077 & CUTE80 datasets which have special character labels, use --data_filtering_off and evaluation with this snippet (only evaluate on alphabets and digits). https://github.com/clovaai/deep-text-recognition-benchmark/blob/master/dataset.py#L186-L188 """ self.filtered_index_list = [] for index in range(self.nSamples): index += 1 # lmdb starts with 1 label_key = 'label-%09d'.encode() % index label = txn.get(label_key).decode('utf-8') if len(label) > self.opt.batch_max_length: # print(f'The length of the label is longer than max_length: length # {len(label)}, {label} in dataset {self.root}') continue # By default, images containing characters which are not in opt.character are filtered. # You can add [UNK] token to `opt.character` in utils.py instead of this filtering. out_of_char = f'[^{self.opt.character}]' if re.search(out_of_char, label.lower()): continue self.filtered_index_list.append(index) self.nSamples = len(self.filtered_index_list) def __len__(self): return self.nSamples def __getitem__(self, index): assert index <= len(self), 'index range error' index = self.filtered_index_list[index] with self.env.begin(write=False) as txn: label_key = 'label-%09d'.encode() % index label = txn.get(label_key).decode('utf-8') img_key = 'image-%09d'.encode() % index imgbuf = txn.get(img_key) buf = six.BytesIO() buf.write(imgbuf) buf.seek(0) try: if self.opt.rgb: img = Image.open(buf).convert('RGB') # for color image else: img = Image.open(buf).convert('L') except IOError: print(f'Corrupted image for {index}') # make dummy image and dummy label for corrupted image. if self.opt.rgb: img = Image.new('RGB', (self.opt.imgW, self.opt.imgH)) else: img = Image.new('L', (self.opt.imgW, self.opt.imgH)) label = '[dummy_label]' if not self.opt.sensitive: label = label.lower() # We only train and evaluate on alphanumerics (or pre-defined character set in train.py) out_of_char = f'[^{self.opt.character}]' label = re.sub(out_of_char, '', label) return (img, label) class RawDataset(Dataset): def __init__(self, root, opt): self.opt = opt self.image_path_list = [] for dirpath, dirnames, filenames in os.walk(root): for name in filenames: _, ext = os.path.splitext(name) ext = ext.lower() if ext == '.jpg' or ext == '.jpeg' or ext == '.png': self.image_path_list.append(os.path.join(dirpath, name)) self.image_path_list = natsorted(self.image_path_list) self.nSamples = len(self.image_path_list) def __len__(self): return self.nSamples def __getitem__(self, index): try: if self.opt.rgb: img = Image.open(self.image_path_list[index]).convert('RGB') # for color image else: img = Image.open(self.image_path_list[index]).convert('L') except IOError: print(f'Corrupted image for {index}') # make dummy image and dummy label for corrupted image. if self.opt.rgb: img = Image.new('RGB', (self.opt.imgW, self.opt.imgH)) else: img = Image.new('L', (self.opt.imgW, self.opt.imgH)) return (img, self.image_path_list[index]) class ResizeNormalize(object): def __init__(self, size, interpolation=Image.BICUBIC): self.size = size self.interpolation = interpolation self.toTensor = transforms.ToTensor() def __call__(self, img): img = img.resize(self.size, self.interpolation) img = self.toTensor(img) img.sub_(0.5).div_(0.5) return img class NormalizePAD(object): def __init__(self, max_size, PAD_type='right'): self.toTensor = transforms.ToTensor() self.max_size = max_size self.max_width_half = math.floor(max_size[2] / 2) self.PAD_type = PAD_type def __call__(self, img): img = self.toTensor(img) img.sub_(0.5).div_(0.5) c, h, w = img.size() Pad_img = torch.FloatTensor(*self.max_size).fill_(0) Pad_img[:, :, :w] = img # right pad if self.max_size[2] != w: # add border Pad Pad_img[:, :, w:] = img[:, :, w - 1].unsqueeze(2).expand(c, h, self.max_size[2] - w) return Pad_img class AlignCollate(object): def __init__(self, imgH=32, imgW=100, keep_ratio_with_pad=False): self.imgH = imgH self.imgW = imgW self.keep_ratio_with_pad = keep_ratio_with_pad def __call__(self, batch): batch = filter(lambda x: x is not None, batch) images, labels = zip(*batch) if self.keep_ratio_with_pad: # same concept with 'Rosetta' paper resized_max_w = self.imgW input_channel = 3 if images[0].mode == 'RGB' else 1 transform = NormalizePAD((input_channel, self.imgH, resized_max_w)) resized_images = [] for image in images: w, h = image.size ratio = w / float(h) if math.ceil(self.imgH * ratio) > self.imgW: resized_w = self.imgW else: resized_w = math.ceil(self.imgH * ratio) resized_image = image.resize((resized_w, self.imgH), Image.BICUBIC) resized_images.append(transform(resized_image)) # resized_image.save('./image_test/%d_test.jpg' % w) image_tensors = torch.cat([t.unsqueeze(0) for t in resized_images], 0) else: transform = ResizeNormalize((self.imgW, self.imgH)) image_tensors = [transform(image) for image in images] image_tensors = torch.cat([t.unsqueeze(0) for t in image_tensors], 0) return image_tensors, labels def tensor2im(image_tensor, imtype=np.uint8): image_numpy = image_tensor.cpu().float().numpy() if image_numpy.shape[0] == 1: image_numpy = np.tile(image_numpy, (3, 1, 1)) image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 return image_numpy.astype(imtype) def save_image(image_numpy, image_path): image_pil = Image.fromarray(image_numpy) image_pil.save(image_path)