Update ocr_engine.py

ocr_engine.py

@@ -21,7 +21,7 @@ def save_debug_image(img, filename_suffix, prefix=""):
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
     filename = os.path.join(DEBUG_DIR, f"{prefix}{timestamp}_{filename_suffix}.png")
     if len(img.shape) == 3:
-        cv2.imwrite(filename, img)
+        cv2.imwrite(filename, cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
     else:
         cv2.imwrite(filename, img)
     logging.info(f"Saved debug image: {filename}")
@@ -32,24 +32,32 @@ def estimate_brightness(img):
     return np.mean(gray)
 
 def preprocess_image(img):
-    """Preprocess image for OCR."""
+    """Preprocess image for OCR with enhanced contrast and noise reduction."""
     gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    denoised = cv2.bilateralFilter(gray, 5, 8, 8)
-    save_debug_image(denoised, "01_preprocess_bilateral")
-    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-    enhanced = clahe.apply(denoised)
+    # Apply Gaussian blur to reduce noise
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    save_debug_image(blurred, "01_preprocess_blur")
+    # Use adaptive histogram equalization for better contrast
+    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
+    enhanced = clahe.apply(blurred)
     save_debug_image(enhanced, "02_preprocess_clahe")
-    return enhanced
+    # Morphological operations to enhance digits
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+    morphed = cv2.morphologyEx(enhanced, cv2.MORPH_CLOSE, kernel)
+    save_debug_image(morphed, "03_preprocess_morph")
+    return morphed
 
 def correct_rotation(img):
-    """Correct image rotation."""
+    """Correct image rotation using edge detection."""
     try:
-        edges = cv2.Canny(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), 50, 150)
+        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+        edges = cv2.Canny(blurred, 50, 150)
         lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=50, minLineLength=30, maxLineGap=10)
         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) > 2:
