Spaces:
Running
Running
File size: 1,819 Bytes
170d9a9 |
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 |
import asyncio
from medrag_multi_modal.document_loader.image_loader import (
FitzPILImageLoader,
PDF2ImageLoader,
PDFPlumberImageLoader,
PyMuPDFImageLoader,
)
URL = "https://archive.org/download/GraysAnatomy41E2015PDF/Grays%20Anatomy-41%20E%20%282015%29%20%5BPDF%5D.pdf"
COLUMN_NAMES = ["page_image", "page_figure_images", "document_name", "page_idx"]
def test_fitzpil_img_loader():
loader = FitzPILImageLoader(
url=URL,
document_name="Gray's Anatomy",
document_file_path="grays_anatomy.pdf",
)
dataset = asyncio.run(loader.load_data(start_page=32, end_page=37))
assert dataset.num_rows == 5
assert dataset.column_names == COLUMN_NAMES
loader.cleanup_image_dir()
def test_pdf2image_img_loader():
loader = PDF2ImageLoader(
url=URL,
document_name="Gray's Anatomy",
document_file_path="grays_anatomy.pdf",
)
dataset = asyncio.run(loader.load_data(start_page=32, end_page=37))
assert dataset.num_rows == 5
assert dataset.column_names == COLUMN_NAMES
loader.cleanup_image_dir()
def test_pdfplumber_img_loader():
loader = PDFPlumberImageLoader(
url=URL,
document_name="Gray's Anatomy",
document_file_path="grays_anatomy.pdf",
)
dataset = asyncio.run(loader.load_data(start_page=32, end_page=37))
assert dataset.num_rows == 5
assert dataset.column_names == COLUMN_NAMES
loader.cleanup_image_dir()
def test_pymupdf_img_loader():
loader = PyMuPDFImageLoader(
url=URL,
document_name="Gray's Anatomy",
document_file_path="grays_anatomy.pdf",
)
dataset = asyncio.run(loader.load_data(start_page=32, end_page=37))
assert dataset.num_rows == 5
assert dataset.column_names == COLUMN_NAMES
loader.cleanup_image_dir()
|