Spaces:
Running
Running
File size: 3,152 Bytes
2e237ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
from pathlib import Path
from os import makedirs
from os.path import join
from pdf_annotate import PdfAnnotator, Location, Appearance
from starlette.responses import FileResponse
from ports.services.visualization_service import VisualizationService
from configuration import ROOT_PATH
DOCLAYNET_COLOR_BY_TYPE = {
"Caption": "#FFC300",
"Footnote": "#581845",
"Formula": "#FF5733",
"List item": "#008B8B",
"Page footer": "#FF5733",
"Page header": "#581845",
"Picture": "#C70039",
"Section header": "#C70039",
"Table": "#FF8C00",
"Text": "#606060",
"Title": "#EED400",
}
class VisualizationServiceAdapter(VisualizationService):
def create_pdf_visualization(self, pdf_path: Path, segment_boxes: list[dict]) -> Path:
pdf_outputs_path = join(ROOT_PATH, "pdf_outputs")
makedirs(pdf_outputs_path, exist_ok=True)
annotator = PdfAnnotator(str(pdf_path))
segment_index = 0
current_page = 1
for segment_box in segment_boxes:
if int(segment_box["page_number"]) != current_page:
segment_index = 0
current_page += 1
page_height = int(segment_box["page_height"])
self._add_prediction_annotation(annotator, segment_box, segment_index, page_height)
segment_index += 1
annotator.write(str(pdf_path))
return pdf_path
def get_visualization_response(self, pdf_path: Path) -> FileResponse:
return FileResponse(path=pdf_path, media_type="application/pdf", filename=pdf_path.name)
def _hex_color_to_rgb(self, color: str) -> tuple:
r, g, b = int(color[1:3], 16), int(color[3:5], 16), int(color[5:7], 16)
alpha = 1
return r / 255, g / 255, b / 255, alpha
def _add_prediction_annotation(
self, annotator: PdfAnnotator, segment_box: dict, segment_index: int, page_height: int
) -> None:
predicted_type = segment_box["type"]
color = DOCLAYNET_COLOR_BY_TYPE[predicted_type]
left, top, right, bottom = (
segment_box["left"],
page_height - segment_box["top"],
segment_box["left"] + segment_box["width"],
page_height - (segment_box["top"] + segment_box["height"]),
)
text_box_size = len(predicted_type) * 8 + 8
annotator.add_annotation(
"square",
Location(x1=left, y1=bottom, x2=right, y2=top, page=int(segment_box["page_number"]) - 1),
Appearance(stroke_color=self._hex_color_to_rgb(color)),
)
annotator.add_annotation(
"square",
Location(x1=left, y1=top, x2=left + text_box_size, y2=top + 10, page=int(segment_box["page_number"]) - 1),
Appearance(fill=self._hex_color_to_rgb(color)),
)
content = predicted_type.capitalize() + f" [{str(segment_index+1)}]"
annotator.add_annotation(
"text",
Location(x1=left, y1=top, x2=left + text_box_size, y2=top + 10, page=int(segment_box["page_number"]) - 1),
Appearance(content=content, font_size=8, fill=(1, 1, 1), stroke_width=3),
)
|