Skip to content

Commit b2a0a75

Browse files
Merge pull request #2 from pythonlessons/feature/captcha_to_text
Feature/captcha to text
2 parents 270e77c + 0a019b6 commit b2a0a75

File tree

16 files changed

+1232
-125
lines changed

16 files changed

+1232
-125
lines changed

CHANGELOG.md

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
1-
## [0.1.3] - 2022-20-12
1+
## [0.1.4] - 2022-12-21
2+
### Added:
3+
- added mltu.augmentors (RandomBrightness, RandomRotate, RandomErodeDilate) - used for simple image augmentation;
4+
5+
## [0.1.3] - 2022-12-20
26

37
Initial release of mltu (Machine Learning Training Utilities)
48

README.md

Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,17 @@
11
# MLTU - Machine Learning Training Utilities (TensorFlow)
22
Machine Learning Training Utilities with TensorFlow 2.* and Python 3
33

4-
## Installation:
5-
Clone the repository and install the requirements:
4+
# Installation:
5+
To use MLTU in your own project, you can install it from PyPI:
66
```bash
7-
git clone https://github.com/pythonlessons/mltu.git
7+
pip install mltu
88
```
9-
cd into the repository
9+
When running tutorials, it's necessary to install mltu for a specific tutorial, for example:
1010
```bash
11-
cd mltu
12-
```
13-
Install the requirements:
14-
```bash
15-
pip install -r requirements.txt
16-
```
17-
18-
Install the mltu package
19-
```bash
20-
pip install .
11+
pip install mltu==0.1.3
2112
```
13+
Each tutorial has its own requirements.txt file for a specific mltu version. As this project progress, the newest versions may have breaking changes, so it's recommended to use the same version as in the tutorial.
2214

