更换文档检测模型

paddle_detection/deploy/serving/cpp/serving_client.py

@@ -0,0 +1,125 @@
+# 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.
+
+import os
+import glob
+import base64
+import argparse
+from paddle_serving_client import Client
+from paddle_serving_client.proto import general_model_config_pb2 as m_config
+import google.protobuf.text_format
+
+parser = argparse.ArgumentParser(description="args for paddleserving")
+parser.add_argument(
+    "--serving_client", type=str, help="the directory of serving_client")
+parser.add_argument("--image_dir", type=str)
+parser.add_argument("--image_file", type=str)
+parser.add_argument("--http_port", type=int, default=9997)
+parser.add_argument(
+    "--threshold", type=float, default=0.5, help="Threshold of score.")
+args = parser.parse_args()
+
+
+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 postprocess(fetch_dict, fetch_vars, draw_threshold=0.5):
+    result = []
+    if "conv2d_441.tmp_1" in fetch_dict:
+        heatmap = fetch_dict["conv2d_441.tmp_1"]
+        print(heatmap)
+        result.append(heatmap)
+    else:
+        bboxes = fetch_dict[fetch_vars[0]]
+        for bbox in bboxes:
+            if bbox[0] > -1 and bbox[1] > draw_threshold:
+                print(f"{int(bbox[0])} {bbox[1]} "
+                      f"{bbox[2]} {bbox[3]} {bbox[4]} {bbox[5]}")
+                result.append(f"{int(bbox[0])} {bbox[1]} "
+                              f"{bbox[2]} {bbox[3]} {bbox[4]} {bbox[5]}")
+    return result
+
+
+def get_model_vars(client_config_dir):
+    # read original serving_client_conf.prototxt
+    client_config_file = os.path.join(client_config_dir,
+                                      "serving_client_conf.prototxt")
+    with open(client_config_file, 'r') as f:
+        model_var = google.protobuf.text_format.Merge(
+            str(f.read()), m_config.GeneralModelConfig())
+    # modify feed_var to run core/general-server/op/
+    [model_var.feed_var.pop() for _ in range(len(model_var.feed_var))]
+    feed_var = m_config.FeedVar()
+    feed_var.name = "input"
+    feed_var.alias_name = "input"
+    feed_var.is_lod_tensor = False
+    feed_var.feed_type = 20
+    feed_var.shape.extend([1])
+    model_var.feed_var.extend([feed_var])
+    with open(
+            os.path.join(client_config_dir, "serving_client_conf_cpp.prototxt"),
+            "w") as f:
+        f.write(str(model_var))
+    # get feed_vars/fetch_vars
+    feed_vars = [var.name for var in model_var.feed_var]
+    fetch_vars = [var.name for var in model_var.fetch_var]
+    return feed_vars, fetch_vars
+
+
+if __name__ == '__main__':
+    url = f"127.0.0.1:{args.http_port}"
+    logid = 10000
+    img_list = get_test_images(args.image_dir, args.image_file)
+    feed_vars, fetch_vars = get_model_vars(args.serving_client)
+
+    client = Client()
+    client.load_client_config(
+        os.path.join(args.serving_client, "serving_client_conf_cpp.prototxt"))
+    client.connect([url])
+
+    for img_file in img_list:
+        with open(img_file, 'rb') as file:
+            image_data = file.read()
+        image = base64.b64encode(image_data).decode('utf8')
+        fetch_dict = client.predict(
+            feed={feed_vars[0]: image}, fetch=fetch_vars)
+        result = postprocess(fetch_dict, fetch_vars, args.threshold)