移动paddle_detection

services/paddle_services/paddle_detection/detector.py

@@ -0,0 +1,49 @@
+import tempfile
+from collections import defaultdict
+
+import cv2
+import numpy as np
+
+from paddle_detection import PADDLE_DET
+from paddle_detection.deploy.third_engine.onnx.infer import PredictConfig
+from paddle_detection.deploy.third_engine.onnx.preprocess import Compose
+from util import image_util, common_util
+
+
+def predict_image(infer_config, predictor, img_path):
+    # load preprocess transforms
+    transforms = Compose(infer_config.preprocess_infos)
+    # predict image
+    inputs = transforms(img_path)
+    inputs["image"] = np.array(inputs["image"]).astype('float32')
+    inputs_name = [var.name for var in predictor.get_inputs()]
+    inputs = {k: inputs[k][None,] for k in inputs_name}
+
+    outputs = predictor.run(output_names=None, input_feed=inputs)
+
+    bboxes = np.array(outputs[0])
+    result = defaultdict(list)
+    for bbox in bboxes:
+        if bbox[0] > -1 and bbox[1] > infer_config.draw_threshold:
+            result[bbox[0]].append({"score": bbox[1], "box": bbox[2:]})
+    return result
+
+
+def detect_image(img_path):
+    infer_cfg = "model/object_det_model/infer_cfg.yml"
+    # load infer config
+    infer_config = PredictConfig(infer_cfg)
+
+    return predict_image(infer_config, PADDLE_DET, img_path)
+
+
+def get_book_areas(image):
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
+        cv2.imwrite(temp_file.name, image)
+    detect_result = detect_image(temp_file.name)
+    common_util.delete_temp_file(temp_file.name)
+    book_areas = detect_result[73]
+    result = []
+    for book_area in book_areas:
+        result.append(image_util.capture(image, book_area["box"]))
+    return result