+            if abs(angle) > 1.5:
                 h, w = img.shape[:2]
                 center = (w // 2, h // 2)
                 M = cv2.getRotationMatrix2D(center, angle, 1.0)
@@ -62,15 +70,20 @@ def correct_rotation(img):
         return img
 
 def detect_roi(img):
-    """Detect region of interest (display)."""
+    """Detect region of interest (display) with refined contour filtering."""
     try:
-        save_debug_image(img, "03_original")
+        save_debug_image(img, "04_original")
         preprocessed = preprocess_image(img)
         brightness_map = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        block_size = max(9, min(31, int(img.shape[0] / 25) * 2 + 1))
+        # Dynamic block size based on image dimensions
+        block_size = max(11, min(31, int(img.shape[0] / 20) * 2 + 1))
         thresh = cv2.adaptiveThreshold(preprocessed, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                        cv2.THRESH_BINARY_INV, block_size, 2)
-        save_debug_image(thresh, "04_roi_threshold")
+        save_debug_image(thresh, "05_roi_threshold")
+        # Morphological operations to connect digit segments
+        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
+        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
+        save_debug_image(thresh, "06_roi_morph")
         contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         
         if contours:
@@ -81,141 +94,105 @@ 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.2 <= aspect_ratio <= 30.0 and w > 30 and h > 10 and roi_brightness > 30):
-                    valid_contours.append((c, roi_brightness))
+                # Relaxed constraints for ROI detection
+                if (100 < area < (img_area * 0.9) and 
+                    0.3 <= aspect_ratio <= 20.0 and w > 40 and h > 15 and roi_brightness > 20):
+                    valid_contours.append((c, area * roi_brightness))
                     logging.debug(f"Contour: 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 = 200
+                # Dynamic padding based on ROI size
+                padding = max(10, min(50, int(min(w, h) * 0.2)))
                 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")
+                save_debug_image(roi_img, "07_detected_roi")
                 logging.info(f"Detected ROI: ({x}, {y}, {w}, {h})")
                 return roi_img, (x, y, w, h)
         
         logging.info("No ROI found, using full image.")
-        save_debug_image(img, "05_no_roi_fallback")
+        save_debug_image(img, "07_no_roi_fallback")
         return img, None
     except Exception as e:
         logging.error(f"ROI detection failed: {str(e)}")
-        save_debug_image(img, "05_roi_error_fallback")
+        save_debug_image(img, "07_roi_error_fallback")
         return img, None
 
-def detect_segments(digit_img, brightness):
-    """Detect seven-segment digits."""
-    h, w = digit_img.shape
-    if h < 5 or w < 3:
-        return None
-
-    segments = {
-        'top': (int(w*0.1), int(w*0.9), 0, int(h*0.25)),
-        '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.75), h),
-        'left_top': (0, int(w*0.3), int(h*0.1), int(h*0.5)),
-        'left_bottom': (0, int(w*0.3), int(h*0.5), int(h*0.9)),
-        'right_top': (int(w*0.7), w, int(h*0.1), int(h*0.5)),
-        'right_bottom': (int(w*0.7), w, int(h*0.5), int(h*0.9))
-    }
-
-    segment_presence = {}
-    for name, (x1, x2, y1, y2) in segments.items():
-        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:
-            segment_presence[name] = False
-            continue
-        pixel_count = np.sum(region == 255)
-        total_pixels = region.size
-        segment_presence[name] = pixel_count / total_pixels > (0.1 if brightness < 80 else 0.25)
-
-    digit_patterns = {
-        '0': ('top', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
-        '1': ('right_top', 'right_bottom'),
-        '2': ('top', 'middle', 'bottom', 'left_bottom', 'right_top'),
-        '3': ('top', 'middle', 'bottom', 'right_top', 'right_bottom'),
-        '4': ('middle', 'left_top', 'right_top', 'right_bottom'),
-        '5': ('top', 'middle', 'bottom', 'left_top', 'right_bottom'),
-        '6': ('top', 'middle', 'bottom', 'left_top', 'left_bottom', 'right_bottom'),
-        '7': ('top', 'right_top', 'right_bottom'),
-        '8': ('top', 'middle', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
-        '9': ('top', 'middle', 'bottom', 'left_top', 'right_top', 'right_bottom')
-    }
-
-    best_match = None
-    max_score = -1
-    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])
-        score = matches - 0.1 * non_matches_penalty
-        if matches >= len(pattern) * 0.55:
-            score += 1.0
-        if score > max_score:
-            max_score = score
-            best_match = digit
-    
-    logging.debug(f"Segment presence: {segment_presence}, Digit: {best_match}")
-    return best_match
-
-def custom_seven_segment_ocr(img, roi_bbox):
-    """Perform OCR for seven-segment displays."""
+def perform_ocr(img, roi_bbox):
+    """Perform OCR optimized for digital displays."""
     try:
         preprocessed = preprocess_image(img)
         brightness = estimate_brightness(img)
-        _, thresh = cv2.threshold(preprocessed, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
-        save_debug_image(thresh, "06_roi_thresh_digits")
+        # Dynamic thresholding based on brightness
+        thresh_value = 0 if brightness < 50 else (127 if brightness < 100 else 200)
+        _, thresh = cv2.threshold(preprocessed, thresh_value, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+        save_debug_image(thresh, "08_ocr_threshold")
+        # Morphological operations to clean up digits
+        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
+        save_debug_image(thresh, "09_ocr_morph")
+        
+        # Optimized EasyOCR parameters for seven-segment displays
         results = easyocr_reader.readtext(thresh, detail=1, paragraph=False, 
-                                          contrast_ths=0.05, adjust_contrast=1.2, 
-                                          text_threshold=0.15, mag_ratio=4.0, 
-                                          allowlist='0123456789.', batch_size=2, y_ths=0.3)
+                                          contrast_ths=0.1, adjust_contrast=1.5, 
+                                          text_threshold=0.2, mag_ratio=3.0, 
+                                          allowlist='0123456789.', batch_size=1, y_ths=0.2)
         
         logging.info(f"EasyOCR results: {results}")
+        if not results:
+            logging.info("No text detected, trying fallback parameters.")
+            results = easyocr_reader.readtext(thresh, detail=1, paragraph=False, 
+                                              contrast_ths=0.05, adjust_contrast=2.0, 
+                                              text_threshold=0.1, mag_ratio=4.0, 
+                                              allowlist='0123456789.', batch_size=1, y_ths=0.2)
+            save_debug_image(thresh, "09_fallback_threshold")
+        
         if not results:
             logging.info("No digits found.")
-            return None
+            return None, 0.0
 
         digits_info = []
         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 (text.isdigit() or text == '.') and h_bbox > 4:
+            if (text.isdigit() or text == '.') and h_bbox > 5 and conf > 0.1:
                 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))
 
+        if not digits_info:
+            logging.info("No valid digits after filtering.")
+            return None, 0.0
+
         digits_info.sort(key=lambda x: x[0])
         recognized_text = ""
