61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
import os.path
|
|
from collections import defaultdict
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
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):
|
|
# 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 = 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 capture(image, rectangle):
|
|
x1, y1, x2, y2 = rectangle
|
|
height, width = image.shape[:2]
|
|
if x1 < 0:
|
|
x1 = 0
|
|
if y1 < 0:
|
|
y1 = 0
|
|
if x2 > width:
|
|
x2 = width
|
|
if y2 > height:
|
|
y2 = height
|
|
return image[int(y1):int(y2), int(x1):int(x2)]
|
|
|
|
|
|
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(capture(image, book_area['box']))
|
|
return result
|