diff --git a/python/test_demo.py b/python/demo.py similarity index 100% rename from python/test_demo.py rename to python/demo.py diff --git a/python/rapid_structure/README.md b/python/rapid_structure/README.md new file mode 100644 index 0000000..990b394 --- /dev/null +++ b/python/rapid_structure/README.md @@ -0,0 +1 @@ +#### Rapid Structure diff --git a/python/rapid_structure/demo.py b/python/rapid_structure/demo.py new file mode 100644 index 0000000..f9d8475 --- /dev/null +++ b/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) diff --git a/python/rapid_structure/models/.gitkeep b/python/rapid_structure/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/python/rapid_structure/rapid_layout/__init__.py b/python/rapid_structure/rapid_layout/__init__.py new file mode 100644 index 0000000..79e4a04 --- /dev/null +++ b/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 diff --git a/python/rapid_structure/rapid_layout/config.yaml b/python/rapid_structure/rapid_layout/config.yaml new file mode 100644 index 0000000..bc86bbf --- /dev/null +++ b/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 \ No newline at end of file diff --git a/python/rapid_structure/rapid_layout/rapid_layout.py b/python/rapid_structure/rapid_layout/rapid_layout.py new file mode 100644 index 0000000..f5afcdb --- /dev/null +++ b/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 diff --git a/python/rapid_structure/rapid_layout/utils.py b/python/rapid_structure/rapid_layout/utils.py new file mode 100644 index 0000000..0df8081 --- /dev/null +++ b/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] diff --git a/python/rapid_structure/test_images/ch.png b/python/rapid_structure/test_images/ch.png new file mode 100644 index 0000000..29a1dc5 Binary files /dev/null and b/python/rapid_structure/test_images/ch.png differ diff --git a/python/rapid_structure/test_images/en.png b/python/rapid_structure/test_images/en.png new file mode 100644 index 0000000..49c259d Binary files /dev/null and b/python/rapid_structure/test_images/en.png differ