pdfparser/main.py at main · purunep/pdfparser · GitHub

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
from fastapi import FastAPI, File, UploadFile, Query
from fastapi.responses import JSONResponse, FileResponse
import tempfile
import shutil
import os
import json
import zipfile
from models import ExtractResponse
from services.pdf_service import PDFTextExtractor
from services.image_service import PDFImageExtractor
from fastapi.middleware.cors import CORSMiddleware
import logging

logging.basicConfig(level=logging.INFO)

app = FastAPI()


app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/extract")
async def extract_pdf(
    file: UploadFile = File(...),
    mode: str = Query("both", regex="^(text|images|both)$", description="Extraction mode: text, images, or both")
):
    logging.info("/extract endpoint called")
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
        shutil.copyfileobj(file.file, tmp_file)
        tmp_path = tmp_file.name
    logging.info(f"Saved uploaded file to {tmp_path}")

    output_dir = tempfile.mkdtemp()
    images_folder = os.path.join(output_dir, "images")
    zip_path = os.path.join(output_dir, "output.zip")
    json_path = os.path.join(output_dir, "output.json")

    try:
        results = []
        if mode in ("text", "both"):
            logging.info("Extracting text chunks...")
            text_chunks = PDFTextExtractor.extract_text_chunks(tmp_path)
            results.extend(text_chunks)
        if mode in ("images", "both"):
            logging.info("Extracting image chunks...")
            image_chunks = PDFImageExtractor.extract_images_from_pdf(tmp_path, images_folder)
            results.extend(image_chunks)
        logging.info("Writing results to JSON...")
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump({"chunks": results}, f, ensure_ascii=False, indent=2)
        logging.info("Creating ZIP file...")
        with zipfile.ZipFile(zip_path, "w") as zipf:
            zipf.write(json_path, arcname="output.json")
            if mode in ("images", "both") and os.path.exists(images_folder):
                for root, _, files in os.walk(images_folder):
                    for img_file in files:
                        img_path = os.path.join(root, img_file)
                        arcname = os.path.relpath(img_path, output_dir)
                        zipf.write(img_path, arcname=arcname)
        logging.info(f"Returning ZIP file: {zip_path}")
        return FileResponse(zip_path, filename="pdfparser_output.zip", media_type="application/zip")
    except Exception as e:
        logging.error(f"Exception occurred: {e}")
        raise
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)