Update ocr_engine.py

ocr_engine.py

@@ -27,73 +27,85 @@ def save_debug_image(img, filename_suffix, prefix=""):
     logging.info(f"Saved debug image: {filename}")
 
 def estimate_brightness(img):
-    """Estimate image brightness to detect illuminated displays"""
+    """Estimate image brightness to detect illuminated displays."""
     gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     return np.mean(gray)
 
+def preprocess_image(img):
+    """Preprocess image for better OCR accuracy."""
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    # Apply Gaussian blur to reduce noise
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    save_debug_image(blurred, "01_preprocess_blurred")
+    # Enhance contrast using CLAHE
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    enhanced = clahe.apply(blurred)
+    save_debug_image(enhanced, "02_preprocess_clahe")
+    return enhanced
+
 def detect_roi(img):
-    """Detect and crop the region of interest (likely the digital display)"""
+    """Detect and crop the region of interest (likely the digital display)."""
     try:
-        save_debug_image(img, "01_original")
-        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        save_debug_image(gray, "02_grayscale")
+        save_debug_image(img, "03_original")
+        preprocessed = preprocess_image(img)
         
-        # Use adaptive thresholding for better robustness
-        thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                       cv2.THRESH_BINARY, 11, 2)
-        save_debug_image(thresh, "03_roi_adaptive_threshold")
+        # Adaptive thresholding with refined parameters
+        thresh = cv2.adaptiveThreshold(preprocessed, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
+                                       cv2.THRESH_BINARY_INV, 15, 5)
+        save_debug_image(thresh, "04_roi_adaptive_threshold")
         
-        kernel = np.ones((7, 7), np.uint8)  # Smaller kernel
-        dilated = cv2.dilate(thresh, kernel, iterations=3)  # Fewer iterations
-        save_debug_image(dilated, "04_roi_dilated")
+        # Morphological operations to connect digits
+        kernel = np.ones((5, 5), np.uint8)
+        dilated = cv2.dilate(thresh, kernel, iterations=2)
+        eroded = cv2.erode(dilated, kernel, iterations=1)
+        save_debug_image(eroded, "05_roi_morphological")
         
-        contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        contours, _ = cv2.findContours(eroded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         
         if contours:
             img_area = img.shape[0] * img.shape[1]
             valid_contours = []
             for c in contours:
                 area = cv2.contourArea(c)
-                # Relaxed area and aspect ratio filters
-                if 500 < area < (img_area * 0.95):
+                if 1000 < area < (img_area * 0.9):
                     x, y, w, h = cv2.boundingRect(c)
                     aspect_ratio = w / h
-                    if 1.5 <= aspect_ratio <= 6.0 and w > 80 and h > 40:
+                    if 1.8 <= aspect_ratio <= 8.0 and w > 100 and h > 50:
                         valid_contours.append(c)
-
+            
             if valid_contours:
-                for contour in sorted(valid_contours, key=cv2.contourArea, reverse=True):
-                    x, y, w, h = cv2.boundingRect(contour)
-                    padding = 60  # Increased 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]
-                    save_debug_image(roi_img, "05_detected_roi")
-                    logging.info(f"Detected ROI with dimensions: ({x}, {y}, {w}, {h})")
-                    return roi_img, (x, y, w, h)
+                contour = max(valid_contours, key=cv2.contourArea)
+                x, y, w, h = cv2.boundingRect(contour)
+                padding = 80
+                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]
+                save_debug_image(roi_img, "06_detected_roi")
+                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 original image.")
-        save_debug_image(img, "05_no_roi_original_fallback")
+        save_debug_image(img, "06_no_roi_original_fallback")
         return img, None
     except Exception as e:
         logging.error(f"ROI detection failed: {str(e)}")
-        save_debug_image(img, "05_roi_detection_error_fallback")
+        save_debug_image(img, "06_roi_detection_error_fallback")
         return img, None
 
 def detect_segments(digit_img):
-    """Detect seven-segment patterns in a digit image"""
+    """Detect seven-segment patterns in a digit image."""
     h, w = digit_img.shape
