Update ocr_engine.py

ocr_engine.py

@@ -6,11 +6,12 @@ import logging
 from datetime import datetime
 import os
 from PIL import Image, ImageEnhance
+from scipy.signal import convolve2d
 
 # Set up logging for detailed debugging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Initialize EasyOCR with English and GPU disabled (enable if you have a compatible GPU)
+# Initialize EasyOCR with English (enable GPU if available)
 easyocr_reader = easyocr.Reader(['en'], gpu=False)
 
 # Directory for debug images
@@ -28,34 +29,67 @@ def save_debug_image(img, filename_suffix, prefix=""):
     logging.debug(f"Saved debug image: {filename}")
 
 def estimate_brightness(img):
-    """Estimate image brightness to detect illuminated displays"""
+    """Estimate image brightness to adjust processing"""
     gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     brightness = np.mean(gray)
     logging.debug(f"Estimated brightness: {brightness}")
     return brightness
 
-def preprocess_image(img):
-    """Enhance contrast, brightness, and reduce noise for better digit detection"""
-    # Convert to PIL for initial enhancement
-    pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-    pil_img = ImageEnhance.Contrast(pil_img).enhance(2.0)  # Stronger contrast
-    pil_img = ImageEnhance.Brightness(pil_img).enhance(1.3)  # Moderate brightness boost
-    img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
-    save_debug_image(img, "00_preprocessed_pil")
-
-    # Apply CLAHE to enhance local contrast
+def deblur_image(img):
+    """Apply deconvolution to reduce blur (approximate Wiener filter)"""
     gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-    enhanced = clahe.apply(gray)
+    # Create a simple point spread function (PSF) for deblurring
+    psf = np.ones((5, 5)) / 25
+    # Normalize image to float32
+    img_float = gray.astype(np.float32) / 255.0
+    # Convolve with PSF (simulate blur)
+    img_blurred = convolve2d(img_float, psf, mode='same')
+    # Avoid division by zero
+    img_blurred = np.where(img_blurred == 0, 1e-10, img_blurred)
+    # Deconvolve
+    img_deblurred = img_float / img_blurred
+    img_deblurred = np.clip(img_deblurred * 255, 0, 255).astype(np.uint8)
+    save_debug_image(img_deblurred, "00_deblurred")
+    return img_deblurred
+
+def preprocess_image(img):
+    """Enhance contrast, brightness, reduce noise, and deblur for digit detection"""
+    # Deblur first
+    deblurred = deblur_image(img)
+    
+    # Convert to PIL for enhancement
+    pil_img = Image.fromarray(deblurred)
+    pil_img = ImageEnhance.Contrast(pil_img).enhance(2.5)  # Aggressive contrast
+    pil_img = ImageEnhance.Brightness(pil_img).enhance(1.5)  # Stronger brightness
+    img_enhanced = np.array(pil_img)
+    save_debug_image(img_enhanced, "00_preprocessed_pil")
+
+    # Apply CLAHE for local contrast enhancement
+    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
+    enhanced = clahe.apply(img_enhanced)
     save_debug_image(enhanced, "00_clahe_enhanced")
 
-    # Apply bilateral filter to reduce noise while preserving edges
-    filtered = cv2.bilateralFilter(enhanced, d=11, sigmaColor=100, sigmaSpace=100)
+    # Aggressive noise reduction
+    filtered = cv2.bilateralFilter(enhanced, d=15, sigmaColor=150, sigmaSpace=150)
     save_debug_image(filtered, "00_bilateral_filtered")
     return filtered
 
+def normalize_image(img):
+    """Resize image to standard dimensions while preserving aspect ratio"""
+    h, w = img.shape[:2]
+    target_height = 720
+    aspect_ratio = w / h
+    target_width = int(target_height * aspect_ratio)
+    if target_width < 320:
+        target_width = 320
+        target_height = int(target_width / aspect_ratio)
+    resized = cv2.resize(img, (target_width, target_height), interpolation=cv2.INTER_CUBIC)
+    save_debug_image(resized, "00_normalized")
+    logging.debug(f"Normalized image to {target_width}x{target_height}")
+    return resized
+
 def detect_roi(img):
