176 lines
5.8 KiB
Python
176 lines
5.8 KiB
Python
import json
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
from operator import itemgetter
|
|
from typing import Literal, SupportsIndex, Tuple
|
|
|
|
import fitz # type: ignore
|
|
import numpy as np
|
|
from funcy import compose, lfilter # type: ignore
|
|
from kn_utils.logging import logger # type: ignore
|
|
from numpy import ndarray as Array
|
|
|
|
BBoxType = tuple[int | float, int | float, int | float, int | float]
|
|
|
|
|
|
@dataclass
|
|
class PageInfo:
|
|
page_num: int
|
|
rotation_matrix: fitz.Matrix
|
|
transformation_matrix: fitz.Matrix
|
|
dpi: int
|
|
width: int | float
|
|
height: int | float
|
|
image_width: int | float
|
|
image_height: int | float
|
|
rotation: int
|
|
|
|
|
|
def transform_image_coordinates_to_pdf_coordinates(
|
|
bbox: BBoxType,
|
|
rotation_matrix: fitz.Matrix,
|
|
transformation_matrix: fitz.Matrix,
|
|
dpi: int | None = None,
|
|
) -> Tuple:
|
|
x1, y1, x2, y2 = map(lambda x: (x / dpi) * 72, bbox) if dpi else bbox # Convert to points, can be done before
|
|
rect = fitz.Rect(x1, y1, x2, y2)
|
|
rect = rect * rotation_matrix * transformation_matrix
|
|
|
|
return rect.x0, rect.y0, rect.x1, rect.y1
|
|
|
|
|
|
def rescale_to_pdf(bbox: BBoxType, page_info: PageInfo) -> BBoxType:
|
|
round3 = lambda x: tuple(map(lambda y: round(y, 3), x))
|
|
|
|
pdf_h, pdf_w = page_info.height, page_info.width
|
|
|
|
if page_info.rotation in {90, 270}:
|
|
pdf_h, pdf_w = pdf_w, pdf_h
|
|
|
|
pix_h, pix_w = page_info.image_height, page_info.image_width
|
|
ratio_h, ratio_w = pdf_h / pix_h, pdf_w / pix_w
|
|
ratio_w, ratio_h, pdf_w, pdf_h, pix_w, pix_h = round3((ratio_w, ratio_h, pdf_w, pdf_h, pix_w, pix_h))
|
|
|
|
return round3((bbox[0] * ratio_w, bbox[1] * ratio_h, bbox[2] * ratio_w, bbox[3] * ratio_h))
|
|
|
|
|
|
def mirror_horizontal(bbox: BBoxType, page_height: int | float) -> BBoxType:
|
|
x0, y0, x1, y1 = bbox
|
|
y0_new = page_height - y0
|
|
y1_new = page_height - y1
|
|
|
|
return x0, y0_new, x1, y1_new
|
|
|
|
|
|
def mirror_vertical(bbox: BBoxType, page_width: int | float) -> BBoxType:
|
|
x0, y0, x1, y1 = bbox
|
|
x0_new = page_width - x1
|
|
x1_new = page_width - x0
|
|
|
|
return x0_new, y0, x1_new, y1
|
|
|
|
|
|
def derotate_image(bbox: BBoxType, page_info: PageInfo) -> BBoxType:
|
|
|
|
logger.debug(f"{page_info.rotation=}")
|
|
match page_info.rotation:
|
|
case 0:
|
|
bbox = mirror_horizontal(bbox, page_info.height)
|
|
case 90:
|
|
pass
|
|
case 180:
|
|
bbox = mirror_vertical(bbox, page_info.height)
|
|
case 270:
|
|
bbox = mirror_vertical(mirror_horizontal(bbox, page_info.height), page_info.height)
|
|
case _:
|
|
logger.warning(f"Unknown rotation: {page_info.rotation}")
|
|
pass
|
|
|
|
return bbox
|
|
|
|
|
|
def transform_table_lines_by_page_info(bboxes: dict, offsets: tuple, page_info: PageInfo) -> dict:
|
|
|
|
transform = partial(rescale_to_pdf, page_info=page_info)
|
|
logger.debug(f"{offsets=}")
|
|
|
|
def apply_offsets(line: tuple) -> tuple:
|
|
x1, y1, x2, y2 = line
|
|
offset_x, offset_y = offsets
|
|
offset_y = page_info.height - offset_y
|
|
logger.debug((f"new offsets: {offset_x}, {offset_y}"))
|
|
|
|
return (x1 + offset_x, y1 + offset_y, x2 + offset_x, y2 + offset_y)
|
|
|
|
derotate = partial(derotate_image, page_info=page_info)
|
|
|
|
unpack = itemgetter("x1", "y1", "x2", "y2")
|
|
pack = lambda x: {"x1": x[0], "y1": x[1], "x2": x[2], "y2": x[3]}
|
|
convert = compose(pack, derotate, apply_offsets, transform, unpack)
|
|
|
|
table_lines = bboxes.get("tableLines", [])
|
|
bboxes["tableLines"] = list(map(convert, table_lines))
|
|
|
|
for i in range(len(table_lines)):
|
|
logger.debug(json.dumps(table_lines[i], indent=4))
|
|
logger.debug(json.dumps(bboxes["tableLines"][i], indent=4))
|
|
|
|
return bboxes
|
|
|
|
|
|
def extract_images_from_pdf(
|
|
pdf_bytes: bytes, vlp_output: dict, dpi: int = 200
|
|
) -> tuple[list[Array], list[dict], list[PageInfo]]:
|
|
with fitz.open(stream=pdf_bytes) as fh:
|
|
table_images = []
|
|
table_info = []
|
|
page_info = []
|
|
|
|
vlp_output = vlp_output["data"] if isinstance(vlp_output, dict) else vlp_output
|
|
|
|
for page_dict in vlp_output:
|
|
page_num = int(page_dict["page_idx"])
|
|
boxes = page_dict["boxes"]
|
|
boxes = filter(lambda box_obj: box_obj["label"] == "table", boxes)
|
|
|
|
page: fitz.Page = fh[page_num]
|
|
page.wrap_contents()
|
|
|
|
page_image = page.get_pixmap(dpi=200)
|
|
current_page_info = PageInfo(
|
|
page_num,
|
|
page.rotation_matrix,
|
|
page.transformation_matrix,
|
|
dpi,
|
|
page.rect[-2],
|
|
page.rect[-1],
|
|
page_image.w,
|
|
page_image.h,
|
|
page.rotation,
|
|
)
|
|
|
|
for box_obj in boxes:
|
|
bbox = box_obj["box"]
|
|
x1, y1, x2, y2 = itemgetter("x1", "y1", "x2", "y2")(bbox)
|
|
rect = fitz.Rect((x1, y1), (x2, y2))
|
|
# FIXME: Check if de-rotation works as intended and is necessary at all.
|
|
# Note that there exists also a derotation_matrix. If changing this, also change the
|
|
# current_page_info object to include the derotation_matrix.
|
|
rect = rect * page.transformation_matrix * page.rotation_matrix
|
|
pixmap = page.get_pixmap(clip=rect, dpi=dpi, colorspace=fitz.csGRAY)
|
|
shape = (pixmap.h, pixmap.w, pixmap.n) if pixmap.n > 1 else (pixmap.h, pixmap.w)
|
|
image = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(*shape)
|
|
|
|
table_images.append(image)
|
|
table_info.append(
|
|
{
|
|
"pageNum": page_num,
|
|
"bbox": bbox,
|
|
"uuid": box_obj["uuid"],
|
|
"label": box_obj["label"],
|
|
}
|
|
)
|
|
page_info.append(current_page_info)
|
|
|
|
return table_images, table_info, page_info
|