boc/iom/visual_geolocation/ocr_pipeline.py

"""
OCR Pipeline
Extract all text from images for semantic geolocation
"""

from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
import numpy as np


@dataclass
class OCRResult:
    """OCR detection result"""
    text: str
    confidence: float
    bbox: List[int]  # [x1, y1, x2, y2]
    language: Optional[str] = None
    text_type: Optional[str] = None  # street_name, business_name, license_plate, etc.


class OCRPipeline:
    """
    Multi-engine OCR pipeline
    
    Supports:
    - Tesseract (local) — requires tesseract binary
    - EasyOCR (local) — requires easyocr package
    - Simulated (demo) — for testing without OCR engine
    """
    
    def __init__(self, engine: str = "auto"):
        self.engine = engine
        self.available = False
        self._init_engine()
    
    def _init_engine(self):
        """Initialize OCR engine"""
        if self.engine == "auto":
            # Try Tesseract first
            if self._check_tesseract():
                self.engine = "tesseract"
                self.available = True
            # Then EasyOCR
            elif self._check_easyocr():
                self.engine = "easyocr"
                self.available = True
            # Fall back to simulated
            else:
                self.engine = "simulated"
                self.available = True
        
        elif self.engine == "tesseract":
            self.available = self._check_tesseract()
        
        elif self.engine == "easyocr":
            self.available = self._check_easyocr()
        
        elif self.engine == "simulated":
            self.available = True
    
    def _check_tesseract(self) -> bool:
        """Check if Tesseract is available"""
        try:
            import pytesseract
            # Try to get version
            pytesseract.get_tesseract_version()
            self.tesseract = pytesseract
            return True
        except:
            return False
    
    def _check_easyocr(self) -> bool:
        """Check if EasyOCR is available"""
        try:
            import easyocr
            self.reader = easyocr.Reader(['en', 'sv', 'th'])
            return True
        except:
            return False
    
    def extract_text(self, image_path: str) -> List[OCRResult]:
        """Extract all text from image"""
        if not self.available:
            return []
        
        if self.engine == "tesseract":
            return self._extract_tesseract(image_path)
        elif self.engine == "easyocr":
            return self._extract_easyocr(image_path)
        elif self.engine == "simulated":
            return self._extract_simulated(image_path)
        
        return []
    
    def _extract_tesseract(self, image_path: str) -> List[OCRResult]:
        """Extract text using Tesseract"""
        from PIL import Image
        
        img = Image.open(image_path)
        
        # Get detailed OCR data
        data = self.tesseract.image_to_data(img, output_type=self.tesseract.Output.DICT)
        
        results = []
        n_boxes = len(data["text"])
        
        for i in range(n_boxes):
            conf = int(data["conf"][i])
            if conf > 30:  # Confidence threshold
                text = data["text"][i].strip()
                if text:
                    x, y, w, h = data["left"][i], data["top"][i], data["width"][i], data["height"][i]
                    
                    result = OCRResult(
                        text=text,
                        confidence=conf / 100.0,
                        bbox=[x, y, x + w, y + h],
                        language=None,
                        text_type=self._classify_text(text)
                    )
                    results.append(result)
        
        return results
    
    def _extract_easyocr(self, image_path: str) -> List[OCRResult]:
        """Extract text using EasyOCR"""
        results = self.reader.readtext(image_path)
        
        ocr_results = []
        for (bbox, text, conf) in results:
            x1, y1 = int(bbox[0][0]), int(bbox[0][1])
            x2, y2 = int(bbox[2][0]), int(bbox[2][1])
            
            result = OCRResult(
                text=text,
                confidence=conf,
                bbox=[x1, y1, x2, y2],
                language=None,
                text_type=self._classify_text(text)
            )
            ocr_results.append(result)
        
        return ocr_results
    
    def _extract_simulated(self, image_path: str) -> List[OCRResult]:
        """Simulated OCR for demo purposes"""
        # In production, this would be replaced with real OCR
        # For demo, return simulated results based on image analysis
        
        from PIL import Image
        import random
        
        img = Image.open(image_path)
        width, height = img.size
        
        # Simulate text detection based on image characteristics
        # In reality, this would use actual OCR
        
        simulated_texts = [
            ("Sukhumvit Road", "street_name", 0.92),
            ("Bangkok", "city_name", 0.88),
            ("10110", "postal_code", 0.95),
            ("Restaurant", "business_name", 0.75),
            ("MASSAGE", "business_name", 0.82),
        ]
        
        results = []
        for text, text_type, conf in simulated_texts:
            # Random position
            x = random.randint(50, width - 200)
            y = random.randint(50, height - 100)
            w = len(text) * 15
            h = 30
            
            result = OCRResult(
                text=text,
                confidence=conf,
                bbox=[x, y, x + w, y + h],
                language="en",
                text_type=text_type
            )
            results.append(result)
        
        return results
    
    def _classify_text(self, text: str) -> Optional[str]:
        """Classify text type (street name, business, etc.)"""
        text_lower = text.lower()
        
        # Street name patterns
        street_keywords = ["street", "st", "avenue", "ave", "road", "rd", 
                          "boulevard", "blvd", "lane", "ln", "gatan", "vägen"]
        if any(kw in text_lower for kw in street_keywords):
            return "street_name"
        
        # Business name patterns
        business_keywords = ["restaurant", "cafe", "shop", "store", "hotel",
                            "massage", "spa", "bar", "pub"]
        if any(kw in text_lower for kw in business_keywords):
            return "business_name"
        
        # License plate patterns (simplified)
        if len(text) <= 10 and any(c.isdigit() for c in text) and any(c.isalpha() for c in text):
            return "license_plate"
        
        # House number
        if text.isdigit() and len(text) <= 4:
            return "house_number"
        
        # Phone number
        if any(c.isdigit() for c in text) and len(text) >= 8:
            return "phone_number"
        
        # Postal code
        if text.isdigit() and len(text) == 5:
            return "postal_code"
        
        return "unknown"
    
    def extract_structured_data(self, image_path: str) -> Dict:
        """Extract structured data from image text"""
        results = self.extract_text(image_path)
        
        structured = {
            "street_names": [],
            "business_names": [],
            "license_plates": [],
            "house_numbers": [],
            "phone_numbers": [],
            "postal_codes": [],
            "other_text": []
        }
        
        for result in results:
            if result.text_type == "street_name":
                structured["street_names"].append(result.text)
            elif result.text_type == "business_name":
                structured["business_names"].append(result.text)
            elif result.text_type == "license_plate":
                structured["license_plates"].append(result.text)
            elif result.text_type == "house_number":
                structured["house_numbers"].append(result.text)
            elif result.text_type == "phone_number":
                structured["phone_numbers"].append(result.text)
            elif result.text_type == "postal_code":
                structured["postal_codes"].append(result.text)
            else:
                structured["other_text"].append(result.text)
        
        return structured


