order_reader.py

#!/usr/bin/env python3
"""
SmartOrderReader - Invoice and Order Data Extraction Tool

This script processes invoice and order images to extract:
1. Order Number / Invoice Number / Order ID
2. Order Date / Creation Date / Invoice Date

It outputs the data in both table and CSV formats.
"""

import os
import sys
import re
from typing import List, Dict, Tuple, Optional
from pathlib import Path
import argparse

try:
    import pytesseract
    from PIL import Image
    import cv2
    import pandas as pd
    from tabulate import tabulate
except ImportError as e:
    print(f"Error: Missing required library. Please install dependencies:")
    print(f"  pip install -r requirements.txt")
    print(f"\nDetails: {e}")
    sys.exit(1)


class OrderDataExtractor:
    """Extracts order information from invoice/order images using OCR."""
    
    # Patterns for order numbers (support various characters)
    ORDER_NUMBER_PATTERNS = [
        r'order\s*#\s*:?\s*([A-Za-z0-9\-._/]+)',
        r'order\s*no\.?\s*:?\s*([A-Za-z0-9\-._/]+)',
        r'order\s*number\s*:?\s*([A-Za-z0-9\-._/]+)',
        r'order\s*id\s*:?\s*([A-Za-z0-9\-._/]+)',
        r'invoice\s*#\s*:?\s*([A-Za-z0-9\-._/]+)',
        r'invoice\s*no\.?\s*:?\s*([A-Za-z0-9\-._/]+)',
        r'invoice\s*number\s*:?\s*([A-Za-z0-9\-._/]+)',
        r'po\s*#\s*:?\s*([A-Za-z0-9\-._/]+)',
        r'reference\s*#\s*:?\s*([A-Za-z0-9\-._/]+)',
        r'ref\s*#\s*:?\s*([A-Za-z0-9\-._/]+)',
        # Fallback patterns with just colon
        r'order\s*:\s*([A-Za-z0-9\-._/]+)',
        r'invoice\s*:\s*([A-Za-z0-9\-._/]+)',
    ]
    
    # Patterns for dates
    DATE_PATTERNS = [
        r'order\s*date\s*:?\s*([0-9]{1,2}[/-][0-9]{1,2}[/-][0-9]{2,4})',
        r'invoice\s*date\s*:?\s*([0-9]{1,2}[/-][0-9]{1,2}[/-][0-9]{2,4})',
        r'date\s*:?\s*([0-9]{1,2}[/-][0-9]{1,2}[/-][0-9]{2,4})',
        r'created\s*on\s*:?\s*([0-9]{1,2}[/-][0-9]{1,2}[/-][0-9]{2,4})',
        r'issue\s*date\s*:?\s*([0-9]{1,2}[/-][0-9]{1,2}[/-][0-9]{2,4})',
        r'transaction\s*date\s*:?\s*([0-9]{1,2}[/-][0-9]{1,2}[/-][0-9]{2,4})',
        # ISO format dates
        r'order\s*date\s*:?\s*([0-9]{4}[/-][0-9]{1,2}[/-][0-9]{1,2})',
        r'invoice\s*date\s*:?\s*([0-9]{4}[/-][0-9]{1,2}[/-][0-9]{1,2})',
        r'date\s*:?\s*([0-9]{4}[/-][0-9]{1,2}[/-][0-9]{1,2})',
        # Month name formats
        r'order\s*date\s*:?\s*([A-Za-z]+\s+[0-9]{1,2},?\s*[0-9]{4})',
        r'invoice\s*date\s*:?\s*([A-Za-z]+\s+[0-9]{1,2},?\s*[0-9]{4})',
        r'date\s*:?\s*([A-Za-z]+\s+[0-9]{1,2},?\s*[0-9]{4})',
    ]
    
    def __init__(self, lang: str = 'eng'):
        """
        Initialize the extractor.
        
        Args:
            lang: Language for OCR. Use 'eng' for English, 'eng+urd' for English+Urdu
        """
        self.lang = lang
    
    def preprocess_image(self, image_path: str) -> Image.Image:
        """
        Preprocess image for better OCR results.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Preprocessed PIL Image
        """
        # Read image with OpenCV
        img = cv2.imread(image_path)
        
        if img is None:
            raise ValueError(f"Could not read image: {image_path}")
        
        # Convert to grayscale
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Apply thresholding to get better contrast
        _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        
        # Convert back to PIL Image
        return Image.fromarray(thresh)
    
    def extract_text(self, image_path: str) -> str:
        """
        Extract text from image using OCR.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Extracted text
        """
        try:
            # Preprocess image
            img = self.preprocess_image(image_path)
            
            # Perform OCR with specified language
            text = pytesseract.image_to_string(img, lang=self.lang)
            
            return text
        except Exception as e:
            print(f"Warning: Error extracting text from {image_path}: {e}")
            return ""
    
    def extract_order_number(self, text: str) -> str:
        """
        Extract order number from text.
        
        Args:
            text: OCR extracted text
            
        Returns:
            Order number or "Not Found"
        """
        # Try each pattern on original text to preserve case
        for pattern in self.ORDER_NUMBER_PATTERNS:
            match = re.search(pattern, text, re.IGNORECASE | re.MULTILINE)
            if match:
                # Preserve original case from the document
                return match.group(1).strip()
        
        return "Not Found"
    
    def extract_order_date(self, text: str) -> str:
        """
        Extract order date from text.
        
        Args:
            text: OCR extracted text
            
        Returns:
            Order date or "Not Found"
        """
        # Try each pattern on the original text to preserve case
        for pattern in self.DATE_PATTERNS:
            match = re.search(pattern, text, re.IGNORECASE | re.MULTILINE)
            if match:
                return match.group(1).strip()
        
        return "Not Found"
    
    def process_image(self, image_path: str) -> Dict[str, str]:
        """
        Process a single image and extract order data.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Dictionary with image filename, order number, and order date
        """
        filename = os.path.basename(image_path)
        
        try:
            # Extract text
            text = self.extract_text(image_path)
            
            if not text.strip():
                print(f"Warning: No text extracted from {filename}")
                return {
                    'Image File Name': filename,
                    'Order Number': 'Not Found',
                    'Order Date': 'Not Found'
                }
            
            # Extract order number and date
            order_number = self.extract_order_number(text)
            order_date = self.extract_order_date(text)
            
            return {
                'Image File Name': filename,
                'Order Number': order_number,
                'Order Date': order_date
            }
            
        except Exception as e:
            print(f"Error processing {filename}: {e}")
            return {
                'Image File Name': filename,
                'Order Number': 'Not Found',
                'Order Date': 'Not Found'
            }
    
    def process_images(self, image_paths: List[str]) -> List[Dict[str, str]]:
        """
        Process multiple images.
        
        Args:
            image_paths: List of image file paths
            
        Returns:
            List of dictionaries with extracted data
        """
        results = []
        
        for image_path in image_paths:
            print(f"Processing: {os.path.basename(image_path)}...")
            result = self.process_image(image_path)
            results.append(result)
        
        return results


