更换文档检测模型

2024-08-27 14:42:45 +08:00
parent aea6f19951
commit 1514e09c40
2072 changed files with 254336 additions and 4967 deletions
--- a/paddle_detection/configs/rotate/tools/onnx_infer.py
+++ b/paddle_detection/configs/rotate/tools/onnx_infer.py
@@ -0,0 +1,302 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import sys
+import six
+import glob
+import copy
+import yaml
+import argparse
+import cv2
+import numpy as np
+from shapely.geometry import Polygon
+from onnxruntime import InferenceSession
+
+
+# preprocess ops
+def decode_image(img_path):
+    with open(img_path, 'rb') as f:
+        im_read = f.read()
+    data = np.frombuffer(im_read, dtype='uint8')
+    im = cv2.imdecode(data, 1)  # BGR mode, but need RGB mode
+    im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
+    img_info = {
+        "im_shape": np.array(
+            im.shape[:2], dtype=np.float32),
+        "scale_factor": np.array(
+            [1., 1.], dtype=np.float32)
+    }
+    return im, img_info
+
+
+class Resize(object):
+    def __init__(self, target_size, keep_ratio=True, interp=cv2.INTER_LINEAR):
+        if isinstance(target_size, int):
+            target_size = [target_size, target_size]
+        self.target_size = target_size
+        self.keep_ratio = keep_ratio
+        self.interp = interp
+
+    def __call__(self, im, im_info):
+        assert len(self.target_size) == 2
+        assert self.target_size[0] > 0 and self.target_size[1] > 0
+        im_channel = im.shape[2]
+        im_scale_y, im_scale_x = self.generate_scale(im)
+        im = cv2.resize(
+            im,
+            None,
+            None,
+            fx=im_scale_x,
+            fy=im_scale_y,
+            interpolation=self.interp)
+        im_info['im_shape'] = np.array(im.shape[:2]).astype('float32')
+        im_info['scale_factor'] = np.array(
+            [im_scale_y, im_scale_x]).astype('float32')
+        return im, im_info
+
+    def generate_scale(self, im):
+        origin_shape = im.shape[:2]
+        im_c = im.shape[2]
+        if self.keep_ratio:
+            im_size_min = np.min(origin_shape)
+            im_size_max = np.max(origin_shape)
+            target_size_min = np.min(self.target_size)
+            target_size_max = np.max(self.target_size)
+            im_scale = float(target_size_min) / float(im_size_min)
+            if np.round(im_scale * im_size_max) > target_size_max:
+                im_scale = float(target_size_max) / float(im_size_max)
+            im_scale_x = im_scale
+            im_scale_y = im_scale
+        else:
+            resize_h, resize_w = self.target_size
+            im_scale_y = resize_h / float(origin_shape[0])
+            im_scale_x = resize_w / float(origin_shape[1])
+        return im_scale_y, im_scale_x
+
+
+class Permute(object):
+    def __init__(self, ):
+        super(Permute, self).__init__()
+
+    def __call__(self, im, im_info):
+        im = im.transpose((2, 0, 1))
+        return im, im_info
+
+
+class NormalizeImage(object):
+    def __init__(self, mean, std, is_scale=True, norm_type='mean_std'):
+        self.mean = mean
+        self.std = std
+        self.is_scale = is_scale
+        self.norm_type = norm_type
+
+    def __call__(self, im, im_info):
+        im = im.astype(np.float32, copy=False)
+        if self.is_scale:
+            scale = 1.0 / 255.0
+            im *= scale
+
+        if self.norm_type == 'mean_std':
+            mean = np.array(self.mean)[np.newaxis, np.newaxis, :]
+            std = np.array(self.std)[np.newaxis, np.newaxis, :]
+            im -= mean
+            im /= std
+        return im, im_info
+
+
+class PadStride(object):
+    def __init__(self, stride=0):
+        self.coarsest_stride = stride
+
+    def __call__(self, im, im_info):
+        coarsest_stride = self.coarsest_stride
+        if coarsest_stride <= 0:
+            return im, im_info
+        im_c, im_h, im_w = im.shape
+        pad_h = int(np.ceil(float(im_h) / coarsest_stride) * coarsest_stride)
+        pad_w = int(np.ceil(float(im_w) / coarsest_stride) * coarsest_stride)
+        padding_im = np.zeros((im_c, pad_h, pad_w), dtype=np.float32)
+        padding_im[:, :im_h, :im_w] = im
+        return padding_im, im_info
+
+
+class Compose:
+    def __init__(self, transforms):
+        self.transforms = []
+        for op_info in transforms:
+            new_op_info = op_info.copy()
+            op_type = new_op_info.pop('type')
+            self.transforms.append(eval(op_type)(**new_op_info))
+
+    def __call__(self, img_path):
+        img, im_info = decode_image(img_path)
+        for t in self.transforms:
+            img, im_info = t(img, im_info)
+        inputs = copy.deepcopy(im_info)
+        inputs['image'] = img
+        return inputs
+
+
+# postprocess
+def rbox_iou(g, p):
+    g = np.array(g)
+    p = np.array(p)
+    g = Polygon(g[:8].reshape((4, 2)))
+    p = Polygon(p[:8].reshape((4, 2)))
+    g = g.buffer(0)
+    p = p.buffer(0)
+    if not g.is_valid or not p.is_valid:
+        return 0
+    inter = Polygon(g).intersection(Polygon(p)).area
+    union = g.area + p.area - inter
+    if union == 0:
+        return 0
+    else:
+        return inter / union
+
+
+def multiclass_nms_rotated(pred_bboxes,
+                           pred_scores,
+                           iou_threshlod=0.1,
+                           score_threshold=0.1):
+    """
+    Args:
+        pred_bboxes (numpy.ndarray): [B, N, 8]
+        pred_scores (numpy.ndarray): [B, C, N]
+    
+    Return:
+        bboxes (numpy.ndarray): [N, 10]
+        bbox_num (numpy.ndarray): [B]
+    """
+    bbox_num = []
+    bboxes = []
+    for bbox_per_img, score_per_img in zip(pred_bboxes, pred_scores):
+        num_per_img = 0
+        for cls_id, score_per_cls in enumerate(score_per_img):
+            keep_mask = score_per_cls > score_threshold
+            bbox = bbox_per_img[keep_mask]
+            score = score_per_cls[keep_mask]
+
+            idx = score.argsort()[::-1]
+            bbox = bbox[idx]
+            score = score[idx]
+            keep_idx = []
+            for i, b in enumerate(bbox):
+                supressed = False
+                for gi in keep_idx:
+                    g = bbox[gi]
+                    if rbox_iou(b, g) > iou_threshlod:
+                        supressed = True
+                        break
+
+                if supressed:
+                    continue
+
+                keep_idx.append(i)
+
+            keep_box = bbox[keep_idx]
+            keep_score = score[keep_idx]
+            keep_cls_ids = np.ones(len(keep_idx)) * cls_id
+            bboxes.append(
+                np.concatenate(
+                    [keep_cls_ids[:, None], keep_score[:, None], keep_box],
+                    axis=-1))
+            num_per_img += len(keep_idx)
+
+        bbox_num.append(num_per_img)
+
+    return np.concatenate(bboxes, axis=0), np.array(bbox_num)
+
+
+def get_test_images(infer_dir, infer_img):
+    """
+    Get image path list in TEST mode
+    """
+    assert infer_img is not None or infer_dir is not None, \
+        "--image_file or --image_dir should be set"
+    assert infer_img is None or os.path.isfile(infer_img), \
+            "{} is not a file".format(infer_img)
+    assert infer_dir is None or os.path.isdir(infer_dir), \
+            "{} is not a directory".format(infer_dir)
+
+    # infer_img has a higher priority
+    if infer_img and os.path.isfile(infer_img):
+        return [infer_img]
+
+    images = set()
+    infer_dir = os.path.abspath(infer_dir)
+    assert os.path.isdir(infer_dir), \
+        "infer_dir {} is not a directory".format(infer_dir)
+    exts = ['jpg', 'jpeg', 'png', 'bmp']
+    exts += [ext.upper() for ext in exts]
+    for ext in exts:
+        images.update(glob.glob('{}/*.{}'.format(infer_dir, ext)))
+    images = list(images)
+
+    assert len(images) > 0, "no image found in {}".format(infer_dir)
+    print("Found {} inference images in total.".format(len(images)))
+
+    return images
+
+
+def predict_image(infer_config, predictor, img_list):
+    # load preprocess transforms
+    transforms = Compose(infer_config['Preprocess'])
+    # predict image
+    for img_path in img_list:
+        inputs = transforms(img_path)
+        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, bbox_num = multiclass_nms_rotated(
+            np.array(outputs[0]), np.array(outputs[1]))
+        print("ONNXRuntime predict: ")
+        for bbox in bboxes:
+            if bbox[0] > -1 and bbox[1] > infer_config['draw_threshold']:
+                print(f"{int(bbox[0])} {bbox[1]} "
+                      f"{bbox[2]} {bbox[3]} {bbox[4]} {bbox[5]}"
+                      f"{bbox[6]} {bbox[7]} {bbox[8]} {bbox[9]}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--infer_cfg", type=str, help="infer_cfg.yml")
+    parser.add_argument(
+        '--onnx_file',
+        type=str,
+        default="model.onnx",
+        help="onnx model file path")
+    parser.add_argument("--image_dir", type=str)
+    parser.add_argument("--image_file", type=str)
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    FLAGS = parse_args()
+    # load image list
+    img_list = get_test_images(FLAGS.image_dir, FLAGS.image_file)
+    # load predictor
+    predictor = InferenceSession(FLAGS.onnx_file)
+    # load infer config
+    with open(FLAGS.infer_cfg) as f:
+        infer_config = yaml.safe_load(f)
+
+    predict_image(infer_config, predictor, img_list)