import time import json import os from PIL import Image # Optional import of Unstract LLMWhisperer client – may not be installed in all environments try: from unstract.llmwhisperer import LLMWhispererClientV2 except ImportError: LLMWhispererClientV2 = None print("[WARN] unstract.llmwhisperer module not found – OCR functionality will be disabled.") from dotenv import load_dotenv # Load environment variables load_dotenv() api_key = os.getenv("UNSTRANCT_API_KEY") # Initialize client only if API key is available client = None if LLMWhispererClientV2 is not None and api_key and api_key != "your_api_key_here": try: client = LLMWhispererClientV2(base_url="https://llmwhisperer-api.us-central.unstract.com/api/v2", api_key=api_key) except Exception as e: print(f"Warning: Failed to initialize OCR client: {e}") client = None else: if LLMWhispererClientV2 is None: print("[WARN] OCR client library unavailable – OCR will be skipped.") else: print("Warning: No valid OCR API key found. OCR will be skipped.") # Create OCR results directory if it doesn't exist results_dir = "data/output/ocr_results" if not os.path.exists(results_dir): os.makedirs(results_dir, exist_ok=True) print(f"Created results directory: {results_dir}") def get_max_confidence_table(tables, table_type): """Get the table with maximum confidence for a given table type""" filtered_tables = [table for table in tables if table["table_type"] == table_type] if not filtered_tables: return None return max(filtered_tables, key=lambda x: x["confidence"]) def crop_table_from_image(image_path, coordinates, margin_ratio: float = 0.10): """Crop table region from image using coordinates with optional margin expansion.""" with Image.open(image_path) as img: # Use full width of image to prevent horizontal clipping of columns x1 = 0 x2 = img.width y1 = int(coordinates["y1"]) y2 = int(coordinates["y2"]) # Expand vertically by margin_ratio height = max(0, y2 - y1) pad_y = int(round(height * margin_ratio)) y1 -= pad_y y2 += pad_y # Ensure coordinates are within image bounds y1 = max(0, min(y1, img.height)) y2 = max(0, min(y2, img.height)) # .copy() detaches the crop from the file handle so it survives the 'with' block cropped_img = img.crop((x1, y1, x2, y2)).copy() return cropped_img def save_cropped_image(cropped_img, filename, table_type): """Save cropped image temporarily for OCR processing""" temp_filename = f"temp_{filename}_{table_type.replace(' ', '_').lower()}.jpg" # Ensure image is in RGB mode for JPEG compatibility if cropped_img.mode != 'RGB': try: rgb_img = cropped_img.convert('RGB') except Exception: # Fallback: create a new RGB image and paste from PIL import Image rgb_img = Image.new('RGB', cropped_img.size, (255, 255, 255)) rgb_img.paste(cropped_img, mask=cropped_img.split()[-1] if cropped_img.mode in ('RGBA', 'LA') else None) rgb_img.save(temp_filename, format='JPEG') rgb_img.close() # Release the intermediate conversion buffer. else: cropped_img.save(temp_filename, format='JPEG') return temp_filename def process_ocr(image_path): """Process OCR on an image and return extracted text""" if not client: return "OCR not available - no valid API key" try: result = client.whisper(file_path=image_path) while True: status = client.whisper_status(whisper_hash=result['whisper_hash']) if status['status'] == 'processed': resultx = client.whisper_retrieve(whisper_hash=result['whisper_hash']) break time.sleep(5) extracted_text = resultx['extraction']['result_text'] return extracted_text except Exception as e: return f"OCR failed: {str(e)}" def process_image_with_tables(image_path, json_path): """Process a single image with its corresponding JSON file""" print(f"Processing: {os.path.basename(image_path)}") # Load JSON data with open(json_path, 'r') as f: data = json.load(f) # Get filename without extension filename = os.path.splitext(os.path.basename(image_path))[0] # Extract table coordinates tables = data["table_detection"]["table_coordinates"] # Get tables with maximum confidence for each type marks_table = get_max_confidence_table(tables, "Marks Table") info_table = get_max_confidence_table(tables, "Information Table") # Process Marks Table print(f" Processing Marks Table...") cropped_img_marks = None if marks_table: print(f" Marks Table detected (confidence: {marks_table['confidence']:.3f})") cropped_img_marks = crop_table_from_image(image_path, marks_table["coordinates"], margin_ratio=0.10) else: print(f" Marks Table NOT detected. Falling back to cropping bottom 65% of the image...") try: with Image.open(image_path) as img: width, height = img.size cropped_img_marks = img.crop((0, int(height * 0.35), width, height)).copy() except Exception as e: print(f" Failed to crop bottom 65% of image: {e}") cropped_img_marks = None if cropped_img_marks: temp_file = save_cropped_image(cropped_img_marks, filename, "marks") cropped_img_marks.close() # No longer needed once the temp file is written. try: marks_text = process_ocr(temp_file) marks_output_path = os.path.join(results_dir, f"{filename}_marks.txt") with open(marks_output_path, 'w', encoding='utf-8') as f: f.write(marks_text) if "OCR not available" in marks_text or "OCR failed" in marks_text: print(f" Warning: {marks_text}") else: print(f" Saved marks table to: {marks_output_path}") except Exception as e: print(f" Error processing marks table: {e}") finally: # Clean up temporary file if os.path.exists(temp_file): os.remove(temp_file) else: print(f" Skipping marks table processing as no image was available") # Process Information Table print(f" Processing Information Table...") cropped_img_info = None if info_table: print(f" Information Table detected (confidence: {info_table['confidence']:.3f})") cropped_img_info = crop_table_from_image(image_path, info_table["coordinates"], margin_ratio=0.15) else: print(f" Information Table NOT detected. Falling back to cropping top 40% of the image...") try: with Image.open(image_path) as img: width, height = img.size cropped_img_info = img.crop((0, 0, width, int(height * 0.40))).copy() except Exception as e: print(f" Failed to crop top 40% of image: {e}") cropped_img_info = None if cropped_img_info: temp_file = save_cropped_image(cropped_img_info, filename, "info") cropped_img_info.close() # No longer needed once the temp file is written. try: info_text = process_ocr(temp_file) info_output_path = os.path.join(results_dir, f"{filename}_info.txt") with open(info_output_path, 'w', encoding='utf-8') as f: f.write(info_text) if "OCR not available" in info_text or "OCR failed" in info_text: print(f" Warning: {info_text}") else: print(f" Saved info table to: {info_output_path}") except Exception as e: print(f" Error processing info table: {e}") finally: # Clean up temporary file if os.path.exists(temp_file): os.remove(temp_file) else: print(f" Skipping info table processing as no image was available") if not marks_table and not info_table: print(f" No tables detected at all by the model.") def list_available_images(): """List all available images in the inputs folder""" inputs_dir = "data/input" if not os.path.exists(inputs_dir): print(f"Error: {inputs_dir} directory not found!") return [] # Get all image files and their corresponding JSON files image_files = [] for file in os.listdir(inputs_dir): if file.lower().endswith(('.jpg', '.jpeg', '.png')): # Extract base name (remove _preprocessed and extension) base_name = file.replace('_preprocessed', '').replace('.jpg', '').replace('.jpeg', '').replace('.png', '') json_file = f"{base_name}_result.json" json_path = os.path.join(inputs_dir, json_file) if os.path.exists(json_path): image_files.append((os.path.join(inputs_dir, file), json_path)) else: print(f"Warning: No JSON file found for {file} (looking for {json_file})") return image_files def process_single_image(image_path, json_path): """Process a single image""" print(f"Processing: {os.path.basename(image_path)}") try: process_image_with_tables(image_path, json_path) print("Processing completed successfully!") except Exception as e: print(f"Error processing {image_path}: {e}") def main(): """Main function with interactive image selection""" print("OCR Table Processing Tool") print("=" * 40) # Get available images image_files = list_available_images() if not image_files: print("No image files with corresponding JSON files found!") return while True: print(f"\nAvailable images ({len(image_files)} found):") print("-" * 30) for i, (image_path, json_path) in enumerate(image_files, 1): filename = os.path.basename(image_path) print(f"{i}. {filename}") print("\nCommands:") print(" - Process specific image (1-{})".format(len(image_files))) print(" all - Process all images") print(" list - Show this list again") print(" quit/exit - Exit the program") choice = input("\nEnter your choice: ").strip().lower() if choice in ['quit', 'exit', 'q']: print("Goodbye!") break elif choice == 'all': print(f"\nProcessing all {len(image_files)} images...") print("=" * 50) for i, (image_path, json_path) in enumerate(image_files, 1): print(f"[{i}/{len(image_files)}] Processing: {os.path.basename(image_path)}") try: process_image_with_tables(image_path, json_path) except Exception as e: print(f"Error processing {image_path}: {e}") print("-" * 30) print("All images processed!") elif choice == 'list': continue elif choice.isdigit(): image_index = int(choice) - 1 if 0 <= image_index < len(image_files): image_path, json_path = image_files[image_index] print(f"\nSelected: {os.path.basename(image_path)}") process_single_image(image_path, json_path) else: print(f"Invalid choice! Please enter a number between 1 and {len(image_files)}") else: print("Invalid choice! Please enter a number, 'all', 'list', or 'quit'") if __name__ == "__main__": main()