-    if h < 15 or w < 10:
+    if h < 20 or w < 15:
         return None
 
     segments = {
-        'top': (int(w*0.15), int(w*0.85), 0, int(h*0.2)),
-        'middle': (int(w*0.15), int(w*0.85), int(h*0.4), int(h*0.6)),
-        'bottom': (int(w*0.15), int(w*0.85), int(h*0.8), h),
-        'left_top': (0, int(w*0.25), int(h*0.05), int(h*0.5)),
-        'left_bottom': (0, int(w*0.25), int(h*0.5), int(h*0.95)),
-        'right_top': (int(w*0.75), w, int(h*0.05), int(h*0.5)),
-        'right_bottom': (int(w*0.75), w, int(h*0.5), int(h*0.95))
+        'top': (int(w*0.1), int(w*0.9), 0, int(h*0.15)),
+        'middle': (int(w*0.1), int(w*0.9), int(h*0.45), int(h*0.55)),
+        'bottom': (int(w*0.1), int(w*0.9), int(h*0.85), h),
+        'left_top': (0, int(w*0.2), int(h*0.1), int(h*0.5)),
+        'left_bottom': (0, int(w*0.2), int(h*0.5), int(h*0.9)),
+        'right_top': (int(w*0.8), w, int(h*0.1), int(h*0.5)),
+        'right_bottom': (int(w*0.8), w, int(h*0.5), int(h*0.9))
     }
 
     segment_presence = {}
@@ -106,7 +118,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.45  # Lowered threshold
+        segment_presence[name] = pixel_count / total_pixels > 0.4
 
     digit_patterns = {
         '0': ('top', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
@@ -126,37 +138,34 @@ def detect_segments(digit_img):
     for digit, pattern in digit_patterns.items():
         matches = sum(1 for segment in pattern if segment_presence.get(segment, False))
         non_matches_penalty = sum(1 for segment in segment_presence if segment not in pattern and segment_presence[segment])
-        current_score = matches - non_matches_penalty
+        score = matches - 0.5 * non_matches_penalty
         if all(segment_presence.get(s, False) for s in pattern):
-            current_score += 0.5
-        if current_score > max_score:
-            max_score = current_score
+            score += 1.0
+        if score > max_score:
+            max_score = score
             best_match = digit
-        elif current_score == max_score and best_match is not None:
-            current_digit_non_matches = sum(1 for segment in segment_presence if segment not in pattern and segment_presence[segment])
-            best_digit_pattern = digit_patterns[best_match]
-            best_digit_non_matches = sum(1 for segment in segment_presence if segment not in best_digit_pattern and segment_presence[segment])
-            if current_digit_non_matches < best_digit_non_matches:
-                best_match = digit
     
     logging.debug(f"Segment presence: {segment_presence}, Detected digit: {best_match}")
     return best_match
 
 def custom_seven_segment_ocr(img, roi_bbox):
-    """Perform custom OCR for seven-segment displays"""
+    """Perform custom OCR for seven-segment displays."""
     try:
-        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        preprocessed = preprocess_image(img)
         brightness = estimate_brightness(img)
-        if brightness > 150:
-            _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-        else:
-            _, thresh = cv2.threshold(gray, 80, 255, cv2.THRESH_BINARY)  # Lower threshold
-        save_debug_image(thresh, "06_roi_thresh_for_digits")
+        thresh_value = 100 if brightness < 100 else 0
+        _, thresh = cv2.threshold(preprocessed, thresh_value, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+        save_debug_image(thresh, "07_roi_thresh_for_digits")
+
+        # Morphological operations to enhance digit segments
+        kernel = np.ones((3, 3), np.uint8)
+        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
+        save_debug_image(thresh, "08_morph_closed")
 
         results = easyocr_reader.readtext(thresh, detail=1, paragraph=False, 
-                                          contrast_ths=0.2, adjust_contrast=0.8, 
-                                          text_threshold=0.7, mag_ratio=2.0, 
-                                          allowlist='0123456789.', y_ths=0.3)
+                                          contrast_ths=0.3, adjust_contrast=1.0, 
+                                          text_threshold=0.6, mag_ratio=3.0, 
+                                          allowlist='0123456789.', y_ths=0.2)
         
         logging.info(f"EasyOCR results: {results}")
         if not results:
@@ -167,7 +176,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 == '.') and h_bbox > 8:
+            if len(text) == 1 and (text.isdigit() or text == '.') and h_bbox > 10:
                 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))
@@ -180,21 +189,18 @@ def custom_seven_segment_ocr(img, roi_bbox):
             if x_max <= x_min or y_max <= y_min:
                 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.9 or easyocr_char == '.' or digit_img_crop.shape[0] < 15 or digit_img_crop.shape[1] < 10:
