diff --git a/README.md b/README.md
index a6958df..5742e6e 100644
--- a/README.md
+++ b/README.md
@@ -39,13 +39,53 @@ pip install layoutparser[ocr]
**For Windows Users:** Please read [installation.md](installation.md) for details about installing Detectron2.
+## **Recent updates**
+
+2021.6.8 Update new layout detection model (PaddleDetection) and ocr model (PaddleOCR). We test Detectron2 and PaddleDetection models on PubLayNet and TableBank datasets, the indicators are as follows:
+
+PubLayNet Dataset:
+
+| Model | mAP | CPU time cost | GPU time cost |
+| :-------------: | :---: | :-----------: | :-----------: |
+| Detectron2 | 88.98 | 16545.5ms | 209.5ms |
+| PaddleDetection | 93.6 | 1713.7ms | 66.6ms |
+
+TableBank Dataset:
+
+| Model | mAP | CPU time cost | GPU time cost |
+| :-------------: | :---: | :-----------: | :-----------: |
+| Detectron2 | 91.26 | 7623.2ms | 104.2.ms |
+| PaddleDetection | 96.2 | 1968.4ms | 65.1ms |
+
+**Envrionment：**
+
+ **CPU：** Intel(R) Xeon(R) CPU E5-2650 v4 @ 2.20GHz，24core
+
+ **GPU：** a single NVIDIA Tesla P40
+
+You can also find detailed installation instructions in [installation.md](installation.md). But generally, it's just `pip install`
+some libraries:
+
+```Python
+# Install PaddlePaddle
+# CUDA10.1
+python -m pip install paddlepaddle-gpu==2.1.0.post101 -f https://paddlepaddle.org.cn/whl/mkl/stable.html
+# CPU
+python -m pip install paddlepaddle -i https://mirror.baidu.com/pypi/simple
+
+# Install the paddle ocr components when necessary
+pip install layoutparser[paddleocr]
+```
+
+For more PaddlePaddle CUDA version or environment to quick install, please refer to the [PaddlePaddle Quick Installation document](https://www.paddlepaddle.org.cn/install/quick)
+
## Quick Start
We provide a series of examples for to help you start using the layout parser library:
1. [Table OCR and Results Parsing](https://github.com/Layout-Parser/layout-parser/blob/master/examples/OCR%20Tables%20and%20Parse%20the%20Output.ipynb): `layoutparser` can be used for conveniently OCR documents and convert the output in to structured data.
-
2. [Deep Layout Parsing Example](https://github.com/Layout-Parser/layout-parser/blob/master/examples/Deep%20Layout%20Parsing.ipynb): With the help of Deep Learning, `layoutparser` supports the analysis very complex documents and processing of the hierarchical structure in the layouts.
+3. [Deep Layout Parsing using Paddle](examples/Deep%20Layout%20Parsing%20using%20Paddle.ipynb): `layoutparser` supports the analysis very complex documents and processing of the hierarchical structure in the layouts Using Paddle models.
## DL Assisted Layout Prediction Example
@@ -54,7 +94,7 @@ We provide a series of examples for to help you start using the layout parser li
*The images shown in the figure above are: a screenshot of [this paper](https://arxiv.org/abs/2004.08686), an image from the [PRIMA Layout Analysis Dataset](https://www.primaresearch.org/dataset/), a screenshot of the [WSJ website](http://wsj.com), and an image from the [HJDataset](https://dell-research-harvard.github.io/HJDataset/).*
-With only 4 lines of code in `layoutparse`, you can unlock the information from complex documents that existing tools could not provide. You can either choose a deep learning model from the [ModelZoo](https://github.com/Layout-Parser/layout-parser/blob/master/docs/notes/modelzoo.md), or load the model that you trained on your own. And use the following code to predict the layout as well as visualize it:
+With only 4 lines of code in `layoutparse`, you can unlock the information from complex documents that existing tools could not provide. You can either choose a deep learning model from the [ModelZoo](docs/notes/modelzoo.md), or load the model that you trained on your own. And use the following code to predict the layout as well as visualize it:
```python
>>> import layoutparser as lp
@@ -63,6 +103,19 @@ With only 4 lines of code in `layoutparse`, you can unlock the information from
>>> lp.draw_box(image, layout,) # With extra configurations
```
+Use PaddleDetection model：
+
+```python
+>>> import layoutparser as lp
+>>> model = lp.PaddleDetectionLayoutModel('lp://PubLayNet/ppyolov2_r50vd_dcn_365e_publaynet/config')
+>>> layout = model.detect(image) # You need to load the image somewhere else, e.g., image = cv2.imread(...)
+>>> lp.draw_box(image, layout,) # With extra configurations
+```
+
+If you want to train a Paddledetection model yourself, please refer to：[Train PaddleDetection model](https://github.com/PaddlePaddle/PaddleDetection/blob/release/2.1/docs/tutorials/GETTING_STARTED.md)
+
+If you want to learn more about PaddleOCR, please refer to: [PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR)、[PaddleOCR infer](https://github.com/PaddlePaddle/PaddleOCR/blob/release/2.1/doc/doc_ch/whl.md)
+
## Contributing
We encourage you to contribute to Layout Parser! Please check out the [Contributing guidelines](.github/CONTRIBUTING.md) for guidelines about how to proceed. Join us!
@@ -78,4 +131,5 @@ If you find `layoutparser` helpful to your work, please consider citing our tool
journal={arXiv preprint arXiv:2103.15348},
year={2021}
}
-```
\ No newline at end of file
+```
+
diff --git a/dev-requirements.txt b/dev-requirements.txt
index d65250c..ed39f39 100644
--- a/dev-requirements.txt
+++ b/dev-requirements.txt
@@ -10,4 +10,6 @@ sphinx_rtd_theme
google-cloud-vision==1
pytesseract
pycocotools
-git+https://github.com/facebookresearch/detectron2.git@v0.4#egg=detectron2
\ No newline at end of file
+git+https://github.com/facebookresearch/detectron2.git@v0.4#egg=detectron2
+paddlepaddle==2.1.0
+paddleocr>=2.0.1
\ No newline at end of file
diff --git a/docs/notes/installation.md b/docs/notes/installation.md
deleted file mode 120000
index b5a6a49..0000000
--- a/docs/notes/installation.md
+++ /dev/null
@@ -1 +0,0 @@
-../../installation.md
\ No newline at end of file
diff --git a/docs/notes/installation.md b/docs/notes/installation.md
new file mode 100644
index 0000000..d8372a9
--- /dev/null
+++ b/docs/notes/installation.md
@@ -0,0 +1,63 @@
+# Installation
+
+## Install Python
+
+Layout Parser is a Python package that requires Python >= 3.6. If you do not have Python installed on your computer, you might want to turn to [the official instruction](https://www.python.org/downloads/) to download and install the appropriate version of Python.
+
+## Install the Layout Parser main library
+
+Installing the Layout Parser library is very straightforward: you just need to run the following command:
+
+```bash
+pip3 install -U layoutparser
+```
+
+## [Optional] Install Detectron2 for Using Layout Models
+
+### For Mac OS and Linux Users
+
+If you would like to use deep learning models for layout detection, you also need to install Detectron2 on your computer. This could be done by running the following command:
+
+```bash
+pip3 install 'git+https://github.com/facebookresearch/detectron2.git@v0.4#egg=detectron2'
+```
+
+This might take some time as the command will *compile* the library. You might also want to install a Detectron2 version
+with GPU support or encounter some issues during the installation process. Please refer to the official Detectron2
+[installation instruction](https://github.com/facebookresearch/detectron2/blob/master/INSTALL.md) for detailed
+information.
+
+### For Windows users
+
+As reported by many users, the installation of Detectron2 can be rather tricky on Windows platforms. In our extensive tests, we find that it is nearly impossible to provide a one-line installation command for Windows users. As a workaround solution, for now we list the possible challenges for installing Detectron2 on Windows, and attach helpful resources for solving them. We are also investigating other possibilities to avoid installing Detectron2 to use pre-trained models. If you have any suggestions or ideas, please feel free to [submit an issue](https://github.com/Layout-Parser/layout-parser/issues) in our repo.
+
+1. Challenges for installing `pycocotools`
+ - You can find detailed instructions on [this post](https://changhsinlee.com/pycocotools/) from Chang Hsin Lee.
+ - Another solution is try to install `pycocotools-windows`, see https://github.com/cocodataset/cocoapi/issues/415.
+2. Challenges for installing `Detectron2`
+ - [@ivanpp](https://github.com/ivanpp) curates a detailed description for installing `Detectron2` on Windows: [Detectron2 walkthrough (Windows)](https://ivanpp.cc/detectron2-walkthrough-windows/#step3installdetectron2)
+ - `Detectron2` maintainers claim that they won't provide official support for Windows (see [1](https://github.com/facebookresearch/detectron2/issues/9#issuecomment-540974288) and [2](https://detectron2.readthedocs.io/en/latest/tutorials/install.html)), but Detectron2 is continuously built on windows with CircleCI (see [3](https://github.com/facebookresearch/detectron2/blob/master/INSTALL.md#common-installation-issues)). Hopefully this situation will be improved in the future.
+
+
+## [Optional] Install OCR utils
+
+Layout Parser also comes with supports for OCR functions. In order to use them, you need to install the OCR utils via:
+
+```bash
+pip3 install -U layoutparser[ocr]
+```
+
+Additionally, if you want to use the Tesseract-OCR engine, you also need to install it on your computer. Please check the
+[official documentation](https://tesseract-ocr.github.io/tessdoc/Installation.html) for detailed installation instructions.
+
+## Known issues
+
+Error: instantiating `lp.GCVAgent.with_credential` returns module 'google.cloud.vision' has no attribute 'types'.
+
+
+In this case, you have a newer version of the google-cloud-vision. Please consider downgrading the API using:
+```bash
+pip install layoutparser[ocr]
+```
+
+Error: instantiating `lp.GCVAgent.with_credential` returns module 'google.cloud.vision' has no attribute 'types'.
diff --git a/setup.py b/setup.py
index 2dac7bf..3aeb31a 100644
--- a/setup.py
+++ b/setup.py
@@ -27,12 +27,17 @@
"torch",
"torchvision",
"iopath",
+ "paddlepaddle==2.1.0",
+ "tqdm",
],
extras_require={
"ocr": [
'google-cloud-vision==1',
'pytesseract'
],
+ "paddleocr": [
+ 'paddleocr>=2.0.1'
+ ],
},
include_package_data=True
)
\ No newline at end of file
diff --git a/src/layoutparser/__init__.py b/src/layoutparser/__init__.py
index ee5b120..cb2eb3c 100644
--- a/src/layoutparser/__init__.py
+++ b/src/layoutparser/__init__.py
@@ -11,11 +11,11 @@
from .ocr import (
GCVFeatureType, GCVAgent,
- TesseractFeatureType, TesseractAgent
+ TesseractFeatureType, TesseractAgent, PaddleocrAgent
)
from .models import (
- Detectron2LayoutModel
+ Detectron2LayoutModel, PaddleDetectionLayoutModel
)
from .io import (
diff --git a/src/layoutparser/models/__init__.py b/src/layoutparser/models/__init__.py
index d7f1d60..51e46cf 100644
--- a/src/layoutparser/models/__init__.py
+++ b/src/layoutparser/models/__init__.py
@@ -1,4 +1,4 @@
from . import catalog as _UNUSED
# A trick learned from
# https://github.com/facebookresearch/detectron2/blob/62cf3a2b6840734d2717abdf96e2dd57ed6612a6/detectron2/checkpoint/__init__.py#L6
-from .layoutmodel import Detectron2LayoutModel
+from .layoutmodel import Detectron2LayoutModel, PaddleDetectionLayoutModel
diff --git a/src/layoutparser/models/catalog.py b/src/layoutparser/models/catalog.py
index 80dd317..86c68bc 100644
--- a/src/layoutparser/models/catalog.py
+++ b/src/layoutparser/models/catalog.py
@@ -36,6 +36,7 @@
"faster_rcnn_R_50_FPN_3x": "https://www.dropbox.com/s/f3b12qc4hc0yh4m/config.yml?dl=1",
"mask_rcnn_R_50_FPN_3x": "https://www.dropbox.com/s/u9wbsfwz4y0ziki/config.yml?dl=1",
"mask_rcnn_X_101_32x8d_FPN_3x": "https://www.dropbox.com/s/nau5ut6zgthunil/config.yaml?dl=1",
+ "ppyolov2_r50vd_dcn_365e_publaynet": "https://paddle-model-ecology.bj.bcebos.com/model/layout-parser/ppyolov2_r50vd_dcn_365e_publaynet.tar",
},
"PrimaLayout": {
"mask_rcnn_R_50_FPN_3x": "https://www.dropbox.com/s/yc92x97k50abynt/config.yaml?dl=1"
@@ -46,6 +47,8 @@
"TableBank": {
"faster_rcnn_R_50_FPN_3x": "https://www.dropbox.com/s/7cqle02do7ah7k4/config.yaml?dl=1",
"faster_rcnn_R_101_FPN_3x": "https://www.dropbox.com/s/h63n6nv51kfl923/config.yaml?dl=1",
+ "ppyolov2_r50vd_dcn_365e_tableBank_word": "https://paddle-model-ecology.bj.bcebos.com/model/layout-parser/ppyolov2_r50vd_dcn_365e_tableBank_word.tar",
+ "ppyolov2_r50vd_dcn_365e_tableBank_latex": "https://paddle-model-ecology.bj.bcebos.com/model/layout-parser/ppyolov2_r50vd_dcn_365e_tableBank_latex.tar",
},
}
diff --git a/src/layoutparser/models/layoutmodel.py b/src/layoutparser/models/layoutmodel.py
index 5544653..57e44ec 100644
--- a/src/layoutparser/models/layoutmodel.py
+++ b/src/layoutparser/models/layoutmodel.py
@@ -1,15 +1,28 @@
from abc import ABC, abstractmethod
import os
+import cv2
import importlib
+import time
+import yaml
+from functools import reduce
from PIL import Image
import numpy as np
import torch
+import tarfile
+import requests
+from tqdm import tqdm
+
+import paddle
+from paddle.inference import Config
+from paddle.inference import create_predictor
+
+from .preprocess import preprocess, Resize, NormalizeImage, Permute, PadStride
from .catalog import PathManager, LABEL_MAP_CATALOG
from ..elements import *
-__all__ = ["Detectron2LayoutModel"]
+__all__ = ["Detectron2LayoutModel", "PaddleDetectionLayoutModel"]
class BaseLayoutModel(ABC):
@@ -112,6 +125,7 @@ def __init__(
cfg = self._config.get_cfg()
config_path = PathManager.get_local_path(config_path)
+
cfg.merge_from_file(config_path)
cfg.merge_from_list(extra_config)
@@ -167,3 +181,257 @@ def detect(self, image):
outputs = self.model(image)
layout = self.gather_output(outputs)
return layout
+
+
+class PaddleDetectionLayoutModel(object):
+ """
+ Args:
+ config (object): config of model, defined by `Config(model_dir)`
+ model_path (str):The path to the saved weights of the model.
+ threshold (float): threshold to reserve the result for output
+ input_shape(list): the image shape after reshape
+ batch_size(int)：test batch size
+ label_map (:obj:`dict`, optional):The map from the model prediction (ids) to realword labels (strings).
+ enforce_cpu (bool): whether use cpu, if false, indicates use GPU
+ enforce_mkldnn(bool): whether use mkldnn to accelerate the computation
+ thread_num(int): the number of threads
+ use_dynamic_shape (bool): use dynamic shape or not
+ trt_min_shape (int): min shape for dynamic shape in trt
+ trt_max_shape (int): max shape for dynamic shape in trt
+ trt_opt_shape (int): opt shape for dynamic shape in trt.
+
+ """
+
+ def __init__(self,
+ config_path=None,
+ model_path=None,
+ threshold=0.5,
+ input_shape=[3,640,640],
+ batch_size=1,
+ label_map=None,
+ enforce_cpu=False,
+ enable_mkldnn=True,
+ thread_num=10,
+ use_dynamic_shape=False,
+ trt_min_shape=1,
+ trt_max_shape=1280,
+ trt_opt_shape=640,
+ min_subgraph_size=3):
+
+ if config_path is not None and config_path.startswith("lp://"):
+ prefix = "lp://"
+ model_name = config_path[len(prefix) :].split('/')[1]
+ url = PathManager.get_local_path(config_path)
+ if label_map is None:
+ dataset_name = config_path.lstrip("lp://").split("/")[0]
+ label_map = LABEL_MAP_CATALOG[dataset_name]
+
+ BASE_DIR = os.path.expanduser("~/.paddledet/")
+ BASE_INFERENCE_MODEL_DIR = os.path.join(BASE_DIR, 'inference_model')
+
+ model_dir = os.path.join(BASE_INFERENCE_MODEL_DIR, model_name, model_name+'_infer')
+ if not os.path.exists(model_dir):
+ os.makedirs(model_dir)
+
+ maybe_download(model_storage_directory=model_dir, url=url)
+
+ if model_path is not None:
+ model_dir = model_path
+
+ self.predictor = self.load_predictor(
+ model_dir,
+ batch_size=batch_size,
+ enforce_cpu=enforce_cpu,
+ enable_mkldnn=enable_mkldnn,
+ thread_num=thread_num,
+ min_subgraph_size=min_subgraph_size,
+ use_dynamic_shape=use_dynamic_shape,
+ trt_min_shape=trt_min_shape,
+ trt_max_shape=trt_max_shape,
+ trt_opt_shape=trt_opt_shape)
+
+ self.threshold = threshold
+ self.input_shape = input_shape
+ self.label_map = label_map
+
+ def load_predictor(self,
+ model_dir,
+ batch_size=1,
+ enforce_cpu=False,
+ enable_mkldnn=True,
+ thread_num=10,
+ min_subgraph_size=3,
+ use_dynamic_shape=False,
+ trt_min_shape=1,
+ trt_max_shape=1280,
+ trt_opt_shape=640):
+ """set AnalysisConfig, generate AnalysisPredictor
+ Args:
+ model_dir (str): root path of __model__ and __params__
+ enforce_cpu (bool): whether use cpu
+ use_dynamic_shape (bool): use dynamic shape or not
+ trt_min_shape (int): min shape for dynamic shape in trt
+ trt_max_shape (int): max shape for dynamic shape in trt
+ trt_opt_shape (int): opt shape for dynamic shape in trt
+ Returns:
+ predictor (PaddlePredictor): AnalysisPredictor
+ Raises:
+ ValueError: predict by TensorRT need enforce_cpu == False.
+ """
+
+ use_calib_mode = False
+ config = Config(
+ os.path.join(model_dir, 'inference.pdmodel'),
+ os.path.join(model_dir, 'inference.pdiparams'))
+
+ if not enforce_cpu:
+ # initial GPU memory(M), device ID
+ config.enable_use_gpu(200, 0)
+ # optimize graph and fuse op
+ config.switch_ir_optim(True)
+ else:
+ config.disable_gpu()
+ config.set_cpu_math_library_num_threads(thread_num)
+ if enable_mkldnn:
+ # cache 10 different shapes for mkldnn to avoid memory leak
+ config.set_mkldnn_cache_capacity(10)
+ config.enable_mkldnn()
+
+
+ # disable print log when predict
+ config.disable_glog_info()
+ # enable shared memory
+ config.enable_memory_optim()
+ # disable feed, fetch OP, needed by zero_copy_run
+ config.switch_use_feed_fetch_ops(False)
+ predictor = create_predictor(config)
+ return predictor
+
+ def create_inputs(self, im, im_info):
+ """generate input for different model type
+ Args:
+ im (np.ndarray): image (np.ndarray)
+ im_info (dict): info of image
+ Returns:
+ inputs (dict): input of model
+ """
+ inputs = {}
+ inputs['image'] = np.array((im, )).astype('float32')
+ inputs['im_shape'] = np.array((im_info['im_shape'], )).astype('float32')
+ inputs['scale_factor'] = np.array(
+ (im_info['scale_factor'], )).astype('float32')
+
+ return inputs
+
+ def preprocess(self, im):
+ im, im_info = preprocess(im, self.input_shape)
+ inputs = self.create_inputs(im, im_info)
+ return inputs
+
+ def postprocess(self, np_boxes, np_masks, inputs):
+ # postprocess output of predictor
+ results = {}
+ results['boxes'] = np_boxes
+ if np_masks is not None:
+ results['masks'] = np_masks
+ return results
+
+ def gather_output(self, results):
+ layout = Layout()
+ np_boxes = results['boxes']
+ expect_boxes = (np_boxes[:, 1] > self.threshold) & (np_boxes[:, 0] > -1)
+ np_boxes = np_boxes[expect_boxes, :]
+
+ for dt in np_boxes:
+ clsid, bbox, score = int(dt[0]), dt[2:], dt[1]
+ x_1, y_1, x_2, y_2 = bbox
+
+ if self.label_map is not None:
+ label = self.label_map[clsid]
+
+ cur_block = TextBlock(
+ Rectangle(x_1, y_1, x_2, y_2), type=label, score=score
+ )
+ layout.append(cur_block)
+
+ return layout
+
+ def detect(self,
+ image):
+ '''
+ Args:
+ image (str/np.ndarray): path of image/ np.ndarray read by cv2
+ Returns:
+ results (dict): include 'boxes': np.ndarray: shape:[N,6], N: number of box,
+ matix element:[class, score, x_min, y_min, x_max, y_max]
+ MaskRCNN's results include 'masks': np.ndarray:
+ shape: [N, im_h, im_w]
+ '''
+ inputs = self.preprocess(image)
+
+ np_boxes, np_masks = None, None
+ input_names = self.predictor.get_input_names()
+ for i in range(len(input_names)):
+ input_tensor = self.predictor.get_input_handle(input_names[i])
+ input_tensor.copy_from_cpu(inputs[input_names[i]])
+
+ self.predictor.run()
+ output_names = self.predictor.get_output_names()
+ boxes_tensor = self.predictor.get_output_handle(output_names[0])
+ np_boxes = boxes_tensor.copy_to_cpu()
+
+ # do not perform postprocess in benchmark mode
+ results = []
+ if reduce(lambda x, y: x * y, np_boxes.shape) < 6:
+ print('[WARNNING] No object detected.')
+ results = {'boxes': np.array([])}
+ else:
+ results = self.postprocess(
+ np_boxes, np_masks, inputs)
+
+ layout = self.gather_output(results)
+ return layout
+
+
+def download_with_progressbar(url, save_path):
+ response = requests.get(url, stream=True)
+ total_size_in_bytes = int(response.headers.get('content-length', 0))
+ block_size = 1024 # 1 Kibibyte
+ progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
+ with open(save_path, 'wb') as file:
+ for data in response.iter_content(block_size):
+ progress_bar.update(len(data))
+ file.write(data)
+ progress_bar.close()
+ if total_size_in_bytes == 0 or progress_bar.n != total_size_in_bytes:
+ raise Exception(
+ "Something went wrong while downloading model/image from {}".
+ format(url))
+
+def maybe_download(model_storage_directory, url):
+ # using custom model
+ tar_file_name_list = [
+ 'inference.pdiparams', 'inference.pdiparams.info', 'inference.pdmodel'
+ ]
+ if not os.path.exists(
+ os.path.join(model_storage_directory, 'inference.pdiparams')
+ ) or not os.path.exists(
+ os.path.join(model_storage_directory, 'inference.pdmodel')):
+ tmp_path = os.path.join(model_storage_directory, url.split('/')[-1])
+ print('download {} to {}'.format(url, tmp_path))
+ os.makedirs(model_storage_directory, exist_ok=True)
+ download_with_progressbar(url, tmp_path)
+ with tarfile.open(tmp_path, 'r') as tarObj:
+ for member in tarObj.getmembers():
+ filename = None
+ for tar_file_name in tar_file_name_list:
+ if tar_file_name in member.name:
+ filename = tar_file_name
+ if filename is None:
+ continue
+ file = tarObj.extractfile(member)
+ with open(
+ os.path.join(model_storage_directory, filename),
+ 'wb') as f:
+ f.write(file.read())
+ os.remove(tmp_path)
\ No newline at end of file
diff --git a/src/layoutparser/models/preprocess.py b/src/layoutparser/models/preprocess.py
new file mode 100644
index 0000000..7ad899f
--- /dev/null
+++ b/src/layoutparser/models/preprocess.py
@@ -0,0 +1,218 @@
+# 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.
+
+from PIL import Image
+import cv2
+import numpy as np
+
+
+def decode_image(im_file, im_info):
+ """read rgb image
+ Args:
+ im_file (str|np.ndarray): input can be image path or np.ndarray
+ im_info (dict): info of image
+ Returns:
+ im (np.ndarray): processed image (np.ndarray)
+ im_info (dict): info of processed image
+ """
+ if isinstance(im_file, str):
+ with open(im_file, 'rb') as f:
+ im_read = f.read()
+ data = np.frombuffer(im_read, dtype='uint8')
+ im = cv2.imdecode(data, 1) # BGR mode, but need RGB mode
+ im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
+ else:
+ im = im_file
+ im_info['im_shape'] = np.array(im.shape[:2], dtype=np.float32)
+ im_info['scale_factor'] = np.array([1., 1.], dtype=np.float32)
+ return im, im_info
+
+
+class Resize(object):
+ """resize image by target_size and max_size
+ Args:
+ target_size (int): the target size of image
+ keep_ratio (bool): whether keep_ratio or not, default true
+ interp (int): method of resize
+ """
+
+ def __init__(
+ self,
+ target_size=[640, 640],
+ keep_ratio=False,
+ interp=2, ):
+ if isinstance(target_size, int):
+ target_size = [target_size, target_size]
+ self.target_size = target_size
+ self.keep_ratio = keep_ratio
+ self.interp = interp
+
+ def __call__(self, im, im_info):
+ """
+ Args:
+ im (np.ndarray): image (np.ndarray)
+ im_info (dict): info of image
+ Returns:
+ im (np.ndarray): processed image (np.ndarray)
+ im_info (dict): info of processed image
+ """
+ assert len(self.target_size) == 2
+ assert self.target_size[0] > 0 and self.target_size[1] > 0
+ im_channel = im.shape[2]
+ im_scale_y, im_scale_x = self.generate_scale(im)
+ # set image_shape
+ im_info['input_shape'][1] = int(im_scale_y * im.shape[0])
+ im_info['input_shape'][2] = int(im_scale_x * im.shape[1])
+ im = cv2.resize(
+ im,
+ None,
+ None,
+ fx=im_scale_x,
+ fy=im_scale_y,
+ interpolation=self.interp)
+ im_info['im_shape'] = np.array(im.shape[:2]).astype('float32')
+ im_info['scale_factor'] = np.array(
+ [im_scale_y, im_scale_x]).astype('float32')
+ return im, im_info
+
+ def generate_scale(self, im):
+ """
+ Args:
+ im (np.ndarray): image (np.ndarray)
+ Returns:
+ im_scale_x: the resize ratio of X
+ im_scale_y: the resize ratio of Y
+ """
+ origin_shape = im.shape[:2]
+ im_c = im.shape[2]
+ if self.keep_ratio:
+ im_size_min = np.min(origin_shape)
+ im_size_max = np.max(origin_shape)
+ target_size_min = np.min(self.target_size)
+ target_size_max = np.max(self.target_size)
+ im_scale = float(target_size_min) / float(im_size_min)
+ if np.round(im_scale * im_size_max) > target_size_max:
+ im_scale = float(target_size_max) / float(im_size_max)
+ im_scale_x = im_scale
+ im_scale_y = im_scale
+ else:
+ resize_h, resize_w = self.target_size
+ im_scale_y = resize_h / float(origin_shape[0])
+ im_scale_x = resize_w / float(origin_shape[1])
+ return im_scale_y, im_scale_x
+
+
+class NormalizeImage(object):
+ """normalize image
+ Args:
+ mean (list): im - mean
+ std (list): im / std
+ is_scale (bool): whether need im / 255
+ is_channel_first (bool): if True: image shape is CHW, else: HWC
+ """
+
+ def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], is_scale=True):
+ self.mean = mean
+ self.std = std
+ self.is_scale = is_scale
+
+ def __call__(self, im, im_info):
+ """
+ Args:
+ im (np.ndarray): image (np.ndarray)
+ im_info (dict): info of image
+ Returns:
+ im (np.ndarray): processed image (np.ndarray)
+ im_info (dict): info of processed image
+ """
+ im = im.astype(np.float32, copy=False)
+ mean = np.array(self.mean)[np.newaxis, np.newaxis, :]
+ std = np.array(self.std)[np.newaxis, np.newaxis, :]
+
+ if self.is_scale:
+ im = im / 255.0
+
+ im -= mean
+ im /= std
+ return im, im_info
+
+
+class Permute(object):
+ """permute image
+ Args:
+ to_bgr (bool): whether convert RGB to BGR
+ channel_first (bool): whether convert HWC to CHW
+ """
+
+ def __init__(self, ):
+ super(Permute, self).__init__()
+
+ def __call__(self, im, im_info):
+ """
+ Args:
+ im (np.ndarray): image (np.ndarray)
+ im_info (dict): info of image
+ Returns:
+ im (np.ndarray): processed image (np.ndarray)
+ im_info (dict): info of processed image
+ """
+ im = im.transpose((2, 0, 1)).copy()
+ return im, im_info
+
+
+class PadStride(object):
+ """ padding image for model with FPN , instead PadBatch(pad_to_stride, pad_gt) in original config
+ Args:
+ stride (bool): model with FPN need image shape % stride == 0
+ """
+
+ def __init__(self, stride=0):
+ self.coarsest_stride = stride
+
+ def __call__(self, im, im_info):
+ """
+ Args:
+ im (np.ndarray): image (np.ndarray)
+ im_info (dict): info of image
+ Returns:
+ im (np.ndarray): processed image (np.ndarray)
+ im_info (dict): info of processed image
+ """
+ coarsest_stride = self.coarsest_stride
+ if coarsest_stride <= 0:
+ return im, im_info
+ im_c, im_h, im_w = im.shape
+ pad_h = int(np.ceil(float(im_h) / coarsest_stride) * coarsest_stride)
+ pad_w = int(np.ceil(float(im_w) / coarsest_stride) * coarsest_stride)
+ padding_im = np.zeros((im_c, pad_h, pad_w), dtype=np.float32)
+ padding_im[:, :im_h, :im_w] = im
+ return padding_im, im_info
+
+
+def preprocess(im, input_shape):
+ # process image by preprocess_ops
+ im_info = {
+ 'scale_factor': np.array(
+ [1., 1.], dtype=np.float32),
+ 'im_shape': None,
+ 'input_shape': input_shape,
+ }
+ im, im_info = decode_image(im, im_info)
+ resize = Resize()
+ permute = Permute()
+ normalizeImage = NormalizeImage()
+ im, im_info = resize(im, im_info)
+ im, im_info = normalizeImage(im, im_info)
+ im, im_info = permute(im, im_info)
+ return im, im_info
\ No newline at end of file
diff --git a/src/layoutparser/ocr.py b/src/layoutparser/ocr.py
index 65f5ccc..0e8c290 100644
--- a/src/layoutparser/ocr.py
+++ b/src/layoutparser/ocr.py
@@ -11,11 +11,12 @@
import numpy as np
import pandas as pd
from cv2 import imencode
+from paddleocr import PaddleOCR
from .elements import *
from .io import load_dataframe
-__all__ = ["GCVFeatureType", "GCVAgent", "TesseractFeatureType", "TesseractAgent"]
+__all__ = ["GCVFeatureType", "GCVAgent", "TesseractFeatureType", "TesseractAgent", "PaddleocrAgent"]
def _cvt_GCV_vertices_to_points(vertices):
@@ -504,3 +505,55 @@ def save_response(res, file_name):
with open(file_name, "wb") as fp:
pickle.dump(res, fp, protocol=pickle.HIGHEST_PROTOCOL)
+
+
+class PaddleocrAgent():
+ """
+ A wrapper for `PaddleOCR `_ Text
+ Detection APIs based on `paddleocr `_.
+ """
+
+ DEPENDENCIES = ["paddleocr"]
+
+ def __init__(self, languages="en", use_gpu=True, use_angle_cls=False, det=True, rec=True, cls=False, **kwargs):
+ """Create a Tesseract OCR Agent.
+
+ Args:
+ languages (:obj:`list` or :obj:`str`, optional):
+ You can specify the language code(s) of the documents to detect to improve
+ accuracy. The supported language and their code can be found on
+ `its github repo `_.
+ It supports llaguages：`ch`, `en`, `french`, `german`, `korean`, `japan`.
+ Defaults to 'eng'.
+ """
+ self.lang = languages
+ self.use_gpu = use_gpu
+ self.use_angle_cls = use_angle_cls
+ self.det = det
+ self.rec = rec
+ self.cls = cls
+ self.configs = kwargs
+ self.ocr = PaddleOCR(use_gpu=self.use_gpu, use_angle_cls=self.use_angle_cls, lang=self.lang)
+
+ def detect(
+ self, image, return_response=False, return_only_text=True, agg_output_level=None
+ ):
+ """Send the input image for OCR.
+
+ Args:
+ image (:obj:`np.ndarray` or :obj:`str`):
+ The input image array or the name of the image file
+ return_response (:obj:`bool`, optional):
+ Whether directly return all output (string and boxes
+ info) from Tesseract.
+ Defaults to `False`.
+ return_only_text (:obj:`bool`, optional):
+ Whether return only the texts in the OCR results.
+ Defaults to `False`.
+ agg_output_level (:obj:`~TesseractFeatureType`, optional):
+ When set, aggregate the GCV output with respect to the
+ specified aggregation level. Defaults to `None`.
+ """
+ result = self.ocr.ocr(image, det=self.det, rec=self.rec, cls=self.cls)
+ txts = '\n'.join(line[1][0] for line in result)
+ return txts