项目架构调整,模型全部采用接口调用
This commit is contained in:
@@ -1,4 +1,8 @@
|
||||
import os
|
||||
|
||||
from onnxruntime import InferenceSession
|
||||
|
||||
PADDLE_DET = InferenceSession("model/object_det_model/ppyoloe_plus_crn_l_80e_coco_w_nms.onnx",
|
||||
providers=["CPUExecutionProvider"], provider_options=[{"device_id": 0}])
|
||||
MODEL_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))),
|
||||
'model', 'object_det_model')
|
||||
PADDLE_DET = InferenceSession(os.path.join(MODEL_DIR, 'ppyoloe_plus_crn_l_80e_coco_w_nms.onnx'),
|
||||
providers=['CPUExecutionProvider'], provider_options=[{'device_id': 0}])
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
import tempfile
|
||||
import os.path
|
||||
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
|
||||
from util import image_util
|
||||
from . import PADDLE_DET, MODEL_DIR
|
||||
from .deploy.third_engine.onnx.infer import PredictConfig
|
||||
from .deploy.third_engine.onnx.preprocess import Compose
|
||||
|
||||
|
||||
def predict_image(infer_config, predictor, img_path):
|
||||
@@ -15,7 +15,7 @@ def predict_image(infer_config, predictor, img_path):
|
||||
transforms = Compose(infer_config.preprocess_infos)
|
||||
# predict image
|
||||
inputs = transforms(img_path)
|
||||
inputs["image"] = np.array(inputs["image"]).astype('float32')
|
||||
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}
|
||||
|
||||
@@ -25,25 +25,23 @@ def predict_image(infer_config, predictor, img_path):
|
||||
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:]})
|
||||
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"
|
||||
infer_cfg = os.path.join(MODEL_DIR, '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)
|
||||
def get_book_areas(img_path):
|
||||
detect_result = detect_image(img_path)
|
||||
book_areas = detect_result[73]
|
||||
result = []
|
||||
image = cv2.imread(img_path)
|
||||
for book_area in book_areas:
|
||||
result.append(image_util.capture(image, book_area["box"]))
|
||||
result.append(image_util.capture(image, book_area['box']))
|
||||
return result
|
||||
|
||||
Reference in New Issue
Block a user