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)