This commit is contained in:
@@ -0,0 +1,225 @@
|
||||
import os
|
||||
import tempfile
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel
|
||||
import pytesseract
|
||||
from PIL import Image
|
||||
from pdf2image import convert_from_path
|
||||
import uvicorn
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
app = FastAPI(
|
||||
title="OCR API",
|
||||
description="API for extracting text from PDF files using OCR",
|
||||
version="1.0.0"
|
||||
)
|
||||
|
||||
class OCRResponse(BaseModel):
|
||||
text: str
|
||||
pages: int
|
||||
processing_time: float
|
||||
confidence: Optional[float] = None
|
||||
|
||||
class ErrorResponse(BaseModel):
|
||||
error: str
|
||||
detail: str
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {"message": "OCR API is running. Use /docs for API documentation."}
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
return {"status": "healthy", "service": "OCR API"}
|
||||
|
||||
async def process_pdf_with_ocr(pdf_path: str) -> OCRResponse:
|
||||
"""
|
||||
Process a PDF file and extract text using OCR
|
||||
"""
|
||||
import time
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Convert PDF to images
|
||||
logger.info(f"Converting PDF to images: {pdf_path}")
|
||||
images = convert_from_path(pdf_path, dpi=300)
|
||||
|
||||
if not images:
|
||||
raise HTTPException(status_code=400, detail="Could not extract images from PDF")
|
||||
|
||||
# Extract text from each page
|
||||
all_text = []
|
||||
total_confidence = 0
|
||||
confidence_count = 0
|
||||
|
||||
for i, image in enumerate(images):
|
||||
logger.info(f"Processing page {i+1}/{len(images)}")
|
||||
|
||||
# Perform OCR on the image
|
||||
ocr_data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
|
||||
|
||||
# Extract text and confidence scores
|
||||
page_text = []
|
||||
for j in range(len(ocr_data['text'])):
|
||||
if int(ocr_data['conf'][j]) > 0: # Filter out low confidence results
|
||||
page_text.append(ocr_data['text'][j])
|
||||
total_confidence += float(ocr_data['conf'][j])
|
||||
confidence_count += 1
|
||||
|
||||
page_text = ' '.join(page_text).strip()
|
||||
if page_text:
|
||||
all_text.append(f"--- Page {i+1} ---\n{page_text}")
|
||||
|
||||
# Calculate average confidence
|
||||
avg_confidence = total_confidence / confidence_count if confidence_count > 0 else None
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
return OCRResponse(
|
||||
text='\n\n'.join(all_text),
|
||||
pages=len(images),
|
||||
processing_time=round(processing_time, 2),
|
||||
confidence=round(avg_confidence, 2) if avg_confidence else None
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing PDF: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail=f"Error processing PDF: {str(e)}")
|
||||
|
||||
@app.post("/ocr/pdf", response_model=OCRResponse)
|
||||
async def extract_text_from_pdf(
|
||||
file: UploadFile = File(...),
|
||||
background_tasks: BackgroundTasks = None
|
||||
):
|
||||
"""
|
||||
Extract text from a PDF file using OCR
|
||||
|
||||
- **file**: PDF file to process
|
||||
- Returns extracted text with metadata
|
||||
"""
|
||||
|
||||
# Validate file type
|
||||
if not file.filename.lower().endswith('.pdf'):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="File must be a PDF"
|
||||
)
|
||||
|
||||
# Check file size (limit to 50MB)
|
||||
if file.size and file.size > 50 * 1024 * 1024:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="File size too large. Maximum size is 50MB"
|
||||
)
|
||||
|
||||
try:
|
||||
# Create temporary file
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_file:
|
||||
# Write uploaded file to temporary file
|
||||
content = await file.read()
|
||||
temp_file.write(content)
|
||||
temp_file_path = temp_file.name
|
||||
|
||||
try:
|
||||
# Process the PDF
|
||||
result = await process_pdf_with_ocr(temp_file_path)
|
||||
|
||||
# Clean up temporary file
|
||||
background_tasks.add_task(os.unlink, temp_file_path)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
# Clean up temporary file on error
|
||||
if os.path.exists(temp_file_path):
|
||||
os.unlink(temp_file_path)
|
||||
raise e
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling file upload: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")
|
||||
|
||||
@app.post("/ocr/image", response_model=OCRResponse)
|
||||
async def extract_text_from_image(
|
||||
file: UploadFile = File(...),
|
||||
background_tasks: BackgroundTasks = None
|
||||
):
|
||||
"""
|
||||
Extract text from an image file using OCR
|
||||
|
||||
- **file**: Image file to process (PNG, JPG, JPEG, etc.)
|
||||
- Returns extracted text with metadata
|
||||
"""
|
||||
|
||||
# Validate file type
|
||||
allowed_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif'}
|
||||
file_extension = Path(file.filename).suffix.lower()
|
||||
|
||||
if file_extension not in allowed_extensions:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Unsupported file type. Allowed types: {', '.join(allowed_extensions)}"
|
||||
)
|
||||
|
||||
try:
|
||||
# Create temporary file
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as temp_file:
|
||||
content = await file.read()
|
||||
temp_file.write(content)
|
||||
temp_file_path = temp_file.name
|
||||
|
||||
try:
|
||||
import time
|
||||
start_time = time.time()
|
||||
|
||||
# Open and process image
|
||||
image = Image.open(temp_file_path)
|
||||
|
||||
# Perform OCR
|
||||
ocr_data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
|
||||
|
||||
# Extract text and confidence
|
||||
text_parts = []
|
||||
total_confidence = 0
|
||||
confidence_count = 0
|
||||
|
||||
for i in range(len(ocr_data['text'])):
|
||||
if int(ocr_data['conf'][i]) > 0:
|
||||
text_parts.append(ocr_data['text'][i])
|
||||
total_confidence += float(ocr_data['conf'][i])
|
||||
confidence_count += 1
|
||||
|
||||
text = ' '.join(text_parts).strip()
|
||||
avg_confidence = total_confidence / confidence_count if confidence_count > 0 else None
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Clean up temporary file
|
||||
background_tasks.add_task(os.unlink, temp_file_path)
|
||||
|
||||
return OCRResponse(
|
||||
text=text,
|
||||
pages=1,
|
||||
processing_time=round(processing_time, 2),
|
||||
confidence=round(avg_confidence, 2) if avg_confidence else None
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
# Clean up temporary file on error
|
||||
if os.path.exists(temp_file_path):
|
||||
os.unlink(temp_file_path)
|
||||
raise e
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing image: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
Reference in New Issue
Block a user