Update ocr_engine.py

ocr_engine.py

@@ -34,12 +34,12 @@ def estimate_brightness(img):
 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")
+    # Apply bilateral filter to preserve edges
+    denoised = cv2.bilateralFilter(gray, 11, 17, 17)
+    save_debug_image(denoised, "01_preprocess_bilateral")
     # Enhance contrast using CLAHE
     clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-    enhanced = clahe.apply(blurred)
+    enhanced = clahe.apply(denoised)
     save_debug_image(enhanced, "02_preprocess_clahe")
     return enhanced
 
@@ -48,15 +48,19 @@ def detect_roi(img):
     try:
         save_debug_image(img, "03_original")
         preprocessed = preprocess_image(img)
+        brightness_map = cv2.GaussianBlur(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), (15, 15), 0)
         
-        # Adaptive thresholding with refined parameters
+        # Dynamic adaptive thresholding
+        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, 15, 5)
-        save_debug_image(thresh, "04_roi_adaptive_threshold")
+                                       cv2.THRESH_BINARY_INV, block_size, 5)
+        _, otsu_thresh = cv2.threshold(preprocessed, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+        combined_thresh = cv2.bitwise_and(thresh, otsu_thresh)
+        save_debug_image(combined_thresh, "04_roi_combined_threshold")
         
         # Morphological operations to connect digits
         kernel = np.ones((5, 5), np.uint8)
-        dilated = cv2.dilate(thresh, kernel, iterations=2)
+        dilated = cv2.dilate(combined_thresh, kernel, iterations=2)
         eroded = cv2.erode(dilated, kernel, iterations=1)
         save_debug_image(eroded, "05_roi_morphological")
         
@@ -67,14 +71,16 @@ def detect_roi(img):
             valid_contours = []
             for c in contours:
                 area = cv2.contourArea(c)
-                if 1000 < area < (img_area * 0.9):
-                    x, y, w, h = cv2.boundingRect(c)
-                    aspect_ratio = w / h
-                    if 1.8 <= aspect_ratio <= 8.0 and w > 100 and h > 50:
-                        valid_contours.append(c)
+                x, y, w, h = cv2.boundingRect(c)
+                roi_brightness = np.mean(brightness_map[y:y+h, x:x+w])
+                aspect_ratio = w / h
+                if (1000 < area < (img_area * 0.9) and 
+                    1.0 <= aspect_ratio <= 10.0 and w > 80 and h > 40 and roi_brightness > 100):
+                    valid_contours.append((c, roi_brightness))
+                    logging.debug(f"Contour: Area={area}, Aspect={aspect_ratio:.2f}, Brightness={roi_brightness:.2f}")
             
             if valid_contours:
-                contour = max(valid_contours, key=cv2.contourArea)
+                contour, _ = max(valid_contours, key=lambda x: x[1])  # Max brightness
                 x, y, w, h = cv2.boundingRect(contour)
                 padding = 80
                 x, y = max(0, x - padding), max(0, y - padding)
@@ -84,6 +90,21 @@ def detect_roi(img):
                 logging.info(f"Detected ROI with dimensions: ({x}, {y}, {w}, {h})")
                 return roi_img, (x, y, w, h)
         
+        logging.info("No suitable ROI found, attempting fallback criteria.")
+        # Fallback with relaxed criteria
+        valid_contours = [c for c in contours if 500 < cv2.contourArea(c) < (img_area * 0.95) and 
+                        0.8 <= cv2.boundingRect(c)[2]/cv2.boundingRect(c)[3] <= 12.0]
+        if valid_contours:
+            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_fallback")
+            logging.info(f"Detected fallback 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, "06_no_roi_original_fallback")
         return img, None
@@ -92,20 +113,20 @@ def detect_roi(img):
         save_debug_image(img, "06_roi_detection_error_fallback")
         return img, None
 
-def detect_segments(digit_img):
+def detect_segments(digit_img, brightness):
     """Detect seven-segment patterns in a digit image."""
     h, w = digit_img.shape
     if h < 20 or w < 15:
         return None
 
     segments = {
-        '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))
+        '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.45), int(h*0.55)),
+        'bottom': (int(w*0.15), int(w*0.85), int(h*0.8), h),
+        'left_top': (0, int(w*0.25), int(h*0.15), int(h*0.5)),
+        'left_bottom': (0, int(w*0.25), int(h*0.5), int(h*0.85)),
+        'right_top': (int(w*0.75), w, int(h*0.15), int(h*0.5)),
+        'right_bottom': (int(w*0.75), w, int(h*0.5), int(h*0.85))
     }
 
     segment_presence = {}
@@ -118,7 +139,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.4
+        segment_presence[name] = pixel_count / total_pixels > (0.3 if brightness < 100 else 0.5)
 
     digit_patterns = {
         '0': ('top', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
@@ -138,8 +159,8 @@ 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])
-        score = matches - 0.5 * non_matches_penalty
-        if all(segment_presence.get(s, False) for s in pattern):
+        score = matches - 0.3 * non_matches_penalty
+        if matches >= len(pattern) * 0.8:
             score += 1.0
         if score > max_score:
             max_score = score
@@ -162,10 +183,11 @@ def custom_seven_segment_ocr(img, roi_bbox):
         thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
         save_debug_image(thresh, "08_morph_closed")
 
+        batch_size = max(4, min(16, int(img.shape[0] * img.shape[1] / 100000)))
         results = easyocr_reader.readtext(thresh, detail=1, paragraph=False, 
                                           contrast_ths=0.3, adjust_contrast=1.0, 
                                           text_threshold=0.6, mag_ratio=3.0, 
-                                          allowlist='0123456789.', y_ths=0.2)
+                                          allowlist='0123456789.', batch_size=batch_size, y_ths=0.2)
         
         logging.info(f"EasyOCR results: {results}")
         if not results:
@@ -193,21 +215,18 @@ def custom_seven_segment_ocr(img, roi_bbox):
             if easyocr_conf > 0.95 or easyocr_char == '.':
                 recognized_text += easyocr_char
             else:
-                digit_from_segments = detect_segments(digit_img_crop)
+                digit_from_segments = detect_segments(digit_img_crop, brightness)
                 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):