+            save_debug_image(digit_img_crop, f"09_digit_crop_{idx}_{easyocr_char}")
+            if easyocr_conf > 0.95 or easyocr_char == '.':
                 recognized_text += easyocr_char
             else:
                 digit_from_segments = detect_segments(digit_img_crop)
-                if digit_from_segments:
-                    recognized_text += digit_from_segments
-                else:
-                    recognized_text += easyocr_char
+                recognized_text += digit_from_segments if digit_from_segments else easyocr_char
         
         logging.info(f"Before validation, recognized_text: {recognized_text}")
         text = re.sub(r"[^\d\.]", "", recognized_text)
         if text.count('.') > 1:
             text = text.replace('.', '', text.count('.') - 1)
-        if text and re.fullmatch(r"^\d*\.?\d*$", text) and len(text) > 0:
+        if text and re.fullmatch(r"^\d*\.?\d+$", text):
             if text.startswith('.'):
                 text = "0" + text
             if text.endswith('.'):
@@ -209,92 +215,56 @@ def custom_seven_segment_ocr(img, roi_bbox):
         return None
 
 def extract_weight_from_image(pil_img):
-    """Extract weight from a PIL image of a digital scale display"""
+    """Extract weight from a PIL image of a digital scale display."""
     try:
         img = np.array(pil_img)
         img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        save_debug_image(img, "00_input_image")  # Log input image
+        save_debug_image(img, "00_input_image")
 
         brightness = estimate_brightness(img)
-        conf_threshold = 0.6 if brightness > 150 else (0.5 if brightness > 80 else 0.4)
+        conf_threshold = 0.7 if brightness > 150 else (0.6 if brightness > 80 else 0.5)
 
         roi_img, roi_bbox = detect_roi(img)
         custom_result = custom_seven_segment_ocr(roi_img, roi_bbox)
         if custom_result:
-            if "." in custom_result:
-                int_part, dec_part = custom_result.split(".")
-                int_part = int_part.lstrip("0") or "0"
-                dec_part = dec_part.rstrip('0')
-                if not dec_part and int_part != "0":
-                    custom_result = int_part
-                elif not dec_part and int_part == "0":
-                    custom_result = "0"
-                else:
-                    custom_result = f"{int_part}.{dec_part}"
-            else:
-                custom_result = custom_result.lstrip('0') or "0"
             try:
-                float(custom_result)
-                logging.info(f"Custom OCR result: {custom_result}, Confidence: 100.0%")
-                return custom_result, 100.0
+                weight = float(custom_result)
+                if 0.1 <= weight <= 300:
+                    logging.info(f"Custom OCR result: {custom_result}, Confidence: 95.0%")
+                    return custom_result, 95.0
+                else:
+                    logging.warning(f"Custom OCR result {custom_result} outside typical weight range.")
             except ValueError:
-                logging.warning(f"Custom OCR result '{custom_result}' is not a valid number, falling back.")
-                custom_result = None
+                logging.warning(f"Custom OCR result '{custom_result}' is not a valid number.")
 
-        logging.info("Custom OCR failed or invalid, falling back to general EasyOCR.")
-        processed_roi_img_gray = cv2.cvtColor(roi_img, cv2.COLOR_BGR2GRAY)
-        kernel_sharpening = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
-        sharpened_roi = cv2.filter2D(processed_roi_img_gray, -1, kernel_sharpening)
-        save_debug_image(sharpened_roi, "08_fallback_sharpened")
-        processed_roi_img_final = cv2.adaptiveThreshold(sharpened_roi, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                                       cv2.THRESH_BINARY, 21, 5)
-        save_debug_image(processed_roi_img_final, "09_fallback_adaptive_thresh")
+        logging.info("Custom OCR failed or invalid, falling back to enhanced EasyOCR.")
+        preprocessed_roi = preprocess_image(roi_img)
+        kernel_sharpening = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
+        sharpened_roi = cv2.filter2D(preprocessed_roi, -1, kernel_sharpening)
+        save_debug_image(sharpened_roi, "10_fallback_sharpened")
+        final_roi = cv2.adaptiveThreshold(sharpened_roi, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
+                                          cv2.THRESH_BINARY_INV, 25, 8)
+        save_debug_image(final_roi, "11_fallback_adaptive_thresh")
 