-        for idx, (x_min, x_max, y_min, y_max, easyocr_char, easyocr_conf) in enumerate(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_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.8 or easyocr_char == '.':
-                recognized_text += easyocr_char
-            else:
-                digit_from_segments = detect_segments(digit_img_crop, brightness)
-                recognized_text += digit_from_segments if digit_from_segments else easyocr_char
+        total_conf = 0.0
+        conf_count = 0
+        for _, _, _, _, char, conf in digits_info:
+            recognized_text += char
+            total_conf += conf
+            conf_count += 1
         
-        logging.info(f"Recognized text: {recognized_text}")
+        avg_conf = total_conf / conf_count if conf_count > 0 else 0.0
+        logging.info(f"Recognized text: {recognized_text}, Average confidence: {avg_conf:.2f}")
+
+        # Validate and clean the recognized text
         text = re.sub(r"[^\d\.]", "", recognized_text)
         if text.count('.') > 1:
             text = text.replace('.', '', text.count('.') - 1)
+        text = text.strip('.')
         if text and re.fullmatch(r"^\d*\.?\d*$", text):
-            text = text.strip('.')
-            if text == '':
-                return None
-            return text.lstrip('0') or '0'
+            text = text.lstrip('0') or '0'
+            if text == '0' and avg_conf < 0.9:
+                avg_conf *= 0.7
+            return text, avg_conf * 100
         logging.info(f"Text '{recognized_text}' failed validation.")
-        return None
+        return None, 0.0
     except Exception as e:
-        logging.error(f"Seven-segment OCR failed: {str(e)}")
-        return None
+        logging.error(f"OCR failed: {str(e)}")
+        return None, 0.0
 
 def extract_weight_from_image(pil_img):
     """Extract weight from a digital scale image."""
@@ -225,148 +202,37 @@ def extract_weight_from_image(pil_img):
         save_debug_image(img, "00_input_image")
         img = correct_rotation(img)
         brightness = estimate_brightness(img)
-        conf_threshold = 0.6 if brightness > 150 else (0.4 if brightness > 80 else 0.2)
+        conf_threshold = 0.5 if brightness > 120 else (0.3 if brightness > 60 else 0.2)
 
         roi_img, roi_bbox = detect_roi(img)
         if roi_bbox:
-            conf_threshold *= 1.05 if (roi_bbox[2] * roi_bbox[3]) > (img.shape[0] * img.shape[1] * 0.5) else 1.0
+            conf_threshold *= 1.1 if (roi_bbox[2] * roi_bbox[3]) > (img.shape[0] * img.shape[1] * 0.4) else 1.0
 
-        custom_result = custom_seven_segment_ocr(roi_img, roi_bbox)
-        if custom_result and custom_result != '0':
+        result, confidence = perform_ocr(roi_img, roi_bbox)
+        if result and confidence >= conf_threshold * 100:
             try:
-                weight = float(custom_result)
+                weight = float(result)
                 if 0.00001 <= weight <= 10000:
-                    logging.info(f"Custom OCR: {custom_result}, Confidence: 90.0%")
-                    return custom_result, 90.0
-                logging.warning(f"Custom OCR {custom_result} out of range.")
+                    logging.info(f"Detected weight: {result} kg, Confidence: {confidence:.2f}%")
+                    return result, confidence
+                logging.warning(f"Weight {result} out of range.")
             except ValueError:
-                logging.warning(f"Custom OCR '{custom_result}' invalid number.")
-
-        logging.info("Custom OCR failed, using EasyOCR fallback.")
-        preprocessed_roi = preprocess_image(roi_img)
-        final_roi = cv2.adaptiveThreshold(preprocessed_roi, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                          cv2.THRESH_BINARY_INV, max(9, min(31, int(roi_img.shape[0] / 25) * 2 + 1)), 2)
-        save_debug_image(final_roi, "08_fallback_thresh")
+                logging.warning(f"Invalid weight format: {result}")
 
-        results = easyocr_reader.readtext(final_roi, detail=1, paragraph=False, 
-                                          contrast_ths=0.05, adjust_contrast=1.2, 
-                                          text_threshold=0.15, mag_ratio=4.0, 
-                                          allowlist='0123456789. kglb', batch_size=2, y_ths=0.3)
-
-        if not results:
-            logging.info("First EasyOCR pass failed, trying fallback.")
-            results = easyocr_reader.readtext(final_roi, detail=1, paragraph=False, 
-                                              contrast_ths=0.02, adjust_contrast=1.5, 
-                                              text_threshold=0.1, mag_ratio=5.0, 
-                                              allowlist='0123456789. kglb', batch_size=2, y_ths=0.3)
-            save_debug_image(final_roi, "08_fallback_thresh_fallback")
-
-        logging.info(f"EasyOCR results: {results}")
-        candidates = []
-        unit = None
-        for (bbox, text, conf) in results:
-            if 'kg' in text.lower():
-                unit = 'kg'
-                continue
-            elif 'g' in text.lower():
-                unit = 'g'
-                continue
-            elif 'lb' in text.lower():
-                unit = 'lb'
-                continue
-            text = re.sub(r"[^\d\.]", "", text)
-            if text.count('.') > 1:
-                text = text.replace('.', '', text.count('.') - 1)
-            text = text.strip('.')
-            if re.fullmatch(r"^\d*\.?\d*$", text):
-                try:
-                    weight = float(text)
-                    if unit == 'g':
-                        weight /= 1000
-                    elif unit == 'lb':
-                        weight *= 0.453592
-                    range_score = 1.5 if 0.00001 <= weight <= 10000 else 0.5
-                    digit_count = len(text.replace('.', ''))
-                    digit_score = 1.4 if 1 <= digit_count <= 8 else 0.6
-                    score = conf * range_score * digit_score
-                    if roi_bbox:
-                        x_roi, y_roi, w_roi, h_roi = roi_bbox
-                        roi_area = w_roi * h_roi
-                        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.02:
-                            score *= 0.4
-                    candidates.append((text, conf, score, unit))
-                    logging.info(f"Candidate: '{text}', Unit: {unit or 'none'}, Conf: {conf}, Score: {score}")
-                except ValueError:
-                    logging.warning(f"Could not convert '{text}' to float.")
-
-        if not candidates and not roi_bbox:
-            logging.info("No candidates, trying full image.")
-            preprocessed_full = preprocess_image(img)
-            final_full = cv2.adaptiveThreshold(preprocessed_full, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                               cv2.THRESH_BINARY_INV, max(9, min(31, int(img.shape[0] / 25) * 2 + 1)), 2)
-            save_debug_image(final_full, "08_fallback_full")
-            results = easyocr_reader.readtext(final_full, detail=1, paragraph=False, 
-                                              contrast_ths=0.05, adjust_contrast=1.5, 
-                                              text_threshold=0.15, mag_ratio=4.0, 
-                                              allowlist='0123456789. kglb', batch_size=2, y_ths=0.3)
-            logging.info(f"Full image EasyOCR: {results}")
-            for (bbox, text, conf) in results:
-                if 'kg' in text.lower():
-                    unit = 'kg'
-                    continue
-                elif 'g' in text.lower():
-                    unit = 'g'
-                    continue
-                elif 'lb' in text.lower():
-                    unit = 'lb'
-                    continue
-                text = re.sub(r"[^\d\.]", "", text)
-                if text.count('.') > 1:
-                    text = text.replace('.', '', text.count('.') - 1)
-                text = text.strip('.')
-                if re.fullmatch(r"^\d*\.?\d*$", text):
-                    try:
-                        weight = float(text)
-                        if unit == 'g':
-                            weight /= 1000
-                        elif unit == 'lb':
-                            weight *= 0.453592
-                        range_score = 1.2 if 0.00001 <= weight <= 10000 else 0.4
-                        digit_count = len(text.replace('.', ''))
-                        digit_score = 1.2 if 1 <= digit_count <= 8 else 0.5
-                        score = conf * range_score * digit_score * 0.7
-                        candidates.append((text, conf, score, unit))
-                        logging.info(f"Full image candidate: '{text}', Unit: {unit or 'none'}, Conf: {conf}, Score: {score}")
-                    except ValueError:
-                        logging.warning(f"Could not convert '{text}' to float (full image).")
-
-        if not candidates:
-            logging.info("No valid weight detected.")
-            return "Not detected", 0.0
-
-        best_weight, best_conf, best_score, best_unit = max(candidates, key=lambda x: x[2])
-        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')
-            best_weight = f"{int_part}.{dec_part}" if dec_part else int_part
-        else:
-            best_weight = best_weight.lstrip('0') or "0"
-
-        try:
-            final_weight = float(best_weight)
-            if final_weight < 0.00001 or final_weight > 10000:
-                best_conf *= 0.4
-            elif final_weight == 0 and best_conf < 0.95:
-                best_conf *= 0.5
-        except ValueError:
-            pass
+        logging.info("Primary OCR failed, using full image fallback.")
+        result, confidence = perform_ocr(img, None)
+        if result and confidence >= conf_threshold * 0.8 * 100:
+            try:
+                weight = float(result)
+                if 0.00001 <= weight <= 10000:
+                    logging.info(f"Full image weight: {result} kg, Confidence: {confidence:.2f}%")
+                    return result, confidence
+                logging.warning(f"Full image weight {result} out of range.")
+            except ValueError:
+                logging.warning(f"Invalid full image weight format: {result}")
 
-        logging.info(f"Final weight: {best_weight} kg, Confidence: {round(best_conf * 100, 2)}%, Unit: {best_unit or 'none'}")
-        return best_weight, round(best_conf * 100, 2)
+        logging.info("No valid weight detected.")
+        return "Not detected", 0.0
     except Exception as e:
         logging.error(f"Weight extraction failed: {str(e)}")
         return "Not detected", 0.0