-    """Detect and crop the region of interest (likely the digital display)"""
+    """Detect the digital display region, with fallback to full image"""
     try:
         save_debug_image(img, "01_original")
         gray = preprocess_image(img)
@@ -63,21 +97,21 @@ def detect_roi(img):
 
         # Try multiple thresholding methods
         brightness = estimate_brightness(img)
-        if brightness > 150:
+        if brightness > 120:
             thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                           cv2.THRESH_BINARY, 31, 5)
+                                           cv2.THRESH_BINARY_INV, 41, 7)  # Inverted for bright displays
             save_debug_image(thresh, "03_roi_adaptive_threshold_high")
         else:
-            _, thresh = cv2.threshold(gray, 40, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-            save_debug_image(thresh, "03_roi_otsu_threshold_low")
+            _, thresh = cv2.threshold(gray, 20, 255, cv2.THRESH_BINARY_INV)  # Low threshold for dim displays
+            save_debug_image(thresh, "03_roi_simple_threshold_low")
 
-        # Morphological operations to clean up noise and connect digits
-        kernel = np.ones((5, 5), np.uint8)
-        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
+        # Morphological operations to connect digits
+        kernel = np.ones((7, 7), np.uint8)
+        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=3)
         save_debug_image(thresh, "03_roi_morph_cleaned")
 
-        kernel = np.ones((11, 11), np.uint8)
-        dilated = cv2.dilate(thresh, kernel, iterations=5)
+        kernel = np.ones((15, 15), np.uint8)
+        dilated = cv2.dilate(thresh, kernel, iterations=6)
         save_debug_image(dilated, "04_roi_dilated")
 
         contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
@@ -87,16 +121,16 @@ def detect_roi(img):
             valid_contours = []
             for c in contours:
                 area = cv2.contourArea(c)
-                if 200 < area < (img_area * 0.99):  # Very relaxed area filter
+                if 100 < area < (img_area * 0.999):  # Extremely relaxed area filter
                     x, y, w, h = cv2.boundingRect(c)
                     aspect_ratio = w / h if h > 0 else 0
-                    if 0.5 <= aspect_ratio <= 10.0 and w > 30 and h > 20:  # Very relaxed filters
+                    if 0.3 <= aspect_ratio <= 15.0 and w > 20 and h > 10:  # Very relaxed filters
                         valid_contours.append(c)
 
             if valid_contours:
-                contour = max(valid_contours, key=cv2.contourArea)  # Largest contour
+                contour = max(valid_contours, key=cv2.contourArea)
                 x, y, w, h = cv2.boundingRect(contour)
-                padding = 100  # Generous padding
+                padding = 120  # Very generous padding
                 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]
@@ -104,8 +138,8 @@ def detect_roi(img):
                 logging.info(f"Detected ROI with dimensions: ({x}, {y}, {w}, {h})")
                 return roi_img, (x, y, w, h)
 
-        logging.info("No suitable ROI found, returning preprocessed image.")
-        save_debug_image(img, "05_no_roi_original_fallback")
+        logging.info("No suitable ROI found, returning full image.")
+        save_debug_image(img, "05_no_roi_full_fallback")
         return img, None
     except Exception as e:
         logging.error(f"ROI detection failed: {str(e)}")
@@ -115,18 +149,18 @@ def detect_roi(img):
 def detect_segments(digit_img):
     """Detect seven-segment patterns in a digit image"""
     h, w = digit_img.shape