-        results = easyocr_reader.readtext(processed_roi_img_final, detail=1, paragraph=False, 
-                                          contrast_ths=0.3, adjust_contrast=0.9, 
-                                          text_threshold=0.5, mag_ratio=2.0, 
-                                          allowlist='0123456789.', batch_size=4, y_ths=0.3)
+        results = easyocr_reader.readtext(final_roi, detail=1, paragraph=False, 
+                                          contrast_ths=0.4, adjust_contrast=1.2, 
+                                          text_threshold=0.5, mag_ratio=4.0, 
+                                          allowlist='0123456789.', batch_size=4, y_ths=0.2)
 
         best_weight = None
         best_conf = 0.0
         best_score = 0.0
         for (bbox, text, conf) in results:
-            text = text.lower().strip()
-            text = text.replace(",", ".").replace(";", ".").replace(":", ".").replace(" ", "")
-            text = text.replace("o", "0").replace("O", "0").replace("q", "0").replace("Q", "0")
-            text = text.replace("s", "5").replace("S", "5")
-            text = text.replace("g", "9").replace("G", "6")
-            text = text.replace("l", "1").replace("I", "1").replace("|", "1")
-            text = text.replace("b", "8").replace("B", "8")
-            text = text.replace("z", "2").replace("Z", "2")
-            text = text.replace("a", "4").replace("A", "4")
-            text = text.replace("e", "3")
-            text = text.replace("t", "7")
-            text = text.replace("~", "").replace("`", "")
-            text = re.sub(r"(kgs|kg|k|lb|g|gr|pounds|lbs)\b", "", text)
             text = re.sub(r"[^\d\.]", "", text)
             if text.count('.') > 1:
-                parts = text.split('.')
-                text = parts[0] + '.' + ''.join(parts[1:])
+                text = text.replace('.', '', text.count('.') - 1)
             text = text.strip('.')
-            if re.fullmatch(r"^\d*\.?\d{0,3}$", text) and len(text.replace('.', '')) > 0:
+            if re.fullmatch(r"^\d*\.?\d+$", text):
                 try:
                     weight = float(text)
-                    range_score = 1.0
-                    if 0.1 <= weight <= 250:
-                        range_score = 1.5
-                    elif weight > 250 and weight <= 500:
-                        range_score = 1.2
-                    elif weight > 500 and weight <= 1000:
-                        range_score = 1.0
-                    else:
-                        range_score = 0.5
+                    range_score = 1.5 if 0.1 <= weight <= 300 else 0.8
                     digit_count = len(text.replace('.', ''))
-                    digit_score = 1.0
-                    if digit_count >= 2 and digit_count <= 5:
-                        digit_score = 1.3
-                    elif digit_count == 1:
-                        digit_score = 0.8
+                    digit_score = 1.3 if 2 <= digit_count <= 5 else 0.9
                     score = conf * range_score * digit_score
                     if roi_bbox:
                         (x_roi, y_roi, w_roi, h_roi) = roi_bbox
@@ -302,11 +272,8 @@ 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.03:
-                            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.2:
-                            score *= 0.7
+                        if roi_area > 0 and bbox_area / roi_area < 0.05:
+                            score *= 0.6
                     if score > best_score and conf > conf_threshold:
                         best_weight = text
                         best_conf = conf
@@ -320,24 +287,19 @@ def extract_weight_from_image(pil_img):
             logging.info("No valid weight detected after all attempts.")
             return "Not detected", 0.0
 
+        # Format the weight
         if "." in best_weight:
             int_part, dec_part = best_weight.split(".")
             int_part = int_part.lstrip("0") or "0"
             dec_part = dec_part.rstrip('0')
-            if not dec_part and int_part != "0":
-                best_weight = int_part
-            elif not dec_part and int_part == "0":
-                best_weight = "0"
-            else:
-                best_weight = f"{int_part}.{dec_part}"
+            best_weight = f"{int_part}.{dec_part}" if dec_part else int_part
         else:
             best_weight = best_weight.lstrip('0') or "0"
 
         try:
-            final_float_weight = float(best_weight)
-            if final_float_weight < 0.01 or final_float_weight > 1000:
-                logging.warning(f"Detected weight {final_float_weight} is outside typical range, reducing confidence.")
-                best_conf *= 0.5
+            final_weight = float(best_weight)
+            if final_weight < 0.1 or final_weight > 300:
+                best_conf *= 0.7
         except ValueError:
             pass