From a7d832c894c24feb7bd7f4d04425b9d613dadfa0 Mon Sep 17 00:00:00 2001 From: Bram Date: Thu, 17 Jul 2025 19:20:55 +0200 Subject: [PATCH] ocr concurrently --- Dockers/ocr-api/main.py | 100 +++++++++++++++++++++++++++++++--------- 1 file changed, 78 insertions(+), 22 deletions(-) diff --git a/Dockers/ocr-api/main.py b/Dockers/ocr-api/main.py index bf377e1..021c70c 100644 --- a/Dockers/ocr-api/main.py +++ b/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)