# Copyright (c) 2021 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 pathlib __dir__ = pathlib.Path(os.path.abspath(__file__)) sys.path.append(str(__dir__)) sys.path.append(str(__dir__.parent.parent)) import cv2 import paddle from paddle import inference import numpy as np from PIL import Image from paddle.vision import transforms from tools.predict import resize_image from post_processing import get_post_processing from utils.util import draw_bbox, save_result class InferenceEngine(object): """InferenceEngine Inference engina class which contains preprocess, run, postprocess """ def __init__(self, args): """ Args: args: Parameters generated using argparser. Returns: None """ super().__init__() self.args = args # init inference engine ( self.predictor, self.config, self.input_tensor, self.output_tensor, ) = self.load_predictor( os.path.join(args.model_dir, "inference.pdmodel"), os.path.join(args.model_dir, "inference.pdiparams"), ) # build transforms self.transforms = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] ), ] ) # wamrup if self.args.warmup > 0: for idx in range(args.warmup): print(idx) x = np.random.rand( 1, 3, self.args.crop_size, self.args.crop_size ).astype("float32") self.input_tensor.copy_from_cpu(x) self.predictor.run() self.output_tensor.copy_to_cpu() self.post_process = get_post_processing( { "type": "SegDetectorRepresenter", "args": { "thresh": 0.3, "box_thresh": 0.7, "max_candidates": 1000, "unclip_ratio": 1.5, }, } ) def load_predictor(self, model_file_path, params_file_path): """load_predictor initialize the inference engine Args: model_file_path: inference model path (*.pdmodel) model_file_path: inference parmaeter path (*.pdiparams) Return: predictor: Predictor created using Paddle Inference. config: Configuration of the predictor. input_tensor: Input tensor of the predictor. output_tensor: Output tensor of the predictor. """ args = self.args config = inference.Config(model_file_path, params_file_path) if args.use_gpu: config.enable_use_gpu(1000, 0) if args.use_tensorrt: config.enable_tensorrt_engine( workspace_size=1 << 30, precision_mode=precision, max_batch_size=args.max_batch_size, min_subgraph_size=args.min_subgraph_size, # skip the minmum trt subgraph use_calib_mode=False, ) # collect shape trt_shape_f = os.path.join(model_dir, "_trt_dynamic_shape.txt") if not os.path.exists(trt_shape_f): config.collect_shape_range_info(trt_shape_f) logger.info(f"collect dynamic shape info into : {trt_shape_f}") try: config.enable_tuned_tensorrt_dynamic_shape(trt_shape_f, True) except Exception as E: logger.info(E) logger.info("Please keep your paddlepaddle-gpu >= 2.3.0!") else: config.disable_gpu() # The thread num should not be greater than the number of cores in the CPU. if args.enable_mkldnn: # cache 10 different shapes for mkldnn to avoid memory leak config.set_mkldnn_cache_capacity(10) config.enable_mkldnn() if args.precision == "fp16": config.enable_mkldnn_bfloat16() if hasattr(args, "cpu_threads"): config.set_cpu_math_library_num_threads(args.cpu_threads) else: # default cpu threads as 10 config.set_cpu_math_library_num_threads(10) # enable memory optim config.enable_memory_optim() config.disable_glog_info() config.switch_use_feed_fetch_ops(False) config.switch_ir_optim(True) # create predictor predictor = inference.create_predictor(config) # get input and json tensor property input_names = predictor.get_input_names() input_tensor = predictor.get_input_handle(input_names[0]) output_names = predictor.get_output_names() output_tensor = predictor.get_output_handle(output_names[0]) return predictor, config, input_tensor, output_tensor def preprocess(self, img_path, short_size): """preprocess Preprocess to the input. Args: img_path: Image path. Returns: Input data after preprocess. """ img = cv2.imread(img_path, 1) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) h, w = img.shape[:2] img = resize_image(img, short_size) img = self.transforms(img) img = np.expand_dims(img, axis=0) shape_info = {"shape": [(h, w)]} return img, shape_info def postprocess(self, x, shape_info, is_output_polygon): """postprocess Postprocess to the inference engine json. Args: x: Inference engine json. Returns: Output data after argmax. """ box_list, score_list = self.post_process( shape_info, x, is_output_polygon=is_output_polygon ) box_list, score_list = box_list[0], score_list[0] if len(box_list) > 0: if is_output_polygon: idx = [x.sum() > 0 for x in box_list] box_list = [box_list[i] for i, v in enumerate(idx) if v] score_list = [score_list[i] for i, v in enumerate(idx) if v] else: idx = ( box_list.reshape(box_list.shape[0], -1).sum(axis=1) > 0 ) # 去掉全为0的框 box_list, score_list = box_list[idx], score_list[idx] else: box_list, score_list = [], [] return box_list, score_list def run(self, x): """run Inference process using inference engine. Args: x: Input data after preprocess. Returns: Inference engine json """ self.input_tensor.copy_from_cpu(x) self.predictor.run() output = self.output_tensor.copy_to_cpu() return output def get_args(add_help=True): """ parse args """ import argparse def str2bool(v): return v.lower() in ("true", "t", "1") parser = argparse.ArgumentParser( description="PaddlePaddle Classification Training", add_help=add_help ) parser.add_argument("--model_dir", default=None, help="inference model dir") parser.add_argument("--batch_size", type=int, default=1) parser.add_argument("--short_size", default=1024, type=int, help="short size") parser.add_argument("--img_path", default="./images/demo.jpg") parser.add_argument("--benchmark", default=False, type=str2bool, help="benchmark") parser.add_argument("--warmup", default=0, type=int, help="warmup iter") parser.add_argument("--polygon", action="store_true", help="json polygon or box") parser.add_argument("--use_gpu", type=str2bool, default=True) parser.add_argument("--use_tensorrt", type=str2bool, default=False) parser.add_argument("--precision", type=str, default="fp32") parser.add_argument("--gpu_mem", type=int, default=500) parser.add_argument("--gpu_id", type=int, default=0) parser.add_argument("--enable_mkldnn", type=str2bool, default=False) parser.add_argument("--cpu_threads", type=int, default=10) args = parser.parse_args() return args def main(args): """ Main inference function. Args: args: Parameters generated using argparser. Returns: class_id: Class index of the input. prob: : Probability of the input. """ inference_engine = InferenceEngine(args) # init benchmark if args.benchmark: import auto_log autolog = auto_log.AutoLogger( model_name="db", batch_size=args.batch_size, inference_config=inference_engine.config, gpu_ids="auto" if args.use_gpu else None, ) # enable benchmark if args.benchmark: autolog.times.start() # preprocess img, shape_info = inference_engine.preprocess(args.img_path, args.short_size) if args.benchmark: autolog.times.stamp() output = inference_engine.run(img) if args.benchmark: autolog.times.stamp() # postprocess box_list, score_list = inference_engine.postprocess( output, shape_info, args.polygon ) if args.benchmark: autolog.times.stamp() autolog.times.end(stamp=True) autolog.report() img = draw_bbox(cv2.imread(args.img_path)[:, :, ::-1], box_list) # 保存结果到路径 os.makedirs("json", exist_ok=True) img_path = pathlib.Path(args.img_path) output_path = os.path.join("json", img_path.stem + "_infer_result.jpg") cv2.imwrite(output_path, img[:, :, ::-1]) save_result( output_path.replace("_infer_result.jpg", ".txt"), box_list, score_list, args.polygon, ) if __name__ == "__main__": args = get_args() main(args)