ocr concurrently

Dockers/ocr-api/main.py

@@ -1,8 +1,11 @@
 import os
 import tempfile
 import logging
+import asyncio
 from typing import List, Optional
 from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
+import time
 
 from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
 from fastapi.responses import JSONResponse
@@ -40,11 +43,51 @@ async def root():
 async def health_check():
     return {"status": "healthy", "service": "OCR API"}
 
-async def process_pdf_with_ocr(pdf_path: str) -> OCRResponse:
+def process_single_page(args):
     """
-    Process a PDF file and extract text using OCR
+    Process a single page with OCR (runs in thread pool)
+    """
+    page_num, image = args
+    try:
+        logger.info(f"Processing page {page_num + 1}")
+
+        # Perform OCR on the image
+        ocr_data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
+
+        # Extract text and confidence scores
+        page_text = []
+        total_confidence = 0
+        confidence_count = 0
+
+        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()
+        avg_confidence = total_confidence / confidence_count if confidence_count > 0 else 0
+
+        return {
+            'page_num': page_num,
+            'text': page_text,
+            'confidence': avg_confidence,
+            'success': True
+        }
+    except Exception as e:
+        logger.error(f"Error processing page {page_num + 1}: {str(e)}")
+        return {
+            'page_num': page_num,
+            'text': '',
+            'confidence': 0,
+            'success': False,
+            'error': str(e)
+        }
+
+async def process_pdf_with_ocr(pdf_path: str, max_workers: int = None) -> OCRResponse:
+    """
+    Process a PDF file and extract text using OCR with concurrent page processing
     """
-    import time
     start_time = time.time()
 
     try:
@@ -55,34 +98,45 @@ async def process_pdf_with_ocr(pdf_path: str) -> OCRResponse:
         if not images:
             raise HTTPException(status_code=400, detail="Could not extract images from PDF")
 
-        # Extract text from each page
+        # Determine optimal number of workers
+        if max_workers is None:
+            max_workers = min(len(images), os.cpu_count() or 4)
+
+        logger.info(f"Processing {len(images)} pages with {max_workers} workers")
+
+        # Process pages concurrently using ThreadPoolExecutor
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            # Create tasks for each page
+            page_tasks = [(i, image) for i, image in enumerate(images)]
+
+            # Submit all tasks and wait for completion
+            results = list(executor.map(process_single_page, page_tasks))
+
+        # Sort results by page number and combine
+        results.sort(key=lambda x: x['page_num'])
+
         all_text = []
         total_confidence = 0
         confidence_count = 0
+        successful_pages = 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])
+        for result in results:
+            if result['success']:
+                successful_pages += 1
+                if result['text']:
+                    all_text.append(f"--- Page {result['page_num'] + 1} ---\n{result['text']}")
+                    total_confidence += result['confidence']
                     confidence_count += 1
-
-            page_text = ' '.join(page_text).strip()
-            if page_text:
-                all_text.append(f"--- Page {i+1} ---\n{page_text}")
+            else:
+                logger.warning(f"Page {result['page_num'] + 1} failed: {result.get('error', 'Unknown error')}")
 
         # Calculate average confidence
         avg_confidence = total_confidence / confidence_count if confidence_count > 0 else None
 
         processing_time = time.time() - start_time
 
+        logger.info(f"Completed processing {successful_pages}/{len(images)} pages in {processing_time:.2f}s")
+
         return OCRResponse(
             text='\n\n'.join(all_text),
             pages=len(images),
@@ -97,12 +151,14 @@ async def process_pdf_with_ocr(pdf_path: str) -> OCRResponse:
 @app.post("/ocr/pdf", response_model=OCRResponse)
 async def extract_text_from_pdf(
     file: UploadFile = File(...),
+    max_workers: int = None,
     background_tasks: BackgroundTasks = None
 ):
     """
-    Extract text from a PDF file using OCR
+    Extract text from a PDF file using OCR with concurrent processing
 
     - **file**: PDF file to process
+    - **max_workers**: Maximum number of concurrent workers (default: min(pages, CPU cores))
     - Returns extracted text with metadata
     """
 
@@ -130,7 +186,7 @@ async def extract_text_from_pdf(
 
         try:
             # Process the PDF
-            result = await process_pdf_with_ocr(temp_file_path)
+            result = await process_pdf_with_ocr(temp_file_path, max_workers)
 
             # Clean up temporary file
             background_tasks.add_task(os.unlink, temp_file_path)