85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
from pathlib import Path
|
|
|
|
import fitz # PyMuPDF
|
|
from loguru import logger
|
|
|
|
import config
|
|
import concurrent.futures
|
|
|
|
|
|
def is_pdf(filename: str) -> bool:
|
|
return Path(filename).suffix.lower() == ".pdf"
|
|
|
|
|
|
def pdf_to_images(pdf_path: str, dpi: int = config.OCR_DPI,
|
|
max_pages: int = config.PDF_MAX_PAGES) -> list[Path]:
|
|
pdf_path = Path(pdf_path)
|
|
output_dir = pdf_path.parent / pdf_path.stem
|
|
output_dir.mkdir(exist_ok=True)
|
|
|
|
ext = config.PDF_IMAGE_FORMAT # "jpeg"
|
|
page_indices = list(range(max_pages))
|
|
|
|
# 이미 변환된 이미지가 모두 있으면 스킵
|
|
cached = [output_dir / f"page_{i+1:03d}.{ext}" for i in page_indices]
|
|
if all(p.exists() for p in cached):
|
|
logger.debug(f"PDF 변환 캐시 사용: {pdf_path.name}")
|
|
return cached
|
|
|
|
pdf_path = Path(pdf_path)
|
|
print(pdf_path.exists())
|
|
print(pdf_path.stat().st_size)
|
|
if not pdf_path.exists() or pdf_path.stat().st_size == 0:
|
|
raise ValueError(f"유효하지 않은 PDF 파일입니다: {pdf_path.name}")
|
|
|
|
matrix = fitz.Matrix(dpi / 72, dpi / 72)
|
|
doc = fitz.open(str(pdf_path))
|
|
|
|
img_paths = []
|
|
for i in range(min(max_pages, len(doc))):
|
|
img_path = output_dir / f"page_{i+1:03d}.{ext}"
|
|
if img_path.exists():
|
|
logger.debug(f"페이지 캐시 사용: {img_path.name}")
|
|
else:
|
|
pixmap = doc[i].get_pixmap(matrix=matrix, alpha=False)
|
|
|
|
# 최대 너비 초과 시 비율 유지하며 축소
|
|
if hasattr(config, 'PDF_MAX_WIDTH') and pixmap.width > config.PDF_MAX_WIDTH:
|
|
scale = config.PDF_MAX_WIDTH / pixmap.width
|
|
pixmap = doc[i].get_pixmap(matrix=fitz.Matrix(dpi / 72 * scale, dpi / 72 * scale), alpha=False)
|
|
|
|
if ext in ("jpeg", "jpg"):
|
|
pixmap.save(str(img_path), jpg_quality=config.PDF_JPEG_QUALITY) # JPEG 저장
|
|
else:
|
|
pixmap.save(str(img_path))
|
|
logger.debug(f"PDF 변환: {pdf_path.name} | 페이지 {i+1}")
|
|
img_paths.append(img_path)
|
|
|
|
doc.close()
|
|
logger.info(f"PDF 변환 완료: {pdf_path.name} | {len(img_paths)}페이지 | {dpi}dpi {ext}")
|
|
return img_paths
|
|
|
|
|
|
def pdf_to_images_batch(pdf_paths: list[str],
|
|
max_workers: int = config.OCR_MAX_WORKERS) -> dict[str, list[Path]]:
|
|
"""
|
|
여러 PDF를 ThreadPoolExecutor로 병렬 변환
|
|
반환: {pdf_path: [img_path, ...]}
|
|
"""
|
|
results = {}
|
|
|
|
def _convert(pdf_path: str):
|
|
return pdf_path, pdf_to_images(pdf_path)
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
futures = {executor.submit(_convert, p): p for p in pdf_paths}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
try:
|
|
path, imgs = future.result()
|
|
results[path] = imgs
|
|
except Exception as e:
|
|
path = futures[future]
|
|
logger.error(f"PDF 변환 실패: {path} | {e}")
|
|
results[path] = []
|
|
|
|
return results |