first commit

benchmark/PaddleOCR_DBNet/tools/__init__.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2019/12/8 13:14
+# @Author  : zhoujun

benchmark/PaddleOCR_DBNet/tools/eval.py

@@ -0,0 +1,93 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2018/6/11 15:54
+# @Author  : zhoujun
+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 argparse
+import time
+import paddle
+from tqdm.auto import tqdm
+
+
+class EVAL:
+    def __init__(self, model_path, gpu_id=0):
+        from models import build_model
+        from data_loader import get_dataloader
+        from post_processing import get_post_processing
+        from utils import get_metric
+
+        self.gpu_id = gpu_id
+        if (
+            self.gpu_id is not None
+            and isinstance(self.gpu_id, int)
+            and paddle.device.is_compiled_with_cuda()
+        ):
+            paddle.device.set_device("gpu:{}".format(self.gpu_id))
+        else:
+            paddle.device.set_device("cpu")
+        checkpoint = paddle.load(model_path)
+        config = checkpoint["config"]
+        config["arch"]["backbone"]["pretrained"] = False
+
+        self.validate_loader = get_dataloader(
+            config["dataset"]["validate"], config["distributed"]
+        )
+
+        self.model = build_model(config["arch"])
+        self.model.set_state_dict(checkpoint["state_dict"])
+
+        self.post_process = get_post_processing(config["post_processing"])
+        self.metric_cls = get_metric(config["metric"])
+
+    def eval(self):
+        self.model.eval()
+        raw_metrics = []
+        total_frame = 0.0
+        total_time = 0.0
+        for i, batch in tqdm(
+            enumerate(self.validate_loader),
+            total=len(self.validate_loader),
+            desc="test model",
+        ):
+            with paddle.no_grad():
+                start = time.time()
+                preds = self.model(batch["img"])
+                boxes, scores = self.post_process(
+                    batch, preds, is_output_polygon=self.metric_cls.is_output_polygon
+                )
+                total_frame += batch["img"].shape[0]
+                total_time += time.time() - start
+                raw_metric = self.metric_cls.validate_measure(batch, (boxes, scores))
+                raw_metrics.append(raw_metric)
+        metrics = self.metric_cls.gather_measure(raw_metrics)
+        print("FPS:{}".format(total_frame / total_time))
+        return {
+            "recall": metrics["recall"].avg,
+            "precision": metrics["precision"].avg,
+            "fmeasure": metrics["fmeasure"].avg,
+        }
+
+
+def init_args():
+    parser = argparse.ArgumentParser(description="DBNet.paddle")
+    parser.add_argument(
+        "--model_path",
+        required=False,
+        default="output/DBNet_resnet18_FPN_DBHead/checkpoint/1.pth",
+        type=str,
+    )
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    args = init_args()
+    eval = EVAL(args.model_path)
+    result = eval.eval()
+    print(result)

benchmark/PaddleOCR_DBNet/tools/export_model.py

@@ -0,0 +1,57 @@
+import os
+import sys
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(__dir__)
+sys.path.insert(0, os.path.abspath(os.path.join(__dir__, "..")))
+
+import argparse
+
+import paddle
+from paddle.jit import to_static
+
+from models import build_model
+from utils import Config, ArgsParser
+
+
+def init_args():
+    parser = ArgsParser()
+    args = parser.parse_args()
+    return args
+
+
+def load_checkpoint(model, checkpoint_path):
+    """
+    load checkpoints
+    :param checkpoint_path: Checkpoint path to be loaded
+    """
+    checkpoint = paddle.load(checkpoint_path)
+    model.set_state_dict(checkpoint["state_dict"])
+    print("load checkpoint from {}".format(checkpoint_path))
+
+
+def main(config):
+    model = build_model(config["arch"])
+    load_checkpoint(model, config["trainer"]["resume_checkpoint"])
+    model.eval()
+
+    save_path = config["trainer"]["output_dir"]
+    save_path = os.path.join(save_path, "inference")
+    infer_shape = [3, -1, -1]
+    model = to_static(
+        model,
+        input_spec=[
+            paddle.static.InputSpec(shape=[None] + infer_shape, dtype="float32")
+        ],
+    )
+
+    paddle.jit.save(model, save_path)
+    print("inference model is saved to {}".format(save_path))
+
+
+if __name__ == "__main__":
+    args = init_args()
+    assert os.path.exists(args.config_file)
+    config = Config(args.config_file)
+    config.merge_dict(args.opt)
+    main(config.cfg)

