File size: 2,523 Bytes
d0ba6da bd1dc81 d0ba6da 44066b7 d0ba6da 44066b7 d0ba6da 44066b7 d0ba6da 44066b7 d0ba6da e11580d d0ba6da 44066b7 d0ba6da |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import os
import cv2
import json
import torch
import logging
import numpy as np
from flask import Flask, jsonify, request
from flask.wrappers import Response
from iam_line_recognition.model_main import CRNN
from iam_line_recognition.utils import ctc_decode
from iam_line_recognition.dataset import HWRecogIAMDataset
app = Flask(__name__)
logging.basicConfig(level=logging.INFO)
file_model_local = "artifacts/crnn_H_32_W_768_E_196.pth"
file_model_cont = "/data/models/crnn_H_32_W_768_E_196.pth"
device = "cpu"
num_classes = len(HWRecogIAMDataset.LABEL_2_CHAR) + 1
image_height = 32
mean_arr = np.array([[0.485, 0.456, 0.406]])
std_arr = np.array([[0.229, 0.224, 0.225]])
hw_recog_model = CRNN(num_classes, image_height)
try:
logging.info(f"loading model from {file_model_local}")
hw_recog_model.load_state_dict(torch.load(file_model_local, map_location=device))
except:
logging.info(f"loading model from {file_model_cont}")
hw_recog_model.load_state_dict(torch.load(file_model_cont, map_location=device))
hw_recog_model.to(device)
hw_recog_model.eval()
def predict_hw(img_test: np.ndarray) -> str:
img_test = np.expand_dims(img_test, 0)
img_test = img_test.astype(np.float32) / 255.0
img_test = (img_test - mean_arr) / std_arr
img_test = np.transpose(img_test, axes=[0, 3, 1, 2])
img_tensor = torch.tensor(img_test).float()
img_tensor = img_tensor.to(device, dtype=torch.float)
log_probs = hw_recog_model(img_tensor)
pred_labels = ctc_decode(log_probs.detach())
str_pred = [HWRecogIAMDataset.LABEL_2_CHAR[i] for i in pred_labels[0]]
str_pred = "".join(str_pred)
return str_pred
@app.route("/predict", methods=["POST"])
def predict() -> Response:
logging.info("IAM Handwriting recognition app")
img_file = request.files["image_file"]
try:
img_arr = np.fromstring(img_file.read(), np.uint8)
except:
img_arr = np.fromstring(img_file.getvalue(), np.uint8)
img_dec = cv2.imdecode(img_arr, cv2.IMREAD_COLOR)
img_dec = cv2.cvtColor(img_dec, cv2.COLOR_BGR2RGB)
img_dec = cv2.resize(img_dec, (768, 32), interpolation=cv2.INTER_LINEAR)
str_pred = predict_hw(img_dec)
dict_pred = {
"file_name": img_file.filename,
"prediction": str_pred,
}
logging.info(dict_pred)
try:
json_pred = jsonify(dict_pred)
except TypeError as e:
json_pred = jsonify({"error": str(e)})
return json_pred
if __name__ == "__main__":
app.run(host="0.0.0.0", debug=True, port=7860)
|