refactoring
This commit is contained in:
@@ -7,6 +7,7 @@ from captcha.DatasetFactory import DatasetFactory
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
from tensorflow import keras
|
from tensorflow import keras
|
||||||
|
|
||||||
|
|
||||||
class CaptchaReader:
|
class CaptchaReader:
|
||||||
|
|
||||||
def __init__(self, modelFilepath, captchaShape):
|
def __init__(self, modelFilepath, captchaShape):
|
||||||
@@ -14,19 +15,17 @@ class CaptchaReader:
|
|||||||
self.captchaShape = captchaShape
|
self.captchaShape = captchaShape
|
||||||
|
|
||||||
def getTextInCaptchaImage(self, captchaImageFile):
|
def getTextInCaptchaImage(self, captchaImageFile):
|
||||||
# FK-TODO: refactor
|
return self._getTextsInCaptchaImage(self._getCaptchaImage(captchaImageFile))[0]
|
||||||
modelDAO = ModelDAO(inColab = False)
|
|
||||||
model = modelDAO.loadModel(self.modelFilepath)
|
def _getCaptchaImage(self, captchaImageFile):
|
||||||
prediction_model = ModelFactory.createPredictionModel(model)
|
return self._asSingleSampleBatch(DatasetFactory.encodeImage(captchaImageFile, self.captchaShape))
|
||||||
charNumConverter = CharNumConverter(CaptchaGenerator.characters)
|
|
||||||
datasetFactory = DatasetFactory(self.captchaShape,charNumConverter.char_to_num, batch_size = 64)
|
|
||||||
batchImages = self._asSingleSampleBatch(datasetFactory._encode_single_sample(captchaImageFile, 'dummy')['image'])
|
|
||||||
preds = prediction_model.predict(batchImages)
|
|
||||||
predictionsDecoder = PredictionsDecoder(CaptchaGenerator.captchaLength, charNumConverter.num_to_char)
|
|
||||||
pred_texts = predictionsDecoder.decode_batch_predictions(preds)
|
|
||||||
return pred_texts[0]
|
|
||||||
|
|
||||||
def _asSingleSampleBatch(self, img):
|
def _asSingleSampleBatch(self, img):
|
||||||
array = keras.utils.img_to_array(img)
|
return np.expand_dims(keras.utils.img_to_array(img), axis=0)
|
||||||
array = np.expand_dims(array, axis=0)
|
|
||||||
return array
|
def _getTextsInCaptchaImage(self, captchaImage):
|
||||||
|
preds = self._createPredictionModel().predict(captchaImage)
|
||||||
|
return PredictionsDecoder(CaptchaGenerator.captchaLength, CharNumConverter(CaptchaGenerator.characters).num_to_char).decode_batch_predictions(preds)
|
||||||
|
|
||||||
|
def _createPredictionModel(self):
|
||||||
|
return ModelFactory.createPredictionModel(ModelDAO(inColab=False).loadModel(self.modelFilepath))
|
||||||
|
|||||||
@@ -10,15 +10,18 @@ class DatasetFactory:
|
|||||||
|
|
||||||
def createDataset(self, x, y):
|
def createDataset(self, x, y):
|
||||||
dataset = tf.data.Dataset.from_tensor_slices((x, y))
|
dataset = tf.data.Dataset.from_tensor_slices((x, y))
|
||||||
dataset = dataset.map(self._encode_single_sample, num_parallel_calls=tf.data.AUTOTUNE)
|
dataset = dataset.map(self._encodeImageAndLabel, num_parallel_calls=tf.data.AUTOTUNE)
|
||||||
dataset = dataset.batch(self.batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)
|
dataset = dataset.batch(self.batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)
|
||||||
return dataset
|
return dataset
|
||||||
|
|
||||||
def _encode_single_sample(self, img_path, label):
|
def _encodeImageAndLabel(self, imageFilename, label):
|
||||||
img = tf.io.read_file(img_path)
|
return {
|
||||||
|
"image": DatasetFactory.encodeImage(imageFilename, self.captchaShape),
|
||||||
|
"label": self.char_to_num(tf.strings.unicode_split(label, input_encoding="UTF-8"))}
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def encodeImage(imageFilename, captchaShape):
|
||||||
|
img = tf.io.read_file(imageFilename)
|
||||||
img = tf.io.decode_jpeg(img, channels=3)
|
img = tf.io.decode_jpeg(img, channels=3)
|
||||||
img = tf.image.resize(img, [self.captchaShape.height, self.captchaShape.width])
|
img = tf.image.resize(img, [captchaShape.height, captchaShape.width])
|
||||||
# Map the characters in label to numbers
|
return img
|
||||||
label = self.char_to_num(tf.strings.unicode_split(label, input_encoding="UTF-8"))
|
|
||||||
# Return a dict as our model is expecting two inputs
|
|
||||||
return {"image": img, "label": label}
|
|
||||||
|
|||||||
Reference in New Issue
Block a user