-    if h < 8 or w < 4:  # Very relaxed size constraints
+    if h < 6 or w < 3:  # Extremely relaxed size constraints
         logging.debug(f"Digit image too small: {w}x{h}")
         return None
 
     segments = {
-        'top': (int(w*0.1), int(w*0.9), 0, int(h*0.25)),
-        'middle': (int(w*0.1), int(w*0.9), int(h*0.35), int(h*0.65)),
-        'bottom': (int(w*0.1), int(w*0.9), int(h*0.75), h),
-        'left_top': (0, int(w*0.3), int(h*0.05), int(h*0.55)),
-        'left_bottom': (0, int(w*0.3), int(h*0.45), int(h*0.95)),
-        'right_top': (int(w*0.7), w, int(h*0.05), int(h*0.55)),
-        'right_bottom': (int(w*0.7), w, int(h*0.45), int(h*0.95))
+        'top': (int(w*0.05), int(w*0.95), 0, int(h*0.3)),
+        'middle': (int(w*0.05), int(w*0.95), int(h*0.35), int(h*0.65)),
+        'bottom': (int(w*0.05), int(w*0.95), int(h*0.7), h),
+        'left_top': (0, int(w*0.35), int(h*0.05), int(h*0.55)),
+        'left_bottom': (0, int(w*0.35), int(h*0.45), int(h*0.95)),
+        'right_top': (int(w*0.65), w, int(h*0.05), int(h*0.55)),
+        'right_bottom': (int(w*0.65), w, int(h*0.45), int(h*0.95))
     }
 
     segment_presence = {}
@@ -139,7 +173,7 @@ def detect_segments(digit_img):
             continue
         pixel_count = np.sum(region == 255)
         total_pixels = region.size
-        segment_presence[name] = pixel_count / total_pixels > 0.3  # Very low threshold
+        segment_presence[name] = pixel_count / total_pixels > 0.25  # Very low threshold
         logging.debug(f"Segment {name}: {pixel_count}/{total_pixels} = {pixel_count/total_pixels:.2f}")
 
     digit_patterns = {
@@ -179,25 +213,25 @@ def detect_segments(digit_img):
 def custom_seven_segment_ocr(img, roi_bbox):
     """Perform custom OCR for seven-segment displays"""
     try:
-        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        gray = preprocess_image(img)
         brightness = estimate_brightness(img)
-        # Try multiple thresholding approaches
-        if brightness > 150:
-            _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+        # Multiple thresholding approaches
+        if brightness > 120:
+            _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
             save_debug_image(thresh, "06_roi_otsu_threshold")
         else:
-            _, thresh = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY)
+            _, thresh = cv2.threshold(gray, 15, 255, cv2.THRESH_BINARY_INV)  # Very low threshold
             save_debug_image(thresh, "06_roi_simple_threshold")
         
         # Morphological cleaning
-        kernel = np.ones((3, 3), np.uint8)
-        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
+        kernel = np.ones((5, 5), np.uint8)
+        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)
         save_debug_image(thresh, "06_roi_morph_cleaned")
 
         results = easyocr_reader.readtext(thresh, detail=1, paragraph=False, 
-                                         contrast_ths=0.1, adjust_contrast=1.0, 
-                                         text_threshold=0.3, mag_ratio=4.0, 
-                                         allowlist='0123456789.-', y_ths=0.6)
+                                         contrast_ths=0.05, adjust_contrast=1.2, 
+                                         text_threshold=0.2, mag_ratio=6.0, 
+                                         allowlist='0123456789.-', y_ths=0.7)
         
         logging.info(f"Custom OCR EasyOCR results: {results}")
         if not results:
@@ -208,7 +242,7 @@ def custom_seven_segment_ocr(img, roi_bbox):
         for (bbox, text, conf) in results:
             (x1, y1), (x2, y2), (x3, y3), (x4, y4) = bbox
             h_bbox = max(y1, y2, y3, y4) - min(y1, y2, y3, y4)
