| from PIL import Image
|
| import fitz
|
|
|
| from docgenie.generation.constants import IMAGE_RENDER_EXT
|
| from docgenie.generation.models._file import SyntheticDatasetFileStructure
|
| from docgenie.generation.models._log import SynDocumentLog
|
|
|
|
|
| def get_pdf_size_pt(docid: str, dsfiles: SyntheticDatasetFileStructure):
|
| pdf_path = dsfiles.final_pdf_directory / f"{docid}.pdf"
|
| doc = fitz.open(pdf_path)
|
| page = doc[0]
|
| width_pt, height_pt = page.rect.width, page.rect.height
|
| width_px = width_pt
|
| height_px = height_pt
|
| doc.close()
|
| return width_px, height_px
|
|
|
| def get_image_size_px(docid: str, dsfiles: SyntheticDatasetFileStructure):
|
|
|
| image_path = dsfiles.img_directory / f"{docid}.{IMAGE_RENDER_EXT}"
|
| img = Image.open(image_path)
|
| width_px, height_px = img.size
|
| return width_px, height_px
|
|
|
| def get_document_size_for_bbox_unnormalization(docid: str, dsfiles: SyntheticDatasetFileStructure):
|
| doclog = SynDocumentLog(document_id=docid, logdir=dsfiles.document_logs_directory)
|
| if doclog.ocr_required:
|
|
|
| return get_image_size_px(docid=docid, dsfiles=dsfiles)
|
| else:
|
|
|
| return get_pdf_size_pt(docid, dsfiles) |
