ocr_engine.py

import easyocr
import numpy as np
import cv2
import re
import logging

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

# Initialize EasyOCR
easyocr_reader = easyocr.Reader(['en'], gpu=False)

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

def detect_roi(img):
    """Detect and crop the region of interest (likely the digital display)"""
    try:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        brightness = estimate_brightness(img)
        thresh_value = 230 if brightness > 100 else 190
        _, thresh = cv2.threshold(gray, thresh_value, 255, cv2.THRESH_BINARY)
        kernel = np.ones((9, 9), np.uint8)
        dilated = cv2.dilate(thresh, kernel, iterations=3)
        contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if contours:
            valid_contours = [c for c in contours if cv2.contourArea(c) > 500]
            if valid_contours:
                for contour in sorted(valid_contours, key=cv2.contourArea, reverse=True):
                    x, y, w, h = cv2.boundingRect(contour)
                    aspect_ratio = w / h
                    if 1.5 <= aspect_ratio <= 4.0 and w > 50 and h > 30:
                        x, y = max(0, x-40), max(0, y-40)
                        w, h = min(w+80, img.shape[1]-x), min(h+80, img.shape[0]-y)
                        return img[y:y+h, x:x+w], (x, y, w, h)
        return img, None
    except Exception as e:
        logging.error(f"ROI detection failed: {str(e)}")
        return img, None