-            if text.startswith('.'):
-                text = "0" + text
-            if text.endswith('.'):
-                text = text.rstrip('.')
-            if text == '.' or text == '':
+        if text and re.fullmatch(r"^\d*\.?\d*$", text):
+            text = text.strip('.')
+            if text == '':
                 return None
-            return text
+            return text.lstrip('0') or '0'
         logging.info(f"Custom OCR text '{recognized_text}' failed validation.")
         return None
     except Exception as e:
@@ -221,15 +240,30 @@ def extract_weight_from_image(pil_img):
         img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
         save_debug_image(img, "00_input_image")
 
+        # Rotation correction
+        edges = cv2.Canny(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), 100, 200)
+        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100, minLineLength=100, maxLineGap=10)
+        if lines is not None:
+            angle = np.mean([np.arctan2(line[0][3] - line[0][1], line[0][2] - line[0][0]) * 180 / np.pi for line in lines])
+            if abs(angle) > 5:
+                (h, w) = img.shape[:2]
+                center = (w // 2, h // 2)
+                M = cv2.getRotationMatrix2D(center, angle, 1.0)
+                img = cv2.warpAffine(img, M, (w, h))
+                save_debug_image(img, "00_rotated_image")
+
         brightness = estimate_brightness(img)
-        conf_threshold = 0.7 if brightness > 150 else (0.6 if brightness > 80 else 0.5)
+        conf_threshold = 0.7 if brightness > 150 else (0.6 if brightness > 80 else 0.4)
+        if roi_bbox := detect_roi(img)[1]:
+            roi_area = roi_bbox[2] * roi_bbox[3]
+            conf_threshold *= 1.2 if roi_area > (img.shape[0] * img.shape[1] * 0.5) else 1.0
 
         roi_img, roi_bbox = detect_roi(img)
         custom_result = custom_seven_segment_ocr(roi_img, roi_bbox)
         if custom_result:
             try:
                 weight = float(custom_result)
-                if 0.1 <= weight <= 300:
+                if 0.01 <= weight <= 500:
                     logging.info(f"Custom OCR result: {custom_result}, Confidence: 95.0%")
                     return custom_result, 95.0
                 else:
@@ -242,70 +276,13 @@ def extract_weight_from_image(pil_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")
+        block_size = max(11, min(31, int(roi_img.shape[0] / 20) * 2 + 1))
         final_roi = cv2.adaptiveThreshold(sharpened_roi, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                          cv2.THRESH_BINARY_INV, 25, 8)
+                                          cv2.THRESH_BINARY_INV, block_size, 8)
         save_debug_image(final_roi, "11_fallback_adaptive_thresh")
 
+        batch_size = max(4, min(16, int(roi_img.shape[0] * roi_img.shape[1] / 100000)))
         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 = 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)
-                    range_score = 1.5 if 0.1 <= weight <= 300 else 0.8
-                    digit_count = len(text.replace('.', ''))
-                    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
-                        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.05:
-                            score *= 0.6
-                    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 during EasyOCR fallback.")
-                    continue
-
-        if not best_weight:
-            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')
-            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.1 or final_weight > 300:
-                best_conf *= 0.7
-        except ValueError:
-            pass
-
-        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:
-        logging.error(f"Weight extraction failed unexpectedly: {str(e)}")
-        return "Not detected", 0.0
+                                          allowlist='0123456789. kglb', batch_size=batch Horrible error: invalid syntax