Update ocr_engine.py

ocr_engine.py

@@ -8,6 +8,8 @@ import logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Initialize EasyOCR
+# Consider using 'en' and potentially 'ch_sim' or other relevant languages if your scales have non-English characters.
+# gpu=True can speed up processing if a compatible GPU is available.
 easyocr_reader = easyocr.Reader(['en'], gpu=False)
 
 def estimate_brightness(img):
@@ -20,21 +22,38 @@ def detect_roi(img):
     try:
         gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
         brightness = estimate_brightness(img)
-        thresh_value = 230 if brightness > 100 else 190
+
+        # Adaptive thresholding based on brightness
+        # For darker images, a lower threshold might be needed.
+        # For very bright images, a higher threshold.
+        thresh_value = 230 if brightness > 180 else (190 if brightness > 100 else 150)
         _, thresh = cv2.threshold(gray, thresh_value, 255, cv2.THRESH_BINARY)
-        kernel = np.ones((9, 9), np.uint8)
-        dilated = cv2.dilate(thresh, kernel, iterations=3)
+        
+        # Increased kernel size for dilation to better connect segments of digits
+        kernel = np.ones((11, 11), np.uint8) 
+        dilated = cv2.dilate(thresh, kernel, iterations=4) # Increased iterations
+        
         contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
         if contours:
-            valid_contours = [c for c in contours if cv2.contourArea(c) > 500]
+            # Filter contours by a more robust area range
+            valid_contours = [c for c in contours if 1000 < cv2.contourArea(c) < (img.shape[0] * img.shape[1] * 0.8)] # Added max area limit
+            
             if valid_contours:
+                # Sort by area descending and iterate
                 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:
-                        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)
+                    
+                    # Tighter aspect ratio and size constraints for typical digital displays
+                    if 1.8 <= aspect_ratio <= 5.0 and w > 80 and h > 40: # Adjusted min w and h
+                        # Expand ROI to ensure full digits are captured
+                        padding = 30 # 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)
                         return img[y:y+h, x:x+w], (x, y, w, h)
+        
+        logging.info("No suitable ROI found, returning original image.")
         return img, None
     except Exception as e:
         logging.error(f"ROI detection failed: {str(e)}")
@@ -43,32 +62,41 @@ def detect_roi(img):
 def detect_segments(digit_img):
     """Detect seven-segment patterns in a digit image"""
     h, w = digit_img.shape
-    if h < 10 or w < 10:
+    if h < 15 or w < 10: # Increased minimum dimensions for a digit
         return None
 
     # Define segment regions (top, middle, bottom, left-top, left-bottom, right-top, right-bottom)
