107 lines
2.5 KiB
Python
107 lines
2.5 KiB
Python
import os
|
|
import uuid
|
|
import tempfile
|
|
import requests
|
|
from minio import Minio
|
|
from pdf2image import convert_from_path
|
|
import pytesseract
|
|
import concurrent.futures
|
|
import PyPDF2
|
|
|
|
# ================= CONFIG =================
|
|
MINIO_ENDPOINT = "192.168.10.76:9000"
|
|
MINIO_ACCESS = "minioadmin"
|
|
MINIO_SECRET = "Admin@!@2025"
|
|
BUCKET_NAME = "pdf-storage"
|
|
|
|
POPPLER_PATH = r"C:\Program Files\poppler\Library\bin" # Ubuntu = None
|
|
MAX_FILE_SIZE_MB = 20
|
|
|
|
# ================= INIT MINIO =================
|
|
minio_client = Minio(
|
|
MINIO_ENDPOINT,
|
|
access_key=MINIO_ACCESS,
|
|
secret_key=MINIO_SECRET,
|
|
secure=False
|
|
)
|
|
|
|
if not minio_client.bucket_exists(BUCKET_NAME):
|
|
minio_client.make_bucket(BUCKET_NAME)
|
|
|
|
|
|
# ================= HELPERS =================
|
|
|
|
def pdf_has_text(file_path):
|
|
reader = PyPDF2.PdfReader(file_path)
|
|
for page in reader.pages:
|
|
if page.extract_text():
|
|
return True
|
|
return False
|
|
|
|
|
|
def ocr_page(page):
|
|
return pytesseract.image_to_string(page, config="--psm 6")
|
|
|
|
|
|
def process_pdf(url):
|
|
response = requests.get(url, timeout=30)
|
|
|
|
if response.status_code != 200:
|
|
raise Exception("Download failed")
|
|
|
|
if "application/pdf" not in response.headers.get("Content-Type", ""):
|
|
raise Exception("Not a valid PDF")
|
|
|
|
if len(response.content) > MAX_FILE_SIZE_MB * 1024 * 1024:
|
|
raise Exception("File too large")
|
|
|
|
return _process_bytes(response.content)
|
|
|
|
|
|
def process_uploaded_pdf(file_bytes):
|
|
return _process_bytes(file_bytes)
|
|
|
|
|
|
def _process_bytes(file_bytes):
|
|
file_id = str(uuid.uuid4())
|
|
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
|
|
tmp.write(file_bytes)
|
|
local_pdf = tmp.name
|
|
|
|
# Upload PDF
|
|
minio_client.fput_object(
|
|
BUCKET_NAME,
|
|
f"{file_id}.pdf",
|
|
local_pdf,
|
|
content_type="application/pdf"
|
|
)
|
|
|
|
# Detect digital PDF
|
|
if pdf_has_text(local_pdf):
|
|
reader = PyPDF2.PdfReader(local_pdf)
|
|
full_text = ""
|
|
for page in reader.pages:
|
|
full_text += page.extract_text() or ""
|
|
source_type = "digital_pdf"
|
|
else:
|
|
pages = convert_from_path(
|
|
local_pdf,
|
|
dpi=300,
|
|
poppler_path=POPPLER_PATH
|
|
)
|
|
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
results = executor.map(ocr_page, pages)
|
|
|
|
full_text = "\n".join(results)
|
|
source_type = "scanned_pdf"
|
|
|
|
os.remove(local_pdf)
|
|
|
|
return {
|
|
"file_id": file_id,
|
|
"type": source_type,
|
|
"text_length": len(full_text),
|
|
"text": full_text
|
|
} |