sneedium/captcha_pixelplanet

1import logging2import re3 4import cv25import lmdb6import six7from fastai.vision import *8from torchvision import transforms9 10from transforms import CVColorJitter, CVDeterioration, CVGeometry11from utils import CharsetMapper, onehot12 13 14class ImageDataset(Dataset):15    "`ImageDataset` read data from LMDB database."16 17    def __init__(self,18                 path:PathOrStr,19                 is_training:bool=True,20                 img_h:int=32,21                 img_w:int=100,22                 max_length:int=25,23                 check_length:bool=True,24                 case_sensitive:bool=False,25                 charset_path:str='data/charset_36.txt',26                 convert_mode:str='RGB',27                 data_aug:bool=True,28                 deteriorate_ratio:float=0.,29                 multiscales:bool=True,30                 one_hot_y:bool=True,31                 return_idx:bool=False,32                 return_raw:bool=False,33                 **kwargs):34        self.path, self.name = Path(path), Path(path).name35        assert self.path.is_dir() and self.path.exists(), f"{path} is not a valid directory."36        self.convert_mode, self.check_length = convert_mode, check_length37        self.img_h, self.img_w = img_h, img_w38        self.max_length, self.one_hot_y = max_length, one_hot_y39        self.return_idx, self.return_raw = return_idx, return_raw40        self.case_sensitive, self.is_training = case_sensitive, is_training41        self.data_aug, self.multiscales = data_aug, multiscales42        self.charset = CharsetMapper(charset_path, max_length=max_length+1)43        self.c = self.charset.num_classes44 45        self.env = lmdb.open(str(path), readonly=True, lock=False, readahead=False, meminit=False)46        assert self.env, f'Cannot open LMDB dataset from {path}.'47        with self.env.begin(write=False) as txn:48            self.length = int(txn.get('num-samples'.encode()))49 50        if self.is_training and self.data_aug:51            self.augment_tfs = transforms.Compose([52                CVGeometry(degrees=45, translate=(0.0, 0.0), scale=(0.5, 2.), shear=(45, 15), distortion=0.5, p=0.5),53                CVDeterioration(var=20, degrees=6, factor=4, p=0.25),54                CVColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.25)55            ])56        self.totensor = transforms.ToTensor()57    58    def __len__(self): return self.length59 60    def _next_image(self, index):61        next_index = random.randint(0, len(self) - 1)62        return self.get(next_index)63 64    def _check_image(self, x, pixels=6):65        if x.size[0] <= pixels or x.size[1] <= pixels: return False66        else: return True67 68    def resize_multiscales(self, img, borderType=cv2.BORDER_CONSTANT): 69        def _resize_ratio(img, ratio, fix_h=True):70            if ratio * self.img_w < self.img_h:71                if fix_h: trg_h = self.img_h72                else: trg_h = int(ratio * self.img_w)73                trg_w = self.img_w74            else: trg_h, trg_w = self.img_h, int(self.img_h / ratio)75            img = cv2.resize(img, (trg_w, trg_h))76            pad_h, pad_w = (self.img_h - trg_h) / 2, (self.img_w - trg_w) / 277            top, bottom = math.ceil(pad_h), math.floor(pad_h)78            left, right = math.ceil(pad_w), math.floor(pad_w)79            img = cv2.copyMakeBorder(img, top, bottom, left, right, borderType)80            return img81        82        if self.is_training: 83            if random.random() < 0.5:84                base, maxh, maxw = self.img_h, self.img_h, self.img_w85                h, w = random.randint(base, maxh), random.randint(base, maxw)86                return _resize_ratio(img, h/w)87            else: return _resize_ratio(img, img.shape[0] / img.shape[1])  # keep aspect ratio88        else:  return _resize_ratio(img, img.shape[0] / img.shape[1])  # keep aspect ratio89 90    def resize(self, img):91        if self.multiscales: return self.resize_multiscales(img, cv2.BORDER_REPLICATE)92        else: return cv2.resize(img, (self.img_w, self.img_h))93         94    def get(self, idx):95        with self.env.begin(write=False) as txn:96            image_key, label_key = f'image-{idx+1:09d}', f'label-{idx+1:09d}'97            try:98                label = str(txn.get(label_key.encode()), 'utf-8')  # label99                label = re.sub('[^0-9a-zA-Z]+', '', label)100                if self.check_length and self.max_length > 0:101                    if len(label) > self.max_length or len(label) <= 0:102                        #logging.info(f'Long or short text image is found: {self.name}, {idx}, {label}, {len(label)}')103                        return self._next_image(idx)104                label = label[:self.max_length]105 106                imgbuf = txn.get(image_key.encode())  # image107                buf = six.BytesIO()108                buf.write(imgbuf)109                buf.seek(0)110                with warnings.catch_warnings():111                    warnings.simplefilter("ignore", UserWarning) # EXIF warning from TiffPlugin112                    image = PIL.Image.open(buf).convert(self.convert_mode)113                if self.is_training and not self._check_image(image):114                    #logging.info(f'Invalid image is found: {self.name}, {idx}, {label}, {len(label)}')115                    return self._next_image(idx)116            except:117                import traceback118                traceback.print_exc()119                logging.info(f'Corrupted image is found: {self.name}, {idx}, {label}, {len(label)}')120                return self._next_image(idx)121            return image, label, idx122 123    def _process_training(self, image):124        if self.data_aug: image = self.augment_tfs(image)125        image = self.resize(np.array(image))126        return image127 128    def _process_test(self, image):129        return self.resize(np.array(image)) # TODO:move is_training to here130 131    def __getitem__(self, idx):132        image, text, idx_new = self.get(idx)133        if not self.is_training: assert idx == idx_new, f'idx {idx} != idx_new {idx_new} during testing.'134 135        if self.is_training: image = self._process_training(image)136        else: image = self._process_test(image)137        if self.return_raw: return image, text138        image = self.totensor(image)139 140        length = tensor(len(text) + 1).to(dtype=torch.long)  # one for end token141        label = self.charset.get_labels(text, case_sensitive=self.case_sensitive)142        label = tensor(label).to(dtype=torch.long)143        if self.one_hot_y: label = onehot(label, self.charset.num_classes)144 145        if self.return_idx: y = [label, length, idx_new]146        else: y = [label, length]147        return image, y148 149 150class TextDataset(Dataset):151    def __init__(self,152                 path:PathOrStr, 153                 delimiter:str='\t',154                 max_length:int=25, 155                 charset_path:str='data/charset_36.txt', 156                 case_sensitive=False, 157                 one_hot_x=True,158                 one_hot_y=True,159                 is_training=True,160                 smooth_label=False,161                 smooth_factor=0.2,162                 use_sm=False,163                 **kwargs):164        self.path = Path(path)165        self.case_sensitive, self.use_sm = case_sensitive, use_sm166        self.smooth_factor, self.smooth_label = smooth_factor, smooth_label167        self.charset = CharsetMapper(charset_path, max_length=max_length+1)168        self.one_hot_x, self.one_hot_y, self.is_training = one_hot_x, one_hot_y, is_training169        if self.is_training and self.use_sm: self.sm = SpellingMutation(charset=self.charset)170 171        dtype = {'inp': str, 'gt': str}172        self.df = pd.read_csv(self.path, dtype=dtype, delimiter=delimiter, na_filter=False)173        self.inp_col, self.gt_col = 0, 1174 175    def __len__(self): return len(self.df)176 177    def __getitem__(self, idx):178        text_x = self.df.iloc[idx, self.inp_col]179        text_x = re.sub('[^0-9a-zA-Z]+', '', text_x)180        if not self.case_sensitive: text_x = text_x.lower()181        if self.is_training and self.use_sm: text_x = self.sm(text_x)182 183        length_x = tensor(len(text_x) + 1).to(dtype=torch.long)  # one for end token184        label_x = self.charset.get_labels(text_x, case_sensitive=self.case_sensitive)185        label_x = tensor(label_x)186        if self.one_hot_x:187            label_x = onehot(label_x, self.charset.num_classes)188            if self.is_training and self.smooth_label: 189                label_x = torch.stack([self.prob_smooth_label(l) for l in label_x])190        x =  [label_x, length_x]191    192        text_y = self.df.iloc[idx, self.gt_col]193        text_y = re.sub('[^0-9a-zA-Z]+', '', text_y)194        if not self.case_sensitive: text_y = text_y.lower()195        length_y = tensor(len(text_y) + 1).to(dtype=torch.long)  # one for end token196        label_y = self.charset.get_labels(text_y, case_sensitive=self.case_sensitive)197        label_y = tensor(label_y)198        if self.one_hot_y: label_y = onehot(label_y, self.charset.num_classes)199        y = [label_y, length_y]200 201        return x, y202 203    def prob_smooth_label(self, one_hot):204        one_hot = one_hot.float()205        delta = torch.rand([]) * self.smooth_factor206        num_classes = len(one_hot)207        noise = torch.rand(num_classes)208        noise = noise / noise.sum() * delta209        one_hot = one_hot * (1 - delta) + noise210        return one_hot211 212 213class SpellingMutation(object):214    def __init__(self, pn0=0.7, pn1=0.85, pn2=0.95, pt0=0.7, pt1=0.85, charset=None):215        """ 216        Args:217            pn0: the prob of not modifying characters is (pn0)218            pn1: the prob of modifying one characters is (pn1 - pn0)219            pn2: the prob of modifying two characters is (pn2 - pn1), 220                 and three (1 - pn2)221            pt0: the prob of replacing operation is pt0.222            pt1: the prob of inserting operation is (pt1 - pt0),223                 and deleting operation is (1 - pt1)224        """225        super().__init__()226        self.pn0, self.pn1, self.pn2 = pn0, pn1, pn2227        self.pt0, self.pt1 = pt0, pt1228        self.charset = charset229        logging.info(f'the probs: pn0={self.pn0}, pn1={self.pn1} ' + 230                     f'pn2={self.pn2}, pt0={self.pt0}, pt1={self.pt1}')231 232    def is_digit(self, text, ratio=0.5):233        length = max(len(text), 1)234        digit_num = sum([t in self.charset.digits for t in text])235        if digit_num / length < ratio: return False236        return True237 238    def is_unk_char(self, char):239        # return char == self.charset.unk_char240        return (char not in self.charset.digits) and (char not in self.charset.alphabets)241 242    def get_num_to_modify(self, length):243        prob = random.random()244        if prob < self.pn0: num_to_modify = 0245        elif prob < self.pn1: num_to_modify = 1246        elif prob < self.pn2: num_to_modify = 2247        else: num_to_modify = 3248        249        if length <= 1: num_to_modify = 0250        elif length >= 2 and length <= 4: num_to_modify = min(num_to_modify, 1)251        else: num_to_modify = min(num_to_modify, length // 2)  # smaller than length // 2252        return num_to_modify253 254    def __call__(self, text, debug=False):255        if self.is_digit(text): return text256        length = len(text)257        num_to_modify = self.get_num_to_modify(length)258        if num_to_modify <= 0: return text259 260        chars = []261        index = np.arange(0, length)262        random.shuffle(index)263        index = index[: num_to_modify]264        if debug: self.index = index265        for i, t in enumerate(text):266            if i not in index: chars.append(t)267            elif self.is_unk_char(t): chars.append(t)268            else:269                prob = random.random()270                if prob < self.pt0: # replace271                    chars.append(random.choice(self.charset.alphabets))272                elif prob < self.pt1: # insert273                    chars.append(random.choice(self.charset.alphabets))274                    chars.append(t)275                else: # delete276                    continue277        new_text = ''.join(chars[: self.charset.max_length-1])278        return new_text if len(new_text) >= 1 else text