-            if len(text) == 1 and (text.isdigit() or text in '.-') and h_bbox > 4:
+            if len(text) <= 2 and any(c in '0123456789.-' for c in text) and h_bbox > 3:
                 x_min, x_max = int(min(x1, x4)), int(max(x2, x3))
                 y_min, y_max = int(min(y1, y2)), int(max(y3, y4))
                 digits_info.append((x_min, x_max, y_min, y_max, text, conf))
@@ -222,7 +256,7 @@ def custom_seven_segment_ocr(img, roi_bbox):
                 continue
             digit_img_crop = thresh[y_min:y_max, x_min:x_max]
             save_debug_image(digit_img_crop, f"07_digit_crop_{idx}_{easyocr_char}")
-            if easyocr_conf > 0.8 or easyocr_char in '.-' or digit_img_crop.shape[0] < 8 or digit_img_crop.shape[1] < 4:
+            if easyocr_conf > 0.7 or easyocr_char in '.-' or digit_img_crop.shape[0] < 6 or digit_img_crop.shape[1] < 3:
                 recognized_text += easyocr_char
             else:
                 digit_from_segments = detect_segments(digit_img_crop)
@@ -232,10 +266,9 @@ def custom_seven_segment_ocr(img, roi_bbox):
                     recognized_text += easyocr_char
         
         logging.info(f"Custom OCR before validation, recognized_text: {recognized_text}")
-        # Relaxed validation for debugging
         if recognized_text:
             return recognized_text
-        logging.info(f"Custom OCR text '{recognized_text}' failed validation.")
+        logging.info(f"Custom OCR text '{recognized_text}' is empty.")
         return None
     except Exception as e:
         logging.error(f"Custom seven-segment OCR failed: {str(e)}")
@@ -248,13 +281,16 @@ def extract_weight_from_image(pil_img):
         img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
         save_debug_image(img, "00_input_image")
 
+        # Normalize image dimensions
+        img = normalize_image(img)
         brightness = estimate_brightness(img)
-        conf_threshold = 0.3 if brightness > 150 else (0.2 if brightness > 80 else 0.1)
+        conf_threshold = 0.2 if brightness > 120 else 0.1
 
         roi_img, roi_bbox = detect_roi(img)
         custom_result = custom_seven_segment_ocr(roi_img, roi_bbox)
         if custom_result:
-            # Basic cleaning
+            logging.info(f"Raw custom OCR result: {custom_result}")
+            # Minimal cleaning
             text = re.sub(r"[^\d\.\-]", "", custom_result)  # Allow negative signs
             if text.count('.') > 1:
                 text = text.replace('.', '', text.count('.') - 1)
@@ -267,34 +303,34 @@ def extract_weight_from_image(pil_img):
                     logging.warning(f"Custom OCR result '{text}' is invalid after cleaning.")
                 else:
                     try:
-                        float(text)
-                        logging.info(f"Custom OCR result: {text}, Confidence: 100.0%")
-                        return text, 100.0
+                        weight = float(text)
+                        logging.info(f"Custom OCR result: {text}, Confidence: 90.0%")
+                        return text, 90.0
                     except ValueError:
                         logging.warning(f"Custom OCR result '{text}' is not a valid number, falling back.")
-            logging.warning(f"Custom OCR result '{custom_result}' failed validation, falling back.")
+            logging.warning(f"Custom OCR result '{custom_result}' failed cleaning, falling back.")
 
         logging.info("Custom OCR failed or invalid, falling back to general EasyOCR.")
         processed_roi_img = preprocess_image(roi_img)
         
-        # Try multiple thresholding approaches
-        if brightness > 150:
+        # Multiple thresholding approaches
+        if brightness > 120:
             thresh = cv2.adaptiveThreshold(processed_roi_img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                           cv2.THRESH_BINARY, 41, 7)
+                                           cv2.THRESH_BINARY_INV, 51, 9)
             save_debug_image(thresh, "09_fallback_adaptive_thresh")
         else:
