Update ocr_engine.py

ocr_engine.py

@@ -21,7 +21,7 @@ def detect_roi(img):
         gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
         # Stricter threshold for bright areas
         brightness = estimate_brightness(img)
-        thresh_value = 220 if brightness > 100 else 180
+        thresh_value = 230 if brightness > 100 else 190
         _, thresh = cv2.threshold(gray, thresh_value, 255, cv2.THRESH_BINARY)
         # Morphological operations to connect digits
         kernel = np.ones((9, 9), np.uint8)
@@ -29,13 +29,13 @@ def detect_roi(img):
         # Find contours
         contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         if contours:
-            # Filter contours by size and aspect ratio (typical for displays)
+            # Filter contours by size and aspect ratio
             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:  # Typical display aspect ratio
+                    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]
@@ -44,16 +44,35 @@ def detect_roi(img):
         logging.error(f"ROI detection failed: {str(e)}")
         return img
 
-def correct_seven_segment(text, bbox):
-    """Correct common seven-segment misreads based on bounding box shape"""
-    if "6" in text:
-        # Check bounding box aspect ratio to differentiate "6" from "2"
+def correct_seven_segment(text, bbox, img):
+    """Correct common seven-segment misreads based on bounding box and pixel distribution"""
+    if "2" in text or "6" in text:
+        # Extract bounding box coordinates
         (x1, y1), (x2, y2), (x3, y3), (x4, y4) = bbox