def detect_segments(digit_img):
    """Detect seven-segment patterns in a digit image"""
    h, w = digit_img.shape
    if h < 10 or w < 10:
        return None

    # Define segment regions (top, middle, bottom, left-top, left-bottom, right-top, right-bottom)
    segments = {
        'top': (0, w, 0, h//5),
        'middle': (0, w, 2*h//5, 3*h//5),
        'bottom': (0, w, 4*h//5, h),
        'left_top': (0, w//5, 0, h//2),
        'left_bottom': (0, w//5, h//2, h),
        'right_top': (4*w//5, w, 0, h//2),
        'right_bottom': (4*w//5, w, h//2, h)
    }

    segment_presence = {}
    for name, (x1, x2, y1, y2) in segments.items():
        region = digit_img[y1:y2, x1:x2]
        if region.size == 0:
            return None
        # Count white pixels in the region
        pixel_count = np.sum(region == 255)
        total_pixels = region.size
        # Segment is present if more than 50% of the region is white
        segment_presence[name] = pixel_count > total_pixels * 0.5

    # Seven-segment digit patterns
    digit_patterns = {
        '0': ('top', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
        '1': ('right_top', 'right_bottom'),
        '2': ('top', 'middle', 'bottom', 'left_bottom', 'right_top'),
        '3': ('top', 'middle', 'bottom', 'right_top', 'right_bottom'),
        '4': ('middle', 'left_top', 'right_top', 'right_bottom'),
        '5': ('top', 'middle', 'bottom', 'left_top', 'right_bottom'),
        '6': ('top', 'middle', 'bottom', 'left_top', 'left_bottom', 'right_bottom'),
        '7': ('top', 'right_top', 'right_bottom'),
        '8': ('top', 'middle', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
        '9': ('top', 'middle', 'bottom', 'left_top', 'right_top', 'right_bottom')
    }

    best_match = None
    max_matches = 0
    for digit, pattern in digit_patterns.items():
        matches = sum(1 for segment in pattern if segment_presence.get(segment, False))
        non_matches = sum(1 for segment in segment_presence if segment not in pattern and segment_presence[segment])
        score = matches - non_matches
        if score > max_matches:
            max_matches = score
            best_match = digit

    return best_match

def custom_seven_segment_ocr(img, roi_bbox):
    """Perform custom OCR for seven-segment displays"""
    try:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        # Use EasyOCR to get bounding boxes for digits
        results = easyocr_reader.readtext(thresh, detail=1, paragraph=False, 
                                        contrast_ths=0.1, adjust_contrast=0.7, 
                                        text_threshold=0.9, mag_ratio=1.5, 
                                        allowlist='0123456789.')

        if not results:
            return None

        # Sort bounding boxes left to right
        digits = []
        for (bbox, _, _) in results:
            (x1, y1), (x2, y2), (x3, y3), (x4, y4) = bbox
            x_min, x_max = min(x1, x4), max(x2, x3)
            y_min, y_max = min(y1, y2), max(y3, y4)
            digits.append((x_min, x_max, y_min, y_max))

        digits.sort(key=lambda x: x[0])  # Sort by x_min (left to right)

        # Extract and recognize each digit
        recognized_text = ""
        for x_min, x_max, y_min, y_max in digits:
            x_min, y_min = max(0, int(x_min)), max(0, int(y_min))
            x_max, y_max = min(thresh.shape[1], int(x_max)), min(thresh.shape[0], int(y_max))
            if x_max <= x_min or y_max <= y_min:
                continue
            digit_img = thresh[y_min:y_max, x_min:x_max]
            digit = detect_segments(digit_img)
            if digit:
                recognized_text += digit

        # Validate the recognized text
        text = recognized_text
        text = re.sub(r"[^\d\.]", "", text)
        if re.fullmatch(r"\d{1,4}(\.\d{0,3})?", text):
            return text
        return None
    except Exception as e:
        logging.error(f"Custom seven-segment OCR failed: {str(e)}")
        return None

def extract_weight_from_image(pil_img):
    try:
        img = np.array(pil_img)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        brightness = estimate_brightness(img)
        conf_threshold = 0.9 if brightness > 100 else 0.7

        # Detect ROI
        roi_img, roi_bbox = detect_roi(img)

        # Try custom seven-segment OCR first
        custom_result = custom_seven_segment_ocr(roi_img, roi_bbox)
        if custom_result:
            # Format the custom result
            if "." in custom_result:
                int_part, dec_part = custom_result.split(".")
                int_part = int_part.lstrip("0") or "0"
                custom_result = f"{int_part}.{dec_part.rstrip('0')}"
            else:
                custom_result = custom_result.lstrip('0') or "0"
            return custom_result, 100.0  # High confidence for custom OCR

        # Fallback to EasyOCR if custom OCR fails
        images_to_process = [
            ("raw", roi_img, {'contrast_ths': 0.1, 'adjust_contrast': 0.7, 'text_threshold': 0.9, 'mag_ratio': 1.5, 'allowlist': '0123456789.'}),
        ]

        best_weight = None
        best_conf = 0.0
        best_score = 0.0

        for mode, proc_img, ocr_params in images_to_process:
            if mode == "raw":
                proc_img = cv2.cvtColor(proc_img, cv2.COLOR_BGR2GRAY)
            results = easyocr_reader.readtext(proc_img, detail=1, paragraph=False, **ocr_params)
            
            for (bbox, text, conf) in results:
                text = text.lower().strip()
                text = text.replace(",", ".").replace(";", ".")
                text = text.replace("o", "0").replace("O", "0")
                text = text.replace("s", "5").replace("S", "5")
                text = text.replace("g", "9").replace("G", "6")
                text = text.replace("l", "1").replace("I", "1")
                text = text.replace("b", "8").replace("B", "8")
                text = text.replace("z", "2").replace("Z", "2")
                text = text.replace("q", "9").replace("Q", "9")
                text = text.replace("kgs", "").replace("kg", "").replace("k", "")
                text = re.sub(r"[^\d\.]", "", text)

                if re.fullmatch(r"\d{1,4}(\.\d{0,3})?", text):
                    try:
                        weight = float(text)
                        range_score = 1.0 if 0.1 <= weight <= 500 else 0.3
                        digit_score = 1.5 if 10 <= weight < 100 else 1.0
                        score = conf * range_score * digit_score
                        if score > best_score and conf > conf_threshold:
                            best_weight = text
                            best_conf = conf
                            best_score = score
                    except ValueError:
                        continue

        if not best_weight:
            logging.info("No valid weight detected")
            return "Not detected", 0.0

        if "." in best_weight:
            int_part, dec_part = best_weight.split(".")
            int_part = int_part.lstrip("0") or "0"
            best_weight = f"{int_part}.{dec_part.rstrip('0')}"
        else:
            best_weight = best_weight.lstrip('0') or "0"

        return best_weight, round(best_conf * 100, 2)

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