FAST-ABINet-OCR / dataset.py
nigger game
gfgf
01e491f
import logging
import re
import cv2
import lmdb
import six
from fastai.vision import *
from torchvision import transforms
from transforms import CVColorJitter, CVDeterioration, CVGeometry
from utils import CharsetMapper, onehot
class ImageDataset(Dataset):
"`ImageDataset` read data from LMDB database."
def __init__(self,
path:PathOrStr,
is_training:bool=True,
img_h:int=32,
img_w:int=100,
max_length:int=25,
check_length:bool=True,
case_sensitive:bool=False,
charset_path:str='data/charset_36.txt',
convert_mode:str='RGB',
data_aug:bool=True,
deteriorate_ratio:float=0.,
multiscales:bool=True,
one_hot_y:bool=True,
return_idx:bool=False,
return_raw:bool=False,
**kwargs):
self.path, self.name = Path(path), Path(path).name
assert self.path.is_dir() and self.path.exists(), f"{path} is not a valid directory."
self.convert_mode, self.check_length = convert_mode, check_length
self.img_h, self.img_w = img_h, img_w
self.max_length, self.one_hot_y = max_length, one_hot_y
self.return_idx, self.return_raw = return_idx, return_raw
self.case_sensitive, self.is_training = case_sensitive, is_training
self.data_aug, self.multiscales = data_aug, multiscales
self.charset = CharsetMapper(charset_path, max_length=max_length+1)
self.c = self.charset.num_classes
self.env = lmdb.open(str(path), readonly=True, lock=False, readahead=False, meminit=False)
assert self.env, f'Cannot open LMDB dataset from {path}.'
with self.env.begin(write=False) as txn:
self.length = int(txn.get('num-samples'.encode()))
if self.is_training and self.data_aug:
self.augment_tfs = transforms.Compose([
CVGeometry(degrees=45, translate=(0.0, 0.0), scale=(0.5, 2.), shear=(45, 15), distortion=0.5, p=0.5),
CVDeterioration(var=20, degrees=6, factor=4, p=0.25),
CVColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.25)
])
self.totensor = transforms.ToTensor()
def __len__(self): return self.length
def _next_image(self, index):
next_index = random.randint(0, len(self) - 1)
return self.get(next_index)
def _check_image(self, x, pixels=6):
if x.size[0] <= pixels or x.size[1] <= pixels: return False
else: return True
def resize_multiscales(self, img, borderType=cv2.BORDER_CONSTANT):
def _resize_ratio(img, ratio, fix_h=True):
if ratio * self.img_w < self.img_h:
if fix_h: trg_h = self.img_h
else: trg_h = int(ratio * self.img_w)
trg_w = self.img_w
else: trg_h, trg_w = self.img_h, int(self.img_h / ratio)
img = cv2.resize(img, (trg_w, trg_h))
pad_h, pad_w = (self.img_h - trg_h) / 2, (self.img_w - trg_w) / 2
top, bottom = math.ceil(pad_h), math.floor(pad_h)
left, right = math.ceil(pad_w), math.floor(pad_w)
img = cv2.copyMakeBorder(img, top, bottom, left, right, borderType)
return img
if self.is_training:
if random.random() < 0.5:
base, maxh, maxw = self.img_h, self.img_h, self.img_w
h, w = random.randint(base, maxh), random.randint(base, maxw)
return _resize_ratio(img, h/w)
else: return _resize_ratio(img, img.shape[0] / img.shape[1]) # keep aspect ratio
else: return _resize_ratio(img, img.shape[0] / img.shape[1]) # keep aspect ratio
def resize(self, img):
if self.multiscales: return self.resize_multiscales(img, cv2.BORDER_REPLICATE)
else: return cv2.resize(img, (self.img_w, self.img_h))
def get(self, idx):
with self.env.begin(write=False) as txn:
image_key, label_key = f'image-{idx+1:09d}', f'label-{idx+1:09d}'
try:
label = str(txn.get(label_key.encode()), 'utf-8') # label
label = re.sub('[^0-9a-zA-Z]+', '', label)
if self.check_length and self.max_length > 0:
if len(label) > self.max_length or len(label) <= 0:
#logging.info(f'Long or short text image is found: {self.name}, {idx}, {label}, {len(label)}')
return self._next_image(idx)
label = label[:self.max_length]
imgbuf = txn.get(image_key.encode()) # image
buf = six.BytesIO()
buf.write(imgbuf)
buf.seek(0)
with warnings.catch_warnings():
warnings.simplefilter("ignore", UserWarning) # EXIF warning from TiffPlugin
image = PIL.Image.open(buf).convert(self.convert_mode)
if self.is_training and not self._check_image(image):
#logging.info(f'Invalid image is found: {self.name}, {idx}, {label}, {len(label)}')
return self._next_image(idx)
except:
import traceback
traceback.print_exc()
logging.info(f'Corrupted image is found: {self.name}, {idx}, {label}, {len(label)}')
return self._next_image(idx)
return image, label, idx
def _process_training(self, image):
if self.data_aug: image = self.augment_tfs(image)
image = self.resize(np.array(image))
return image
def _process_test(self, image):
return self.resize(np.array(image)) # TODO:move is_training to here
def __getitem__(self, idx):
image, text, idx_new = self.get(idx)
if not self.is_training: assert idx == idx_new, f'idx {idx} != idx_new {idx_new} during testing.'
if self.is_training: image = self._process_training(image)
else: image = self._process_test(image)
if self.return_raw: return image, text
image = self.totensor(image)
length = tensor(len(text) + 1).to(dtype=torch.long) # one for end token
label = self.charset.get_labels(text, case_sensitive=self.case_sensitive)
label = tensor(label).to(dtype=torch.long)
if self.one_hot_y: label = onehot(label, self.charset.num_classes)
if self.return_idx: y = [label, length, idx_new]
else: y = [label, length]
return image, y
class TextDataset(Dataset):
def __init__(self,
path:PathOrStr,
delimiter:str='\t',
max_length:int=25,
charset_path:str='data/charset_36.txt',
case_sensitive=False,
one_hot_x=True,
one_hot_y=True,
is_training=True,
smooth_label=False,
smooth_factor=0.2,
use_sm=False,
**kwargs):
self.path = Path(path)
self.case_sensitive, self.use_sm = case_sensitive, use_sm
self.smooth_factor, self.smooth_label = smooth_factor, smooth_label
self.charset = CharsetMapper(charset_path, max_length=max_length+1)
self.one_hot_x, self.one_hot_y, self.is_training = one_hot_x, one_hot_y, is_training
if self.is_training and self.use_sm: self.sm = SpellingMutation(charset=self.charset)
dtype = {'inp': str, 'gt': str}
self.df = pd.read_csv(self.path, dtype=dtype, delimiter=delimiter, na_filter=False)
self.inp_col, self.gt_col = 0, 1
def __len__(self): return len(self.df)
def __getitem__(self, idx):
text_x = self.df.iloc[idx, self.inp_col]
text_x = re.sub('[^0-9a-zA-Z]+', '', text_x)
if not self.case_sensitive: text_x = text_x.lower()
if self.is_training and self.use_sm: text_x = self.sm(text_x)
length_x = tensor(len(text_x) + 1).to(dtype=torch.long) # one for end token
label_x = self.charset.get_labels(text_x, case_sensitive=self.case_sensitive)
label_x = tensor(label_x)
if self.one_hot_x:
label_x = onehot(label_x, self.charset.num_classes)
if self.is_training and self.smooth_label:
label_x = torch.stack([self.prob_smooth_label(l) for l in label_x])
x = [label_x, length_x]
text_y = self.df.iloc[idx, self.gt_col]
text_y = re.sub('[^0-9a-zA-Z]+', '', text_y)
if not self.case_sensitive: text_y = text_y.lower()
length_y = tensor(len(text_y) + 1).to(dtype=torch.long) # one for end token
label_y = self.charset.get_labels(text_y, case_sensitive=self.case_sensitive)
label_y = tensor(label_y)
if self.one_hot_y: label_y = onehot(label_y, self.charset.num_classes)
y = [label_y, length_y]
return x, y
def prob_smooth_label(self, one_hot):
one_hot = one_hot.float()
delta = torch.rand([]) * self.smooth_factor
num_classes = len(one_hot)
noise = torch.rand(num_classes)
noise = noise / noise.sum() * delta
one_hot = one_hot * (1 - delta) + noise
return one_hot
class SpellingMutation(object):
def __init__(self, pn0=0.7, pn1=0.85, pn2=0.95, pt0=0.7, pt1=0.85, charset=None):
"""
Args:
pn0: the prob of not modifying characters is (pn0)
pn1: the prob of modifying one characters is (pn1 - pn0)
pn2: the prob of modifying two characters is (pn2 - pn1),
and three (1 - pn2)
pt0: the prob of replacing operation is pt0.
pt1: the prob of inserting operation is (pt1 - pt0),
and deleting operation is (1 - pt1)
"""
super().__init__()
self.pn0, self.pn1, self.pn2 = pn0, pn1, pn2
self.pt0, self.pt1 = pt0, pt1
self.charset = charset
logging.info(f'the probs: pn0={self.pn0}, pn1={self.pn1} ' +
f'pn2={self.pn2}, pt0={self.pt0}, pt1={self.pt1}')
def is_digit(self, text, ratio=0.5):
length = max(len(text), 1)
digit_num = sum([t in self.charset.digits for t in text])
if digit_num / length < ratio: return False
return True
def is_unk_char(self, char):
# return char == self.charset.unk_char
return (char not in self.charset.digits) and (char not in self.charset.alphabets)
def get_num_to_modify(self, length):
prob = random.random()
if prob < self.pn0: num_to_modify = 0
elif prob < self.pn1: num_to_modify = 1
elif prob < self.pn2: num_to_modify = 2
else: num_to_modify = 3
if length <= 1: num_to_modify = 0
elif length >= 2 and length <= 4: num_to_modify = min(num_to_modify, 1)
else: num_to_modify = min(num_to_modify, length // 2) # smaller than length // 2
return num_to_modify
def __call__(self, text, debug=False):
if self.is_digit(text): return text
length = len(text)
num_to_modify = self.get_num_to_modify(length)
if num_to_modify <= 0: return text
chars = []
index = np.arange(0, length)
random.shuffle(index)
index = index[: num_to_modify]
if debug: self.index = index
for i, t in enumerate(text):
if i not in index: chars.append(t)
elif self.is_unk_char(t): chars.append(t)
else:
prob = random.random()
if prob < self.pt0: # replace
chars.append(random.choice(self.charset.alphabets))
elif prob < self.pt1: # insert
chars.append(random.choice(self.charset.alphabets))
chars.append(t)
else: # delete
continue
new_text = ''.join(chars[: self.charset.max_length-1])
return new_text if len(new_text) >= 1 else text