Update ocr_engine.py

ocr_engine.py

@@ -10,39 +10,36 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 # Initialize EasyOCR
 easyocr_reader = easyocr.Reader(['en'], gpu=False)
 
-def estimate_blur(img):
-    """Estimate image blur using Laplacian variance"""
-    try:
-        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        return cv2.Laplacian(gray, cv2.CV_64F).var()
-    except Exception as e:
-        logging.error(f"Blur estimation failed: {str(e)}")
-        return 100  # Default value for fallback
+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)
-        # Adaptive thresholding to handle varying lighting
-        thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                      cv2.THRESH_BINARY_INV, 11, 2)
-        # Dilate to connect text regions
-        kernel = np.ones((5, 5), np.uint8)
-        dilated = cv2.dilate(thresh, kernel, iterations=1)
+        # Threshold to isolate bright areas (like illuminated displays)
+        brightness = estimate_brightness(img)
+        thresh_value = 200 if brightness > 100 else 150  # Adjust based on brightness
+        _, thresh = cv2.threshold(gray, thresh_value, 255, cv2.THRESH_BINARY)
+        # Dilate to connect digits
+        kernel = np.ones((7, 7), np.uint8)
+        dilated = cv2.dilate(thresh, kernel, iterations=2)
         # Find contours
         contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         if contours:
             # Get the largest contour with reasonable size
-            valid_contours = [c for c in contours if cv2.contourArea(c) > 1000]
+            valid_contours = [c for c in contours if cv2.contourArea(c) > 500]
             if valid_contours:
                 largest_contour = max(valid_contours, key=cv2.contourArea)
                 x, y, w, h = cv2.boundingRect(largest_contour)
-                # Add padding and ensure bounds
-                x, y = max(0, x-20), max(0, y-20)
-                w, h = min(w+40, img.shape[1]-x), min(h+40, img.shape[0]-y)
-                if w > 50 and h > 30:  # Minimum size for valid ROI
+                # Add more padding and ensure bounds
+                x, y = max(0, x-30), max(0, y-30)
+                w, h = min(w+60, img.shape[1]-x), min(h+60, img.shape[0]-y)
+                if w > 50 and h > 30:
                     return img[y:y+h, x:x+w]
-        return img  # Fallback to original image
+        return img
     except Exception as e:
         logging.error(f"ROI detection failed: {str(e)}")
         return img
@@ -52,32 +49,34 @@ def enhance_image(img, mode="standard"):
     try:
         gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
 
-        if mode == "high_contrast":
-            # Stronger denoising and contrast for blurry images
+        if mode == "seven_segment":
+            # Gentle preprocessing for seven-segment displays
+            denoised = cv2.GaussianBlur(gray, (5, 5), 0)
+            _, thresh = cv2.threshold(denoised, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+        elif mode == "high_contrast":
             denoised = cv2.bilateralFilter(gray, d=11, sigmaColor=100, sigmaSpace=100)
             clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
+            thresh = clahe.apply(denoised)
         elif mode == "low_noise":
-            # Gentle denoising for clear but noisy images
             denoised = cv2.bilateralFilter(gray, d=7, sigmaColor=50, sigmaSpace=50)
             clahe = cv2.createCLAHE(clipLimit=1.5, tileGridSize=(8, 8))
+            thresh = clahe.apply(denoised)
         else:
-            # Standard preprocessing
             denoised = cv2.bilateralFilter(gray, d=9, sigmaColor=75, sigmaSpace=75)
             clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+            thresh = clahe.apply(denoised)
 
-        contrast = clahe.apply(denoised)
-
-        # Adaptive thresholding
-        thresh = cv2.adaptiveThreshold(contrast, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                      cv2.THRESH_BINARY, 11, 2)
+        if mode != "seven_segment":
+            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)
 
-        # Adaptive sharpening
-        blur_score = estimate_blur(img)
-        sharpen_strength = 5 if blur_score < 100 else 3
+        # Reduced sharpening for seven-segment displays
+        brightness = estimate_brightness(img)
+        sharpen_strength = 3 if mode == "seven_segment" or brightness > 100 else 5
         sharpen_kernel = np.array([[0, -1, 0], [-1, sharpen_strength, -1], [0, -1, 0]])
         sharpened = cv2.filter2D(morphed, -1, sharpen_kernel)
 
@@ -99,19 +98,20 @@ def extract_weight_from_image(pil_img):
         img = np.array(pil_img)
         img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
 
-        # Estimate blur for adaptive thresholding
-        blur_score = estimate_blur(img)
-        conf_threshold = 0.35 if blur_score < 100 else 0.55  # Slightly stricter thresholds
+        # Estimate brightness for adaptive thresholding
+        brightness = estimate_brightness(img)
+        conf_threshold = 0.5 if brightness > 100 else 0.4  # Stricter for bright displays
 
         # Detect ROI
         roi_img = detect_roi(img)
 
         # Process multiple image versions
         images_to_process = [
-            ("standard", enhance_image(roi_img, mode="standard"), {}),
-            ("high_contrast", enhance_image(roi_img, mode="high_contrast"), {}),
-            ("low_noise", enhance_image(roi_img, mode="low_noise"), {}),
-            ("original", roi_img, {'allowlist': '0123456789.'})  # Restrict to digits and decimal
+            ("seven_segment", enhance_image(roi_img, mode="seven_segment"), {'contrast_ths': 0.3, 'allowlist': '0123456789.'}),
+            ("standard", enhance_image(roi_img, mode="standard"), {'contrast_ths': 0.1}),
+            ("high_contrast", enhance_image(roi_img, mode="high_contrast"), {'contrast_ths': 0.1}),
+            ("low_noise", enhance_image(roi_img, mode="low_noise"), {'contrast_ths': 0.1}),
+            ("original", roi_img, {'contrast_ths': 0.3, 'allowlist': '0123456789.'})
         ]
 
         best_weight = None
@@ -135,6 +135,7 @@ def extract_weight_from_image(pil_img):
                 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("6", "2").replace("9", "2")  # Specific correction for seven-segment
                 text = text.replace("kgs", "").replace("kg", "").replace("k", "")
                 text = re.sub(r"[^\d\.]", "", text)
 
@@ -144,7 +145,9 @@ 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
-                        score = conf * range_score
+                        # Prefer two-digit weights for scales
+                        digit_score = 1.1 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