diff --git a/clas_api.py b/clas_api.py new file mode 100644 index 0000000..f50b9b0 --- /dev/null +++ b/clas_api.py @@ -0,0 +1,26 @@ +from flask import Flask, request +from paddleclas import PaddleClas + +from util.common_util import process_request + +app = Flask(__name__) +CLAS = PaddleClas(model_name="text_image_orientation") + + +@app.route('/clas/orientation', methods=['POST']) +@process_request +def orientation(): + """ + 判断图片旋转角度，逆时针旋转该角度后为正。可能值["0", "90", "180", "270"] + :return: 最有可能的两个角度 + """ + img_path = request.form.get('img_path') + clas_result = CLAS.predict(input_data=img_path) + clas_result = next(clas_result)[0] + if clas_result["scores"][0] < 0.5: + return ['0', '90'] + return clas_result["label_names"] + + +if __name__ == '__main__': + app.run('0.0.0.0', 5002) diff --git a/det_api.py b/det_api.py index 9084945..abd8e2c 100644 --- a/det_api.py +++ b/det_api.py @@ -1,24 +1,23 @@ import os.path import cv2 -from flask import Flask, request, Blueprint +from flask import Flask, request from paddle_detection import detector +from util import image_util from util.common_util import process_request app = Flask(__name__) -det_bp = Blueprint('det_bp', __name__) -@det_bp.route('/books', methods=['POST']) +@app.route('/det/books', methods=['POST']) @process_request def books(): - img_path = request.form['img_path'] + img_path = request.form.get('img_path') image = cv2.imread(img_path) result = detector.get_book_areas(image) - dirname = os.path.dirname(img_path) - img_name, ext = os.path.basename(img_path).rsplit('.', 1) + dirname, img_name, ext = image_util.parse_path(img_path) books_path = [] for i in range(len(result)): save_path = os.path.join(dirname, img_name + '_book_' + str(i) + '.' + ext) @@ -27,7 +26,5 @@ def books(): return books_path -app.register_blueprint(det_bp, url_prefix='/det') - if __name__ == '__main__': app.run('0.0.0.0', 5000) diff --git a/dewarp_api.py b/dewarp_api.py new file mode 100644 index 0000000..f755463 --- /dev/null +++ b/dewarp_api.py @@ -0,0 +1,26 @@ +import os + +import cv2 +from flask import Flask, request + +from doc_dewarp import dewarp +from util import image_util +from util.common_util import process_request + +app = Flask(__name__) + + +@app.route('/dewarp', methods=['POST']) +@process_request +def dewarp(): + img_path = request.form.get('img_path') + img = cv2.imread(img_path) + dewarped_img = dewarp.dewarp_image(img) + dirname, img_name, ext = image_util.parse_path(img_path) + save_path = os.path.join(dirname, img_name + '_dewarped.' + ext) + cv2.imwrite(save_path, dewarped_img) + return save_path + + +if __name__ == '__main__': + app.run('0.0.0.0', 5001) diff --git a/photo_review/auto_photo_review.py b/photo_review/auto_photo_review.py index 6664117..8ecec4f 100644 --- a/photo_review/auto_photo_review.py +++ b/photo_review/auto_photo_review.py @@ -182,15 +182,11 @@ def information_extraction(ie, phrecs, identity): result = merge_result(result, ie_result['result']) else: - target_images = [] - target_images += model_util.request_book_areas(img_path) # 识别文档区域并裁剪 - if not target_images: - target_images.append(image) # 识别失败 + target_images = model_util.request_book_areas(img_path) # 识别文档区域并裁剪 angle_count = defaultdict(int, {'0': 0}) # 分割后图片的最优角度统计 for target_image in target_images: - # dewarped_image = dewarp.dewarp_image(target_image) # 去扭曲 - dewarped_image = target_image - angles = image_util.parse_rotation_angles(dewarped_image) + dewarped_image = model_util.request_dewarped_image(target_image) # 去扭曲 + angles = model_util.request_image_orientation(dewarped_image) split_results = image_util.split(dewarped_image) for split_result in split_results: diff --git a/util/image_util.py b/util/image_util.py index 5a9d4bb..8c417dc 100644 --- a/util/image_util.py +++ b/util/image_util.py @@ -283,3 +283,9 @@ def save_to_local(img_url, save_path=None): file.write(response.content) return save_path + + +def parse_path(img_path): + dirname = os.path.dirname(img_path) + img_name, ext = os.path.basename(img_path).rsplit('.', 1) + return dirname, img_name, ext diff --git a/util/model_util.py b/util/model_util.py index 547f525..b3291db 100644 --- a/util/model_util.py +++ b/util/model_util.py @@ -1,6 +1,5 @@ import logging -import cv2 import requests from tenacity import retry, stop_after_attempt, wait_random @@ -8,13 +7,46 @@ from tenacity import retry, stop_after_attempt, wait_random @retry(stop=stop_after_attempt(3), wait=wait_random(1, 3), reraise=True, after=lambda x: logging.warning('获取文档区域失败！')) def request_book_areas(img_path): + """ + 请求文档区域识别接口 + :param img_path: 待识别图片路径 + :return: 文档图片路径列表 + """ url = 'http://det_api:5000/det/books' response = requests.post(url, {'img_path': img_path}) if response.status_code == 200: - response_data = response.json() - books = [] - for books_path in response_data: - books.append(cv2.imread(books_path)) - return books + return response.json() else: - return [] + return [img_path] + + +@retry(stop=stop_after_attempt(3), wait=wait_random(1, 3), reraise=True, + after=lambda x: logging.warning('矫正扭曲图片失败！')) +def request_dewarped_image(img_path): + """ + 请求矫正图片接口 + :param img_path: 待矫正图片路径 + :return: 矫正后的图片路径 + """ + url = 'http://det_api:5001/dewarp' + response = requests.post(url, {'img_path': img_path}) + if response.status_code == 200: + return response.json() + else: + return img_path + + +@retry(stop=stop_after_attempt(3), wait=wait_random(1, 3), reraise=True, + after=lambda x: logging.warning('获取图片方向失败！')) +def request_image_orientation(img_path): + """ + 请求图片方向分类接口 + :param img_path: 待分类图片路径 + :return: 最有可能的两个图片方向 + """ + url = 'http://det_api:5002/clas/orientation' + response = requests.post(url, {'img_path': img_path}) + if response.status_code == 200: + return response.json() + else: + return ['0', '90']