修改为东南天坐标系
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,244 @@
|
||||
"""
|
||||
This module contains variables that can permitted to be tweaked by the system environment. For
|
||||
example, model parameters that changes the output of an inference call. Constants do NOT belong in
|
||||
this module. Constants are values that are usually names for common options (e.g., color names) or
|
||||
settings that should not be altered without making a code change (e.g., definition of 1Gb of memory
|
||||
in bytes). Constants should go into `./constants.py`
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from unstructured.partition.utils.constants import OCR_AGENT_TESSERACT
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_tempdir(dir: str) -> str:
|
||||
tempdir = Path(dir) / f"tmp/{os.getpgid(0)}"
|
||||
return str(tempdir)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ENVConfig:
|
||||
"""class for configuring enviorment parameters"""
|
||||
|
||||
def __post_init__(self):
|
||||
if self.GLOBAL_WORKING_DIR_ENABLED:
|
||||
self._setup_tmpdir(self.GLOBAL_WORKING_PROCESS_DIR)
|
||||
|
||||
def _get_string(self, var: str, default_value: str = "") -> str:
|
||||
"""attempt to get the value of var from the os environment; if not present return the
|
||||
default_value"""
|
||||
return os.environ.get(var, default_value)
|
||||
|
||||
def _get_int(self, var: str, default_value: int) -> int:
|
||||
if value := self._get_string(var):
|
||||
return int(value)
|
||||
return default_value
|
||||
|
||||
def _get_float(self, var: str, default_value: float) -> float:
|
||||
if value := self._get_string(var):
|
||||
return float(value)
|
||||
return default_value
|
||||
|
||||
def _get_bool(self, var: str, default_value: bool) -> bool:
|
||||
if value := self._get_string(var):
|
||||
return value.lower() in ("true", "1", "t")
|
||||
return default_value
|
||||
|
||||
def _setup_tmpdir(self, tmpdir: str) -> None:
|
||||
Path(tmpdir).mkdir(parents=True, exist_ok=True)
|
||||
tempfile.tempdir = tmpdir
|
||||
|
||||
@property
|
||||
def IMAGE_CROP_PAD(self) -> int:
|
||||
"""extra image content to add around an identified element region; measured in pixels"""
|
||||
return self._get_int("IMAGE_CROP_PAD", 0)
|
||||
|
||||
@property
|
||||
def TABLE_IMAGE_CROP_PAD(self) -> int:
|
||||
"""extra image content to add around an identified table region; measured in pixels
|
||||
|
||||
The padding adds image data around an identified table bounding box for downstream table
|
||||
structure detection model use as input
|
||||
"""
|
||||
return self._get_int("TABLE_IMAGE_CROP_PAD", 0)
|
||||
|
||||
@property
|
||||
def TESSERACT_TEXT_HEIGHT_QUANTILE(self) -> float:
|
||||
"""the quantile to check for text height"""
|
||||
return self._get_float("TESSERACT_TEXT_HEIGHT_QUANTILE", 0.5)
|
||||
|
||||
@property
|
||||
def TESSERACT_MIN_TEXT_HEIGHT(self) -> int:
|
||||
"""minimum text height acceptable from tesseract OCR results
|
||||
|
||||
if estimated text height from tesseract OCR results is lower than this value the image is
|
||||
scaled up to be processed again
|
||||
"""
|
||||
return self._get_int("TESSERACT_MIN_TEXT_HEIGHT", 12)
|
||||
|
||||
@property
|
||||
def TESSERACT_MAX_TEXT_HEIGHT(self) -> int:
|
||||
"""maximum text height acceptable from tesseract OCR results
|
||||
|
||||
if estimated text height from tesseract OCR results is higher than this value the image is
|
||||
scaled down to be processed again
|
||||
"""
|
||||
return self._get_int("TESSERACT_MAX_TEXT_HEIGHT", 100)
|
||||
|
||||
@property
|
||||
def TESSERACT_OPTIMUM_TEXT_HEIGHT(self) -> int:
|
||||
"""optimum text height for tesseract OCR"""
|
||||
return self._get_int("TESSERACT_OPTIMUM_TEXT_HEIGHT", 20)
|
||||
|
||||
@property
|
||||
def TESSERACT_CHARACTER_CONFIDENCE_THRESHOLD(self) -> int:
|
||||
"""Tesseract predictions with confidence below this threshold are ignored"""
|
||||
return self._get_float("TESSERACT_CHARACTER_CONFIDENCE_THRESHOLD", 0.0)
|
||||
|
||||
@property
|
||||
def GOOGLEVISION_API_ENDPOINT(self) -> str:
|
||||
"""API endpoint to use for Google Vision"""
|
||||
return self._get_string("GOOGLEVISION_API_ENDPOINT", "")
|
||||
|
||||
@property
|
||||
def OCR_AGENT(self) -> str:
|
||||
"""OCR Agent to use"""
|
||||
return self._get_string("OCR_AGENT", OCR_AGENT_TESSERACT)
|
||||
|
||||
@property
|
||||
def OCR_AGENT_CACHE_SIZE(self) -> int:
|
||||
"""Maximum number of OCR agents to cache per process"""
|
||||
return self._get_int("OCR_AGENT_CACHE_SIZE", 1)
|
||||
|
||||
@property
|
||||
def EXTRACT_IMAGE_BLOCK_CROP_HORIZONTAL_PAD(self) -> int:
|
||||
"""extra image block content to add around an identified element(`Image`, `Table`) region
|
||||
horizontally; measured in pixels
|
||||
"""
|
||||
return self._get_int("EXTRACT_IMAGE_BLOCK_CROP_HORIZONTAL_PAD", 0)
|
||||
|
||||
@property
|
||||
def EXTRACT_IMAGE_BLOCK_CROP_VERTICAL_PAD(self) -> int:
|
||||
"""extra image block content to add around an identified element(`Image`, `Table`) region
|
||||
vertically; measured in pixels
|
||||
"""
|
||||
return self._get_int("EXTRACT_IMAGE_BLOCK_CROP_VERTICAL_PAD", 0)
|
||||
|
||||
@property
|
||||
def EXTRACT_TABLE_AS_CELLS(self) -> bool:
|
||||
"""adds `table_as_cells` to a Table element's metadata when it is True"""
|
||||
return self._get_bool("EXTRACT_TABLE_AS_CELLS", False)
|
||||
|
||||
@property
|
||||
def OCR_LAYOUT_SUBREGION_THRESHOLD(self) -> float:
|
||||
"""threshold to determine if an OCR region is a sub-region of a given block
|
||||
when aggregating the text from OCR'd elements that lie within the given block
|
||||
|
||||
When the intersection region area divided by self area is larger than this threshold self is
|
||||
considered a subregion of the other
|
||||
"""
|
||||
return self._get_float("OCR_LAYOUT_SUBREGION_THRESHOLD", 0.5)
|
||||
|
||||
@property
|
||||
def EMBEDDED_IMAGE_SAME_REGION_THRESHOLD(self) -> float:
|
||||
"""threshold to consider the bounding boxes of two embedded images as the same region"""
|
||||
return self._get_float("EMBEDDED_IMAGE_SAME_REGION_THRESHOLD", 0.6)
|
||||
|
||||
@property
|
||||
def EMBEDDED_TEXT_AGGREGATION_SUBREGION_THRESHOLD(self) -> float:
|
||||
"""threshold to determine if an embedded region is a sub-region of a given block
|
||||
when aggregating the text from embedded elements that lie within the given block
|
||||
|
||||
When the intersection region area divided by self area is larger than this threshold self is
|
||||
considered a subregion of the other
|
||||
"""
|
||||
return self._get_float("EMBEDDED_TEXT_AGGREGATION_SUBREGION_THRESHOLD", 0.99)
|
||||
|
||||
@property
|
||||
def EMBEDDED_TEXT_SAME_REGION_THRESHOLD(self) -> float:
|
||||
"""threshold to consider the bounding boxes of two embedded images as the same region"""
|
||||
return self._get_float("EMBEDDED_TEXT_SAME_REGION_THRESHOLD", 0.9)
|
||||
|
||||
@property
|
||||
def PDF_ANNOTATION_THRESHOLD(self) -> float:
|
||||
"""The threshold value (between 0.0 and 1.0) that determines the minimum overlap required
|
||||
for an annotation to be considered within the element.
|
||||
"""
|
||||
|
||||
return self._get_float("PDF_ANNOTATION_THRESHOLD", 0.9)
|
||||
|
||||
@property
|
||||
def PDF_MAX_EMBED_INVISIBLE_TEXT_RATIO(self) -> float:
|
||||
"""maximum ratio of invisible text for a text to be considered embedded text"""
|
||||
return self._get_float("PDF_MAX_EMBED_INVISIBLE_TEXT_RATIO", 0.1)
|
||||
|
||||
@property
|
||||
def GLOBAL_WORKING_DIR_ENABLED(self) -> bool:
|
||||
"""Enable usage of GLOBAL_WORKING_DIR and GLOBAL_WORKING_PROCESS_DIR."""
|
||||
return self._get_bool("GLOBAL_WORKING_DIR_ENABLED", False)
|
||||
|
||||
@property
|
||||
def GLOBAL_WORKING_DIR(self) -> str:
|
||||
"""Path to Unstructured cache directory."""
|
||||
return self._get_string("GLOBAL_WORKING_DIR", str(Path.home() / ".cache/unstructured"))
|
||||
|
||||
@property
|
||||
def GLOBAL_WORKING_PROCESS_DIR(self) -> str:
|
||||
"""Path to Unstructured cache tempdir. Overrides TMPDIR, TEMP and TMP.
|
||||
Defaults to '{GLOBAL_WORKING_DIR}/tmp/{os.getpgid(0)}'.
|
||||
"""
|
||||
default_tmpdir = get_tempdir(dir=self.GLOBAL_WORKING_DIR)
|
||||
tmpdir = self._get_string("GLOBAL_WORKING_PROCESS_DIR", default_tmpdir)
|
||||
if tmpdir == "":
|
||||
tmpdir = default_tmpdir
|
||||
if self.GLOBAL_WORKING_DIR_ENABLED:
|
||||
self._setup_tmpdir(tmpdir)
|
||||
return tmpdir
|
||||
|
||||
@property
|
||||
def ANALYSIS_DUMP_OD_SKIP(self) -> bool:
|
||||
"""Analysis dump object detection skip flag."""
|
||||
return self._get_bool("ANALYSIS_DUMP_OD_SKIP", False)
|
||||
|
||||
@property
|
||||
def ANALYSIS_BBOX_SKIP(self) -> bool:
|
||||
"""Analysis draw bboxes on pages skip flag."""
|
||||
return self._get_bool("ANALYSIS_BBOX_SKIP", False)
|
||||
|
||||
@property
|
||||
def ANALYSIS_BBOX_DRAW_GRID(self) -> bool:
|
||||
"""Flag for drawing the analysis bboxes on a single image (as grid)"""
|
||||
return self._get_bool("ANALYSIS_BBOX_DRAW_GRID", False)
|
||||
|
||||
@property
|
||||
def ANALYSIS_BBOX_DRAW_CAPTION(self) -> bool:
|
||||
"""Flag for drawing the caption above the analysed page (for e.g. layout source)"""
|
||||
return self._get_bool("ANALYSIS_BBOX_DRAW_CAPTION", True)
|
||||
|
||||
@property
|
||||
def ANALYSIS_BBOX_RESIZE(self) -> Optional[float]:
|
||||
"""Analaysis bbox resize value"""
|
||||
resize = self._get_float("ANALYSIS_BBOX_RESIZE", -1.0)
|
||||
if resize == -1.0:
|
||||
return None
|
||||
return resize
|
||||
|
||||
@property
|
||||
def ANALYSIS_BBOX_FORMAT(self) -> str:
|
||||
"""The format for analysed pages with bboxes drawn on them. Default is 'png'."""
|
||||
return self._get_string("ANALYSIS_BBOX_FORMAT", "png")
|
||||
|
||||
@property
|
||||
def TEXT_COVERAGE_THRESHOLD(self) -> float:
|
||||
"""the minimum iou between extracted text bboxes and their target inferred element bbox for
|
||||
the inferred element to be considered contaning extracted text"""
|
||||
return self._get_float("TEXT_COVERAGE_THRESHOLD", 0.25)
|
||||
|
||||
|
||||
env_config = ENVConfig()
|
||||
@@ -0,0 +1,207 @@
|
||||
import os
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class Source(Enum):
|
||||
PDFMINER = "pdfminer"
|
||||
OCR_TESSERACT = "ocr_tesseract"
|
||||
OCR_PADDLE = "ocr_paddle"
|
||||
OCR_GOOGLEVISION = "ocr_googlevision"
|
||||
|
||||
|
||||
class OCRMode(Enum):
|
||||
INDIVIDUAL_BLOCKS = "individual_blocks"
|
||||
FULL_PAGE = "entire_page"
|
||||
|
||||
|
||||
class PartitionStrategy:
|
||||
AUTO = "auto"
|
||||
FAST = "fast"
|
||||
OCR_ONLY = "ocr_only"
|
||||
HI_RES = "hi_res"
|
||||
|
||||
|
||||
SORT_MODE_XY_CUT = "xy-cut"
|
||||
SORT_MODE_BASIC = "basic"
|
||||
SORT_MODE_DONT = "dont"
|
||||
|
||||
OCR_AGENT_TESSERACT_OLD = "tesseract"
|
||||
OCR_AGENT_PADDLE_OLD = "paddle"
|
||||
|
||||
OCR_AGENT_TESSERACT = "unstructured.partition.utils.ocr_models.tesseract_ocr.OCRAgentTesseract"
|
||||
OCR_AGENT_PADDLE = "unstructured.partition.utils.ocr_models.paddle_ocr.OCRAgentPaddle"
|
||||
OCR_AGENT_GOOGLEVISION = (
|
||||
"unstructured.partition.utils.ocr_models.google_vision_ocr.OCRAgentGoogleVision"
|
||||
)
|
||||
|
||||
OCR_AGENT_MODULES_WHITELIST = os.getenv(
|
||||
"OCR_AGENT_MODULES_WHITELIST",
|
||||
"unstructured.partition.utils.ocr_models.tesseract_ocr,"
|
||||
"unstructured.partition.utils.ocr_models.paddle_ocr,"
|
||||
"unstructured.partition.utils.ocr_models.google_vision_ocr",
|
||||
).split(",")
|
||||
|
||||
UNSTRUCTURED_INCLUDE_DEBUG_METADATA = os.getenv("UNSTRUCTURED_INCLUDE_DEBUG_METADATA", False)
|
||||
|
||||
# this field is defined by unstructured_pytesseract
|
||||
TESSERACT_TEXT_HEIGHT = "height"
|
||||
|
||||
TESSERACT_LANGUAGES_SPLITTER = "+"
|
||||
|
||||
# source: https://tesseract-ocr.github.io/tessdoc/Data-Files-in-different-versions.html
|
||||
# All languages have been changed to lowercase and have been altered to remove dates and "(contrib)"
|
||||
# Ex: "Greek, Ancient (to 1453) (contrib)" -> "greek, ancient"
|
||||
# Where it seemed appropriate, languages have been split into multiple keys with the same value.
|
||||
# Ex: "greek, modern":"ell", "greek":"ell", "chinese - simplified":"chi_sim", "chinese":"chi_sim",
|
||||
# On tesseract-ocr.github.io, "Spanish" matches with both "spa_old" and "spa".
|
||||
# Here, it only matches with "spa" and "spanish - old":"spa_old" was added.
|
||||
TESSERACT_LANGUAGES_AND_CODES = {
|
||||
"afrikaans": "afr",
|
||||
"amharic": "amh",
|
||||
"arabic": "ara",
|
||||
"assamese": "asm",
|
||||
"azerbaijani": "aze",
|
||||
"azerbaijani - cyrilic": "aze_cyrl",
|
||||
"belarusian": "bel",
|
||||
"bengali": "ben",
|
||||
"tibetan": "bod",
|
||||
"bosnian": "bos",
|
||||
"breton": "bre",
|
||||
"bulgarian": "bul",
|
||||
"catalan; Valencian": "cat",
|
||||
"cebuano": "ceb",
|
||||
"czech": "ces",
|
||||
"chinese - simplified": "chi_sim",
|
||||
"chinese": "chi_sim",
|
||||
"chinese - traditional": "chi_tra",
|
||||
"cherokee": "chr",
|
||||
"corsican": "cos",
|
||||
"welsh": "cym",
|
||||
"danish": "dan",
|
||||
"danish - fraktur": "dan_frak",
|
||||
"german": "deu",
|
||||
"german - fraktur (contrib)": "deu_frak", # "contrib" not removed because it would repeat key
|
||||
"dzongkha": "dzo",
|
||||
"greek, modern": "ell",
|
||||
"greek": "ell",
|
||||
"english": "eng",
|
||||
"english, middle": "enm",
|
||||
"esperanto": "epo",
|
||||
"math / equation detection module": "equ",
|
||||
"estonian": "est",
|
||||
"basque": "eus",
|
||||
"faroese": "fao",
|
||||
"persian": "fas",
|
||||
"filipino (old - tagalog)": "fil",
|
||||
"filipino": "fil",
|
||||
"finnish": "fin",
|
||||
"french": "fra",
|
||||
"german - fraktur": "frk",
|
||||
"french, middle": "frm",
|
||||
"western frisian": "fry",
|
||||
"scottish gaelic": "gla",
|
||||
"irish": "gle",
|
||||
"galician": "glg",
|
||||
"greek, ancient": "grc",
|
||||
"gujarati": "guj",
|
||||
"haitian": "hat",
|
||||
"haitian creole": "hat",
|
||||
"hebrew": "heb",
|
||||
"hindi": "hin",
|
||||
"croatian": "hrv",
|
||||
"hungarian": "hun",
|
||||
"armenian": "hye",
|
||||
"inuktitut": "iku",
|
||||
"indonesian": "ind",
|
||||
"icelandic": "isl",
|
||||
"italian": "ita",
|
||||
"italian - old": "ita_old",
|
||||
"javanese": "jav",
|
||||
"japanese": "jpn",
|
||||
"kannada": "kan",
|
||||
"georgian": "kat",
|
||||
"georgian - old": "kat_old",
|
||||
"kazakh": "kaz",
|
||||
"central khmer": "khm",
|
||||
"kirghiz": "kir",
|
||||
"kyrgyz": "kir",
|
||||
"kurmanji (kurdish - latin script)": "kmr",
|
||||
"korean": "kor",
|
||||
"korean (vertical)": "kor_vert",
|
||||
"kurdish (arabic script)": "kur",
|
||||
"lao": "lao",
|
||||
"latin": "lat",
|
||||
"latvian": "lav",
|
||||
"lithuanian": "lit",
|
||||
"luxembourgish": "ltz",
|
||||
"malayalam": "mal",
|
||||
"marathi": "mar",
|
||||
"macedonian": "mkd",
|
||||
"maltese": "mlt",
|
||||
"mongolian": "mon",
|
||||
"maori": "mri",
|
||||
"malay": "msa",
|
||||
"burmese": "mya",
|
||||
"nepali": "nep",
|
||||
"dutch": "nld",
|
||||
"flemish": "nld",
|
||||
"norwegian": "nor",
|
||||
"occitan": "oci",
|
||||
"oriya": "ori",
|
||||
"orientation and script detection module": "osd",
|
||||
"panjabi": "pan",
|
||||
"punjabi": "pan",
|
||||
"polish": "pol",
|
||||
"portuguese": "por",
|
||||
"pushto": "pus",
|
||||
"pashto": "pus",
|
||||
"quechua": "que",
|
||||
"romanian": "ron",
|
||||
"moldavian": "ron",
|
||||
"moldovan": "ron",
|
||||
"russian": "rus",
|
||||
"sanskrit": "san",
|
||||
"sinhala": "sin",
|
||||
"sinhalese": "sin",
|
||||
"slovak": "slk",
|
||||
"slovak - fraktur": "slk_frak",
|
||||
"slovenian": "slv",
|
||||
"sindhi": "snd",
|
||||
"spanish": "spa",
|
||||
"castilian": "spa",
|
||||
"spanish - old": "spa_old",
|
||||
"castilian - old": "spa_old",
|
||||
"albanian": "sqi",
|
||||
"serbian": "srp",
|
||||
"serbian - latin": "srp_latn",
|
||||
"sundanese": "sun",
|
||||
"swahili": "swa",
|
||||
"swedish": "swe",
|
||||
"syriac": "syr",
|
||||
"tamil": "tam",
|
||||
"tatar": "tat",
|
||||
"telugu": "tel",
|
||||
"tajik": "tgk",
|
||||
"tagalog": "tgl",
|
||||
"thai": "tha",
|
||||
"tigrinya": "tir",
|
||||
"tonga": "ton",
|
||||
"turkish": "tur",
|
||||
"uighur": "uig",
|
||||
"uyghur": "uig",
|
||||
"ukrainian": "ukr",
|
||||
"urdu": "urd",
|
||||
"uzbek": "uzb",
|
||||
"uzbek - cyrilic": "uzb_cyrl",
|
||||
"vietnamese": "vie",
|
||||
"yiddish": "yid",
|
||||
"yoruba": "yor",
|
||||
}
|
||||
|
||||
# 2 ** 31 - 1, max byte size for image data
|
||||
TESSERACT_MAX_SIZE = 2147483647
|
||||
|
||||
# default image colors
|
||||
IMAGE_COLOR_DEPTH = 32
|
||||
|
||||
HTML_MAX_PREDECESSOR_LEN = 15
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,116 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from io import BytesIO
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from google.cloud.vision import Image, ImageAnnotatorClient, ImageContext, Paragraph, TextAnnotation
|
||||
|
||||
from unstructured.logger import logger, trace_logger
|
||||
from unstructured.partition.utils.config import env_config
|
||||
from unstructured.partition.utils.constants import Source
|
||||
from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from PIL import Image as PILImage
|
||||
from unstructured_inference.inference.elements import TextRegion, TextRegions
|
||||
from unstructured_inference.inference.layoutelement import LayoutElements
|
||||
|
||||
|
||||
class OCRAgentGoogleVision(OCRAgent):
|
||||
"""OCR service implementation for Google Vision API."""
|
||||
|
||||
def __init__(self, language: Optional[str] = None) -> None:
|
||||
self.language = language
|
||||
client_options = {}
|
||||
api_endpoint = env_config.GOOGLEVISION_API_ENDPOINT
|
||||
if api_endpoint:
|
||||
logger.info(f"Using Google Vision OCR with endpoint {api_endpoint}")
|
||||
client_options["api_endpoint"] = api_endpoint
|
||||
else:
|
||||
logger.info("Using Google Vision OCR with default endpoint")
|
||||
self.client = ImageAnnotatorClient(client_options=client_options)
|
||||
|
||||
def is_text_sorted(self) -> bool:
|
||||
return True
|
||||
|
||||
def get_text_from_image(self, image: PILImage.Image) -> str:
|
||||
image_context = ImageContext(language_hints=[self.language]) if self.language else None
|
||||
with BytesIO() as buffer:
|
||||
image.save(buffer, format="PNG")
|
||||
response = self.client.document_text_detection(
|
||||
image=Image(content=buffer.getvalue()), image_context=image_context
|
||||
)
|
||||
document = response.full_text_annotation
|
||||
assert isinstance(document, TextAnnotation)
|
||||
return document.text
|
||||
|
||||
def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
|
||||
trace_logger.detail("Processing entire page OCR with Google Vision API...")
|
||||
image_context = ImageContext(language_hints=[self.language]) if self.language else None
|
||||
with BytesIO() as buffer:
|
||||
image.save(buffer, format="PNG")
|
||||
response = self.client.document_text_detection(
|
||||
image=Image(content=buffer.getvalue()), image_context=image_context
|
||||
)
|
||||
document = response.full_text_annotation
|
||||
assert isinstance(document, TextAnnotation)
|
||||
regions = self._parse_regions(document)
|
||||
return regions
|
||||
|
||||
def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
|
||||
from unstructured.partition.pdf_image.inference_utils import (
|
||||
build_layout_elements_from_ocr_regions,
|
||||
)
|
||||
|
||||
ocr_regions = self.get_layout_from_image(
|
||||
image,
|
||||
)
|
||||
ocr_text = self.get_text_from_image(
|
||||
image,
|
||||
)
|
||||
return build_layout_elements_from_ocr_regions(
|
||||
ocr_regions=ocr_regions,
|
||||
ocr_text=ocr_text,
|
||||
group_by_ocr_text=False,
|
||||
)
|
||||
|
||||
def _parse_regions(self, ocr_data: TextAnnotation) -> TextRegions:
|
||||
from unstructured_inference.inference.elements import TextRegions
|
||||
|
||||
from unstructured.partition.pdf_image.inference_utils import build_text_region_from_coords
|
||||
|
||||
text_regions: list[TextRegion] = []
|
||||
for page_idx, page in enumerate(ocr_data.pages):
|
||||
for block in page.blocks:
|
||||
for paragraph in block.paragraphs:
|
||||
vertices = paragraph.bounding_box.vertices
|
||||
x1, y1 = vertices[0].x, vertices[0].y
|
||||
x2, y2 = vertices[2].x, vertices[2].y
|
||||
text_region = build_text_region_from_coords(
|
||||
x1,
|
||||
y1,
|
||||
x2,
|
||||
y2,
|
||||
text=self._get_text_from_paragraph(paragraph),
|
||||
source=Source.OCR_GOOGLEVISION,
|
||||
)
|
||||
text_regions.append(text_region)
|
||||
return TextRegions.from_list(text_regions)
|
||||
|
||||
def _get_text_from_paragraph(self, paragraph: Paragraph) -> str:
|
||||
breaks = TextAnnotation.DetectedBreak.BreakType
|
||||
para = ""
|
||||
line = ""
|
||||
for word in paragraph.words:
|
||||
for symbol in word.symbols:
|
||||
line += symbol.text
|
||||
if symbol.property.detected_break.type_ == breaks.SPACE:
|
||||
line += " "
|
||||
if symbol.property.detected_break.type_ == breaks.EOL_SURE_SPACE:
|
||||
line += " "
|
||||
para += line
|
||||
line = ""
|
||||
if symbol.property.detected_break.type_ == breaks.LINE_BREAK:
|
||||
para += line
|
||||
line = ""
|
||||
return para
|
||||
@@ -0,0 +1,96 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import importlib
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from unstructured.logger import logger
|
||||
from unstructured.partition.utils.config import env_config
|
||||
from unstructured.partition.utils.constants import (
|
||||
OCR_AGENT_MODULES_WHITELIST,
|
||||
OCR_AGENT_PADDLE,
|
||||
OCR_AGENT_PADDLE_OLD,
|
||||
OCR_AGENT_TESSERACT,
|
||||
OCR_AGENT_TESSERACT_OLD,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from PIL import Image as PILImage
|
||||
from unstructured_inference.inference.elements import TextRegions
|
||||
from unstructured_inference.inference.layoutelement import LayoutElements
|
||||
|
||||
|
||||
class OCRAgent(ABC):
|
||||
"""Defines the interface for an Optical Character Recognition (OCR) service."""
|
||||
|
||||
@classmethod
|
||||
def get_agent(cls, language: str) -> OCRAgent:
|
||||
"""Get the configured OCRAgent instance.
|
||||
|
||||
The OCR package used by the agent is determined by the `OCR_AGENT` environment variable.
|
||||
"""
|
||||
ocr_agent_cls_qname = cls._get_ocr_agent_cls_qname()
|
||||
return cls.get_instance(ocr_agent_cls_qname, language)
|
||||
|
||||
@staticmethod
|
||||
@functools.lru_cache(maxsize=env_config.OCR_AGENT_CACHE_SIZE)
|
||||
def get_instance(ocr_agent_module: str, language: str) -> "OCRAgent":
|
||||
module_name, class_name = ocr_agent_module.rsplit(".", 1)
|
||||
if module_name not in OCR_AGENT_MODULES_WHITELIST:
|
||||
raise ValueError(
|
||||
f"Environment variable OCR_AGENT module name {module_name} must be set to a "
|
||||
f"whitelisted module part of {OCR_AGENT_MODULES_WHITELIST}."
|
||||
)
|
||||
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
loaded_class = getattr(module, class_name)
|
||||
return loaded_class(language)
|
||||
except (ImportError, AttributeError) as e:
|
||||
logger.error(f"Failed to get OCRAgent instance: {e}")
|
||||
raise RuntimeError(
|
||||
"Could not get the OCRAgent instance. Please check the OCR package and the "
|
||||
"OCR_AGENT environment variable."
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_text_from_image(self, image: PILImage.Image) -> str:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def is_text_sorted(self) -> bool:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _get_ocr_agent_cls_qname() -> str:
|
||||
"""Get the fully-qualified class name of the configured OCR agent.
|
||||
|
||||
The qualified name (qname) looks like:
|
||||
"unstructured.partition.utils.ocr_models.tesseract_ocr.OCRAgentTesseract"
|
||||
|
||||
The qname provides the full module address and class name of the OCR agent.
|
||||
"""
|
||||
ocr_agent_qname = env_config.OCR_AGENT
|
||||
|
||||
# -- map legacy method of setting OCR agent by key-name to full qname --
|
||||
qnames_by_keyname = {
|
||||
OCR_AGENT_TESSERACT_OLD: OCR_AGENT_TESSERACT,
|
||||
OCR_AGENT_PADDLE_OLD: OCR_AGENT_PADDLE,
|
||||
}
|
||||
if qname_mapped_from_keyname := qnames_by_keyname.get(ocr_agent_qname.lower()):
|
||||
logger.warning(
|
||||
f"OCR agent name {ocr_agent_qname} is outdated and will be removed in a future"
|
||||
f" release; please use {qname_mapped_from_keyname} instead"
|
||||
)
|
||||
return qname_mapped_from_keyname
|
||||
|
||||
return ocr_agent_qname
|
||||
@@ -0,0 +1,144 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image as PILImage
|
||||
|
||||
from unstructured.documents.elements import ElementType
|
||||
from unstructured.logger import logger, trace_logger
|
||||
from unstructured.partition.utils.constants import Source
|
||||
from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from unstructured_inference.inference.elements import TextRegion, TextRegions
|
||||
from unstructured_inference.inference.layoutelement import LayoutElements
|
||||
|
||||
|
||||
class OCRAgentPaddle(OCRAgent):
|
||||
"""OCR service implementation for PaddleOCR."""
|
||||
|
||||
def __init__(self, language: str = "en"):
|
||||
self.agent = self.load_agent(language)
|
||||
|
||||
def load_agent(self, language: str):
|
||||
"""Loads the PaddleOCR agent as a global variable to ensure that we only load it once."""
|
||||
|
||||
import paddle
|
||||
from unstructured_paddleocr import PaddleOCR
|
||||
|
||||
# Disable signal handlers at C++ level upon failing
|
||||
# ref: https://www.paddlepaddle.org.cn/documentation/docs/en/api/paddle/
|
||||
# disable_signal_handler_en.html#disable-signal-handler
|
||||
paddle.disable_signal_handler()
|
||||
# Use paddlepaddle-gpu if there is gpu device available
|
||||
gpu_available = paddle.device.cuda.device_count() > 0
|
||||
if gpu_available:
|
||||
logger.info(f"Loading paddle with GPU on language={language}...")
|
||||
else:
|
||||
logger.info(f"Loading paddle with CPU on language={language}...")
|
||||
try:
|
||||
# Enable MKL-DNN for paddle to speed up OCR if OS supports it
|
||||
# ref: https://paddle-inference.readthedocs.io/en/master/
|
||||
# api_reference/cxx_api_doc/Config/CPUConfig.html
|
||||
paddle_ocr = PaddleOCR(
|
||||
use_angle_cls=True,
|
||||
use_gpu=gpu_available,
|
||||
lang=language,
|
||||
enable_mkldnn=True,
|
||||
show_log=False,
|
||||
)
|
||||
except AttributeError:
|
||||
paddle_ocr = PaddleOCR(
|
||||
use_angle_cls=True,
|
||||
use_gpu=gpu_available,
|
||||
lang=language,
|
||||
enable_mkldnn=False,
|
||||
show_log=False,
|
||||
)
|
||||
return paddle_ocr
|
||||
|
||||
def get_text_from_image(self, image: PILImage.Image) -> str:
|
||||
ocr_regions = self.get_layout_from_image(image)
|
||||
return "\n\n".join(ocr_regions.texts)
|
||||
|
||||
def is_text_sorted(self):
|
||||
return False
|
||||
|
||||
def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
|
||||
"""Get the OCR regions from image as a list of text regions with paddle."""
|
||||
|
||||
trace_logger.detail("Processing entire page OCR with paddle...")
|
||||
|
||||
# TODO(yuming): pass in language parameter once we
|
||||
# have the mapping for paddle lang code
|
||||
# see CORE-2034
|
||||
ocr_data = self.agent.ocr(np.array(image), cls=True)
|
||||
ocr_regions = self.parse_data(ocr_data)
|
||||
|
||||
return ocr_regions
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
|
||||
ocr_regions = self.get_layout_from_image(image)
|
||||
|
||||
# NOTE(christine): For paddle, there is no difference in `ocr_layout` and `ocr_text` in
|
||||
# terms of grouping because we get ocr_text from `ocr_layout, so the first two grouping
|
||||
# and merging steps are not necessary.
|
||||
return LayoutElements(
|
||||
element_coords=ocr_regions.element_coords,
|
||||
texts=ocr_regions.texts,
|
||||
element_class_ids=np.zeros(ocr_regions.texts.shape),
|
||||
element_class_id_map={0: ElementType.UNCATEGORIZED_TEXT},
|
||||
)
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def parse_data(self, ocr_data: list[Any]) -> TextRegions:
|
||||
"""Parse the OCR result data to extract a list of TextRegion objects from paddle.
|
||||
|
||||
The function processes the OCR result dictionary, looking for bounding
|
||||
box information and associated text to create instances of the TextRegion
|
||||
class, which are then appended to a list.
|
||||
|
||||
Parameters:
|
||||
- ocr_data (list): A list containing the OCR result data
|
||||
|
||||
Returns:
|
||||
- TextRegions:
|
||||
TextRegions object, containing data from all text regions in numpy arrays; each row
|
||||
represents a detected text region within the OCR-ed image.
|
||||
|
||||
Note:
|
||||
- An empty string or a None value for the 'text' key in the input
|
||||
dictionary will result in its associated bounding box being ignored.
|
||||
"""
|
||||
|
||||
from unstructured_inference.inference.elements import TextRegions
|
||||
|
||||
from unstructured.partition.pdf_image.inference_utils import build_text_region_from_coords
|
||||
|
||||
text_regions: list[TextRegion] = []
|
||||
for idx in range(len(ocr_data)):
|
||||
res = ocr_data[idx]
|
||||
if not res:
|
||||
continue
|
||||
|
||||
for line in res:
|
||||
x1 = min([i[0] for i in line[0]])
|
||||
y1 = min([i[1] for i in line[0]])
|
||||
x2 = max([i[0] for i in line[0]])
|
||||
y2 = max([i[1] for i in line[0]])
|
||||
text = line[1][0]
|
||||
if not text:
|
||||
continue
|
||||
cleaned_text = text.strip()
|
||||
if cleaned_text:
|
||||
text_region = build_text_region_from_coords(
|
||||
x1, y1, x2, y2, text=cleaned_text, source=Source.OCR_PADDLE
|
||||
)
|
||||
text_regions.append(text_region)
|
||||
|
||||
# FIXME (yao): find out if paddle supports a vectorized output format so we can skip the
|
||||
# step of parsing a list
|
||||
return TextRegions.from_list(text_regions)
|
||||
@@ -0,0 +1,260 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import unstructured_pytesseract
|
||||
from lxml import etree
|
||||
from PIL import Image as PILImage
|
||||
|
||||
from unstructured.logger import trace_logger
|
||||
from unstructured.partition.utils.config import env_config
|
||||
from unstructured.partition.utils.constants import (
|
||||
IMAGE_COLOR_DEPTH,
|
||||
TESSERACT_MAX_SIZE,
|
||||
TESSERACT_TEXT_HEIGHT,
|
||||
Source,
|
||||
)
|
||||
from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from unstructured_inference.inference.elements import TextRegions
|
||||
from unstructured_inference.inference.layoutelement import LayoutElements
|
||||
|
||||
_RE_X_CONF = re.compile(r"x_conf (\d+\.\d+)")
|
||||
|
||||
# -- force tesseract to be single threaded, otherwise we see major performance problems --
|
||||
if "OMP_THREAD_LIMIT" not in os.environ:
|
||||
os.environ["OMP_THREAD_LIMIT"] = "1"
|
||||
|
||||
|
||||
class OCRAgentTesseract(OCRAgent):
|
||||
"""OCR service implementation for Tesseract."""
|
||||
|
||||
hocr_namespace = {"h": "http://www.w3.org/1999/xhtml"}
|
||||
|
||||
def __init__(self, language: str = "eng"):
|
||||
self.language = language
|
||||
|
||||
def is_text_sorted(self):
|
||||
return True
|
||||
|
||||
def get_text_from_image(self, image: PILImage.Image) -> str:
|
||||
return unstructured_pytesseract.image_to_string(np.array(image), lang=self.language)
|
||||
|
||||
def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
|
||||
"""Get the OCR regions from image as a list of text regions with tesseract."""
|
||||
|
||||
trace_logger.detail("Processing entire page OCR with tesseract...")
|
||||
zoom = 1
|
||||
ocr_df: pd.DataFrame = self.image_to_data_with_character_confidence_filter(
|
||||
np.array(image),
|
||||
lang=self.language,
|
||||
character_confidence_threshold=env_config.TESSERACT_CHARACTER_CONFIDENCE_THRESHOLD,
|
||||
)
|
||||
ocr_df = ocr_df.dropna()
|
||||
|
||||
# tesseract performance degrades when the text height is out of the preferred zone so we
|
||||
# zoom the image (in or out depending on estimated text height) for optimum OCR results
|
||||
# but this needs to be evaluated based on actual use case as the optimum scaling also
|
||||
# depend on type of characters (font, language, etc); be careful about this
|
||||
# functionality
|
||||
text_height = ocr_df[TESSERACT_TEXT_HEIGHT].quantile(
|
||||
env_config.TESSERACT_TEXT_HEIGHT_QUANTILE
|
||||
)
|
||||
if (
|
||||
text_height < env_config.TESSERACT_MIN_TEXT_HEIGHT
|
||||
or text_height > env_config.TESSERACT_MAX_TEXT_HEIGHT
|
||||
):
|
||||
max_zoom = max(
|
||||
0,
|
||||
np.round(np.sqrt(TESSERACT_MAX_SIZE / np.prod(image.size) / IMAGE_COLOR_DEPTH), 1),
|
||||
)
|
||||
# rounding avoids unnecessary precision and potential numerical issues associated
|
||||
# with numbers very close to 1 inside cv2 image processing
|
||||
zoom = min(
|
||||
np.round(env_config.TESSERACT_OPTIMUM_TEXT_HEIGHT / text_height, 1),
|
||||
max_zoom,
|
||||
)
|
||||
ocr_df = self.image_to_data_with_character_confidence_filter(
|
||||
np.array(zoom_image(image, zoom)),
|
||||
lang=self.language,
|
||||
character_confidence_threshold=env_config.TESSERACT_CHARACTER_CONFIDENCE_THRESHOLD,
|
||||
)
|
||||
ocr_df = ocr_df.dropna()
|
||||
ocr_regions = self.parse_data(ocr_df, zoom=zoom)
|
||||
|
||||
return ocr_regions
|
||||
|
||||
def image_to_data_with_character_confidence_filter(
|
||||
self,
|
||||
image: np.ndarray,
|
||||
lang: str = "eng",
|
||||
config: str = "",
|
||||
character_confidence_threshold: float = 0.0,
|
||||
) -> pd.DataFrame:
|
||||
hocr: str = unstructured_pytesseract.image_to_pdf_or_hocr(
|
||||
image,
|
||||
lang=lang,
|
||||
config="-c hocr_char_boxes=1 " + config,
|
||||
extension="hocr",
|
||||
)
|
||||
ocr_df = self.hocr_to_dataframe(hocr, character_confidence_threshold)
|
||||
return ocr_df
|
||||
|
||||
def hocr_to_dataframe(
|
||||
self, hocr: str, character_confidence_threshold: float = 0.0
|
||||
) -> pd.DataFrame:
|
||||
df_entries = []
|
||||
|
||||
if not hocr:
|
||||
return pd.DataFrame(df_entries, columns=["left", "top", "width", "height", "text"])
|
||||
|
||||
root = etree.fromstring(hocr)
|
||||
word_spans = root.findall('.//h:span[@class="ocrx_word"]', self.hocr_namespace)
|
||||
|
||||
for word_span in word_spans:
|
||||
word_title = word_span.get("title", "")
|
||||
bbox_match = re.search(r"bbox (\d+) (\d+) (\d+) (\d+)", word_title)
|
||||
|
||||
text = self.extract_word_from_hocr(
|
||||
word=word_span, character_confidence_threshold=character_confidence_threshold
|
||||
)
|
||||
if text and bbox_match:
|
||||
word_bbox = list(map(int, bbox_match.groups()))
|
||||
left, top, right, bottom = word_bbox
|
||||
df_entries.append(
|
||||
{
|
||||
"left": left,
|
||||
"top": top,
|
||||
"right": right,
|
||||
"bottom": bottom,
|
||||
"text": text,
|
||||
}
|
||||
)
|
||||
ocr_df = pd.DataFrame(df_entries, columns=["left", "top", "right", "bottom", "text"])
|
||||
|
||||
ocr_df["width"] = ocr_df["right"] - ocr_df["left"]
|
||||
ocr_df["height"] = ocr_df["bottom"] - ocr_df["top"]
|
||||
|
||||
ocr_df = ocr_df.drop(columns=["right", "bottom"])
|
||||
return ocr_df
|
||||
|
||||
def extract_word_from_hocr(
|
||||
self, word: etree.Element, character_confidence_threshold: float = 0.0
|
||||
) -> str:
|
||||
"""Extracts a word from an hOCR word tag, filtering out characters with low confidence."""
|
||||
|
||||
character_spans = word.findall('.//h:span[@class="ocrx_cinfo"]', self.hocr_namespace)
|
||||
if len(character_spans) == 0:
|
||||
return ""
|
||||
|
||||
chars = []
|
||||
for character_span in character_spans:
|
||||
char = character_span.text
|
||||
|
||||
char_title = character_span.get("title", "")
|
||||
conf_match = _RE_X_CONF.search(char_title)
|
||||
|
||||
if not (char and conf_match):
|
||||
continue
|
||||
|
||||
character_probability = float(conf_match.group(1)) / 100
|
||||
|
||||
if character_probability >= character_confidence_threshold:
|
||||
chars.append(char)
|
||||
|
||||
return "".join(chars)
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
|
||||
from unstructured.partition.pdf_image.inference_utils import (
|
||||
build_layout_elements_from_ocr_regions,
|
||||
)
|
||||
|
||||
ocr_regions = self.get_layout_from_image(image)
|
||||
|
||||
# NOTE(christine): For tesseract, the ocr_text returned by
|
||||
# `unstructured_pytesseract.image_to_string()` doesn't contain bounding box data but is
|
||||
# well grouped. Conversely, the ocr_layout returned by parsing
|
||||
# `unstructured_pytesseract.image_to_data()` contains bounding box data but is not well
|
||||
# grouped. Therefore, we need to first group the `ocr_layout` by `ocr_text` and then merge
|
||||
# the text regions in each group to create a list of layout elements.
|
||||
|
||||
ocr_text = self.get_text_from_image(image)
|
||||
|
||||
return build_layout_elements_from_ocr_regions(
|
||||
ocr_regions=ocr_regions,
|
||||
ocr_text=ocr_text,
|
||||
group_by_ocr_text=True,
|
||||
)
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def parse_data(self, ocr_data: pd.DataFrame, zoom: float = 1) -> TextRegions:
|
||||
"""Parse the OCR result data to extract a list of TextRegion objects from tesseract.
|
||||
|
||||
The function processes the OCR result data frame, looking for bounding
|
||||
box information and associated text to create instances of the TextRegion
|
||||
class, which are then appended to a list.
|
||||
|
||||
Parameters:
|
||||
- ocr_data (pd.DataFrame):
|
||||
A Pandas DataFrame containing the OCR result data.
|
||||
It should have columns like 'text', 'left', 'top', 'width', and 'height'.
|
||||
|
||||
- zoom (float, optional):
|
||||
A zoom factor to scale the coordinates of the bounding boxes from image scaling.
|
||||
Default is 1.
|
||||
|
||||
Returns:
|
||||
- TextRegions:
|
||||
TextRegions object, containing data from all text regions in numpy arrays; each row
|
||||
represents a detected text region within the OCR-ed image.
|
||||
|
||||
Note:
|
||||
- An empty string or a None value for the 'text' key in the input
|
||||
data frame will result in its associated bounding box being ignored.
|
||||
"""
|
||||
|
||||
from unstructured_inference.inference.elements import TextRegions
|
||||
|
||||
if zoom <= 0:
|
||||
zoom = 1
|
||||
|
||||
texts = ocr_data.text.apply(
|
||||
lambda text: str(text) if not isinstance(text, str) else text.strip()
|
||||
).values
|
||||
mask = texts != ""
|
||||
element_coords = ocr_data[["left", "top", "width", "height"]].values
|
||||
element_coords[:, 2] += element_coords[:, 0]
|
||||
element_coords[:, 3] += element_coords[:, 1]
|
||||
element_coords = element_coords.astype(float) / zoom
|
||||
return TextRegions(
|
||||
element_coords=element_coords[mask],
|
||||
texts=texts[mask],
|
||||
sources=np.array([Source.OCR_TESSERACT] * mask.sum()),
|
||||
)
|
||||
|
||||
|
||||
def zoom_image(image: PILImage.Image, zoom: float = 1) -> PILImage.Image:
|
||||
"""scale an image based on the zoom factor using cv2; the scaled image is post processed by
|
||||
dilation then erosion to improve edge sharpness for OCR tasks"""
|
||||
if zoom <= 0:
|
||||
# no zoom but still does dilation and erosion
|
||||
zoom = 1
|
||||
new_image = cv2.resize(
|
||||
cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR),
|
||||
None,
|
||||
fx=zoom,
|
||||
fy=zoom,
|
||||
interpolation=cv2.INTER_CUBIC,
|
||||
)
|
||||
|
||||
# Skip dilation and erosion for 1x1 kernel as they are no-ops
|
||||
|
||||
return PILImage.fromarray(new_image)
|
||||
@@ -0,0 +1,268 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from unstructured.documents.elements import CoordinatesMetadata, Element
|
||||
from unstructured.logger import trace_logger
|
||||
from unstructured.partition.utils.constants import SORT_MODE_BASIC, SORT_MODE_XY_CUT
|
||||
from unstructured.partition.utils.xycut import recursive_xy_cut, recursive_xy_cut_swapped
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from unstructured_inference.inference.elements import TextRegions
|
||||
|
||||
|
||||
def coordinates_to_bbox(coordinates: CoordinatesMetadata) -> tuple[int, int, int, int]:
|
||||
"""
|
||||
Convert coordinates to a bounding box representation.
|
||||
|
||||
Parameters:
|
||||
coordinates (CoordinatesMetadata): Metadata containing points to represent the bounding box.
|
||||
|
||||
Returns:
|
||||
tuple[int, int, int, int]: A tuple representing the bounding box in the format
|
||||
(left, top, right, bottom).
|
||||
"""
|
||||
|
||||
points = coordinates.points
|
||||
left, top = points[0]
|
||||
right, bottom = points[2]
|
||||
return int(left), int(top), int(right), int(bottom)
|
||||
|
||||
|
||||
def shrink_bbox(bbox: tuple[int, int, int, int], shrink_factor) -> tuple[int, int, int, int]:
|
||||
"""
|
||||
Shrink a bounding box by a given shrink factor while maintaining its top and left.
|
||||
|
||||
Parameters:
|
||||
bbox (tuple[int, int, int, int]): The original bounding box represented by
|
||||
(left, top, right, bottom).
|
||||
shrink_factor (float): The factor by which to shrink the bounding box (0.0 to 1.0).
|
||||
|
||||
Returns:
|
||||
tuple[int, int, int, int]: The shrunken bounding box represented by
|
||||
(left, top, right, bottom).
|
||||
"""
|
||||
|
||||
left, top, right, bottom = bbox
|
||||
width = right - left
|
||||
height = bottom - top
|
||||
new_width = width * shrink_factor
|
||||
new_height = height * shrink_factor
|
||||
dw = width - new_width
|
||||
dh = height - new_height
|
||||
|
||||
new_right = right - dw
|
||||
new_bottom = bottom - dh
|
||||
return int(left), int(top), int(new_right), int(new_bottom)
|
||||
|
||||
|
||||
def coord_has_valid_points(coordinates: CoordinatesMetadata) -> bool:
|
||||
"""
|
||||
Verifies all 4 points in a coordinate exist and are positive.
|
||||
"""
|
||||
if not coordinates:
|
||||
return False
|
||||
if len(coordinates.points) != 4:
|
||||
return False
|
||||
for point in coordinates.points:
|
||||
if len(point) != 2:
|
||||
return False
|
||||
try:
|
||||
if point[0] < 0 or point[1] < 0:
|
||||
return False
|
||||
except TypeError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def bbox_is_valid(bbox: Any) -> bool:
|
||||
"""
|
||||
Verifies all 4 values in a bounding box exist and are positive.
|
||||
"""
|
||||
|
||||
if not bbox:
|
||||
return False
|
||||
if len(bbox) != 4:
|
||||
return False
|
||||
for v in bbox:
|
||||
try:
|
||||
if v < 0:
|
||||
return False
|
||||
except TypeError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def sort_page_elements(
|
||||
page_elements: list[Element],
|
||||
sort_mode: str = SORT_MODE_XY_CUT,
|
||||
shrink_factor: float = 0.9,
|
||||
xy_cut_primary_direction: str = "x",
|
||||
) -> list[Element]:
|
||||
"""
|
||||
Sorts a list of page elements based on the specified sorting mode.
|
||||
|
||||
Parameters:
|
||||
- page_elements (list[Element]): A list of elements representing parts of a page. Each element
|
||||
should have metadata containing coordinates.
|
||||
- sort_mode (str, optional): The mode by which the elements will be sorted. Default is
|
||||
SORT_MODE_XY_CUT.
|
||||
- SORT_MODE_XY_CUT: Sorts elements based on XY-cut sorting approach. Requires the
|
||||
recursive_xy_cut function and coordinates_to_bbox function to be defined. And requires all
|
||||
elements to have valid cooridnates
|
||||
- SORT_MODE_BASIC: Sorts elements based on their coordinates. Elements without coordinates
|
||||
will be pushed to the end.
|
||||
- If an unrecognized sort_mode is provided, the function returns the elements as-is.
|
||||
|
||||
Returns:
|
||||
- list[Element]: A list of sorted page elements.
|
||||
"""
|
||||
|
||||
shrink_factor = float(
|
||||
os.environ.get("UNSTRUCTURED_XY_CUT_BBOX_SHRINK_FACTOR", shrink_factor),
|
||||
)
|
||||
|
||||
xy_cut_primary_direction = os.environ.get(
|
||||
"UNSTRUCTURED_XY_CUT_PRIMARY_DIRECTION",
|
||||
xy_cut_primary_direction,
|
||||
)
|
||||
|
||||
if not page_elements:
|
||||
return []
|
||||
|
||||
coordinates_list = [el.metadata.coordinates for el in page_elements]
|
||||
|
||||
def _coords_ok(strict_points: bool):
|
||||
warned = False
|
||||
|
||||
for coord in coordinates_list:
|
||||
if coord is None or not coord.points:
|
||||
trace_logger.detail( # type: ignore
|
||||
"some or all elements are missing coordinates, skipping sort",
|
||||
)
|
||||
return False
|
||||
elif not coord_has_valid_points(coord):
|
||||
if not warned:
|
||||
trace_logger.detail(f"coord {coord} does not have valid points") # type: ignore
|
||||
warned = True
|
||||
if strict_points:
|
||||
return False
|
||||
return True
|
||||
|
||||
if sort_mode == SORT_MODE_XY_CUT:
|
||||
if not _coords_ok(strict_points=True):
|
||||
return page_elements
|
||||
shrunken_bboxes = []
|
||||
for coords in coordinates_list:
|
||||
bbox = coordinates_to_bbox(coords)
|
||||
shrunken_bbox = shrink_bbox(bbox, shrink_factor)
|
||||
shrunken_bboxes.append(shrunken_bbox)
|
||||
|
||||
res: list[int] = []
|
||||
xy_cut_sorting_func = (
|
||||
recursive_xy_cut_swapped if xy_cut_primary_direction == "x" else recursive_xy_cut
|
||||
)
|
||||
xy_cut_sorting_func(
|
||||
np.asarray(shrunken_bboxes).astype(int),
|
||||
np.arange(len(shrunken_bboxes)),
|
||||
res,
|
||||
)
|
||||
sorted_page_elements = [page_elements[i] for i in res]
|
||||
elif sort_mode == SORT_MODE_BASIC:
|
||||
if not _coords_ok(strict_points=False):
|
||||
return page_elements
|
||||
sorted_page_elements = sorted(
|
||||
page_elements,
|
||||
key=lambda el: (
|
||||
el.metadata.coordinates.points[0][1] if el.metadata.coordinates else float("inf"),
|
||||
el.metadata.coordinates.points[0][0] if el.metadata.coordinates else float("inf"),
|
||||
),
|
||||
)
|
||||
else:
|
||||
sorted_page_elements = page_elements
|
||||
|
||||
return sorted_page_elements
|
||||
|
||||
|
||||
def sort_bboxes_by_xy_cut(
|
||||
bboxes,
|
||||
shrink_factor: float = 0.9,
|
||||
xy_cut_primary_direction: str = "x",
|
||||
):
|
||||
"""Sort bounding boxes using XY-cut algorithm."""
|
||||
|
||||
shrunken_bboxes = []
|
||||
for bbox in bboxes:
|
||||
shrunken_bbox = shrink_bbox(bbox, shrink_factor)
|
||||
shrunken_bboxes.append(shrunken_bbox)
|
||||
|
||||
res: list[int] = []
|
||||
xy_cut_sorting_func = (
|
||||
recursive_xy_cut_swapped if xy_cut_primary_direction == "x" else recursive_xy_cut
|
||||
)
|
||||
xy_cut_sorting_func(
|
||||
np.asarray(shrunken_bboxes).astype(int),
|
||||
np.arange(len(shrunken_bboxes)),
|
||||
res,
|
||||
)
|
||||
return res
|
||||
|
||||
|
||||
def sort_text_regions(
|
||||
elements: TextRegions,
|
||||
sort_mode: str = SORT_MODE_XY_CUT,
|
||||
shrink_factor: float = 0.9,
|
||||
xy_cut_primary_direction: str = "x",
|
||||
) -> TextRegions:
|
||||
"""Sort a list of TextRegion elements based on the specified sorting mode."""
|
||||
|
||||
if not elements:
|
||||
return elements
|
||||
|
||||
bboxes = elements.element_coords
|
||||
|
||||
def _bboxes_ok(strict_points: bool):
|
||||
if np.isnan(bboxes).any():
|
||||
trace_logger.detail( # type: ignore
|
||||
"some or all elements are missing bboxes, skipping sort",
|
||||
)
|
||||
return False
|
||||
|
||||
if bboxes.shape[1] != 4 or np.where(bboxes < 0)[0].size:
|
||||
trace_logger.detail("at least one bbox contains invalid values") # type: ignore
|
||||
if strict_points:
|
||||
return False
|
||||
return True
|
||||
|
||||
if sort_mode == SORT_MODE_XY_CUT:
|
||||
if not _bboxes_ok(strict_points=True):
|
||||
return elements
|
||||
|
||||
shrink_factor = float(
|
||||
os.environ.get("UNSTRUCTURED_XY_CUT_BBOX_SHRINK_FACTOR", shrink_factor),
|
||||
)
|
||||
|
||||
xy_cut_primary_direction = os.environ.get(
|
||||
"UNSTRUCTURED_XY_CUT_PRIMARY_DIRECTION",
|
||||
xy_cut_primary_direction,
|
||||
)
|
||||
|
||||
res = sort_bboxes_by_xy_cut(
|
||||
bboxes=bboxes,
|
||||
shrink_factor=shrink_factor,
|
||||
xy_cut_primary_direction=xy_cut_primary_direction,
|
||||
)
|
||||
sorted_elements = elements.slice(res)
|
||||
elif sort_mode == SORT_MODE_BASIC:
|
||||
# NOTE (yao): lexsort order is revese from the input sequence; so below is first sort by y1,
|
||||
# then x1, then y2, lastly x2
|
||||
sorted_elements = elements.slice(
|
||||
np.lexsort((elements.x2, elements.y2, elements.x1, elements.y1))
|
||||
)
|
||||
else:
|
||||
sorted_elements = elements
|
||||
|
||||
return sorted_elements
|
||||
@@ -0,0 +1,327 @@
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
"""
|
||||
|
||||
This module contains the implementation of the XY-Cut sorting approach
|
||||
from: https://github.com/Sanster/xy-cut
|
||||
|
||||
"""
|
||||
|
||||
|
||||
def projection_by_bboxes(boxes: np.ndarray, axis: int) -> np.ndarray:
|
||||
"""
|
||||
Obtain the projection histogram through a set of bboxes and finally output it in per-pixel form
|
||||
|
||||
Args:
|
||||
boxes: [N, 4]
|
||||
axis: 0 - x coordinates are projected in the horizontal direction, 1 - y coordinates
|
||||
are projected in the vertical direction
|
||||
|
||||
Returns:
|
||||
1D projection histogram, the length is the maximum value of the projection direction
|
||||
coordinate (we don’t need the actual side length of the picture because we just
|
||||
want to find the interval of the text box)
|
||||
"""
|
||||
|
||||
assert axis in [0, 1]
|
||||
length = np.max(boxes[:, axis::2])
|
||||
res = np.zeros(length, dtype=int)
|
||||
# TODO: how to remove for loop?
|
||||
for start, end in boxes[:, axis::2]:
|
||||
res[start:end] += 1
|
||||
return res
|
||||
|
||||
|
||||
# from: https://dothinking.github.io/2021-06-19-%E9%80%92%E5%BD%92%E6%8A%95%E5%BD%B1
|
||||
# %E5%88%86%E5%89%B2%E7%AE%97%E6%B3%95/#:~:text=%E9%80%92%E5%BD%92%E6%8A%95%E5%BD%B1
|
||||
# %E5%88%86%E5%89%B2%EF%BC%88Recursive%20XY,%EF%BC%8C%E5%8F%AF%E4%BB%A5%E5%88%92
|
||||
# %E5%88%86%E6%AE%B5%E8%90%BD%E3%80%81%E8%A1%8C%E3%80%82
|
||||
def split_projection_profile(arr_values: np.ndarray, min_value: float, min_gap: float):
|
||||
"""Split projection profile:
|
||||
|
||||
```
|
||||
┌──┐
|
||||
arr_values │ │ ┌─┐───
|
||||
┌──┐ │ │ │ │ |
|
||||
│ │ │ │ ┌───┐ │ │min_value
|
||||
│ │<- min_gap ->│ │ │ │ │ │ |
|
||||
────┴──┴─────────────┴──┴─┴───┴─┴─┴─┴───
|
||||
0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
|
||||
```
|
||||
|
||||
Args:
|
||||
arr_values (np.array): 1-d array representing the projection profile.
|
||||
min_value (float): Ignore the profile if `arr_value` is less than `min_value`.
|
||||
min_gap (float): Ignore the gap if less than this value.
|
||||
|
||||
Returns:
|
||||
tuple: Start indexes and end indexes of split groups.
|
||||
"""
|
||||
# all indexes with projection height exceeding the threshold
|
||||
arr_index = np.where(arr_values > min_value)[0]
|
||||
if not len(arr_index):
|
||||
return
|
||||
|
||||
# find zero intervals between adjacent projections
|
||||
# | | ||
|
||||
# ||||<- zero-interval -> |||||
|
||||
arr_diff = arr_index[1:] - arr_index[0:-1]
|
||||
arr_diff_index = np.where(arr_diff > min_gap)[0]
|
||||
arr_zero_intvl_start = arr_index[arr_diff_index]
|
||||
arr_zero_intvl_end = arr_index[arr_diff_index + 1]
|
||||
|
||||
# convert to index of projection range:
|
||||
# the start index of zero interval is the end index of projection
|
||||
arr_start = np.insert(arr_zero_intvl_end, 0, arr_index[0])
|
||||
arr_end = np.append(arr_zero_intvl_start, arr_index[-1])
|
||||
arr_end += 1 # end index will be excluded as index slice
|
||||
|
||||
return arr_start, arr_end
|
||||
|
||||
|
||||
def recursive_xy_cut(boxes: np.ndarray, indices: np.ndarray, res: List[int]):
|
||||
"""
|
||||
|
||||
Args:
|
||||
boxes: (N, 4)
|
||||
indices: during the recursion process, the index of box in the original data
|
||||
is always represented.
|
||||
res: save output
|
||||
|
||||
"""
|
||||
# project to the y-axis
|
||||
assert len(boxes) == len(indices)
|
||||
|
||||
_indices = boxes[:, 1].argsort()
|
||||
y_sorted_boxes = boxes[_indices]
|
||||
y_sorted_indices = indices[_indices]
|
||||
|
||||
# debug_vis(y_sorted_boxes, y_sorted_indices)
|
||||
|
||||
y_projection = projection_by_bboxes(boxes=y_sorted_boxes, axis=1)
|
||||
pos_y = split_projection_profile(y_projection, 0, 1)
|
||||
if not pos_y:
|
||||
return
|
||||
|
||||
arr_y0, arr_y1 = pos_y
|
||||
for r0, r1 in zip(arr_y0, arr_y1):
|
||||
# [r0, r1] means that the areas with bbox will be divided horizontally, and these areas
|
||||
# will be divided vertically.
|
||||
_indices = (r0 <= y_sorted_boxes[:, 1]) & (y_sorted_boxes[:, 1] < r1)
|
||||
|
||||
y_sorted_boxes_chunk = y_sorted_boxes[_indices]
|
||||
y_sorted_indices_chunk = y_sorted_indices[_indices]
|
||||
|
||||
_indices = y_sorted_boxes_chunk[:, 0].argsort()
|
||||
x_sorted_boxes_chunk = y_sorted_boxes_chunk[_indices]
|
||||
x_sorted_indices_chunk = y_sorted_indices_chunk[_indices]
|
||||
|
||||
# project in the x direction
|
||||
x_projection = projection_by_bboxes(boxes=x_sorted_boxes_chunk, axis=0)
|
||||
pos_x = split_projection_profile(x_projection, 0, 1)
|
||||
if not pos_x:
|
||||
continue
|
||||
|
||||
arr_x0, arr_x1 = pos_x
|
||||
if len(arr_x0) == 1:
|
||||
# x-direction cannot be divided
|
||||
res.extend(x_sorted_indices_chunk)
|
||||
continue
|
||||
|
||||
# can be separated in the x-direction and continue to call recursively
|
||||
for c0, c1 in zip(arr_x0, arr_x1):
|
||||
_indices = (c0 <= x_sorted_boxes_chunk[:, 0]) & (x_sorted_boxes_chunk[:, 0] < c1)
|
||||
recursive_xy_cut(
|
||||
x_sorted_boxes_chunk[_indices],
|
||||
x_sorted_indices_chunk[_indices],
|
||||
res,
|
||||
)
|
||||
|
||||
|
||||
def recursive_xy_cut_swapped(boxes: np.ndarray, indices: np.ndarray, res: List[int]):
|
||||
"""
|
||||
Args:
|
||||
boxes: (N, 4) - Numpy array representing bounding boxes with shape (N, 4)
|
||||
where each row is (left, top, right, bottom)
|
||||
indices: An array representing indices that correspond to boxes in the original data
|
||||
res: A list to save the output results
|
||||
"""
|
||||
|
||||
# Sort the bounding boxes based on x-coordinates (flipped)
|
||||
assert len(boxes) == len(indices)
|
||||
_indices = boxes[:, 0].argsort()
|
||||
x_sorted_boxes = boxes[_indices]
|
||||
x_sorted_indices = indices[_indices]
|
||||
|
||||
# Project the boxes onto the x-axis and split the projection profile
|
||||
x_projection = projection_by_bboxes(boxes=x_sorted_boxes, axis=0)
|
||||
pos_x = split_projection_profile(x_projection, 0, 1)
|
||||
|
||||
if not pos_x:
|
||||
return
|
||||
|
||||
arr_x0, arr_x1 = pos_x
|
||||
|
||||
# Loop over the segments obtained from the x-axis projection
|
||||
for c0, c1 in zip(arr_x0, arr_x1):
|
||||
# Obtain sub-boxes in the x-axis segment
|
||||
_indices = (c0 <= x_sorted_boxes[:, 0]) & (x_sorted_boxes[:, 0] < c1)
|
||||
x_sorted_boxes_chunk = x_sorted_boxes[_indices]
|
||||
x_sorted_indices_chunk = x_sorted_indices[_indices]
|
||||
|
||||
# Sort the sub-boxes based on y-coordinates (flipped)
|
||||
_indices = x_sorted_boxes_chunk[:, 1].argsort()
|
||||
y_sorted_boxes_chunk = x_sorted_boxes_chunk[_indices]
|
||||
y_sorted_indices_chunk = x_sorted_indices_chunk[_indices]
|
||||
|
||||
# Project the sub-boxes onto the y-axis and split the projection profile
|
||||
y_projection = projection_by_bboxes(boxes=y_sorted_boxes_chunk, axis=1)
|
||||
pos_y = split_projection_profile(y_projection, 0, 1)
|
||||
|
||||
if not pos_y:
|
||||
continue
|
||||
|
||||
arr_y0, arr_y1 = pos_y
|
||||
|
||||
if len(arr_y0) == 1:
|
||||
# If there's no splitting along the y-axis, add the indices to the result
|
||||
res.extend(y_sorted_indices_chunk)
|
||||
continue
|
||||
|
||||
# Recursive call for sub-boxes along the y-axis segments
|
||||
for r0, r1 in zip(arr_y0, arr_y1):
|
||||
_indices = (r0 <= y_sorted_boxes_chunk[:, 1]) & (y_sorted_boxes_chunk[:, 1] < r1)
|
||||
recursive_xy_cut_swapped(
|
||||
y_sorted_boxes_chunk[_indices],
|
||||
y_sorted_indices_chunk[_indices],
|
||||
res,
|
||||
)
|
||||
|
||||
|
||||
def points_to_bbox(points):
|
||||
assert len(points) == 8
|
||||
|
||||
# [x1,y1,x2,y2,x3,y3,x4,y4]
|
||||
left = min(points[::2])
|
||||
right = max(points[::2])
|
||||
top = min(points[1::2])
|
||||
bottom = max(points[1::2])
|
||||
|
||||
left = max(left, 0)
|
||||
top = max(top, 0)
|
||||
right = max(right, 0)
|
||||
bottom = max(bottom, 0)
|
||||
return [left, top, right, bottom]
|
||||
|
||||
|
||||
def bbox2points(bbox):
|
||||
left, top, right, bottom = bbox
|
||||
return [left, top, right, top, right, bottom, left, bottom]
|
||||
|
||||
|
||||
@requires_dependencies("cv2")
|
||||
def vis_polygon(img, points, thickness=2, color=None):
|
||||
import cv2
|
||||
|
||||
br2bl_color = color
|
||||
tl2tr_color = color
|
||||
tr2br_color = color
|
||||
bl2tl_color = color
|
||||
cv2.line(
|
||||
img,
|
||||
(points[0][0], points[0][1]),
|
||||
(points[1][0], points[1][1]),
|
||||
color=tl2tr_color,
|
||||
thickness=thickness,
|
||||
)
|
||||
|
||||
cv2.line(
|
||||
img,
|
||||
(points[1][0], points[1][1]),
|
||||
(points[2][0], points[2][1]),
|
||||
color=tr2br_color,
|
||||
thickness=thickness,
|
||||
)
|
||||
|
||||
cv2.line(
|
||||
img,
|
||||
(points[2][0], points[2][1]),
|
||||
(points[3][0], points[3][1]),
|
||||
color=br2bl_color,
|
||||
thickness=thickness,
|
||||
)
|
||||
|
||||
cv2.line(
|
||||
img,
|
||||
(points[3][0], points[3][1]),
|
||||
(points[0][0], points[0][1]),
|
||||
color=bl2tl_color,
|
||||
thickness=thickness,
|
||||
)
|
||||
return img
|
||||
|
||||
|
||||
@requires_dependencies("cv2")
|
||||
def vis_points(
|
||||
img: np.ndarray,
|
||||
points,
|
||||
texts: List[str],
|
||||
color=(0, 200, 0),
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
|
||||
Args:
|
||||
img:
|
||||
points: [N, 8] 8: x1,y1,x2,y2,x3,y3,x4,y4
|
||||
texts:
|
||||
color:
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
import cv2
|
||||
|
||||
points = np.array(points)
|
||||
assert len(texts) == points.shape[0]
|
||||
|
||||
for i, _points in enumerate(points):
|
||||
vis_polygon(img, _points.reshape(-1, 2), thickness=2, color=color)
|
||||
bbox = points_to_bbox(_points)
|
||||
left, top, right, bottom = bbox
|
||||
cx = (left + right) // 2
|
||||
cy = (top + bottom) // 2
|
||||
|
||||
txt = texts[i]
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
cat_size = cv2.getTextSize(txt, font, 0.5, 2)[0]
|
||||
|
||||
img = cv2.rectangle(
|
||||
img,
|
||||
(cx - 5 * len(txt), cy - cat_size[1] - 5),
|
||||
(cx - 5 * len(txt) + cat_size[0], cy - 5),
|
||||
color,
|
||||
-1,
|
||||
)
|
||||
|
||||
img = cv2.putText(
|
||||
img,
|
||||
txt,
|
||||
(cx - 5 * len(txt), cy - 5),
|
||||
font,
|
||||
0.5,
|
||||
(255, 255, 255),
|
||||
thickness=1,
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def vis_polygons_with_index(image, points):
|
||||
texts = [str(i) for i in range(len(points))]
|
||||
res_img = vis_points(image.copy(), points, texts)
|
||||
return res_img
|
||||
Reference in New Issue
Block a user