#!/usr/bin/env python3
"""OCR template."""
from __future__ import annotations
import logging
import os
from enum import Enum
from pathlib import Path
from queue import SimpleQueue
from typing import Any, Final, Iterable, Optional, TypeAlias
import cv2 as cv2
import numpy as np
import numpy.typing as npt
import pandas as pd # type: ignore
import pdf2image # type: ignore
import rerun as rr # pip install rerun-sdk
import rerun.blueprint as rrb
from paddleocr import PPStructure # type: ignore
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes # type: ignore
EXAMPLE_DIR: Final = Path(os.path.dirname(__file__))
DATASET_DIR: Final = EXAMPLE_DIR / "dataset"
SAMPLE_IMAGE_URLs = ["https://storage.googleapis.com/rerun-example-datasets/ocr/paper.png"]
LayoutStructure: TypeAlias = tuple[
list[str], list[str], list[rrb.Spatial2DView], list[rrb.Spatial2DView], list[rrb.Spatial2DView]
]
# Supportive Classes
class Color:
Red = (255, 0, 0)
Green = (0, 255, 0)
Blue = (0, 0, 255)
Yellow = (255, 255, 0)
Cyan = (0, 255, 255)
Magenta = (255, 0, 255)
Purple = (128, 0, 128)
Orange = (255, 165, 0)
"""
LayoutType:
Defines an enumeration for different types of document layout elements, each associated with a unique number, name,
and color. Types:
- UNKNOWN: Default type for undefined or unrecognized elements, represented by purple.
- TITLE: Represents the title of a document, represented by red.
- TEXT: Represents plain text content within the document, represented by green.
- FIGURE: Represents graphical or image content, represented by blue.
- FIGURE_CAPTION: Represents captions for figures, represented by yellow.
- TABLE: Represents tabular data, represented by cyan.
- TABLE_CAPTION: Represents captions for tables, represented by magenta.
- REFERENCE: Represents citation references within the document, also represented by purple.
- Footer: Represents footer of the document, represented as orange.
"""
class LayoutType(Enum):
UNKNOWN = (0, "unknown", Color.Purple)
TITLE = (1, "title", Color.Red)
TEXT = (2, "text", Color.Green)
FIGURE = (3, "figure", Color.Blue)
FIGURE_CAPTION = (4, "figure_caption", Color.Yellow)
TABLE = (5, "table", Color.Cyan)
TABLE_CAPTION = (6, "table_caption", Color.Magenta)
REFERENCE = (7, "reference", Color.Purple)
FOOTER = (8, "footer", Color.Orange)
def __str__(self) -> str:
return str(self.value[1]) # Returns the string part (type)
@property
def number(self) -> int:
return self.value[0] # Returns the numerical identifier
@property
def type(self) -> str:
return self.value[1] # Returns the type
@property
def color(self) -> tuple[int, int, int]:
return self.value[2] # Returns the color
@staticmethod
def get_class_id(text: str) -> int:
try:
return LayoutType[text.upper()].number
except KeyError:
logging.warning(f"Invalid layout type {text}")
return 0
@staticmethod
def get_type(text: str) -> LayoutType:
try:
return LayoutType[text.upper()]
except KeyError:
logging.warning(f"Invalid layout type {text}")
return LayoutType.UNKNOWN
@classmethod
def get_annotation(cls) -> list[tuple[int, str, tuple[int, int, int]]]:
return [(layout.number, layout.type, layout.color) for layout in cls]
"""
Layout Class:
The main purpose of this class is to:
1. Keep track of the layout types (including type, numbering)
2. Save the detections for each layout (text, img or table)
3. Save the bounding box of each detected layout
4. Generate the recovery text document
"""
class Layout:
def __init__(self, show_unknown: bool = False):
self.counts = {layout_type: 0 for layout_type in LayoutType}
self.records: dict[LayoutType, Any] = {layout_type: [] for layout_type in LayoutType}
self.recovery = """"""
self.show_unknown = show_unknown
def add(
self,
layout_type: LayoutType,
bounding_box: list[int],
detections: Optional[Iterable[dict[str, Any]]] = None,
table: Optional[str] = None,
figure: Optional[dict[str, Any]] = None,
) -> None:
if layout_type in LayoutType:
self.counts[layout_type] += 1
name = f"{layout_type}{self.counts[layout_type]}"
logging.info(f"Saved layout type {layout_type} with name: {name}")
self.records[layout_type].append({
"type": layout_type,
"name": name,
"bounding_box": bounding_box,
"detections": detections,
"table": table,
})
if layout_type != LayoutType.UNKNOWN or self.show_unknown: # Discards the unknown layout types detections
path = f"recording://Image/{layout_type.type.title()}/{name.title()}"
self.recovery += f"\n\n## [{name.title()}]({path})\n\n" # Log Type as Heading
# Enhancement - Logged image for Figure type TODO(#6517)
if layout_type == LayoutType.TABLE:
if table:
self.recovery += table # Log details (table)
elif detections:
for index, detection in enumerate(detections):
path_text = f"recording://Image/{layout_type.type.title()}/{name.title()}/Detections/{index}"
self.recovery += f' [{detection["text"]}]({path_text})' # Log details (text)
else:
logging.warning(f"Invalid layout type detected: {layout_type}")
def get_count(self, layout_type: LayoutType) -> int:
if layout_type in LayoutType:
return self.counts[layout_type]
else:
raise ValueError("Invalid layout type")
def get_records(self) -> dict[LayoutType, list[dict[str, Any]]]:
return self.records
def save_all_layouts(self, results: list[dict[str, Any]]) -> None:
for line in results:
self.save_layout_data(line)
for layout_type in LayoutType:
logging.info(f"Number of detections for type {layout_type}: {self.counts[layout_type]}")
def save_layout_data(self, line: dict[str, Any]) -> None:
type = line.get("type", "empty")
box = line.get("bbox", [0, 0, 0, 0])
layout_type = LayoutType.get_type(type)
detections, table, img = [], None, None
if layout_type == LayoutType.TABLE:
table = self.get_table_markdown(line)
elif layout_type == LayoutType.FIGURE:
detections = self.get_detections(line)
img = line.get("img") # Currently not in use
else:
detections = self.get_detections(line)
self.add(layout_type, box, detections=detections, table=table, figure=img)
@staticmethod
def get_detections(line: dict[str, Any]) -> list[dict[str, Any]]:
detections = []
results = line.get("res")
if results is not None:
for i, result in enumerate(results):
text = result.get("text")
confidence = result.get("confidence")
box = result.get("text_region")
x_min, y_min = box[0]
x_max, y_max = box[2]
new_box = [x_min, y_min, x_max, y_max]
detections.append({"id": i, "text": text, "confidence": confidence, "box": new_box})
return detections
# Safely attempt to extract the HTML table from the results
@staticmethod
def get_table_markdown(line: dict[str, Any]) -> str:
try:
html_table = line.get("res", {}).get("html")
if not html_table:
return "No table found."
dataframes = pd.read_html(html_table)
if not dataframes:
return "No data extracted from the table."
markdown_table = dataframes[0].to_markdown()
return markdown_table # type: ignore[no-any-return]
except Exception as e:
return f"Error processing the table: {str(e)}"
def process_layout_records(log_queue: SimpleQueue[Any], layout: Layout, page_path: str) -> LayoutStructure:
paths, detections_paths = [], []
zoom_paths: list[rrb.Spatial2DView] = []
zoom_paths_figures: list[rrb.Spatial2DView] = []
zoom_paths_tables: list[rrb.Spatial2DView] = []
zoom_paths_texts: list[rrb.Spatial2DView] = []
for layout_type in LayoutType:
for record in layout.records[layout_type]:
record_name = record["name"].title()
record_base_path = f"{page_path}/Image/{layout_type.type.title()}/{record_name}"
paths.append(f"-{record_base_path}/**")
detections_paths.append(f"-{record_base_path}/Detections/**")
# Log bounding box
log_queue.put([
"log",
record_base_path,
[
rr.Boxes2D(
array=record["bounding_box"],
array_format=rr.Box2DFormat.XYXY,
labels=[str(layout_type.type)],
class_ids=[str(layout_type.number)],
),
rr.AnyValues(name=record_name),
],
])
log_detections(log_queue, layout_type, record, record_base_path)
# Prepare zoom path views
update_zoom_paths(
layout,
layout_type,
record,
paths,
page_path,
zoom_paths,
zoom_paths_figures,
zoom_paths_tables,
zoom_paths_texts,
)
return paths, detections_paths, zoom_paths_figures, zoom_paths_tables, zoom_paths_texts
def log_detections(log_queue: SimpleQueue, layout_type: LayoutType, record: dict[str, Any], page_path: str) -> None:
if layout_type == LayoutType.TABLE:
log_queue.put([
"log",
f"Extracted{record['name']}",
[rr.TextDocument(record["table"], media_type=rr.MediaType.MARKDOWN)],
])
else:
for detection in record.get("detections", []):
log_queue.put([
"log",
f"{page_path}/Detections/{detection['id']}",
[
rr.Boxes2D(
array=detection["box"], array_format=rr.Box2DFormat.XYXY, class_ids=[str(layout_type.number)]
),
rr.AnyValues(
DetectionID=detection["id"], Text=detection["text"], Confidence=detection["confidence"]
),
],
])
def update_zoom_paths(
layout: Layout,
layout_type: LayoutType,
record: dict[str, Any],
paths: list[str],
page_path: str,
zoom_paths: list[rrb.Spatial2DView],
zoom_paths_figures: list[rrb.Spatial2DView],
zoom_paths_tables: list[rrb.Spatial2DView],
zoom_paths_texts: list[rrb.Spatial2DView],
) -> None:
if layout_type in [LayoutType.FIGURE, LayoutType.TABLE, LayoutType.TEXT]:
current_paths = paths.copy()
current_paths.remove(f"-{page_path}/Image/{layout_type.type.title()}/{record['name'].title()}/**")
bounds = rrb.VisualBounds2D(
x_range=[record["bounding_box"][0] - 10, record["bounding_box"][2] + 10],
y_range=[record["bounding_box"][1] - 10, record["bounding_box"][3] + 10],
)
# Add to zoom paths
view = rrb.Spatial2DView(
name=record["name"].title(), contents=[f"{page_path}/Image/**"] + current_paths, visual_bounds=bounds
)
zoom_paths.append(view)
# Add to type-specific zoom paths
if layout_type == LayoutType.FIGURE:
zoom_paths_figures.append(view)
elif layout_type == LayoutType.TABLE:
zoom_paths_tables.append(view)
elif layout_type != LayoutType.UNKNOWN or layout.show_unknown:
zoom_paths_texts.append(view)
def generate_blueprint(
layouts: list[Layout],
page_paths: list[str],
processed_layouts: list[LayoutStructure],
) -> rrb.Blueprint:
page_tabs = []
for layout, (page_path, processed_layout) in zip(layouts, zip(page_paths, processed_layouts)):
paths, detections_paths, zoom_paths_figures, zoom_paths_tables, zoom_paths_texts = processed_layout
section_tabs = []
content_data: dict[str, Any] = {
"Figures": zoom_paths_figures,
"Tables": zoom_paths_tables,
"Texts": zoom_paths_texts,
}
for name, paths in content_data.items():
if paths:
section_tabs.append(rrb.Tabs(*paths, name=name)) # type: ignore[arg-type]
page_tabs.append(
rrb.Vertical(
rrb.Horizontal(
rrb.Spatial2DView(
name="Layout",
origin=f"{page_path}/Image/",
contents=[f"{page_path}/Image/**"] + detections_paths,
),
rrb.Spatial2DView(name="Detections", contents=[f"{page_path}/Image/**"]),
rrb.TextDocumentView(name="Recovery", contents=f"{page_path}/Recovery"),
),
rrb.Horizontal(*section_tabs),
name=page_path,
row_shares=[4, 3],
)
)
return rrb.Blueprint(
rrb.Tabs(*page_tabs),
collapse_panels=True,
)
def detect_and_log_layouts(log_queue: SimpleQueue[Any], file_path: str) -> None:
images: list[npt.NDArray[np.uint8]] = []
if file_path.endswith(".pdf"):
# convert pdf to images
images.extend(np.array(img, dtype=np.uint8) for img in pdf2image.convert_from_path(file_path))
else:
# read image
img = cv2.imread(file_path)
coloured_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
images.append(coloured_image.astype(np.uint8))
# Extracte the layout from each image
layouts: list[Layout] = []
page_paths = [f"page_{i + 1}" for i in range(len(images))]
processed_layouts: list[LayoutStructure] = []
for i, (image, page_path) in enumerate(zip(images, page_paths)):
layouts.append(detect_and_log_layout(log_queue, image, page_path))
# Generate and send a blueprint based on the detected layouts
processed_layouts.append(
process_layout_records(
log_queue,
layouts[-1],
page_path,
)
)
logging.info("Sending blueprint...")
blueprint = generate_blueprint(layouts, page_paths, processed_layouts)
log_queue.put(["blueprint", blueprint])
logging.info("Blueprint sent...")
def detect_and_log_layout(log_queue: SimpleQueue, coloured_image: npt.NDArray[np.uint8], page_path: str = "") -> Layout:
# Layout Object - This will contain the detected layouts and their detections
layout = Layout()
# Log Image and add Annotation Context
log_queue.put([
"log",
f"{page_path}/Image",
[rr.Image(coloured_image)],
])
log_queue.put([
"log",
f"{page_path}/Image",
# The annotation is defined in the Layout class based on its properties
[rr.AnnotationContext(LayoutType.get_annotation())],
{
"static": True,
},
])
# Paddle Model - Getting Predictions
logging.info("Start detection... (It usually takes more than 10-20 seconds per page)")
ocr_model_pp = PPStructure(show_log=False, recovery=True)
logging.info("model loaded")
result_pp = ocr_model_pp(coloured_image)
_, w, _ = coloured_image.shape
result_pp = sorted_layout_boxes(result_pp, w)
logging.info("Detection finished...")
# Add results to the layout
layout.save_all_layouts(result_pp)
logging.info("All results are saved...")
# Recovery Text Document for the detected text
log_queue.put([
"log",
f"{page_path}/Recovery",
[rr.TextDocument(layout.recovery, media_type=rr.MediaType.MARKDOWN)],
])
return layout