Update ocr_engine.py

ocr_engine.py

@@ -32,30 +32,30 @@ def estimate_brightness(img):
     return np.mean(gray)
 
 def preprocess_image(img):
-    """Preprocess image with simplified, robust contrast enhancement."""
+    """Preprocess image with enhanced contrast and adaptive thresholding."""
     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
+    # Apply CLAHE with dynamic clip limit
+    clahe_clip = 10.0 if brightness < 80 else 5.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)
+    # Stronger blur to reduce noise
+    blurred = cv2.GaussianBlur(enhanced, (7, 7), 1.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))
+    # Adaptive thresholding with larger block size
+    block_size = max(11, min(41, int(img.shape[0] / 15) * 2 + 1))
     thresh = cv2.adaptiveThreshold(
         blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-        cv2.THRESH_BINARY_INV, block_size, 3
+        cv2.THRESH_BINARY_INV, block_size, 5
     )
     
-    # Minimal morphological operations
-    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
-    thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
+    # Morphological operations for better digit separation
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
+    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
     save_debug_image(thresh, "03_preprocess_morph")
     return thresh, enhanced
 
@@ -63,12 +63,12 @@ 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)
+        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
+        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=20, minLineLength=10, maxLineGap=5)
         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:
+            if abs(angle) > 0.5:
                 h, w = img.shape[:2]
                 center = (w // 2, h // 2)
                 M = cv2.getRotationMatrix2D(center, angle, 1.0)
@@ -81,21 +81,21 @@ def correct_rotation(img):
         return img
 
 def detect_roi(img):
-    """Detect region of interest with broader contour analysis."""
+    """Detect region of interest with relaxed 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]]
+        block_sizes = [max(11, min(41, int(img.shape[0] / s) * 2 + 1)) for s in [5, 10, 15]]
         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
+                cv2.THRESH_BINARY_INV, block_size, 5
             )
-            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
             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)
@@ -105,15 +105,15 @@ def detect_roi(img):
                 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):
+                if (30 < area < (img_area * 0.98) and 
+                    0.02 <= aspect_ratio <= 25.0 and w > 15 and h > 5 and roi_brightness > 10):
                     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)))
+            padding = max(5, min(25, int(min(w, h) * 0.5)))
             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]
@@ -130,183 +130,61 @@ def detect_roi(img):
         return img, None
 
 def detect_digit_template(digit_img, brightness):
-    """Digit recognition with expanded template matching."""
+    """Digit recognition with adjusted 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)
-            ]
+            '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)],
+            '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)],
+            '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)],
+            '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)],
+            '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)],
+            '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)],
+            '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)],
+            '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)],
+            '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)],
+            '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([[0, 0, 0], [0, 1, 0], [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
+            digit_img_resized = (digit_img_resized > 90).astype(np.float32)  # Adjusted binarization threshold
             
             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
+                    if max_val > 0.50 and max_val > best_score:  # 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
+        return best_match if best_score > 0.50 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."""
+    """Perform OCR with enhanced Tesseract and 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
+        # Enhanced 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
+            r'--oem 3 --psm 6 -c tessedit_char_whitelist=0123456789.',  # Block of text
+            r'--oem 3 --psm 10 -c tessedit_char_whitelist=0123456789.'  # Single character
         ]
         for config in configs:
             text = pytesseract.image_to_string(pil_img, config=config)
@@ -321,13 +199,13 @@ def perform_ocr(img, roi_bbox):
                 logging.info(f"Validated Tesseract text: {text}, Confidence: {confidence:.2f}%")
                 return text, confidence
 
-        # Fallback to template-based detection
+        # Enhanced 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:
+            if w > 3 and h > 4 and 0.02 <= w/h <= 5.0:
                 digits_info.append((x, x+w, y, y+h))
         
         if digits_info:
@@ -344,7 +222,7 @@ def perform_ocr(img, roi_bbox):
                 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'):
+                elif x_min - prev_x_max < 15 and prev_x_max != -float('inf'):
                     recognized_text += '.'
                 prev_x_max = x_max
             
@@ -365,19 +243,19 @@ def perform_ocr(img, roi_bbox):
         return None, 0.0
 
 def extract_weight_from_image(pil_img):
-    """Extract weight from any digital scale image."""
+    """Extract weight from any digital scale image with adjusted thresholds."""
     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
+        conf_threshold = 0.60 if brightness > 70 else 0.40  # Lowered threshold
 
         # 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
+            conf_threshold *= 1.2 if (roi_bbox[2] * roi_bbox[3]) > (img.shape[0] * img.shape[1] * 0.03) else 1.0
 
         result, confidence = perform_ocr(roi_img, roi_bbox)
         if result and confidence >= conf_threshold * 100:
@@ -390,10 +268,10 @@ def extract_weight_from_image(pil_img):
             except ValueError:
                 logging.warning(f"Invalid weight format: {result}")
 
-        # Full image fallback
+        # Full image fallback with relaxed threshold
         logging.info("Primary OCR failed, using full image fallback.")
         result, confidence = perform_ocr(img, None)
-        if result and confidence >= conf_threshold * 0.85 * 100:
+        if result and confidence >= conf_threshold * 0.80 * 100:
             try:
                 weight = float(result)
                 if 0.001 <= weight <= 5000: