From 916cec34c6a5be1bd2822094f275e17814b098dd Mon Sep 17 00:00:00 2001
From: littletomatodonkey <dazhiningsibuqu@163.com>
Date: Fri, 9 Apr 2021 11:50:01 +0000
Subject: [PATCH] unify predict mp and sp

---
 tools/infer/predict_system.py    | 17 ++++++++++-
 tools/infer/predict_system_mp.py | 48 --------------------------------
 tools/infer/utility.py           |  1 +
 3 files changed, 17 insertions(+), 49 deletions(-)
 delete mode 100755 tools/infer/predict_system_mp.py

diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py
index 1014475d..b77c9443 100755
--- a/tools/infer/predict_system.py
+++ b/tools/infer/predict_system.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 import os
 import sys
+import subprocess
 
 __dir__ = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(__dir__)
@@ -185,4 +186,18 @@ def main(args):
 
 
 if __name__ == "__main__":
-    main(utility.parse_args())
+    args = utility.parse_args()
+    if args.use_mp:
+        p_list = []
+        total_process_num = args.total_process_num
+        for process_id in range(total_process_num):
+            cmd = [sys.executable, "-u"] + sys.argv + [
+                "--process_id={} ".format(process_id),
+                "--use_mp={} ".format(args.use_mp)
+            ]
+            p = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stdout)
+            p_list.append(p)
+        for p in p_list:
+            p.wait()
+    else:
+        main(args)
diff --git a/tools/infer/predict_system_mp.py b/tools/infer/predict_system_mp.py
deleted file mode 100755
index a4c3bf39..00000000
--- a/tools/infer/predict_system_mp.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# Copyright (c) 2020 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 sys
-import paddle
-import subprocess
-
-__dir__ = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(__dir__)
-sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
-
-import tools.infer.utility as utility
-
-if __name__ == "__main__":
-    args = utility.parse_args()
-
-    p_list = []
-
-    inference_dir = "inference_results"
-    if not os.path.exists(inference_dir):
-        os.makedirs(inference_dir)
-
-    total_process_num = args.total_process_num
-    for process_id in range(total_process_num):
-        cmd = [sys.executable, "-u", "tools/infer/predict_system.py"
-               ] + sys.argv[1:] + ["--process_id={}".format(process_id)]
-
-        with open("{}/results.{}".format(inference_dir, process_id),
-                  "w") as fin:
-            # p = subprocess.Popen(cmd, stdout=fin, stderr=fin)
-            # if you want to print results in the screen, you can use the following command
-            p = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stdout)
-            p_list.append(p)
-    for p in p_list:
-        p.wait()
-    print("all inference results are save in directory: {}".format(
-        inference_dir))
diff --git a/tools/infer/utility.py b/tools/infer/utility.py
index fc416fd8..b273eaf3 100755
--- a/tools/infer/utility.py
+++ b/tools/infer/utility.py
@@ -98,6 +98,7 @@ def parse_args():
     parser.add_argument("--enable_mkldnn", type=str2bool, default=False)
     parser.add_argument("--use_pdserving", type=str2bool, default=False)
 
+    parser.add_argument("--use_mp", type=str2bool, default=False)
     parser.add_argument("--total_process_num", type=int, default=1)
     parser.add_argument("--process_id", type=int, default=0)