2315
# Tutorials and Examples:
24-
...
16+
1. [Text Recognition With TensorFlow and CTC network](https://pylessons.com/ctc-text-recognition), code in ```Tutorials\01_image_to_word``` folder;
17+
2. [TensorFlow OCR model for reading Captchas](https://pylessons.com/tensorflow-ocr-captcha), code in ```Tutorials\02_captcha_to_text``` folder;

Tutorials/01_image_to_word/README.md

Lines changed: 505 additions & 1 deletion
Large diffs are not rendered by default.
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
mltu==0.1.3

Tutorials/01_image_to_word/train.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,3 @@
1-
# https://github.com/rajesh-bhat/spark-ai-summit-2020-text-extraction
2-
# https://github.com/tensorflow/benchmarks/blob/master/scripts/tf_cnn_benchmarks/models/experimental/deepspeech.py
3-
# https://www.robots.ox.ac.uk/~vgg/data/text/#sec-chars
41
import stow
52
from tqdm import tqdm
63
import tensorflow as tf

Tutorials/02_captcha_to_text/README.md

Lines changed: 336 additions & 0 deletions
Large diffs are not rendered by default.
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
import stow
2+
from datetime import datetime
3+
4+
from mltu.configs import BaseModelConfigs
5+
6+
class ModelConfigs(BaseModelConfigs):
7+
def __init__(self):
8+
super().__init__()
9+
self.model_path = stow.join('Models/02_captcha_to_text', datetime.strftime(datetime.now(), "%Y%m%d%H%M"))
10+
self.vocab = ''
11+
self.height = 50
12+
self.width = 200
13+
self.max_text_length = 0
14+
self.batch_size = 64
15+
self.learning_rate = 1e-3
16+
self.train_epochs = 1000
17+
self.train_workers = 20
Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import cv2
2+
import typing
3+
import numpy as np
4+
5+
from mltu.inferenceModel import OnnxInferenceModel
6+
from mltu.utils.text_utils import ctc_decoder, get_cer
7+
8+
class ImageToWordModel(OnnxInferenceModel):
9+
def __init__(self, char_list: typing.Union[str, list], *args, **kwargs):
10+
super().__init__(*args, **kwargs)
11+
self.char_list = char_list
12+
13+
def predict(self, image: np.ndarray):
14+
image = cv2.resize(image, self.input_shape[:2][::-1])
15+
16+
image_pred = np.expand_dims(image, axis=0).astype(np.float32)
17+
18+
preds = self.model.run(None, {self.input_name: image_pred})[0]
19+
20+
text = ctc_decoder(preds, self.char_list)[0]
21+
22+
return text
23+
24+
if __name__ == "__main__":
25+
import pandas as pd
26+
from tqdm import tqdm
27+
from mltu.configs import BaseModelConfigs
28+
29+
configs = BaseModelConfigs.load("Models/02_captcha_to_text/202212211205/configs.yaml")
30+
31+
model = ImageToWordModel(model_path=configs.model_path, char_list=configs.vocab)
32+
33+
df = pd.read_csv("Models/02_captcha_to_text/202212211205/val.csv").values.tolist()
34+
35+
accum_cer = []
36+
for image_path, label in tqdm(df):
37+
image = cv2.imread(image_path)
38+
39+
prediction_text = model.predict(image)
40+
41+
cer = get_cer(prediction_text, label)
42+
print(f"Image: {image_path}, Label: {label}, Prediction: {prediction_text}, CER: {cer}")
43+
44+
accum_cer.append(cer)
45+
46+
print(f"Average CER: {np.average(accum_cer)}")
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
from keras import layers
2+
from keras.models import Model
3+
4+
from mltu.model_utils import residual_block
5+
6+
def train_model(input_dim, output_dim, activation='leaky_relu', dropout=0.2):
7+
8+
inputs = layers.Input(shape=input_dim, name="input")
9+
10+
# normalize images here instead in preprocessing step
11+
input = layers.Lambda(lambda x: x / 255)(inputs)
12+
13+
x1 = residual_block(input, 16, activation=activation, skip_conv=True, strides=1, dropout=dropout)
14+
15+
x2 = residual_block(x1, 16, activation=activation, skip_conv=True, strides=2, dropout=dropout)
16+
x3 = residual_block(x2, 16, activation=activation, skip_conv=False, strides=1, dropout=dropout)
17+
18+
x4 = residual_block(x3, 32, activation=activation, skip_conv=True, strides=2, dropout=dropout)
19+
x5 = residual_block(x4, 32, activation=activation, skip_conv=False, strides=1, dropout=dropout)
20+
21+
x6 = residual_block(x5, 64, activation=activation, skip_conv=True, strides=2, dropout=dropout)
22+
x7 = residual_block(x6, 32, activation=activation, skip_conv=True, strides=1, dropout=dropout)
23+
24+
x8 = residual_block(x7, 64, activation=activation, skip_conv=True, strides=2, dropout=dropout)
25+
x9 = residual_block(x8, 64, activation=activation, skip_conv=False, strides=1, dropout=dropout)
26+
27+
squeezed = layers.Reshape((x9.shape[-3] * x9.shape[-2], x9.shape[-1]))(x9)
28+
29+
blstm = layers.Bidirectional(layers.LSTM(128, return_sequences=True))(squeezed)
30+
blstm = layers.Dropout(dropout)(blstm)
31+
32+
output = layers.Dense(output_dim + 1, activation='softmax', name="output")(blstm)
33+
34+
model = Model(inputs=inputs, outputs=output)
35+
return model
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
import tensorflow as tf
2+
try: [tf.config.experimental.set_memory_growth(gpu, True) for gpu in tf.config.experimental.list_physical_devices('GPU')]
3+
except: pass
4+
5+
from keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, TensorBoard
6+
7+
from mltu.dataProvider import DataProvider
8+
from mltu.preprocessors import ImageReader
9+
from mltu.transformers import ImageResizer, LabelIndexer, LabelPadding
10+
from mltu.augmentors import RandomBrightness, RandomRotate, RandomErodeDilate
11+
from mltu.losses import CTCloss
12+
from mltu.callbacks import Model2onnx, TrainLogger
13+
from mltu.metrics import CWERMetric
14+
15+
from model import train_model
16+
from configs import ModelConfigs
17+
18+
import stow
19+
from urllib.request import urlopen
20+
from io import BytesIO
21+
from zipfile import ZipFile
22+
23+
def download_and_unzip(url, extract_to='Datasets'):
24+
http_response = urlopen(url)
25+
zipfile = ZipFile(BytesIO(http_response.read()))
26+
zipfile.extractall(path=extract_to)
27+
28+
if not stow.exists(stow.join('Datasets', 'captcha_images_v2')):
29+
download_and_unzip('https://github.com/AakashKumarNain/CaptchaCracker/raw/master/captcha_images_v2.zip', extract_to='Datasets')
30+
31+
# Create a list of all the images and labels in the dataset
32+
dataset, vocab, max_len = [], set(), 0
33+
for file in stow.ls(stow.join('Datasets', 'captcha_images_v2')):
34+
dataset.append([stow.relpath(file), file.name])
35+
vocab.update(list(file.name))
36+
max_len = max(max_len, len(file.name))
37+
38+
configs = ModelConfigs()
39+
40+
# Save vocab and maximum text length to configs
41+
configs.vocab = "".join(vocab)
42+
configs.max_text_length = max_len
43+
configs.save()
44+
45+
# Create a data provider for the dataset
46+
data_provider = DataProvider(
47+
dataset=dataset,
48+
skip_validation=True,
49+
batch_size=configs.batch_size,
50+
data_preprocessors=[ImageReader()],
51+
transformers=[
52+
ImageResizer(configs.width, configs.height),
53+
LabelIndexer(configs.vocab),
54+
LabelPadding(max_word_length=configs.max_text_length, padding_value=len(configs.vocab))
55+
],
56+
)
57+
# Split the dataset into training and validation sets
58+
train_data_provider, val_data_provider = data_provider.split(split = 0.9)
59+
60+
# Augment training data with random brightness, rotation and erode/dilate
61+
train_data_provider.augmentors = [RandomBrightness(), RandomRotate(), RandomErodeDilate()]
62+
63+
model = train_model(
64+
input_dim = (configs.height, configs.width, 3),
65+
output_dim = len(configs.vocab),
66+
)
67+
68+
# Compile the model and print summary
69+
model.compile(
70+
optimizer=tf.keras.optimizers.Adam(learning_rate=configs.learning_rate),
71+
loss=CTCloss(),
72+
metrics=[CWERMetric()],
73+
run_eagerly=False
74+
)
75+
model.summary(line_length=110)
76+
# Define path to save the model
77+
stow.mkdir(configs.model_path)
78+
79+
# Define callbacks
80+
earlystopper = EarlyStopping(monitor='val_CER', patience=40, verbose=1)
81+
checkpoint = ModelCheckpoint(f"{configs.model_path}/model.h5", monitor='val_CER', verbose=1, save_best_only=True, mode='min')
82+
trainLogger = TrainLogger(configs.model_path)
83+
tb_callback = TensorBoard(f'{configs.model_path}/logs', update_freq=1)
84+
reduceLROnPlat = ReduceLROnPlateau(monitor='val_CER', factor=0.9, min_delta=1e-10, patience=20, verbose=1, mode='auto')
85+
model2onnx = Model2onnx(f"{configs.model_path}/model.h5")
86+
87+
# Train the model
88+
model.fit(
89+
train_data_provider,
90+
validation_data=val_data_provider,
91+
epochs=configs.train_epochs,
92+
callbacks=[earlystopper, checkpoint, trainLogger, reduceLROnPlat, tb_callback, model2onnx],
93+
workers=configs.train_workers
94+
)
95+
96+
# Save training and validation datasets as csv files
97+
train_data_provider.to_csv(stow.join(configs.model_path, 'train.csv'))
98+
val_data_provider.to_csv(stow.join(configs.model_path, 'val.csv'))

0 commit comments

Comments
 (0)