def format_output(results: List[Dict[str, str]]) -> str:
    """
    Format results as table and CSV.
    
    Args:
        results: List of extraction results
        
    Returns:
        Formatted output string
    """
    output = []
    
    # Create DataFrame
    df = pd.DataFrame(results)
    
    # Table format
    output.append("\n" + "="*80)
    output.append("EXTRACTED ORDER DATA")
    output.append("="*80 + "\n")
    
    table = tabulate(df, headers='keys', tablefmt='grid', showindex=False)
    output.append(table)
    
    # CSV format
    output.append("\n" + "="*80)
    output.append("CSV FORMAT (Copy & Paste Ready)")
    output.append("="*80 + "\n")
    
    csv_output = df.to_csv(index=False)
    output.append(csv_output)
    
    # Summary statistics
    output.append("="*80)
    output.append("SUMMARY")
    output.append("="*80 + "\n")
    
    total_images = len(results)
    missing_order_number = sum(1 for r in results if r['Order Number'] == 'Not Found')
    missing_order_date = sum(1 for r in results if r['Order Date'] == 'Not Found')
    
    output.append(f"Total images processed: {total_images}")
    output.append(f"Images with missing Order Number: {missing_order_number}")
    output.append(f"Images with missing Order Date: {missing_order_date}")
    
    if missing_order_number > 0 or missing_order_date > 0:
        output.append(f"\nNote: {missing_order_number + missing_order_date} field(s) could not be extracted.")
    else:
        output.append("\nAll data extracted successfully!")
    
    return "\n".join(output)


def save_csv(results: List[Dict[str, str]], output_file: str = "order_data.csv"):
    """
    Save results to CSV file.
    
    Args:
        results: List of extraction results
        output_file: Output CSV filename
    """
    df = pd.DataFrame(results)
    df.to_csv(output_file, index=False)
    print(f"\nCSV file saved: {output_file}")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description='SmartOrderReader - Extract order data from invoice/order images',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s image1.jpg image2.png
  %(prog)s *.jpg --output results.csv
  %(prog)s invoice.png --lang eng+urd
        """
    )
    
    parser.add_argument(
        'images',
        nargs='+',
        help='Image files to process (supports jpg, png, jpeg, bmp, tiff)'
    )
    
    parser.add_argument(
        '-o', '--output',
        default='order_data.csv',
        help='Output CSV filename (default: order_data.csv)'
    )
    
    parser.add_argument(
        '-l', '--lang',
        default='eng',
        help='OCR language (default: eng, use eng+urd for English+Urdu)'
    )
    
    args = parser.parse_args()
    
    # Validate image files
    valid_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
    image_paths = []
    
    for img in args.images:
        if not os.path.exists(img):
            print(f"Warning: File not found: {img}")
            continue
        
        if Path(img).suffix.lower() not in valid_extensions:
            print(f"Warning: Unsupported file format: {img}")
            continue
        
        image_paths.append(img)
    
    if not image_paths:
        print("Error: No valid image files provided.")
        sys.exit(1)
    
    print(f"\nSmartOrderReader - Processing {len(image_paths)} image(s)...\n")
    
    # Initialize extractor
    extractor = OrderDataExtractor(lang=args.lang)
    
    # Process images
    results = extractor.process_images(image_paths)
    
    # Display results
    output = format_output(results)
    print(output)
    
    # Save to CSV
    save_csv(results, args.output)


if __name__ == '__main__':
    main()