-        width = abs(x2 - x1)
-        height = abs(y2 - y1)
-        aspect_ratio = width / height if height > 0 else 1.0
-        # "2" typically has a more rectangular shape in seven-segment
-        if aspect_ratio > 0.5:  # Adjust based on typical "2" vs "6" shapes
+        x_min, x_max = min(x1, x4), max(x2, x3)
+        y_min, y_max = min(y1, y2), max(y3, y4)
+        # Ensure bounds are within image
+        x_min, y_min = max(0, int(x_min)), max(0, int(y_min))
+        x_max, y_max = min(img.shape[1], int(x_max)), min(img.shape[0], int(y_max))
+        if x_max <= x_min or y_max <= y_min:
+            return text
+        # Crop the digit area
+        digit_area = img[y_min:y_max, x_min:x_max]
+        if digit_area.size == 0:
+            return text
+        # Convert to grayscale and threshold
+        gray = cv2.cvtColor(digit_area, cv2.COLOR_BGR2GRAY)
+        _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+        # Calculate pixel distribution in upper vs lower half
+        h, w = thresh.shape
+        upper_half = thresh[:h//2, :]
+        lower_half = thresh[h//2:, :]
+        upper_pixels = np.sum(upper_half == 255)
+        lower_pixels = np.sum(lower_half == 255)
+        # "6" has more pixels in the lower half due to the loop, "2" is more balanced
+        if lower_pixels > upper_pixels * 1.5:
+            text = text.replace("2", "6")
+        else:
             text = text.replace("6", "2")
     return text
 
@@ -63,14 +82,16 @@ def enhance_image(img, mode="standard"):
         gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
 
         if mode == "seven_segment":
-            # Minimal preprocessing for seven-segment displays
+            # Extremely minimal preprocessing for seven-segment displays
             _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
             return thresh
         elif mode == "minimal":
-            # Very light preprocessing
-            denoised = cv2.GaussianBlur(gray, (3, 3), 0)
-            _, thresh = cv2.threshold(denoised, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+            # No blurring, just threshold
+            _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
             return thresh
+        elif mode == "raw":
+            # No preprocessing, just convert to grayscale
+            return gray
         elif mode == "high_contrast":
             denoised = cv2.bilateralFilter(gray, d=11, sigmaColor=100, sigmaSpace=100)
             clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
@@ -84,16 +105,19 @@ def enhance_image(img, mode="standard"):
             clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
             thresh = clahe.apply(denoised)
 
-        if mode not in ["seven_segment", "minimal"]:
+        if mode not in ["seven_segment", "minimal", "raw"]:
             thresh = cv2.adaptiveThreshold(thresh, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                           cv2.THRESH_BINARY, 11, 2)
 
-        # Morphological operations
-        kernel = np.ones((3, 3), np.uint8)
-        morphed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=1)
+        # Morphological operations only for non-seven-segment modes
+        if mode not in ["seven_segment", "minimal", "raw"]:
+            kernel = np.ones((3, 3), np.uint8)
+            morphed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=1)
+        else:
+            morphed = thresh
 
-        # Skip sharpening for seven-segment and minimal modes
-        if mode not in ["seven_segment", "minimal"]:
+        # Skip sharpening for seven-segment, minimal, and raw modes
+        if mode not in ["seven_segment", "minimal", "raw"]:
             brightness = estimate_brightness(img)
             sharpen_strength = 3 if brightness > 100 else 5
             sharpen_kernel = np.array([[0, -1, 0], [-1, sharpen_strength, -1], [0, -1, 0]])
@@ -119,19 +143,19 @@ def extract_weight_from_image(pil_img):
 
         # Estimate brightness for adaptive thresholding
         brightness = estimate_brightness(img)
-        conf_threshold = 0.7 if brightness > 100 else 0.5  # Stricter for bright displays
+        conf_threshold = 0.8 if brightness > 100 else 0.6  # Stricter for bright displays
 
         # Detect ROI
         roi_img = detect_roi(img)
 
         # Process multiple image versions
         images_to_process = [
-            ("seven_segment", enhance_image(roi_img, mode="seven_segment"), {'contrast_ths': 0.2, 'adjust_contrast': 0.5, 'allowlist': '0123456789.'}),
-            ("minimal", enhance_image(roi_img, mode="minimal"), {'contrast_ths': 0.2, 'adjust_contrast': 0.5, 'allowlist': '0123456789.'}),
-            ("standard", enhance_image(roi_img, mode="standard"), {'contrast_ths': 0.1, 'adjust_contrast': 0.5}),
-            ("high_contrast", enhance_image(roi_img, mode="high_contrast"), {'contrast_ths': 0.1, 'adjust_contrast': 0.5}),
-            ("low_noise", enhance_image(roi_img, mode="low_noise"), {'contrast_ths': 0.1, 'adjust_contrast': 0.5}),
-            ("original", roi_img, {'contrast_ths': 0.2, 'adjust_contrast': 0.5, 'allowlist': '0123456789.'})
+            ("seven_segment", enhance_image(roi_img, mode="seven_segment"), {'contrast_ths': 0.15, 'adjust_contrast': 0.7, 'text_threshold': 0.8, 'allowlist': '0123456789.'}),
+            ("minimal", enhance_image(roi_img, mode="minimal"), {'contrast_ths': 0.15, 'adjust_contrast': 0.7, 'text_threshold': 0.8, 'allowlist': '0123456789.'}),
+            ("raw", enhance_image(roi_img, mode="raw"), {'contrast_ths': 0.15, 'adjust_contrast': 0.7, 'text_threshold': 0.8, 'allowlist': '0123456789.'}),
+            ("standard", enhance_image(roi_img, mode="standard"), {'contrast_ths': 0.1, 'adjust_contrast': 0.5, 'text_threshold': 0.7}),
+            ("high_contrast", enhance_image(roi_img, mode="high_contrast"), {'contrast_ths': 0.1, 'adjust_contrast': 0.5, 'text_threshold': 0.7}),
+            ("low_noise", enhance_image(roi_img, mode="low_noise"), {'contrast_ths': 0.1, 'adjust_contrast': 0.5, 'text_threshold': 0.7}),
         ]
 
         best_weight = None
@@ -144,7 +168,7 @@ def extract_weight_from_image(pil_img):
             
             for (bbox, text, conf) in results:
                 # Apply seven-segment correction
-                text = correct_seven_segment(text, bbox)
+                text = correct_seven_segment(text, bbox, roi_img)
                 original_text = text
                 text = text.lower().strip()
 
@@ -166,7 +190,7 @@ def extract_weight_from_image(pil_img):
                         weight = float(text)
                         # Score based on realistic weight range (0.1–500 kg)
                         range_score = 1.0 if 0.1 <= weight <= 500 else 0.3
-                        # Strongly prefer two-digit weights for scales
+                        # Prefer two-digit weights for scales
                         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: