def __init__(self, model_path):
alphabet_unicode = config.alphabet_v2
self.alphabet = ''.join([chr(uni) for uni in alphabet_unicode])
# print(len(self.alphabet))
self.nclass = len(self.alphabet) + 1
self.model = CRNN(config.imgH, 1, self.nclass, 256)
self.cuda = False
if torch.cuda.is_available():
self.cuda = True
self.model.cuda()
self.model.load_state_dict({k.replace('module.', ''): v \
for k, v in torch.load(model_path).items()})
else:
# self.model = nn.DataParallel(self.model)
self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
self.model.eval()
self.converter = strLabelConverter(self.alphabet)
class PytorchOcr():
def __init__(self, model_path):
alphabet_unicode = config.alphabet_v2
self.alphabet = ''.join([chr(uni) for uni in alphabet_unicode])
# print(len(self.alphabet))
self.nclass = len(self.alphabet) + 1
self.model = CRNN(config.imgH, 1, self.nclass, 256)
self.cuda = False
if torch.cuda.is_available():
self.cuda = True
self.model.cuda()
self.model.load_state_dict({k.replace('module.', ''): v \
for k, v in torch.load(model_path).items()})
else:
# self.model = nn.DataParallel(self.model)
self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
self.model.eval()
self.converter = strLabelConverter(self.alphabet)
def recognize(self, img):
h,w = img.shape[:2]
if len(img.shape) == 3:
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
image = Image.fromarray(img)
transformer = resizeNormalize((int(w/h*32), 32))
image = transformer(image)
image = image.view(1, *image.size())
image = Variable(image)
if self.cuda:
image = image.cuda()
preds = self.model(image)
_, preds = preds.max(2)
preds = preds.transpose(1, 0).contiguous().view(-1)
preds_size = Variable(torch.IntTensor([preds.size(0)]))
txt = self.converter.decode(preds.data, preds_size.data, raw=False).strip()
return txt
if __name__ == '__main__':
# GPU settings
gpus = tf.config.list_physical_devices("GPU")
if gpus:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
# dataset
dataset = Dataset()
num_classes, blank_index = get_num_classes_and_blank_index()
train_set, train_size = dataset.train_dataset()
valid_set, valid_size = dataset.valid_dataset()
# model
crnn_model = CRNN(num_classes)
# loss
loss = CTCLoss()
# optimizer
optimizer = tf.optimizers.Adadelta()
# metrics
train_loss_metric = tf.metrics.Mean()
valid_loss_metric = tf.metrics.Mean()
accuracy = Accuracy(blank_index)
train_accuracy = tf.metrics.Mean()
valid_accuracy = tf.metrics.Mean()
def train_step(batch_images, batch_labels):
from configuration import BATCH_SIZE, EPOCHS, save_model_dir, save_frequency
if __name__ == '__main__':
# GPU settings
gpus = tf.config.list_physical_devices("GPU")
if gpus:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
# dataset
dataset = Dataset()
train_set, train_size = dataset.train_dataset()
valid_set, valid_size = dataset.valid_dataset()
# model
crnn_model = CRNN()
# loss
loss = CTCLoss()
# optimizer
optimizer = tf.optimizers.Adadelta()
# metrics
train_loss_metric = tf.metrics.Mean()
valid_loss_metric = tf.metrics.Mean()
accuracy = Accuracy()
accuracy_metric = tf.metrics.Mean()
def train_step(batch_images, batch_labels):
with tf.GradientTape() as tape:
from core.utils import get_num_classes_and_blank_index
if __name__ == '__main__':
# GPU settings
gpus = tf.config.list_physical_devices("GPU")
if gpus:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
# dataset
dataset = Dataset()
num_classes, blank_index = get_num_classes_and_blank_index()
test_set, test_size = dataset.test_dataset()
# model
crnn_model = CRNN(num_classes)
crnn_model.load_weights(filepath=Config.save_model_dir+"saved_model")
# loss and metrics
loss_object = CTCLoss()
loss_metric = tf.metrics.Mean()
test_accuracy = Accuracy(blank_index)
test_accuracy_metric = tf.metrics.Mean()
def test_step(batch_images, batch_labels):
pred = crnn_model(batch_images, training=False)
loss_value = loss_object(y_true=batch_labels, y_pred=pred)
acc = test_accuracy(decoded_text=predict_text(pred, blank_index=blank_index), true_label=batch_labels)
loss_metric.update_state(values=loss_value)
test_accuracy_metric.update_state(values=acc)
from core.predict import predict_text
from configuration import save_model_dir
if __name__ == '__main__':
# GPU settings
gpus = tf.config.list_physical_devices("GPU")
if gpus:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
# dataset
dataset = Dataset()
test_set, test_size = dataset.test_dataset()
# model
crnn_model = CRNN()
crnn_model.load_weights(filepath=save_model_dir + "saved_model")
# loss and metrics
loss_object = CTCLoss()
loss_metric = tf.metrics.Mean()
test_accuracy = Accuracy()
test_accuracy_metric = tf.metrics.Mean()
def test_step(batch_images, batch_labels, labels_list):
pred = crnn_model(batch_images, training=False)
loss_value = loss_object(y_true=batch_labels, y_pred=pred)
acc = test_accuracy(decoded_text=predict_text(pred),
true_label=labels_list)
loss_metric.update_state(values=loss_value)
test_accuracy_metric.update_state(values=acc)