Spaces:
Running
Running
File size: 2,073 Bytes
9002555 |
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 |
from typing import Optional
from llama_index.core.readers.base import BaseReader
from llama_index.core.schema import Document
from fastapi import UploadFile
from typing import List
from PyPDF2 import PdfReader
from io import BytesIO
import fitz # PyMuPDF
class Reader(BaseReader):
async def read_from_uploadfile(self, file: UploadFile) -> List[Document]:
try:
# Read the file content asynchronously
file_content = await file.read()
# Initialize PyMuPDF document with file content
pdf_document = fitz.open(stream=file_content, filetype="pdf")
# Extract text and images from each page
pages = []
for page_num in range(len(pdf_document)):
page = pdf_document.load_page(page_num)
# Extract text
text = page.get_text().strip()
if text:
pages.append(Document(text=text, metadata={"page": page_num + 1}))
# Extract images
for img_index, img in enumerate(page.get_images(full=True)):
xref = img[0]
base_image = pdf_document.extract_image(xref)
image_bytes = base_image["image"]
image_stream = BytesIO(image_bytes)
# Store the image as a Document (or any other structure you need)
pages.append(
Document(
text=f"Image {img_index + 1} on page {page_num + 1}",
metadata={
"page": page_num + 1,
"image_index": img_index + 1,
"image": image_stream,
},
)
)
return pages
except Exception as e:
# Handle exceptions more granularly if needed
print(f"Error reading PDF file: {e}")
raise RuntimeError(f"Failed to process the uploaded file: {e}")
|