#!/usr/bin/env python3 """ Image processing script for OCR and entity extraction using OpenAI-compatible API. Processes images from Downloads folder and extracts structured data. """ import os import json import re import base64 from pathlib import Path from typing import Dict, List, Optional import concurrent.futures from dataclasses import dataclass, asdict from openai import OpenAI from tqdm import tqdm import argparse from dotenv import load_dotenv @dataclass class ProcessingResult: """Structure for processing results""" filename: str success: bool data: Optional[Dict] = None error: Optional[str] = None class ImageProcessor: """Process images using OpenAI-compatible vision API""" def __init__(self, api_url: str, api_key: str, model: str = "gpt-4o", index_file: str = "processing_index.json", downloads_dir: Optional[str] = None): self.client = OpenAI(api_key=api_key, base_url=api_url) self.model = model self.downloads_dir = Path(downloads_dir) if downloads_dir else Path.home() / "Downloads" self.index_file = index_file self.processed_files = self.load_index() def load_index(self) -> set: """Load the index of already processed files""" if os.path.exists(self.index_file): try: with open(self.index_file, 'r') as f: data = json.load(f) return set(data.get('processed_files', [])) except Exception as e: print(f"⚠️ Warning: Could not load index file: {e}") return set() return set() def save_index(self, failed_files=None): """Save the current index of processed files""" data = { 'processed_files': sorted(list(self.processed_files)), 'last_updated': str(Path.cwd()) } if failed_files: data['failed_files'] = failed_files with open(self.index_file, 'w') as f: json.dump(data, f, indent=2) def mark_processed(self, filename: str): """Mark a file as processed and update index""" self.processed_files.add(filename) self.save_index() def get_image_files(self) -> List[Path]: """Get all image files from Downloads folder (recursively)""" image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'} image_files = [] for ext in image_extensions: image_files.extend(self.downloads_dir.glob(f'**/*{ext}')) image_files.extend(self.downloads_dir.glob(f'**/*{ext.upper()}')) return sorted(image_files) def get_relative_path(self, file_path: Path) -> str: """Get relative path from downloads directory for unique indexing""" try: return str(file_path.relative_to(self.downloads_dir)) except ValueError: # If file is not relative to downloads_dir, use full path return str(file_path) def get_unprocessed_files(self) -> List[Path]: """Get only files that haven't been processed yet""" all_files = self.get_image_files() unprocessed = [f for f in all_files if self.get_relative_path(f) not in self.processed_files] return unprocessed def encode_image(self, image_path: Path) -> str: """Encode image to base64""" with open(image_path, 'rb') as f: return base64.b64encode(f.read()).decode('utf-8') def get_system_prompt(self) -> str: """Get the system prompt for structured extraction""" return """You are an expert OCR and document analysis system. Extract ALL text from the image in READING ORDER to create a digital twin of the document. IMPORTANT: Transcribe text exactly as it appears on the page, from top to bottom, left to right, including: - All printed text - All handwritten text (inline where it appears) - Stamps and annotations (inline where they appear) - Signatures (note location) Preserve the natural reading flow. Mix printed and handwritten text together in the order they appear. Return ONLY valid JSON in this exact structure: { "document_metadata": { "page_number": "string or null", "document_number": "string or null", "date": "string or null", "document_type": "string or null", "has_handwriting": true/false, "has_stamps": true/false }, "full_text": "Complete text transcription in reading order. Include ALL text - printed, handwritten, stamps, etc. - exactly as it appears from top to bottom.", "text_blocks": [ { "type": "printed|handwritten|stamp|signature|other", "content": "text content", "position": "top|middle|bottom|header|footer|margin" } ], "entities": { "people": ["list of person names"], "organizations": ["list of organizations"], "locations": ["list of locations"], "dates": ["list of dates found"], "reference_numbers": ["list of any reference/ID numbers"] }, "additional_notes": "Any observations about document quality, redactions, damage, etc." }""" def fix_json_with_llm(self, base64_image: str, broken_json: str, error_msg: str) -> dict: """Ask the LLM to fix its own broken JSON""" response = self.client.chat.completions.create( model=self.model, messages=[ { "role": "system", "content": self.get_system_prompt() }, { "role": "user", "content": [ { "type": "text", "text": "Extract all text and entities from this image. Return only valid JSON." }, { "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{base64_image}" } } ] }, { "role": "assistant", "content": broken_json }, { "role": "user", "content": f"Your JSON response has an error: {error_msg}\n\nPlease fix the JSON and return ONLY the corrected valid JSON. Do not explain, just return the fixed JSON." } ], max_tokens=4096, temperature=0.1 ) content = response.choices[0].message.content.strip() # Extract JSON using same logic json_match = re.search(r'```(?:json)?\s*\n(.*?)\n```', content, re.DOTALL) if json_match: content = json_match.group(1).strip() else: json_match = re.search(r'\{.*\}', content, re.DOTALL) if json_match: content = json_match.group(0).strip() return json.loads(content) def process_image(self, image_path: Path) -> ProcessingResult: """Process a single image through the API""" try: # Encode image base64_image = self.encode_image(image_path) # Make API call using OpenAI client response = self.client.chat.completions.create( model=self.model, messages=[ { "role": "system", "content": self.get_system_prompt() }, { "role": "user", "content": [ { "type": "text", "text": "Extract all text and entities from this image. Return only valid JSON." }, { "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{base64_image}" } } ] } ], max_tokens=4096, temperature=0.1 ) # Parse response content = response.choices[0].message.content original_content = content # Keep original for retry # Robust JSON extraction content = content.strip() # 1. Try to find JSON between markdown code fences json_match = re.search(r'```(?:json)?\s*\n(.*?)\n```', content, re.DOTALL) if json_match: content = json_match.group(1).strip() else: # 2. Try to find JSON between curly braces json_match = re.search(r'\{.*\}', content, re.DOTALL) if json_match: content = json_match.group(0).strip() else: # 3. Strip markdown manually if content.startswith('```json'): content = content[7:] elif content.startswith('```'): content = content[3:] if content.endswith('```'): content = content[:-3] content = content.strip() # Try to parse JSON try: extracted_data = json.loads(content) except json.JSONDecodeError as e: # Try to salvage by finding the first complete JSON object try: # Find first { and matching } start = content.find('{') if start == -1: raise ValueError("No JSON object found") brace_count = 0 end = start for i in range(start, len(content)): if content[i] == '{': brace_count += 1 elif content[i] == '}': brace_count -= 1 if brace_count == 0: end = i + 1 break if end > start: content = content[start:end] extracted_data = json.loads(content) else: raise ValueError("Could not find complete JSON object") except Exception: # Last resort: Ask LLM to fix its JSON try: extracted_data = self.fix_json_with_llm(base64_image, original_content, str(e)) except Exception: # Save ORIGINAL LLM response to errors directory (not our extracted version) self.save_broken_json(self.get_relative_path(image_path), original_content) # If even that fails, raise the original error raise e return ProcessingResult( filename=self.get_relative_path(image_path), success=True, data=extracted_data ) except Exception as e: return ProcessingResult( filename=self.get_relative_path(image_path), success=False, error=str(e) ) def process_all(self, max_workers: int = 5, limit: Optional[int] = None, resume: bool = True) -> List[ProcessingResult]: """Process all images with parallel processing""" if resume: image_files = self.get_unprocessed_files() total_files = len(self.get_image_files()) already_processed = len(self.processed_files) print(f"Found {total_files} total image files") print(f"Already processed: {already_processed}") print(f"Remaining to process: {len(image_files)}") else: image_files = self.get_image_files() print(f"Found {len(image_files)} image files to process") if limit: image_files = image_files[:limit] print(f"Limited to {limit} files for this run") if not image_files: print("No files to process!") return [] results = [] failed_files = [] with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: futures = {executor.submit(self.process_image, img): img for img in image_files} with tqdm(total=len(image_files), desc="Processing images") as pbar: for future in concurrent.futures.as_completed(futures): result = future.result() results.append(result) # Save individual result to file if result.success: self.save_individual_result(result) tqdm.write(f"✅ Processed: {result.filename}") else: # Track failed files failed_files.append({ 'filename': result.filename, 'error': result.error }) tqdm.write(f"❌ Failed: {result.filename} - {result.error}") # Mark as processed regardless of success/failure self.mark_processed(result.filename) pbar.update(1) # Save failed files to index for reference if failed_files: self.save_index(failed_files=failed_files) print(f"\n⚠️ {len(failed_files)} files failed - logged in {self.index_file}") return results def save_individual_result(self, result: ProcessingResult): """Save individual result to ./results/folder/imagename.json""" # Create output path mirroring the source structure result_path = Path("./results") / result.filename result_path = result_path.with_suffix('.json') # Create parent directories result_path.parent.mkdir(parents=True, exist_ok=True) # Save the extracted data with open(result_path, 'w', encoding='utf-8') as f: json.dump(result.data, f, indent=2, ensure_ascii=False) def save_broken_json(self, filename: str, broken_content: str): """Save broken JSON to errors directory""" error_path = Path("./errors") / filename error_path = error_path.with_suffix('.json') # Create parent directories error_path.parent.mkdir(parents=True, exist_ok=True) # Save the broken content as-is with open(error_path, 'w', encoding='utf-8') as f: f.write(broken_content) def save_results(self, results: List[ProcessingResult], output_file: str = "processed_results.json"): """Save summary results to JSON file""" output_data = { "total_processed": len(results), "successful": sum(1 for r in results if r.success), "failed": sum(1 for r in results if not r.success), "results": [asdict(r) for r in results] } with open(output_file, 'w', encoding='utf-8') as f: json.dump(output_data, f, indent=2, ensure_ascii=False) print(f"\n✅ Summary saved to {output_file}") print(f" Individual results saved to ./results/") print(f" Successful: {output_data['successful']}") print(f" Failed: {output_data['failed']}") def main(): # Load environment variables load_dotenv() parser = argparse.ArgumentParser(description="Process images with OCR and entity extraction") parser.add_argument("--api-url", help="OpenAI-compatible API base URL (default: from .env or OPENAI_API_URL)") parser.add_argument("--api-key", help="API key (default: from .env or OPENAI_API_KEY)") parser.add_argument("--model", help="Model name (default: from .env, OPENAI_MODEL, or meta-llama/Llama-4-Maverick-17B-128E-Instruct)") parser.add_argument("--workers", type=int, default=5, help="Number of parallel workers (default: 5)") parser.add_argument("--limit", type=int, help="Limit number of images to process (for testing)") parser.add_argument("--output", default="processed_results.json", help="Output JSON file") parser.add_argument("--index", default="processing_index.json", help="Index file to track processed files") parser.add_argument("--downloads-dir", default="./downloads", help="Directory containing images (default: ./downloads)") parser.add_argument("--no-resume", action="store_true", help="Process all files, ignoring index") args = parser.parse_args() # Get values from args or environment variables api_url = args.api_url or os.getenv("OPENAI_API_URL", "http://...") api_key = args.api_key or os.getenv("OPENAI_API_KEY", "abcd1234") model = args.model or os.getenv("OPENAI_MODEL", "meta-llama/Llama-4-Maverick-17B-128E-Instruct") processor = ImageProcessor( api_url=api_url, api_key=api_key, model=model, index_file=args.index, downloads_dir=args.downloads_dir ) results = processor.process_all( max_workers=args.workers, limit=args.limit, resume=not args.no_resume ) processor.save_results(results, args.output) if __name__ == "__main__": main()