Source code for layoutparser.ocr.gcv_agent

# Copyright 2021 The Layout Parser team. 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

import io
import os
import json
import warnings

import numpy as np
from cv2 import imencode

from .base import BaseOCRAgent, BaseOCRElementType
from ..elements import Layout, TextBlock, Quadrilateral, TextBlock
from ..file_utils import is_gcv_available

if is_gcv_available():
    import google.protobuf.json_format as _json_format
    import as _vision

def _cvt_GCV_vertices_to_points(vertices):
    return np.array([[vertex.x, vertex.y] for vertex in vertices])

[docs]class GCVFeatureType(BaseOCRElementType): """ The element types from Google Cloud Vision API """ PAGE = 0 BLOCK = 1 PARA = 2 WORD = 3 SYMBOL = 4 @property def attr_name(self): name_cvt = { GCVFeatureType.PAGE: "pages", GCVFeatureType.BLOCK: "blocks", GCVFeatureType.PARA: "paragraphs", GCVFeatureType.WORD: "words", GCVFeatureType.SYMBOL: "symbols", } return name_cvt[self] @property def child_level(self): child_cvt = { GCVFeatureType.PAGE: GCVFeatureType.BLOCK, GCVFeatureType.BLOCK: GCVFeatureType.PARA, GCVFeatureType.PARA: GCVFeatureType.WORD, GCVFeatureType.WORD: GCVFeatureType.SYMBOL, GCVFeatureType.SYMBOL: None, } return child_cvt[self]
[docs]class GCVAgent(BaseOCRAgent): """A wrapper for `Google Cloud Vision (GCV) <>`_ Text Detection APIs. Note: Google Cloud Vision API returns the output text in two types: * `text_annotations`: In this format, GCV automatically find the best aggregation level for the text, and return the results in a list. We use :obj:`~gather_text_annotations` to reterive this type of information. * `full_text_annotation`: To support better user control, GCV also provides the `full_text_annotation` output, where it returns the hierarchical structure of the output text. To process this output, we provide the :obj:`~gather_full_text_annotation` function to aggregate the texts of the given aggregation level. """ DEPENDENCIES = ["google-cloud-vision"] def __init__(self, languages=None, ocr_image_decode_type=".png"): """Create a Google Cloud Vision OCR Agent. Args: languages (:obj:`list`, optional): You can specify the language code of the documents to detect to improve accuracy. The supported language and their code can be found on `this page <>`_. Defaults to None. ocr_image_decode_type (:obj:`str`, optional): The format to convert the input image to before sending for GCV OCR. Defaults to `".png"`. * `".png"` is suggested as it does not compress the image. * But `".jpg"` could also be a good choice if the input image is very large. """ try: self._client = _vision.ImageAnnotatorClient() except: warnings.warn( "The GCV credential has not been set. You could not run the detect command." ) self._context = _vision.types.ImageContext(language_hints=languages) self.ocr_image_decode_type = ocr_image_decode_type
[docs] @classmethod def with_credential(cls, credential_path, **kwargs): """Specifiy the credential to use for the GCV OCR API. Args: credential_path (:obj:`str`): The path to the credential file """ os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credential_path return cls(**kwargs)
def _detect(self, img_content): img_content = _vision.types.Image(content=img_content) response = self._client.document_text_detection( image=img_content, image_context=self._context ) return response
[docs] def detect( self, image, return_response=False, return_only_text=False, 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 the google cloud response. 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:`~GCVFeatureType`, optional): When set, aggregate the GCV output with respect to the specified aggregation level. Defaults to `None`. """ if isinstance(image, np.ndarray): img_content = imencode(self.ocr_image_decode_type, image)[1].tostring() elif isinstance(image, str): with, "rb") as image_file: img_content = res = self._detect(img_content) if return_response: return res if return_only_text: return res.full_text_annotation.text if agg_output_level is not None: return self.gather_full_text_annotation(res, agg_output_level) return self.gather_text_annotations(res)
[docs] @staticmethod def gather_text_annotations(response): """Convert the text_annotations from GCV output to an :obj:`Layout` object. Args: response (:obj:`AnnotateImageResponse`): The returned Google Cloud Vision AnnotateImageResponse object. Returns: :obj:`Layout`: The reterived layout from the response. """ # The 0th element contains all texts doc = response.text_annotations[1:] gathered_text = Layout() for i, text_comp in enumerate(doc): points = _cvt_GCV_vertices_to_points(text_comp.bounding_poly.vertices) gathered_text.append( TextBlock(block=Quadrilateral(points), text=text_comp.description, id=i) ) return gathered_text
[docs] @staticmethod def gather_full_text_annotation(response, agg_level): """Convert the full_text_annotation from GCV output to an :obj:`Layout` object. Args: response (:obj:`AnnotateImageResponse`): The returned Google Cloud Vision AnnotateImageResponse object. agg_level (:obj:`~GCVFeatureType`): The layout level to aggregate the text in full_text_annotation. Returns: :obj:`Layout`: The reterived layout from the response. """ def iter_level( iter, agg_level=None, text_blocks=None, texts=None, cur_level=GCVFeatureType.PAGE, ): for item in getattr(iter, cur_level.attr_name): if cur_level == agg_level: texts = [] # Go down levels to fetch the texts if cur_level == GCVFeatureType.SYMBOL: texts.append(item.text) elif ( cur_level == GCVFeatureType.WORD and agg_level != GCVFeatureType.SYMBOL ): chars = [] iter_level( item, agg_level, text_blocks, chars, cur_level.child_level ) texts.append("".join(chars)) else: iter_level( item, agg_level, text_blocks, texts, cur_level.child_level ) if cur_level == agg_level: nonlocal element_id points = _cvt_GCV_vertices_to_points(item.bounding_box.vertices) text_block = TextBlock( block=Quadrilateral(points), text=" ".join(texts), score=item.confidence, id=element_id, ) text_blocks.append(text_block) element_id += 1 if agg_level == GCVFeatureType.PAGE: doc = response.text_annotations[0] points = _cvt_GCV_vertices_to_points(doc.bounding_poly.vertices) text_blocks = [TextBlock(block=Quadrilateral(points), text=doc.description)] else: doc = response.full_text_annotation text_blocks = [] element_id = 0 iter_level(doc, agg_level, text_blocks) return Layout(text_blocks)
[docs] def load_response(self, filename): with open(filename, "r") as f: data = return _json_format.Parse( data, _vision.types.AnnotateImageResponse(), ignore_unknown_fields=True )
[docs] def save_response(self, res, file_name): res = _json_format.MessageToJson(res) with open(file_name, "w") as f: json_file = json.loads(res) json.dump(json_file, f)