[python] Add layout module

python/rapid_structure/README.md

@@ -0,0 +1 @@
+#### Rapid Structure

python/rapid_structure/demo.py

@@ -0,0 +1,45 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import copy
+from pathlib import Path
+
+import cv2
+import numpy as np
+from rapid_layout import RapidLayout
+
+
+def vis_layout(img, layout_res: list):
+    tmp_img = copy.deepcopy(img)
+    for v in layout_res:
+        bbox = np.round(v['bbox']).astype(np.int32)
+        label = v['label']
+
+        start_point = (bbox[0], bbox[1])
+        end_point = (bbox[2], bbox[3])
+
+        cv2.rectangle(tmp_img, start_point, end_point, (0, 255, 0), 2)
+        cv2.putText(tmp_img, label, start_point,
+                    cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
+
+    draw_img_save = Path("./inference_results/")
+    if not draw_img_save.exists():
+        draw_img_save.mkdir(parents=True, exist_ok=True)
+
+    image_save = str(draw_img_save / 'layout_result.jpg')
+    cv2.imwrite(image_save, tmp_img)
+    print(f'The infer result has saved in {image_save}')
+
+
+if __name__ == '__main__':
+    model_path = 'models/layout_cdla.onnx'
+    rapid_layout = RapidLayout(model_path)
+
+    img_path = 'test_images/ch.png'
+
+    img = cv2.imread(img_path)
+
+    layout_res, elapse = rapid_layout(img)
+
+    vis_layout(img, layout_res)
+    print(layout_res)

python/rapid_structure/rapid_layout/__init__.py

@@ -0,0 +1,14 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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 .rapid_layout import RapidLayout

python/rapid_structure/rapid_layout/config.yaml

@@ -0,0 +1,24 @@
+model_path: models/layout_table.onnx
+
+use_cuda: false
+CUDAExecutionProvider:
+    device_id: 0
+    arena_extend_strategy: kNextPowerOfTwo
+    cudnn_conv_algo_search: EXHAUSTIVE
+    do_copy_in_default_stream: true
+
+pre_process:
+    Resize:
+        size: [800, 608]
+    NormalizeImage:
+        std: [0.229, 0.224, 0.225]
+        mean: [0.485, 0.456, 0.406]
+        scale: 1./255.
+        order: hwc
+    ToCHWImage:
+    KeepKeys:
+        keep_keys: ['image']
+
+post_process:
+    score_threshold: 0.5
+    nms_threshold: 0.5

python/rapid_structure/rapid_layout/rapid_layout.py

@@ -0,0 +1,66 @@
+# 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 time
+from pathlib import Path
+import numpy as np
+
+from .utils import (OrtInferSession, PicoDetPostProcess, create_operators,
+                    read_yaml, transform)
+
+root_dir = Path(__file__).resolve().parent
+
+
+class RapidLayout(object):
+    def __init__(self, model_path: str):
+        config_path = str(root_dir / 'config.yaml')
+        config = read_yaml(config_path)
+        config['model_path'] = model_path
+
+        session_instance = OrtInferSession(config)
+        self.predictor = session_instance.session
+        self.input_name = session_instance.get_input_name()
+        labels = session_instance.get_metadata()['character'].splitlines()
+
+        self.preprocess_op = create_operators(config['pre_process'])
+        self.postprocess_op = PicoDetPostProcess(labels,
+                                                 **config['post_process'])
+
+    def __call__(self, img):
+        ori_im = img.copy()
+        data = {'image': img}
+        data = transform(data, self.preprocess_op)
+        img = data[0]
+
+        if img is None:
+            return None, 0
+
+        img = np.expand_dims(img, axis=0)
+        img = img.copy()
+
+        preds, elapse = 0, 1
+        starttime = time.time()
+
+        input_dict = {self.input_name: img}
+        preds_onnx = self.predictor.run(None, input_dict)
+
+        np_score_list, np_boxes_list = [], []
+        num_outs = int(len(preds_onnx) / 2)
+        for out_idx in range(num_outs):
+            np_score_list.append(preds_onnx[out_idx])
+            np_boxes_list.append(preds_onnx[out_idx + num_outs])
+        preds = dict(boxes=np_score_list, boxes_num=np_boxes_list)
+
+        post_preds = self.postprocess_op(ori_im, img, preds)
+        elapse = time.time() - starttime
+        return post_preds, elapse

python/rapid_structure/rapid_layout/utils.py

@@ -0,0 +1,418 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import warnings
+
+import yaml
+import cv2
+import numpy as np
+from onnxruntime import (get_available_providers, get_device,
+                         SessionOptions, InferenceSession,
+                         GraphOptimizationLevel)
+
+
+class OrtInferSession(object):
+    def __init__(self, config):
+        sess_opt = SessionOptions()
+        sess_opt.log_severity_level = 4
+        sess_opt.enable_cpu_mem_arena = False
+        sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        cuda_ep = 'CUDAExecutionProvider'
+        cpu_ep = 'CPUExecutionProvider'
+        cpu_provider_options = {"arena_extend_strategy": "kSameAsRequested",}
+
+        EP_list = []
+        if config['use_cuda'] and get_device() == 'GPU' \
+                and cuda_ep in get_available_providers():
+            EP_list = [(cuda_ep, config[cuda_ep])]
+        EP_list.append((cpu_ep, cpu_provider_options))
+
+        self.session = InferenceSession(config['model_path'],
+                                        sess_options=sess_opt,
+                                        providers=EP_list)
+
+        if config['use_cuda'] and cuda_ep not in self.session.get_providers():
+            warnings.warn(f'{cuda_ep} is not avaiable for current env, the inference part is automatically shifted to be executed under {cpu_ep}.\n'
+                          'Please ensure the installed onnxruntime-gpu version matches your cuda and cudnn version, '
+                          'you can check their relations from the offical web site: '
+                          'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html',
+                          RuntimeWarning)
+
+    def get_input_name(self, input_idx=0):
+        return self.session.get_inputs()[input_idx].name
+
+    def get_output_name(self, output_idx=0):
+        return self.session.get_outputs()[output_idx].name
+
+    def get_metadata(self):
+        meta_dict = self.session.get_modelmeta().custom_metadata_map
+        return meta_dict
+
+
+def transform(data, ops=None):
+    """ transform """
+    if ops is None:
+        ops = []
+
+    for op in ops:
+        data = op(data)
+        if data is None:
+            return None
+    return data
+
+
+def create_operators(op_param_dict):
+    """
+    create operators based on the config
+    """
+    ops = []
+    for op_name, param in op_param_dict.items():
+        if param is None:
+            param = {}
+        op = eval(op_name)(**param)
+        ops.append(op)
+    return ops
+
+
+class Resize(object):
+    def __init__(self, size=(640, 640), **kwargs):
+        self.size = size
+
+    def resize_image(self, img):
+        resize_h, resize_w = self.size
+        ori_h, ori_w = img.shape[:2]  # (h, w, c)
+        ratio_h = float(resize_h) / ori_h
+        ratio_w = float(resize_w) / ori_w
+        img = cv2.resize(img, (int(resize_w), int(resize_h)))
+        return img, [ratio_h, ratio_w]
+
+    def __call__(self, data):
+        img = data['image']
+        if 'polys' in data:
+            text_polys = data['polys']
+
+        img_resize, [ratio_h, ratio_w] = self.resize_image(img)
+        if 'polys' in data:
+            new_boxes = []
+            for box in text_polys:
+                new_box = []
+                for cord in box:
+                    new_box.append([cord[0] * ratio_w, cord[1] * ratio_h])
+                new_boxes.append(new_box)
+            data['polys'] = np.array(new_boxes, dtype=np.float32)
+        data['image'] = img_resize
+        return data
+
+
+class NormalizeImage(object):
+    """ normalize image such as substract mean, divide std
+    """
+
+    def __init__(self, scale=None, mean=None, std=None, order='chw', **kwargs):
+        if isinstance(scale, str):
+            scale = eval(scale)
+        self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
+        mean = mean if mean is not None else [0.485, 0.456, 0.406]
+        std = std if std is not None else [0.229, 0.224, 0.225]
+
+        shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
+        self.mean = np.array(mean).reshape(shape).astype('float32')
+        self.std = np.array(std).reshape(shape).astype('float32')
+
+    def __call__(self, data):
+        img = data['image']
+        from PIL import Image
+        if isinstance(img, Image.Image):
+            img = np.array(img)
+        assert isinstance(img,
+                          np.ndarray), "invalid input 'img' in NormalizeImage"
+        data['image'] = (
+            img.astype('float32') * self.scale - self.mean) / self.std
+        return data
+
+
+class ToCHWImage(object):
+    """ convert hwc image to chw image
+    """
+
+    def __init__(self, **kwargs):
+        pass
+
+    def __call__(self, data):
+        img = data['image']
+        from PIL import Image
+        if isinstance(img, Image.Image):
+            img = np.array(img)
+        data['image'] = img.transpose((2, 0, 1))
+        return data
+
+
+class KeepKeys(object):
+    def __init__(self, keep_keys, **kwargs):
+        self.keep_keys = keep_keys
+
+    def __call__(self, data):
+        data_list = []
+        for key in self.keep_keys:
+            data_list.append(data[key])
+        return data_list
+
+
+def read_yaml(yaml_path):
+    with open(yaml_path, 'rb') as f:
+        data = yaml.load(f, Loader=yaml.Loader)
+    return data
+
+
+class PicoDetPostProcess(object):
+    """
+    Args:
+        input_shape (int): network input image size
+        ori_shape (int): ori image shape of before padding
+        scale_factor (float): scale factor of ori image
+        enable_mkldnn (bool): whether to open MKLDNN
+    """
+    def __init__(self,
+                 labels,
+                 strides=[8, 16, 32, 64],
+                 score_threshold=0.4,
+                 nms_threshold=0.5,
+                 nms_top_k=1000,
+                 keep_top_k=100):
+        self.labels = labels
+        self.strides = strides
+        self.score_threshold = score_threshold
+        self.nms_threshold = nms_threshold
+        self.nms_top_k = nms_top_k
+        self.keep_top_k = keep_top_k
+
+    def __call__(self, ori_img, img, preds):
+        scores, raw_boxes = preds['boxes'], preds['boxes_num']
+        batch_size = raw_boxes[0].shape[0]
+        reg_max = int(raw_boxes[0].shape[-1] / 4 - 1)
+        out_boxes_num = []
+        out_boxes_list = []
+        results = []
+        ori_shape, input_shape, scale_factor = self.img_info(ori_img, img)
+
+        for batch_id in range(batch_size):
+            # generate centers
+            decode_boxes = []
+            select_scores = []
+            for stride, box_distribute, score in zip(self.strides, raw_boxes,
+                                                     scores):
+                box_distribute = box_distribute[batch_id]
+                score = score[batch_id]
+                # centers
+                fm_h = input_shape[0] / stride
+                fm_w = input_shape[1] / stride
+                h_range = np.arange(fm_h)
+                w_range = np.arange(fm_w)
+                ww, hh = np.meshgrid(w_range, h_range)
+                ct_row = (hh.flatten() + 0.5) * stride
+                ct_col = (ww.flatten() + 0.5) * stride
+                center = np.stack((ct_col, ct_row, ct_col, ct_row), axis=1)
+
+                # box distribution to distance
+                reg_range = np.arange(reg_max + 1)
+                box_distance = box_distribute.reshape((-1, reg_max + 1))
+                box_distance = self.softmax(box_distance, axis=1)
+                box_distance = box_distance * np.expand_dims(reg_range, axis=0)
+                box_distance = np.sum(box_distance, axis=1).reshape((-1, 4))
+                box_distance = box_distance * stride
+
+                # top K candidate
+                topk_idx = np.argsort(score.max(axis=1))[::-1]
+                topk_idx = topk_idx[:self.nms_top_k]
+                center = center[topk_idx]
+                score = score[topk_idx]
+                box_distance = box_distance[topk_idx]
+
+                # decode box
+                decode_box = center + [-1, -1, 1, 1] * box_distance
+
+                select_scores.append(score)
+                decode_boxes.append(decode_box)
+
+            # nms
+            bboxes = np.concatenate(decode_boxes, axis=0)
+            confidences = np.concatenate(select_scores, axis=0)
+            picked_box_probs = []
+            picked_labels = []
+            for class_index in range(0, confidences.shape[1]):
+                probs = confidences[:, class_index]
+                mask = probs > self.score_threshold
+                probs = probs[mask]
+                if probs.shape[0] == 0:
+                    continue
+                subset_boxes = bboxes[mask, :]
+                box_probs = np.concatenate(
+                    [subset_boxes, probs.reshape(-1, 1)], axis=1)
+                box_probs = self.hard_nms(
+                    box_probs,
+                    iou_threshold=self.nms_threshold,
+                    top_k=self.keep_top_k, )
+                picked_box_probs.append(box_probs)
+                picked_labels.extend([class_index] * box_probs.shape[0])
+
+            if len(picked_box_probs) == 0:
+                out_boxes_list.append(np.empty((0, 4)))
+                out_boxes_num.append(0)
+
+            else:
+                picked_box_probs = np.concatenate(picked_box_probs)
+
+                # resize output boxes
+                picked_box_probs[:, :4] = self.warp_boxes(
+                    picked_box_probs[:, :4], ori_shape[batch_id])
+                im_scale = np.concatenate([
+                    scale_factor[batch_id][::-1], scale_factor[batch_id][::-1]
+                ])
+                picked_box_probs[:, :4] /= im_scale
+                # clas score box
+                out_boxes_list.append(
+                    np.concatenate(
+                        [
+                            np.expand_dims(
+                                np.array(picked_labels),
+                                axis=-1), np.expand_dims(
+                                    picked_box_probs[:, 4], axis=-1),
+                            picked_box_probs[:, :4]
+                        ],
+                        axis=1))
+                out_boxes_num.append(len(picked_labels))
+
+        out_boxes_list = np.concatenate(out_boxes_list, axis=0)
+        out_boxes_num = np.asarray(out_boxes_num).astype(np.int32)
+
+        for dt in out_boxes_list:
+            clsid, bbox, score = int(dt[0]), dt[2:], dt[1]
+            label = self.labels[clsid]
+            result = {'bbox': bbox, 'label': label}
+            results.append(result)
+        return results
+
+    def load_layout_dict(self, layout_dict_path):
+        with open(layout_dict_path, 'r', encoding='utf-8') as fp:
+            labels = fp.readlines()
+        return [label.strip('\n') for label in labels]
+
+    def warp_boxes(self, boxes, ori_shape):
+        """Apply transform to boxes
+        """
+        width, height = ori_shape[1], ori_shape[0]
+        n = len(boxes)
+        if n:
+            # warp points
+            xy = np.ones((n * 4, 3))
+            xy[:, :2] = boxes[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(
+                n * 4, 2)  # x1y1, x2y2, x1y2, x2y1
+            # xy = xy @ M.T  # transform
+            xy = (xy[:, :2] / xy[:, 2:3]).reshape(n, 8)  # rescale
+            # create new boxes
+            x = xy[:, [0, 2, 4, 6]]
+            y = xy[:, [1, 3, 5, 7]]
+            xy = np.concatenate(
+                (x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
+            # clip boxes
+            xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)
+            xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)
+            return xy.astype(np.float32)
+        else:
+            return boxes
+
+    def img_info(self, ori_img, img):
+        origin_shape = ori_img.shape
+        resize_shape = img.shape
+        im_scale_y = resize_shape[2] / float(origin_shape[0])
+        im_scale_x = resize_shape[3] / float(origin_shape[1])
+        scale_factor = np.array([im_scale_y, im_scale_x], dtype=np.float32)
+        img_shape = np.array(img.shape[2:], dtype=np.float32)
+
+        input_shape = np.array(img).astype('float32').shape[2:]
+        ori_shape = np.array((img_shape, )).astype('float32')
+        scale_factor = np.array((scale_factor, )).astype('float32')
+        return ori_shape, input_shape, scale_factor
+
+    @staticmethod
+    def softmax(x, axis=None):
+        def logsumexp(a, axis=None, b=None, keepdims=False):
+            a_max = np.amax(a, axis=axis, keepdims=True)
+
+            if a_max.ndim > 0:
+                a_max[~np.isfinite(a_max)] = 0
+            elif not np.isfinite(a_max):
+                a_max = 0
+
+            tmp = np.exp(a - a_max)
+
+            # suppress warnings about log of zero
+            with np.errstate(divide='ignore'):
+                s = np.sum(tmp, axis=axis, keepdims=keepdims)
+                out = np.log(s)
+
+            if not keepdims:
+                a_max = np.squeeze(a_max, axis=axis)
+            out += a_max
+            return out
+
+        return np.exp(x - logsumexp(x, axis=axis, keepdims=True))
+
+    def hard_nms(self, box_scores, iou_threshold, top_k=-1,
+                 candidate_size=200):
+        """
+        Args:
+            box_scores (N, 5): boxes in corner-form and probabilities.
+            iou_threshold: intersection over union threshold.
+            top_k: keep top_k results. If k <= 0, keep all the results.
+            candidate_size: only consider the candidates with the highest scores.
+        Returns:
+            picked: a list of indexes of the kept boxes
+        """
+        scores = box_scores[:, -1]
+        boxes = box_scores[:, :-1]
+        picked = []
+        indexes = np.argsort(scores)
+        indexes = indexes[-candidate_size:]
+        while len(indexes) > 0:
+            current = indexes[-1]
+            picked.append(current)
+            if 0 < top_k == len(picked) or len(indexes) == 1:
+                break
+            current_box = boxes[current, :]
+            indexes = indexes[:-1]
+            rest_boxes = boxes[indexes, :]
+            iou = self.iou_of(rest_boxes, np.expand_dims(current_box, axis=0),)
+            indexes = indexes[iou <= iou_threshold]
+
+        return box_scores[picked, :]
+
+    def iou_of(self, boxes0, boxes1, eps=1e-5):
+        """Return intersection-over-union (Jaccard index) of boxes.
+        Args:
+            boxes0 (N, 4): ground truth boxes.
+            boxes1 (N or 1, 4): predicted boxes.
+            eps: a small number to avoid 0 as denominator.
+        Returns:
+            iou (N): IoU values.
+        """
+        overlap_left_top = np.maximum(boxes0[..., :2], boxes1[..., :2])
+        overlap_right_bottom = np.minimum(boxes0[..., 2:], boxes1[..., 2:])
+
+        overlap_area = self.area_of(overlap_left_top, overlap_right_bottom)
+        area0 = self.area_of(boxes0[..., :2], boxes0[..., 2:])
+        area1 = self.area_of(boxes1[..., :2], boxes1[..., 2:])
+        return overlap_area / (area0 + area1 - overlap_area + eps)
+
+    @staticmethod
+    def area_of(left_top, right_bottom):
+        """Compute the areas of rectangles given two corners.
+        Args:
+            left_top (N, 2): left top corner.
+            right_bottom (N, 2): right bottom corner.
+        Returns:
+            area (N): return the area.
+        """
+        hw = np.clip(right_bottom - left_top, 0.0, None)
+        return hw[..., 0] * hw[..., 1]