Update ocr_engine.py

ocr_engine.py

@@ -2,69 +2,129 @@ import easyocr
 import numpy as np
 import cv2
 import re
-
-reader = easyocr.Reader(['en'], gpu=False)
+import logging
+from mmocr.utils.ocr import MMOCR
+
+# Set up logging for debugging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Initialize OCR engines
+easyocr_reader = easyocr.Reader(['en'], gpu=False)
+try:
+    mmocr_reader = MMOCR(det='DB_r18', recog='CRNN')
+except:
+    mmocr_reader = None
+    logging.warning("MMOCR initialization failed, falling back to EasyOCR only")
+
+def estimate_blur(img):
+    """Estimate image blur using Laplacian variance"""
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    return cv2.Laplacian(gray, cv2.CV_64F).var()
 
 def enhance_image(img):
-    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    try:
+        # Convert to grayscale
+        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
 
-    # CLAHE (adaptive histogram equalization for better contrast)
-    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-    contrast = clahe.apply(gray)
+        # Bilateral filter for noise reduction while preserving edges
+        denoised = cv2.bilateralFilter(gray, d=9, sigmaColor=75, sigmaSpace=75)
 
-    # Sharpen image
-    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
-    sharpened = cv2.filter2D(contrast, -1, kernel)
+        # CLAHE for contrast enhancement
+        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+        contrast = clahe.apply(denoised)
 
-    # Resize if very small
-    h, w = sharpened.shape
-    if max(h, w) < 500:
-        sharpened = cv2.resize(sharpened, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)
+        # Adaptive thresholding for uneven lighting
+        thresh = cv2.adaptiveThreshold(contrast, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
+                                      cv2.THRESH_BINARY, 11, 2)
 
-    return sharpened
+        # Morphological operations to enhance text
+        kernel = np.ones((3, 3), np.uint8)
+        morphed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=1)
+
+        # Sharpen image
+        sharpen_kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
+        sharpened = cv2.filter2D(morphed, -1, sharpen_kernel)
+
+        # Dynamic resizing
+        h, w = sharpened.shape
+        target_size = 800  # Target max dimension for OCR
+        scale_factor = min(target_size / max(h, w), 2.0) if max(h, w) < 300 else min(target_size / max(h, w), 1.0)
+        if scale_factor != 1.0:
+            sharpened = cv2.resize(sharpened, None, fx=scale_factor, fy=scale_factor, 
+                                 interpolation=cv2.INTER_CUBIC if scale_factor > 1 else cv2.INTER_AREA)
+
+        return sharpened
+    except Exception as e:
+        logging.error(f"Image enhancement failed: {str(e)}")
+        return img  # Return original image as fallback
 
 def extract_weight_from_image(pil_img):
     try:
         img = np.array(pil_img)
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)  # convert from PIL to OpenCV BGR
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+        # Estimate blur to adjust confidence threshold
+        blur_score = estimate_blur(img)
+        conf_threshold = 0.3 if blur_score < 100 else 0.5  # Lower threshold for blurry images
 
         # Preprocess image
         processed = enhance_image(img)
 
-        results = reader.readtext(processed)
-
+        # Initialize results
         best_weight = None
         best_conf = 0.0
 
-        for (bbox, text, conf) in results:
+        # EasyOCR detection
+        easyocr_results = easyocr_reader.readtext(processed, detail=1, paragraph=False)
+        if not easyocr_results:  # Fallback to original image if no results
+            easyocr_results = easyocr_reader.readtext(img, detail=1, paragraph=False)
+
+        # MMOCR detection (if available)
+        mmocr_results = []
+        if mmocr_reader:
+            try:
+                mmocr_result = mmocr_reader.readtext(processed)
+                mmocr_results = [(bbox, text, score) for bbox, text, score in mmocr_result]
+            except:
+                logging.warning("MMOCR processing failed, using EasyOCR results only")
+
+        # Combine results
+        all_results = easyocr_results + mmocr_results
+
+        for (bbox, text, conf) in all_results:
             original_text = text
             text = text.lower().strip()
 
             # Fix common OCR errors
-            text = text.replace(",", ".")
+            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("kgs", "").replace("kg", "")
+            text = text.replace("l", "1").replace("I", "1")
+            text = text.replace("b", "8").replace("B", "8")
+            text = text.replace("kgs", "").replace("kg", "").replace("k9", "").replace("k", "")
             text = re.sub(r"[^\d\.]", "", text)
 
-            if re.fullmatch(r"\d{1,4}(\.\d{1,3})?", text):
-                if conf > best_conf:
+            # Stricter regex for weight (0.0 to 9999.999)
+            if re.fullmatch(r"\d{1,4}(\.\d{0,3})?", text):
+                if conf > best_conf and conf > conf_threshold:
                     best_weight = text
                     best_conf = conf
 
         if not best_weight:
+            logging.info("No valid weight detected")
             return "Not detected", 0.0
 
         # Format output
         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}"
+            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:
-        return f"Error: {str(e)}", 0.0
+        logging.error(f"Weight extraction failed: {str(e)}")
+        return "Not detected", 0.0