-
Notifications
You must be signed in to change notification settings - Fork 0
/
api.py
96 lines (79 loc) · 2.62 KB
/
api.py
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import base64
import io
import json
import logging
import os
import time
import traceback
import cv2
import numpy as np
from flask import Flask, jsonify, request
from flask.json import jsonify
from flask_cors import CORS
from PIL import Image
from main import SceneTextExtractor
app = Flask(__name__)
CORS(app)
logging.getLogger().setLevel(logging.INFO)
models = SceneTextExtractor()
def read_image(img_bytes):
return cv2.imdecode(np.asarray(bytearray(img_bytes.read()), dtype="uint8"), cv2.IMREAD_COLOR)
@app.route("/", methods=["POST"])
def predict():
start_time = time.time()
response = {"success": False}
try:
image = request.files.get("image", None)
if image is not None:
image = read_image(image)
elif request.data is not None:
image = cv2.imdecode(np.fromstring(
request.data, np.uint8), cv2.IMREAD_COLOR)
res = models.process(image)
response["prediction"] = res
response["success"] = True
except Exception as ex:
response["ex"] = ex
print(traceback.format_exc())
response['run_time'] = "%.2f" % (time.time() - start_time)
return jsonify(response)
@app.route("/layout", methods=["POST"])
def predict_layout():
start_time = time.time()
response = {"success": False}
try:
image = request.files.get("image", None)
if image is not None:
image = read_image(image)
elif request.data is not None:
image = cv2.imdecode(np.fromstring(
request.data, np.uint8), cv2.IMREAD_COLOR)
res = models.layout_model.process(image)
response["prediction"] = res
response["success"] = True
except Exception as ex:
response["ex"] = ex
print(traceback.format_exc())
response['run_time'] = "%.2f" % (time.time() - start_time)
return jsonify(response)
@app.route("/ocr", methods=["POST"])
def predict_ocr():
start_time = time.time()
response = {"success": False}
try:
image = request.files.get("image", None)
if image is not None:
image = read_image(image)
elif request.data is not None:
image = cv2.imdecode(np.fromstring(
request.data, np.uint8), cv2.IMREAD_COLOR)
res = models.ocr_model.predict([image])
response["prediction"] = res
response["success"] = True
except Exception as ex:
response["ex"] = ex
print(traceback.format_exc())
response['run_time'] = "%.2f" % (time.time() - start_time)
return jsonify(response)
if __name__ == '__main__':
app.run(debug=True, host='0.0.0.0', port=8080)