projects/Dockers/ocr-api/main.py

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)