ocr_engine.py

import pytesseract
import numpy as np
import cv2
import re
import logging
from datetime import datetime
import os
from PIL import Image

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Directory for debug images
DEBUG_DIR = "debug_images"
os.makedirs(DEBUG_DIR, exist_ok=True)

def save_debug_image(img, filename_suffix, prefix=""):
    """Save image to debug directory with timestamp."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    filename = os.path.join(DEBUG_DIR, f"{prefix}{timestamp}_{filename_suffix}.png")
    if isinstance(img, Image.Image):
        img.save(filename)
    elif len(img.shape) == 3:
        cv2.imwrite(filename, cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    else:
        cv2.imwrite(filename, img)
    logging.info(f"Saved debug image: {filename}")

def estimate_brightness(img):
    """Estimate image brightness."""
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return np.mean(gray)

def preprocess_image(img):
    """Preprocess image with simplified, robust contrast enhancement."""
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    brightness = estimate_brightness(img)
    
    # Apply mild CLAHE for contrast
    clahe_clip = 8.0 if brightness < 90 else 4.0
    clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=(8, 8))
    enhanced = clahe.apply(gray)
    save_debug_image(enhanced, "01_preprocess_clahe")
    
    # Light blur to reduce noise
    blurred = cv2.GaussianBlur(enhanced, (5, 5), 0)
    save_debug_image(blurred, "02_preprocess_blur")
    
    # Dynamic thresholding with larger block size for small displays
    block_size = max(7, min(31, int(img.shape[0] / 20) * 2 + 1))
    thresh = cv2.adaptiveThreshold(
        blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
        cv2.THRESH_BINARY_INV, block_size, 3
    )
    
    # Minimal morphological operations
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
    save_debug_image(thresh, "03_preprocess_morph")
    return thresh, enhanced

def correct_rotation(img):
    """Correct image rotation using edge detection."""
    try:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 30, 100, apertureSize=3)
        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=25, minLineLength=15, maxLineGap=10)
        if lines is not None:
            angles = [np.arctan2(line[0][3] - line[0][1], line[0][2] - line[0][0]) * 180 / np.pi for line in lines]
            angle = np.median(angles)
            if abs(angle) > 0.3:
                h, w = img.shape[:2]
                center = (w // 2, h // 2)
                M = cv2.getRotationMatrix2D(center, angle, 1.0)
                img = cv2.warpAffine(img, M, (w, h))
                save_debug_image(img, "00_rotated_image")
                logging.info(f"Applied rotation: {angle:.2f} degrees")
        return img
    except Exception as e:
        logging.error(f"Rotation correction failed: {str(e)}")
        return img

def detect_roi(img):
    """Detect region of interest with broader contour analysis."""
    try:
        save_debug_image(img, "04_original")
        thresh, enhanced = preprocess_image(img)
        brightness_map = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        block_sizes = [max(7, min(31, int(img.shape[0] / s) * 2 + 1)) for s in [5, 10, 20]]
        valid_contours = []
        img_area = img.shape[0] * img.shape[1]
        
        for block_size in block_sizes:
            temp_thresh = cv2.adaptiveThreshold(
                enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                cv2.THRESH_BINARY_INV, block_size, 3
            )
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
            temp_thresh = cv2.morphologyEx(temp_thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
            save_debug_image(temp_thresh, f"05_roi_threshold_block{block_size}")
            contours, _ = cv2.findContours(temp_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            for c in contours:
                area = cv2.contourArea(c)
                x, y, w, h = cv2.boundingRect(c)
                roi_brightness = np.mean(brightness_map[y:y+h, x:x+w])
                aspect_ratio = w / h
                if (50 < area < (img_area * 0.95) and 
                    0.05 <= aspect_ratio <= 20.0 and w > 20 and h > 8 and roi_brightness > 15):
                    valid_contours.append((c, area * roi_brightness))
                    logging.debug(f"Contour (block {block_size}): Area={area}, Aspect={aspect_ratio:.2f}, Brightness={roi_brightness:.2f}")
        
        if valid_contours:
            contour, _ = max(valid_contours, key=lambda x: x[1])
            x, y, w, h = cv2.boundingRect(contour)
            padding = max(5, min(20, int(min(w, h) * 0.4)))
            x, y = max(0, x - padding), max(0, y - padding)
            w, h = min(w + 2 * padding, img.shape[1] - x), min(h + 2 * padding, img.shape[0] - y)
            roi_img = img[y:y+h, x:x+w]
            save_debug_image(roi_img, "06_detected_roi")
            logging.info(f"Detected ROI: ({x}, {y}, {w}, {h})")
            return roi_img, (x, y, w, h)
        
        logging.info("No ROI found, using full image.")
        save_debug_image(img, "06_no_roi_fallback")
        return img, None
    except Exception as e:
        logging.error(f"ROI detection failed: {str(e)}")
        save_debug_image(img, "06_roi_error_fallback")
        return img, None

def detect_digit_template(digit_img, brightness):
    """Digit recognition with expanded template matching."""
    try:
        h, w = digit_img.shape
        if h < 5 or w < 2:
            logging.debug("Digit image too small for template matching.")
            return None

        # Expanded digit templates for seven-segment display variations
        digit_templates = {
            '0': [
                np.array([[1, 1, 1, 1, 1],
                          [1, 0, 0, 0, 1],
                          [1, 0, 0, 0, 1],
                          [1, 0, 0, 0, 1],
                          [1, 1, 1, 1, 1]], dtype=np.float32),
                np.array([[1, 1, 1, 1],
                          [1, 0, 0, 1],
                          [1, 0, 0, 1],
                          [1, 0, 0, 1],
                          [1, 1, 1, 1]], dtype=np.float32)
            ],
            '1': [
                np.array([[0, 0, 1, 0, 0],
                          [0, 0, 1, 0, 0],
                          [0, 0, 1, 0, 0],
                          [0, 0, 1, 0, 0],
                          [0, 0, 1, 0, 0]], dtype=np.float32),
                np.array([[0, 1, 0],
                          [0, 1, 0],
                          [0, 1, 0],
                          [0, 1, 0],
                          [0, 1, 0]], dtype=np.float32)
            ],
            '2': [
                np.array([[1, 1, 1, 1, 1],
                          [0, 0, 0, 1, 1],
                          [1, 1, 1, 1, 1],
                          [1, 1, 0, 0, 0],
                          [1, 1, 1, 1, 1]], dtype=np.float32),
                np.array([[1, 1, 1, 1],
                          [0, 0, 1, 1],
                          [1, 1, 1, 1],
                          [1, 1, 0, 0],
                          [1, 1, 1, 1]], dtype=np.float32)
            ],
            '3': [
                np.array([[1, 1, 1, 1, 1],
                          [0, 0, 0, 1, 1],
                          [1, 1, 1, 1, 1],
                          [0, 0, 0, 1, 1],
                          [1, 1, 1, 1, 1]], dtype=np.float32),
                np.array([[1, 1, 1, 1],
                          [0, 0, 1, 1],
                          [1, 1, 1, 1],
                          [0, 0, 1, 1],
                          [1, 1, 1, 1]], dtype=np.float32)
            ],
            '4': [
                np.array([[1, 1, 0, 0, 1],
                          [1, 1, 0, 0, 1],
                          [1, 1, 1, 1, 1],
                          [0, 0, 0, 0, 1],
                          [0, 0, 0, 0, 1]], dtype=np.float32),
                np.array([[1, 0, 0, 1],
                          [1, 0, 0, 1],
                          [1, 1, 1, 1],
                          [0, 0, 0, 1],
                          [0, 0, 0, 1]], dtype=np.float32)
            ],
            '5': [
                np.array([[1, 1, 1, 1, 1],
                          [1, 1, 0, 0, 0],
                          [1, 1, 1, 1, 1],
                          [0, 0, 0, 1, 1],
                          [1, 1, 1, 1, 1]], dtype=np.float32),
                np.array([[1, 1, 1, 1],
                          [1, 1, 0, 0],
                          [1, 1, 1, 1],
                          [0, 0, 1, 1],
                          [1, 1, 1, 1]], dtype=np.float32)
            ],
            '6': [
                np.array([[1, 1, 1, 1, 1],
                          [1, 1, 0, 0, 0],
                          [1, 1, 1, 1, 1],
                          [1, 0, 0, 1, 1],
                          [1, 1, 1, 1, 1]], dtype=np.float32),
                np.array([[1, 1, 1, 1],
                          [1, 1, 0, 0],
                          [1, 1, 1, 1],
                          [1, 0, 1, 1],
                          [1, 1, 1, 1]], dtype=np.float32)
            ],
            '7': [
                np.array([[1, 1, 1, 1, 1],
                          [0, 0, 0, 0, 1],
                          [0, 0, 0, 0, 1],
                          [0, 0, 0, 0, 1],
                          [0, 0, 0, 0, 1]], dtype=np.float32),
                np.array([[1, 1, 1, 1],
                          [0, 0, 0, 1],
                          [0, 0, 0, 1],
                          [0, 0, 0, 1],
                          [0, 0, 0, 1]], dtype=np.float32)
            ],
            '8': [
                np.array([[1, 1, 1, 1, 1],
                          [1, 0, 0, 0, 1],
                          [1, 1, 1, 1, 1],
                          [1, 0, 0, 0, 1],
                          [1, 1, 1, 1, 1]], dtype=np.float32),
                np.array([[1, 1, 1, 1],
                          [1, 0, 0, 1],
                          [1, 1, 1, 1],
                          [1, 0, 0, 1],
                          [1, 1, 1, 1]], dtype=np.float32)
            ],
            '9': [
                np.array([[1, 1, 1, 1, 1],
                          [1, 0, 0, 0, 1],
                          [1, 1, 1, 1, 1],
                          [0, 0, 0, 1, 1],
                          [1, 1, 1, 1, 1]], dtype=np.float32),
                np.array([[1, 1, 1, 1],
                          [1, 0, 0, 1],
                          [1, 1, 1, 1],
                          [0, 0, 1, 1],
                          [1, 1, 1, 1]], dtype=np.float32)
            ],
            '.': [
                np.array([[0, 0, 0],
                          [0, 1, 0],
                          [0, 0, 0]], dtype=np.float32),
                np.array([[0, 0],
                          [1, 0],
                          [0, 0]], dtype=np.float32)
            ]
        }

        # Try multiple sizes for digit image
        sizes = [(5, 5), (4, 4), (3, 3)] if h > w else [(3, 3), (2, 2)]
        best_match, best_score = None, -1
        for size in sizes:
            digit_img_resized = cv2.resize(digit_img, size, interpolation=cv2.INTER_AREA)
            digit_img_resized = (digit_img_resized > 100).astype(np.float32)  # Binarize
            
            for digit, templates in digit_templates.items():
                for template in templates:
                    if digit == '.' and size[0] > 3:
                        continue
                    if digit != '.' and size[0] <= 3:
                        continue
                    if template.shape[0] != size[0] or template.shape[1] != size[1]:
                        continue
                    result = cv2.matchTemplate(digit_img_resized, template, cv2.TM_CCOEFF_NORMED)
                    _, max_val, _, _ = cv2.minMaxLoc(result)
                    if max_val > 0.55 and max_val > best_score:  # Further lowered threshold
                        best_score = max_val
                        best_match = digit
        logging.debug(f"Template match: {best_match}, Score: {best_score:.2f}")
        return best_match if best_score > 0.55 else None
    except Exception as e:
        logging.error(f"Template digit detection failed: {str(e)}")
        return None

def perform_ocr(img, roi_bbox):
    """Perform OCR with Tesseract and robust template fallback."""
    try:
        thresh, enhanced = preprocess_image(img)
        brightness = estimate_brightness(img)
        pil_img = Image.fromarray(enhanced)
        save_debug_image(pil_img, "07_ocr_input")
        
        # Try multiple Tesseract configurations
        configs = [
            r'--oem 3 --psm 7 -c tessedit_char_whitelist=0123456789.',  # Single line
            r'--oem 3 --psm 6 -c tessedit_char_whitelist=0123456789.'   # Block of text
        ]
        for config in configs:
            text = pytesseract.image_to_string(pil_img, config=config)
            logging.info(f"Tesseract raw output (config {config}): {text}")
            text = re.sub(r"[^\d\.]", "", text)
            if text.count('.') > 1:
                text = text.replace('.', '', text.count('.') - 1)
            text = text.strip('.')
            if text and re.fullmatch(r"^\d*\.?\d*$", text):
                text = text.lstrip('0') or '0'
                confidence = 95.0 if len(text.replace('.', '')) >= 3 else 90.0
                logging.info(f"Validated Tesseract text: {text}, Confidence: {confidence:.2f}%")
                return text, confidence

        # Fallback to template-based detection
        logging.info("Tesseract failed, using template-based detection.")
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        digits_info = []
        for c in contours:
            x, y, w, h = cv2.boundingRect(c)
            if w > 4 and h > 5 and 0.03 <= w/h <= 4.0:
                digits_info.append((x, x+w, y, y+h))
        
        if digits_info:
            digits_info.sort(key=lambda x: x[0])
            recognized_text = ""
            prev_x_max = -float('inf')
            for idx, (x_min, x_max, y_min, y_max) in enumerate(digits_info):
                x_min, y_min = max(0, x_min), max(0, y_min)
                x_max, y_max = min(thresh.shape[1], x_max), min(thresh.shape[0], y_max)
                if x_max <= x_min or y_max <= y_min:
                    continue
                digit_crop = thresh[y_min:y_max, x_min:x_max]
                save_debug_image(digit_crop, f"08_digit_crop_{idx}")
                digit = detect_digit_template(digit_crop, brightness)
                if digit:
                    recognized_text += digit
                elif x_min - prev_x_max < 10 and prev_x_max != -float('inf'):
                    recognized_text += '.'
                prev_x_max = x_max
            
            text = re.sub(r"[^\d\.]", "", recognized_text)
            if text.count('.') > 1:
                text = text.replace('.', '', text.count('.') - 1)
            text = text.strip('.')
            if text and re.fullmatch(r"^\d*\.?\d*$", text):
                text = text.lstrip('0') or '0'
                confidence = 90.0 if len(text.replace('.', '')) >= 3 else 85.0
                logging.info(f"Validated template text: {text}, Confidence: {confidence:.2f}%")
                return text, confidence
        
        logging.info("No valid digits detected.")
        return None, 0.0
    except Exception as e:
        logging.error(f"OCR failed: {str(e)}")
        return None, 0.0

def extract_weight_from_image(pil_img):
    """Extract weight from any digital scale image."""
    try:
        img = np.array(pil_img)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        save_debug_image(img, "00_input_image")
        img = correct_rotation(img)
        brightness = estimate_brightness(img)
        conf_threshold = 0.65 if brightness > 70 else 0.45

        # Try ROI-based detection
        roi_img, roi_bbox = detect_roi(img)
        if roi_bbox:
            conf_threshold *= 1.15 if (roi_bbox[2] * roi_bbox[3]) > (img.shape[0] * img.shape[1] * 0.05) else 1.0

        result, confidence = perform_ocr(roi_img, roi_bbox)
        if result and confidence >= conf_threshold * 100:
            try:
                weight = float(result)
                if 0.001 <= weight <= 5000:
                    logging.info(f"Detected weight: {result} kg, Confidence: {confidence:.2f}%")
                    return result, confidence
                logging.warning(f"Weight {result} out of range.")
            except ValueError:
                logging.warning(f"Invalid weight format: {result}")

        # Full image fallback
        logging.info("Primary OCR failed, using full image fallback.")
        result, confidence = perform_ocr(img, None)
        if result and confidence >= conf_threshold * 0.85 * 100:
            try:
                weight = float(result)
                if 0.001 <= weight <= 5000:
                    logging.info(f"Full image weight: {result} kg, Confidence: {confidence:.2f}%")
                    return result, confidence
                logging.warning(f"Full image weight {result} out of range.")
            except ValueError:
                logging.warning(f"Invalid full image weight format: {result}")

        logging.info("No valid weight detected.")
        return "Not detected", 0.0
    except Exception as e:
        logging.error(f"Weight extraction failed: {str(e)}")
        return "Not detected", 0.0