benchmark/PaddleOCR_DBNet/tools/infer.py

@@ -0,0 +1,315 @@
+# 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 engine 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 parameter 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 minimum 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 output 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 output.
+        Args:
+            x: Inference engine output.
+        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 output
+        """
+        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="output 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("output", exist_ok=True)
+    img_path = pathlib.Path(args.img_path)
+    output_path = os.path.join("output", 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)

benchmark/PaddleOCR_DBNet/tools/predict.py

@@ -0,0 +1,175 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2019/8/24 12:06
+# @Author  : zhoujun
+
+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 time
+import cv2
+import paddle
+
+from data_loader import get_transforms
+from models import build_model
+from post_processing import get_post_processing
+
+
+def resize_image(img, short_size):
+    height, width, _ = img.shape
+    if height < width:
+        new_height = short_size
+        new_width = new_height / height * width
+    else:
+        new_width = short_size
+        new_height = new_width / width * height
+    new_height = int(round(new_height / 32) * 32)
+    new_width = int(round(new_width / 32) * 32)
+    resized_img = cv2.resize(img, (new_width, new_height))
+    return resized_img
+
+
+class PaddleModel:
+    def __init__(self, model_path, post_p_thre=0.7, gpu_id=None):
+        """
+        初始化模型
+        :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件)
+        :param gpu_id: 在哪一块gpu上运行
+        """
+        self.gpu_id = gpu_id
+
+        if (
+            self.gpu_id is not None
+            and isinstance(self.gpu_id, int)
+            and paddle.device.is_compiled_with_cuda()
+        ):
+            paddle.device.set_device("gpu:{}".format(self.gpu_id))
+        else:
+            paddle.device.set_device("cpu")
+        checkpoint = paddle.load(model_path)
+
+        config = checkpoint["config"]
+        config["arch"]["backbone"]["pretrained"] = False
+        self.model = build_model(config["arch"])
+        self.post_process = get_post_processing(config["post_processing"])
+        self.post_process.box_thresh = post_p_thre
+        self.img_mode = config["dataset"]["train"]["dataset"]["args"]["img_mode"]
+        self.model.set_state_dict(checkpoint["state_dict"])
+        self.model.eval()
+
+        self.transform = []
+        for t in config["dataset"]["train"]["dataset"]["args"]["transforms"]:
+            if t["type"] in ["ToTensor", "Normalize"]:
+                self.transform.append(t)
+        self.transform = get_transforms(self.transform)
+
+    def predict(self, img_path: str, is_output_polygon=False, short_size: int = 1024):
+        """
+        对传入的图像进行预测，支持图像地址,opencv 读取图片，偏慢
+        :param img_path: 图像地址
+        :param is_numpy:
+        :return:
+        """
+        assert os.path.exists(img_path), "file is not exists"
+        img = cv2.imread(img_path, 1 if self.img_mode != "GRAY" else 0)
+        if self.img_mode == "RGB":
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        h, w = img.shape[:2]
+        img = resize_image(img, short_size)
+        # 将图片由(w,h)变为(1,img_channel,h,w)
+        tensor = self.transform(img)
+        tensor = tensor.unsqueeze_(0)
+
+        batch = {"shape": [(h, w)]}
+        with paddle.no_grad():
+            start = time.time()
+            preds = self.model(tensor)
+            box_list, score_list = self.post_process(
+                batch, preds, 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 = [], []
+            t = time.time() - start
+        return preds[0, 0, :, :].detach().cpu().numpy(), box_list, score_list, t
+
+
+def save_depoly(net, input, save_path):
+    input_spec = [paddle.static.InputSpec(shape=[None, 3, None, None], dtype="float32")]
+    net = paddle.jit.to_static(net, input_spec=input_spec)
+
+    # save static model for inference directly
+    paddle.jit.save(net, save_path)
+
+
+def init_args():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="DBNet.paddle")
+    parser.add_argument("--model_path", default=r"model_best.pth", type=str)
+    parser.add_argument(
+        "--input_folder", default="./test/input", type=str, help="img path for predict"
+    )
+    parser.add_argument(
+        "--output_folder", default="./test/output", type=str, help="img path for output"
+    )
+    parser.add_argument("--gpu", default=0, type=int, help="gpu for inference")
+    parser.add_argument(
+        "--thre", default=0.3, type=float, help="the thresh of post_processing"
+    )
+    parser.add_argument("--polygon", action="store_true", help="output polygon or box")
+    parser.add_argument("--show", action="store_true", help="show result")
+    parser.add_argument(
+        "--save_result", action="store_true", help="save box and score to txt file"
+    )
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    import pathlib
+    from tqdm import tqdm
+    import matplotlib.pyplot as plt
+    from utils.util import show_img, draw_bbox, save_result, get_image_file_list
+
+    args = init_args()
+    print(args)
+    # 初始化网络
+    model = PaddleModel(args.model_path, post_p_thre=args.thre, gpu_id=args.gpu)
+    img_folder = pathlib.Path(args.input_folder)
+    for img_path in tqdm(get_image_file_list(args.input_folder)):
+        preds, boxes_list, score_list, t = model.predict(
+            img_path, is_output_polygon=args.polygon
+        )
+        img = draw_bbox(cv2.imread(img_path)[:, :, ::-1], boxes_list)
+        if args.show:
+            show_img(preds)
+            show_img(img, title=os.path.basename(img_path))
+            plt.show()
+        # 保存结果到路径
+        os.makedirs(args.output_folder, exist_ok=True)
+        img_path = pathlib.Path(img_path)
+        output_path = os.path.join(args.output_folder, img_path.stem + "_result.jpg")
+        pred_path = os.path.join(args.output_folder, img_path.stem + "_pred.jpg")
+        cv2.imwrite(output_path, img[:, :, ::-1])
+        cv2.imwrite(pred_path, preds * 255)
+        save_result(
+            output_path.replace("_result.jpg", ".txt"),
+            boxes_list,
+            score_list,
+            args.polygon,
+        )

benchmark/PaddleOCR_DBNet/tools/train.py

@@ -0,0 +1,64 @@
+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 paddle
+import paddle.distributed as dist
+from utils import Config, ArgsParser
+
+
+def init_args():
+    parser = ArgsParser()
+    args = parser.parse_args()
+    return args
+
+
+def main(config, profiler_options):
+    from models import build_model, build_loss
+    from data_loader import get_dataloader
+    from trainer import Trainer
+    from post_processing import get_post_processing
+    from utils import get_metric
+
+    if paddle.device.cuda.device_count() > 1:
+        dist.init_parallel_env()
+        config["distributed"] = True
+    else:
+        config["distributed"] = False
+    train_loader = get_dataloader(config["dataset"]["train"], config["distributed"])
+    assert train_loader is not None
+    if "validate" in config["dataset"]:
+        validate_loader = get_dataloader(config["dataset"]["validate"], False)
+    else:
+        validate_loader = None
+    criterion = build_loss(config["loss"])
+    config["arch"]["backbone"]["in_channels"] = (
+        3 if config["dataset"]["train"]["dataset"]["args"]["img_mode"] != "GRAY" else 1
+    )
+    model = build_model(config["arch"])
+    # set @to_static for benchmark, skip this by default.
+    post_p = get_post_processing(config["post_processing"])
+    metric = get_metric(config["metric"])
+    trainer = Trainer(
+        config=config,
+        model=model,
+        criterion=criterion,
+        train_loader=train_loader,
+        post_process=post_p,
+        metric_cls=metric,
+        validate_loader=validate_loader,
+        profiler_options=profiler_options,
+    )
+    trainer.train()
+
+
+if __name__ == "__main__":
+    args = init_args()
+    assert os.path.exists(args.config_file)
+    config = Config(args.config_file)
+    config.merge_dict(args.opt)
+    main(config.cfg, args.profiler_options)