+    # Adjusted segment proportions for better robustness
     segments = {
-        'top': (0, w, 0, h//5),
-        'middle': (0, w, 2*h//5, 3*h//5),
-        'bottom': (0, w, 4*h//5, h),
-        'left_top': (0, w//5, 0, h//2),
-        'left_bottom': (0, w//5, h//2, h),
-        'right_top': (4*w//5, w, 0, h//2),
-        'right_bottom': (4*w//5, w, h//2, h)
+        'top': (int(w*0.1), int(w*0.9), 0, int(h*0.2)),
+        'middle': (int(w*0.1), int(w*0.9), int(h*0.4), int(h*0.6)),
+        'bottom': (int(w*0.1), int(w*0.9), int(h*0.8), h),
+        'left_top': (0, int(w*0.2), int(h*0.05), int(h*0.5)),
+        'left_bottom': (0, int(w*0.2), int(h*0.5), int(h*0.95)),
+        'right_top': (int(w*0.8), w, int(h*0.05), int(h*0.5)),
+        'right_bottom': (int(w*0.8), w, int(h*0.5), int(h*0.95))
     }
 
     segment_presence = {}
     for name, (x1, x2, y1, y2) in segments.items():
+        # Ensure coordinates are within bounds
+        x1, y1 = max(0, x1), max(0, y1)
+        x2, y2 = min(w, x2), min(h, y2)
+        
         region = digit_img[y1:y2, x1:x2]
         if region.size == 0:
-            return None
+            segment_presence[name] = False
+            continue
+        
         # Count white pixels in the region
         pixel_count = np.sum(region == 255)
         total_pixels = region.size
-        # Segment is present if more than 50% of the region is white
-        segment_presence[name] = pixel_count > total_pixels * 0.5
+        
+        # Segment is present if a significant portion of the region is white
+        # Adjusted threshold for segment presence
+        segment_presence[name] = pixel_count / total_pixels > 0.4 # Increased sensitivity
 
-    # Seven-segment digit patterns
+    # Seven-segment digit patterns - remain the same
     digit_patterns = {
         '0': ('top', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
         '1': ('right_top', 'right_bottom'),
@@ -83,58 +111,105 @@ def detect_segments(digit_img):
     }
 
     best_match = None
-    max_matches = 0
+    max_score = -1 # Initialize with a lower value
+
     for digit, pattern in digit_patterns.items():
         matches = sum(1 for segment in pattern if segment_presence.get(segment, False))
-        non_matches = sum(1 for segment in segment_presence if segment not in pattern and segment_presence[segment])
-        score = matches - non_matches
-        if score > max_matches:
-            max_matches = score
+        
+        # Penalize for segments that should NOT be present but are
+        non_matches_penalty = sum(1 for segment in segment_presence if segment not in pattern and segment_presence[segment])
+        
+        # Prioritize digits with more matched segments and fewer incorrect segments
+        current_score = matches - non_matches_penalty
+        
+        # Add a small bonus for matching exactly all required segments for the digit
+        if all(segment_presence.get(s, False) for s in pattern):
+            current_score += 0.5 
+
+        if current_score > max_score:
+            max_score = current_score
             best_match = digit
+        elif current_score == max_score and best_match is not None:
+            # Tie-breaking: prefer digits with fewer "extra" segments when scores are equal
+            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
 
     return best_match
 
+
 def custom_seven_segment_ocr(img, roi_bbox):
     """Perform custom OCR for seven-segment displays"""
     try:
         gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-
+        
+        # Adaptive thresholding for digits within ROI
+        # Using OTSU for automatic thresholding or a fixed value depending on brightness
+        brightness = estimate_brightness(img)
+        if brightness > 150:
+            _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+        else:
+            _, thresh = cv2.threshold(gray, 120, 255, cv2.THRESH_BINARY) # Adjust threshold for darker displays
+        
         # Use EasyOCR to get bounding boxes for digits
+        # Increased text_threshold for more confident digit detection
+        # Adjusted mag_ratio for better handling of digit sizes
         results = easyocr_reader.readtext(thresh, detail=1, paragraph=False, 
-                                        contrast_ths=0.1, adjust_contrast=0.7, 
-                                        text_threshold=0.9, mag_ratio=1.5, 
+                                        contrast_ths=0.2, adjust_contrast=0.8, # Slightly more contrast adjustment
+                                        text_threshold=0.85, mag_ratio=1.2, # Reduced mag_ratio for potentially closer digits
                                         allowlist='0123456789.')
 
         if not results:
+            logging.info("EasyOCR found no digits for custom seven-segment OCR.")
             return None
 
         # Sort bounding boxes left to right
-        digits = []
-        for (bbox, _, _) in results:
-            (x1, y1), (x2, y2), (x3, y3), (x4, y4) = bbox
-            x_min, x_max = min(x1, x4), max(x2, x3)
-            y_min, y_max = min(y1, y2), max(y3, y4)
-            digits.append((x_min, x_max, y_min, y_max))
+        digits_info = []
+        for (bbox, text, conf) in results:
+            # Ensure the text found by EasyOCR is a single digit or a decimal point
+            if len(text) == 1 and (text.isdigit() or text == '.'):
+                (x1, y1), (x2, y2), (x3, y3), (x4, y4) = bbox
+                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))
 
-        digits.sort(key=lambda x: x[0])  # Sort by x_min (left to right)
+        # Sort by x_min (left to right)
+        digits_info.sort(key=lambda x: x[0]) 
 
-        # Extract and recognize each digit
         recognized_text = ""
-        for x_min, x_max, y_min, y_max in digits:
-            x_min, y_min = max(0, int(x_min)), max(0, int(y_min))
-            x_max, y_max = min(thresh.shape[1], int(x_max)), min(thresh.shape[0], int(y_max))
+        for x_min, x_max, y_min, y_max, easyocr_char, easyocr_conf in digits_info:
+            x_min, y_min = max(0, x_min), max(0, y_min)
+            x_max, y_max = min(thresh.shape[1], x_max), min(thresh.shape[0], y_max)
+            
             if x_max <= x_min or y_max <= y_min:
                 continue
-            digit_img = thresh[y_min:y_max, x_min:x_max]
-            digit = detect_segments(digit_img)
-            if digit:
-                recognized_text += digit
-
+            
+            digit_img_crop = thresh[y_min:y_max, x_min:x_max]
+            
+            # If EasyOCR is very confident about a digit or it's a decimal, use its result directly
+            if easyocr_conf > 0.95 or easyocr_char == '.':
+                recognized_text += easyocr_char
+            else:
+                # Otherwise, try the segment detection
+                digit_from_segments = detect_segments(digit_img_crop)
+                if digit_from_segments:
+                    recognized_text += digit_from_segments
+                else:
+                    # If segment detection also fails, fall back to EasyOCR's less confident result
+                    recognized_text += easyocr_char
+            
         # Validate the recognized text
         text = recognized_text
-        text = re.sub(r"[^\d\.]", "", text)
-        if re.fullmatch(r"\d{1,4}(\.\d{0,3})?", text):
+        text = re.sub(r"[^\d\.]", "", text) # Remove any non-digit/non-dot characters
+        
+        # Ensure there's at most one decimal point
+        if text.count('.') > 1:
+            text = text.replace('.', '', text.count('.') - 1) # Remove extra decimal points
+        
+        # Basic validation for common weight formats
+        if re.fullmatch(r"^\d+(\.\d+)?$", text) and len(text) > 0: # Ensures it starts with digit and has optional decimal
             return text
         return None
     except Exception as e:
@@ -147,74 +222,141 @@ def extract_weight_from_image(pil_img):
         img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
 
         brightness = estimate_brightness(img)
-        conf_threshold = 0.9 if brightness > 100 else 0.7
+        # Adjust confidence threshold more dynamically
+        conf_threshold = 0.9 if brightness > 150 else (0.75 if brightness > 80 else 0.6)
 
         # Detect ROI
         roi_img, roi_bbox = detect_roi(img)
+        
+        # Convert ROI to RGB for display purposes if needed later
+        # roi_img_rgb = cv2.cvtColor(roi_img, cv2.COLOR_BGR2RGB) # For debugging or display
 
         # Try custom seven-segment OCR first
         custom_result = custom_seven_segment_ocr(roi_img, roi_bbox)
         if custom_result:
-            # Format the custom result
+            # Format the custom result: remove leading zeros (unless it's "0" or "0.x") and trailing zeros after decimal
             if "." in custom_result:
                 int_part, dec_part = custom_result.split(".")
                 int_part = int_part.lstrip("0") or "0"
                 custom_result = f"{int_part}.{dec_part.rstrip('0')}"
             else:
                 custom_result = custom_result.lstrip('0') or "0"
+            
+            # Additional validation for custom result
+            if custom_result == "0." or custom_result == ".": # Handle cases like "0." or just "."
+                return "Not detected", 0.0
+
+            logging.info(f"Custom OCR result: {custom_result}, Confidence: 100.0%")
             return custom_result, 100.0  # High confidence for custom OCR
 
         # Fallback to EasyOCR if custom OCR fails
-        images_to_process = [
-            ("raw", roi_img, {'contrast_ths': 0.1, 'adjust_contrast': 0.7, 'text_threshold': 0.9, 'mag_ratio': 1.5, 'allowlist': '0123456789.'}),
-        ]
+        logging.info("Custom OCR failed, falling back to general EasyOCR.")
+        
+        # Apply more aggressive image processing for EasyOCR if custom OCR failed
+        # This could involve different thresholds or contrast adjustments
+        processed_roi_img_gray = cv2.cvtColor(roi_img, cv2.COLOR_BGR2GRAY)
+        
+        # Sharpening
+        kernel_sharpening = np.array([[-1,-1,-1], 
+                                      [-1,9,-1], 
+                                      [-1,-1,-1]])
+        sharpened_roi = cv2.filter2D(processed_roi_img_gray, -1, kernel_sharpening)
+
+        # Apply adaptive thresholding to the sharpened image for better digit isolation
+        processed_roi_img_final = cv2.adaptiveThreshold(sharpened_roi, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
+                                                        cv2.THRESH_BINARY, 11, 2)
 
+        # EasyOCR parameters for general text
+        # Adjusted parameters for better digit recognition
+        # added batch_size for potentially better performance on multiple texts
+        results = easyocr_reader.readtext(processed_roi_img_final, detail=1, paragraph=False, 
+                                        contrast_ths=0.3, adjust_contrast=0.9, 
+                                        text_threshold=0.7, mag_ratio=1.8, # Increased mag_ratio for potentially larger digits
+                                        allowlist='0123456789.', batch_size=4) # Added batch_size
+        
         best_weight = None
         best_conf = 0.0
         best_score = 0.0
 
-        for mode, proc_img, ocr_params in images_to_process:
-            if mode == "raw":
-                proc_img = cv2.cvtColor(proc_img, cv2.COLOR_BGR2GRAY)
-            results = easyocr_reader.readtext(proc_img, detail=1, paragraph=False, **ocr_params)
+        for (bbox, text, conf) in results:
+            text = text.lower().strip()
             
-            for (bbox, text, conf) in results:
-                text = text.lower().strip()
-                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("l", "1").replace("I", "1")
-                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("kgs", "").replace("kg", "").replace("k", "")
-                text = re.sub(r"[^\d\.]", "", text)
-
-                if re.fullmatch(r"\d{1,4}(\.\d{0,3})?", text):
-                    try:
-                        weight = float(text)
-                        range_score = 1.0 if 0.1 <= weight <= 500 else 0.3
-                        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:
-                            best_weight = text
-                            best_conf = conf
-                            best_score = score
-                    except ValueError:
-                        continue
+            # More robust character replacements
+            text = text.replace(",", ".").replace(";", ".").replace(":", ".")
+            text = text.replace("o", "0").replace("O", "0").replace("q", "0") # 'q' can look like 0
+            text = text.replace("s", "5").replace("S", "5")
+            text = text.replace("g", "9").replace("G", "6") # Be careful with G to 6 conversion
+            text = text.replace("l", "1").replace("I", "1").replace("|", "1") # Added | to 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") # 'a' can look like 4
+            text = text.replace("e", "3") # 'e' can look like 3
+
+            # Remove common weight units and other non-numeric characters
+            text = re.sub(r"(kgs|kg|k|lb|g|gr|pounds)\b", "", text) # Use word boundary \b
+            text = re.sub(r"[^\d\.]", "", text)
+
+            # Handle multiple decimal points (keep only the first one)
+            if text.count('.') > 1:
+                parts = text.split('.')
+                text = parts[0] + '.' + ''.join(parts[1:])
+            
+            # Validate the final text format
+            if re.fullmatch(r"^\d{1,4}(\.\d{0,3})?$", text): # Adjusted regex for more flexible digits
+                try:
+                    weight = float(text)
+                    # Refined scoring for weights within a reasonable range
+                    range_score = 1.0
+                    if 0.01 <= weight <= 300: # Typical personal scale range
+                        range_score = 1.2
+                    elif weight > 300 and weight <= 1000: # Larger scales
+                        range_score = 1.1
+                    else: # Very small or very large weights
+                        range_score = 0.8
+                    
+                    digit_count = len(text.replace('.', ''))
+                    digit_score = 1.0
+                    if digit_count >= 3 and digit_count <= 5: # Prefer weights with 3-5 digits (e.g., 50.5, 123.4)
+                        digit_score = 1.3
+                    
+                    score = conf * range_score * digit_score
+                    
+                    # Also consider area of the bounding box relative to ROI for confidence
+                    bbox_area = (bbox[1][0] - bbox[0][0]) * (bbox[2][1] - bbox[1][1])
+                    if roi_bbox:
+                        roi_area = roi_bbox[2] * roi_bbox[3]
+                        if roi_area > 0 and bbox_area / roi_area < 0.05: # Small bounding boxes might be noise
+                            score *= 0.5 
+
+                    if score > best_score and conf > conf_threshold:
+                        best_weight = text
+                        best_conf = conf
+                        best_score = score
+                        logging.info(f"Candidate EasyOCR weight: {text}, Conf: {conf}, Score: {score}")
+
+                except ValueError:
+                    logging.warning(f"Could not convert '{text}' to float.")
+                    continue
 
         if not best_weight:
-            logging.info("No valid weight detected")
+            logging.info("No valid weight detected after all attempts.")
             return "Not detected", 0.0
 
+        # Final formatting of the best detected weight
         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.rstrip('0')}"
+            int_part = int_part.lstrip("0") or "0" # Remove leading zeros, keep "0" for 0.x
+            dec_part = dec_part.rstrip('0') # Remove trailing zeros after decimal
+            if not dec_part and int_part != "0": # If decimal part is empty (e.g., "50."), remove the dot
+                best_weight = int_part
+            elif not dec_part and int_part == "0": # if it's "0." keep it as "0"
+                best_weight = "0"
+            else:
+                best_weight = f"{int_part}.{dec_part}"
         else:
-            best_weight = best_weight.lstrip('0') or "0"
+            best_weight = best_weight.lstrip('0') or "0" # Remove leading zeros, keep "0"
 
+        logging.info(f"Final detected weight: {best_weight}, Confidence: {round(best_conf * 100, 2)}%")
         return best_weight, round(best_conf * 100, 2)
 
     except Exception as e: