Spaces:
Sleeping
Sleeping
File size: 5,586 Bytes
7a11b03 | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """PDF upload validation and text extraction."""
from pathlib import Path
from pypdf import PdfReader
from src.config import DEFAULT_OCR_ENGINE, DEFAULT_OCR_MODE, OCR_ENGINE_NAMES, OCR_MODE_DPI
from src.ocr_service import extract_text_from_pdf_page_with_selected_ocr
def validate_pdf_file(file_name: str) -> None:
"""Raise a helpful error when the uploaded file is not a PDF."""
if not file_name.lower().endswith(".pdf"):
raise ValueError(f"{file_name} is not a PDF file. Please upload only PDF files.")
def save_pdf_bytes(file_name: str, file_bytes: bytes, upload_dir: Path) -> Path:
"""Save uploaded PDF bytes into the local uploads folder."""
validate_pdf_file(file_name)
if not file_bytes:
raise ValueError(f"{file_name} is empty. Please upload a valid PDF file.")
upload_dir.mkdir(parents=True, exist_ok=True)
safe_file_name = Path(file_name).name
saved_path = upload_dir / safe_file_name
with saved_path.open("wb") as pdf_file:
pdf_file.write(file_bytes)
return saved_path
def count_pdf_pages(pdf_path: Path) -> int:
"""Return how many pages a PDF has."""
try:
reader = PdfReader(str(pdf_path))
except Exception as error:
raise ValueError(f"Could not read {pdf_path.name}. The PDF may be damaged.") from error
return len(reader.pages)
def validate_page_range(start_page: int | None, end_page: int | None, total_pages: int) -> tuple[int, int]:
"""Convert a 1-based page range into safe 0-based bounds."""
start = start_page or 1
end = end_page or total_pages
if start < 1:
raise ValueError("Start page must be 1 or greater.")
if end < start:
raise ValueError("End page must be greater than or equal to start page.")
if start > total_pages:
raise ValueError(f"Start page {start} is greater than total pages {total_pages}.")
end = min(end, total_pages)
return start - 1, end
def validate_ocr_mode(ocr_mode: str) -> str:
"""Return a supported OCR mode or raise a beginner-friendly error."""
normalized_mode = (ocr_mode or DEFAULT_OCR_MODE).strip().lower()
if normalized_mode not in OCR_MODE_DPI:
allowed_modes = ", ".join(OCR_MODE_DPI)
raise ValueError(f"Unsupported OCR mode '{ocr_mode}'. Use one of: {allowed_modes}.")
return normalized_mode
def validate_ocr_engine(ocr_engine: str) -> str:
"""Return a supported OCR engine or raise a helpful error."""
normalized_engine = (ocr_engine or DEFAULT_OCR_ENGINE).strip().lower()
if normalized_engine not in OCR_ENGINE_NAMES:
allowed_engines = ", ".join(sorted(OCR_ENGINE_NAMES))
raise ValueError(f"Unsupported OCR engine '{ocr_engine}'. Use one of: {allowed_engines}.")
return normalized_engine
def extract_pages_from_pdf(
pdf_path: Path,
start_page: int | None = None,
end_page: int | None = None,
ocr_mode: str = DEFAULT_OCR_MODE,
ocr_engine: str = DEFAULT_OCR_ENGINE,
) -> list[dict]:
"""Extract text page by page, using the selected OCR engine for scanned pages."""
pages = []
ocr_mode = validate_ocr_mode(ocr_mode)
ocr_engine = validate_ocr_engine(ocr_engine)
try:
reader = PdfReader(str(pdf_path))
except Exception as error:
raise ValueError(f"Could not read {pdf_path.name}. The PDF may be damaged.") from error
total_pages = len(reader.pages)
start_index, stop_index = validate_page_range(start_page, end_page, total_pages)
for page_index in range(start_index, stop_index):
page = reader.pages[page_index]
page_text = page.extract_text() or ""
cleaned_text = " ".join(page_text.split())
extraction_method = "embedded_text"
if not cleaned_text:
cleaned_text, extraction_method = extract_text_from_pdf_page_with_selected_ocr(
pdf_path=pdf_path,
page_index=page_index,
ocr_mode=ocr_mode,
ocr_engine=ocr_engine,
)
if cleaned_text:
used_ocr = extraction_method != "embedded_text"
pages.append(
{
"text": cleaned_text,
"page_number": page_index + 1,
"source_file": pdf_path.name,
"extraction_method": extraction_method,
"ocr_mode": ocr_mode if used_ocr else "not_used",
"ocr_engine": ocr_engine if used_ocr else "not_used",
"total_pages": total_pages,
}
)
if not pages:
raise ValueError(
f"No readable text was found in {pdf_path.name}. "
"The selected OCR engine could not read enough text from this PDF. "
"Try a clearer scan, better lighting, or typed notes."
)
return pages
def load_uploaded_pdfs(
saved_pdf_paths: list[Path],
start_page: int | None = None,
end_page: int | None = None,
ocr_mode: str = DEFAULT_OCR_MODE,
ocr_engine: str = DEFAULT_OCR_ENGINE,
) -> list[dict]:
"""Extract readable text from all uploaded PDFs."""
all_pages = []
for pdf_path in saved_pdf_paths:
all_pages.extend(
extract_pages_from_pdf(
pdf_path=pdf_path,
start_page=start_page,
end_page=end_page,
ocr_mode=ocr_mode,
ocr_engine=ocr_engine,
)
)
if not all_pages:
raise ValueError("No readable text was found in the uploaded PDFs.")
return all_pages
|