From a427ad01b91a915fa576931e2ed380e23483764f Mon Sep 17 00:00:00 2001
From: liuyebo <1515783401@qq.com>
Date: Wed, 28 Aug 2024 17:14:43 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=96=87=E6=A1=A3=E6=A3=80?=
 =?UTF-8?q?=E6=B5=8B=E4=B8=BA=E6=8E=A5=E5=8F=A3?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 det_api.py                        | 32 +++++++++++++++++++++++++++++++
 docker-compose.yml                | 20 +++++++++++++++++--
 paddle_detection/detector.py      | 24 +++++++++++++++++++++++
 photo_review/auto_photo_review.py |  4 +++-
 4 files changed, 77 insertions(+), 3 deletions(-)
 create mode 100644 det_api.py

diff --git a/det_api.py b/det_api.py
new file mode 100644
index 0000000..b0ee859
--- /dev/null
+++ b/det_api.py
@@ -0,0 +1,32 @@
+import base64
+
+import cv2
+import numpy as np
+from flask import Flask, request, jsonify
+
+from paddle_detection import detector
+
+app = Flask(__name__)
+
+
+@app.route("/det/detect_books", methods=['POST'])
+def detect_books():
+    try:
+        file = request.files['image']
+        image_data = file.read()
+        nparr = np.frombuffer(image_data, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        result = detector.get_book_areas(image)
+        encoded_images = []
+        for i in result:
+            _, encoded_image = cv2.imencode('.jpg', i)
+            byte_stream = encoded_image.tobytes()
+            img_str = base64.b64encode(byte_stream).decode('utf-8')
+            encoded_images.append(img_str)
+        return jsonify(encoded_images), 200
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+
+if __name__ == '__main__':
+    app.run()
diff --git a/docker-compose.yml b/docker-compose.yml
index 475b544..848cb09 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -1,13 +1,29 @@
 x-env:
   &template
-  image: fcb_photo_review:1.13.7
+  image: fcb_photo_review:1.13.8
   restart: always
 
 services:
-  photo_review_1:
+  det_api:
     <<: *template
     build:
       context: .
+    container_name: det_api
+    hostname: det_api
+    volumes:
+      - ./log:/app/log
+      - ./model:/app/model
+    command: [ "det_api.py" ]
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - device_ids: [ "0" ]
+              capabilities: [ "gpu" ]
+              driver: "nvidia"
+
+  photo_review_1:
+    <<: *template
     container_name: photo_review_1
     hostname: photo_review_1
     volumes:
diff --git a/paddle_detection/detector.py b/paddle_detection/detector.py
index c49d877..845a1d0 100644
--- a/paddle_detection/detector.py
+++ b/paddle_detection/detector.py
@@ -1,8 +1,12 @@
+import base64
+import logging
 import tempfile
 from collections import defaultdict
 
 import cv2
 import numpy as np
+import requests
+from tenacity import retry, stop_after_attempt, wait_random
 
 from paddle_detection import PADDLE_DET
 from paddle_detection.deploy.third_engine.onnx.infer import PredictConfig
@@ -47,3 +51,23 @@ def get_book_areas(image):
     for book_area in book_areas:
         result.append(image_util.capture(image, book_area["box"]))
     return result
+
+
+@retry(stop=stop_after_attempt(3), wait=wait_random(1, 3), reraise=True,
+       after=lambda x: logging.warning("获取文档区域失败！"))
+def request_book_areas(image):
+    url = "http://det_api:5000/det/detect_books"
+    _, encoded_image = cv2.imencode('.jpg', image)
+    byte_stream = encoded_image.tobytes()
+    files = {"image": ("image.jpg", byte_stream)}
+    response = requests.post(url, files=files)
+    if response.status_code == 200:
+        img_str_list = response.json()
+        result = []
+        for img_str in img_str_list:
+            img_data = base64.b64decode(img_str)
+            np_array = np.frombuffer(img_data, np.uint8)
+            result.append(cv2.imdecode(np_array, cv2.IMREAD_COLOR))
+        return result
+    else:
+        return []
diff --git a/photo_review/auto_photo_review.py b/photo_review/auto_photo_review.py
index 7596e57..6d2f89f 100644
--- a/photo_review/auto_photo_review.py
+++ b/photo_review/auto_photo_review.py
@@ -82,7 +82,9 @@ def information_extraction(ie, phrecs, identity):
 
         image = image_util.read(img_path)
         target_images = []
-        target_images += detector.get_book_areas(image)  # 识别文档区域并裁剪
+        det_time = time.time()
+        target_images += detector.request_book_areas(image)  # 识别文档区域并裁剪
+        logging.info(f"检测目标耗时{time.time() - det_time}秒")
         if not target_images:
             target_images.append(image)  # 识别失败
         angle_count = defaultdict(int, {"0": 0})  # 分割后图片的最优角度统计