Spaces:
Sleeping
Sleeping
from fastapi import FastAPI, File, UploadFile | |
from fastapi.responses import HTMLResponse | |
import fitz # PyMuPDF | |
import pytesseract | |
from PIL import Image | |
import io | |
app = FastAPI() | |
async def home(): | |
html_content = """ | |
<!DOCTYPE html> | |
<html> | |
<head> | |
<title>PDF Text Extraction API</title> | |
<style> | |
body { font-family: Arial, sans-serif; padding: 2rem; background: #f0f0f0; } | |
.container { max-width: 600px; margin: auto; background: white; padding: 2rem; border-radius: 8px; box-shadow: 0 0 10px rgba(0,0,0,0.1);} | |
h1 { color: #333; } | |
p { color: #555; } | |
a { color: #1a73e8; text-decoration: none; } | |
a:hover { text-decoration: underline; } | |
</style> | |
</head> | |
<body> | |
<div class="container"> | |
<h1>Welcome to PDF Text Extraction API</h1> | |
<p>This API allows you to upload PDFs and extract text β including optional OCR for images.</p> | |
<h2>Available endpoints:</h2> | |
<ul> | |
<li><b>POST /extract-text</b> - Extract text from PDF pages.</li> | |
<li><b>POST /extract-text-ocr</b> - Extract text including OCR from images inside PDFs.</li> | |
</ul> | |
<p>Use a tool like <a href="https://www.postman.com/" target="_blank">Postman</a> or write your own client to send PDF files to the endpoints.</p> | |
</div> | |
</body> | |
</html> | |
""" | |
return HTMLResponse(content=html_content, status_code=200) | |
async def extract_text(file: UploadFile = File(...)): | |
try: | |
contents = await file.read() | |
doc = fitz.open(stream=contents, filetype="pdf") | |
extracted_text = "" | |
for i, page in enumerate(doc): | |
extracted_text += f"\n\n--- Page {i + 1} ---\n\n" + page.get_text() | |
return {"filename": file.filename, "text": extracted_text} | |
except Exception as e: | |
return {"error": str(e)} | |
async def extract_text_ocr(file: UploadFile = File(...)): | |
try: | |
contents = await file.read() | |
doc = fitz.open(stream=contents, filetype="pdf") | |
full_text = "" | |
for i in range(len(doc)): | |
page = doc.load_page(i) | |
# Normal text | |
text = page.get_text() | |
# Render page to an image | |
pix = page.get_pixmap() | |
img = Image.open(io.BytesIO(pix.tobytes())) | |
# OCR text | |
ocr_text = pytesseract.image_to_string(img) | |
full_text += f"\n\n--- Page {i + 1} ---\n\n" | |
full_text += text + "\n" | |
full_text += "[OCR Text]\n" + ocr_text | |
return {"filename": file.filename, "text": full_text} | |
except Exception as e: | |
return {"error": str(e)} | |