-            _, thresh = cv2.threshold(processed_roi_img, 30, 255, cv2.THRESH_BINARY)
+            _, thresh = cv2.threshold(processed_roi_img, 15, 255, cv2.THRESH_BINARY_INV)
             save_debug_image(thresh, "09_fallback_simple_thresh")
 
         # Morphological cleaning
-        kernel = np.ones((3, 3), np.uint8)
-        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
+        kernel = np.ones((5, 5), np.uint8)
+        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)
         save_debug_image(thresh, "09_fallback_morph_cleaned")
 
         results = easyocr_reader.readtext(thresh, detail=1, paragraph=False, 
-                                         contrast_ths=0.1, adjust_contrast=1.0, 
-                                         text_threshold=0.2, mag_ratio=5.0, 
-                                         allowlist='0123456789.-', batch_size=4, y_ths=0.6)
+                                         contrast_ths=0.05, adjust_contrast=1.2, 
+                                         text_threshold=0.1, mag_ratio=7.0, 
+                                         allowlist='0123456789.-', batch_size=4, y_ths=0.8)
 
         best_weight = None
         best_conf = 0.0
@@ -319,21 +355,19 @@ def extract_weight_from_image(pil_img):
                 parts = text.split('.')
                 text = parts[0] + '.' + ''.join(parts[1:])
             text = text.strip('.')
-            if len(text.replace('.', '').replace('-', '')) > 0:  # Allow negative weights
+            if len(text.replace('.', '').replace('-', '')) > 0:
                 try:
                     weight = float(text)
                     range_score = 1.0
-                    if 0.0 <= weight <= 250:
+                    if -1000 <= weight <= 1000:  # Allow negative weights
                         range_score = 1.5
-                    elif weight > 250 and weight <= 500:
-                        range_score = 1.2
-                    elif weight > 500 and weight <= 1000:
+                    elif weight > 1000 and weight <= 2000:
                         range_score = 1.0
                     else:
                         range_score = 0.5
                     digit_count = len(text.replace('.', '').replace('-', ''))
                     digit_score = 1.0
-                    if digit_count >= 2 and digit_count <= 5:
+                    if digit_count >= 2 and digit_count <= 6:
                         digit_score = 1.3
                     elif digit_count == 1:
                         digit_score = 0.8
@@ -344,10 +378,10 @@ def extract_weight_from_image(pil_img):
                         x_min, y_min = int(min(b[0] for b in bbox)), int(min(b[1] for b in bbox))
                         x_max, y_max = int(max(b[0] for b in bbox)), int(max(b[1] for b in bbox))
                         bbox_area = (x_max - x_min) * (y_max - y_min)
-                        if roi_area > 0 and bbox_area / roi_area < 0.02:
+                        if roi_area > 0 and bbox_area / roi_area < 0.01:
                             score *= 0.5
                         bbox_aspect_ratio = (x_max - x_min) / (y_max - y_min) if (y_max - y_min) > 0 else 0
-                        if bbox_aspect_ratio < 0.1:
+                        if bbox_aspect_ratio < 0.05:
                             score *= 0.7
                     if score > best_score and conf > conf_threshold:
                         best_weight = text
@@ -375,15 +409,6 @@ def extract_weight_from_image(pil_img):
         else:
             best_weight = best_weight.lstrip('0') or "0"
 
-        try:
-            final_float_weight = float(best_weight)
-            if final_float_weight < 0.0 or final_float_weight > 1000:
-                logging.warning(f"Detected weight {final_float_weight} is outside typical range, reducing confidence.")
-                best_conf *= 0.5
-        except ValueError:
-            logging.warning(f"Final weight '{best_weight}' is not a valid number.")
-            best_conf *= 0.5
-
         logging.info(f"Final detected weight: {best_weight}, Confidence: {round(best_conf * 100, 2)}%")
         return best_weight, round(best_conf * 100, 2)