# Example usage
def test_ocr():
    """Test OCR pipeline"""
    print("=== Testing OCR Pipeline ===\n")
    
    ocr = OCRPipeline(engine="simulated")
    
    print(f"Using engine: {ocr.engine}")
    print(f"Available: {ocr.available}\n")
    
    # Create test image
    from PIL import Image, ImageDraw, ImageFont
    
    img = Image.new('RGB', (640, 480), color='white')
    draw = ImageDraw.Draw(img)
    
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 30)
    except:
        font = ImageFont.load_default()
    
    draw.text((50, 50), "Sukhumvit Road", fill='black', font=font)
    draw.text((50, 150), "Bangkok 10110", fill='black', font=font)
    draw.text((50, 250), "Sainokuni Restaurant", fill='black', font=font)
    
    img.save('/tmp/test_ocr.jpg')
    
    # Extract text
    results = ocr.extract_text('/tmp/test_ocr.jpg')
    
    print(f"Detected {len(results)} text regions:")
    for result in results:
        print(f"  - '{result.text}' ({result.text_type}, conf: {result.confidence:.2f})")
    
    # Extract structured data
    structured = ocr.extract_structured_data('/tmp/test_ocr.jpg')
    
    print(f"\nStructured data:")
    for key, values in structured.items():
        if values:
            print(f"  {key}: {values}")
    
    return results


if __name__ == '__main__':
    test_ocr()