Update ocr_engine.py

ocr_engine.py

@@ -38,15 +38,39 @@ def preprocess_image(img):
     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))
+    clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8))
     enhanced = clahe.apply(denoised)
     save_debug_image(enhanced, "02_preprocess_clahe")
-    return enhanced
+    # Sharpen the image
+    kernel_sharpening = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
+    sharpened = cv2.filter2D(enhanced, -1, kernel_sharpening)
+    save_debug_image(sharpened, "03_preprocess_sharpened")
+    return sharpened
+
+def correct_rotation(img):
+    """Correct image rotation using Hough Transform."""
+    try:
+        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:
+            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)  # Use median for robustness
+            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")
+                logging.info(f"Applied rotation correction: {angle:.2f} degrees")
+        return img
+    except Exception as e:
+        logging.error(f"Rotation correction failed: {str(e)}")
+        return img
 
 def detect_roi(img):
     """Detect and crop the region of interest (likely the digital display)."""
     try:
-        save_debug_image(img, "03_original")
+        save_debug_image(img, "04_original")
         preprocessed = preprocess_image(img)
         brightness_map = cv2.GaussianBlur(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), (15, 15), 0)
         
@@ -56,13 +80,13 @@ def detect_roi(img):
                                        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")
+        save_debug_image(combined_thresh, "05_roi_combined_threshold")
         
         # Morphological operations to connect digits
         kernel = np.ones((5, 5), np.uint8)
         dilated = cv2.dilate(combined_thresh, kernel, iterations=2)
         eroded = cv2.erode(dilated, kernel, iterations=1)
-        save_debug_image(eroded, "05_roi_morphological")
+        save_debug_image(eroded, "06_roi_morphological")
         
         contours, _ = cv2.findContours(eroded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         
@@ -74,49 +98,49 @@ 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 (1000 < area < (img_area * 0.9) and 
-                    1.0 <= aspect_ratio <= 10.0 and w > 80 and h > 40 and roi_brightness > 100):
+                if (500 < area < (img_area * 0.9) and 
+                    0.8 <= aspect_ratio <= 12.0 and w > 60 and h > 30 and roi_brightness > 80):
                     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=lambda x: x[1])  # Max brightness
                 x, y, w, h = cv2.boundingRect(contour)
-                padding = 80
+                padding = 100
                 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")
+                save_debug_image(roi_img, "07_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, 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]
+        valid_contours = [c for c in contours if 300 < cv2.contourArea(c) < (img_area * 0.95) and 
+                          0.5 <= cv2.boundingRect(c)[2]/cv2.boundingRect(c)[3] <= 15.0]
         if valid_contours:
             contour = max(valid_contours, key=cv2.contourArea)
             x, y, w, h = cv2.boundingRect(contour)
-            padding = 80
+            padding = 100
             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")
+            save_debug_image(roi_img, "07_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")
+        save_debug_image(img, "07_no_roi_original_fallback")
         return img, None
     except Exception as e:
         logging.error(f"ROI detection failed: {str(e)}")
-        save_debug_image(img, "06_roi_detection_error_fallback")
+        save_debug_image(img, "07_roi_detection_error_fallback")
         return img, None
 
 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:
+    if h < 15 or w < 10:
         return None
 
     segments = {
@@ -139,7 +163,7 @@ def detect_segments(digit_img, brightness):
             continue
         pixel_count = np.sum(region == 255)
         total_pixels = region.size
-        segment_presence[name] = pixel_count / total_pixels > (0.3 if brightness < 100 else 0.5)
+        segment_presence[name] = pixel_count / total_pixels > (0.25 if brightness < 100 else 0.45)
 
     digit_patterns = {
         '0': ('top', 'bottom', 'left_top', 'left_bottom', 'right_top', 'right_bottom'),
@@ -159,8 +183,8 @@ def detect_segments(digit_img, brightness):
     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.3 * non_matches_penalty
-        if matches >= len(pattern) * 0.8:
+        score = matches - 0.2 * non_matches_penalty
+        if matches >= len(pattern) * 0.75:
             score += 1.0
         if score > max_score:
             max_score = score
@@ -176,12 +200,12 @@ def custom_seven_segment_ocr(img, roi_bbox):
         brightness = estimate_brightness(img)
         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")
+        save_debug_image(thresh, "08_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")
+        save_debug_image(thresh, "09_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, 
@@ -198,7 +222,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 > 10:
+            if len(text) == 1 and (text.isdigit() or text == '.') and h_bbox > 8:
                 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))
@@ -211,7 +235,7 @@ 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"09_digit_crop_{idx}_{easyocr_char}")
+            save_debug_image(digit_img_crop, f"10_digit_crop_{idx}_{easyocr_char}")
             if easyocr_conf > 0.95 or easyocr_char == '.':
                 recognized_text += easyocr_char
             else:
@@ -240,20 +264,12 @@ 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")
+        # Apply rotation correction
+        img = correct_rotation(img)
 
         brightness = estimate_brightness(img)
         conf_threshold = 0.7 if brightness > 150 else (0.6 if brightness > 80 else 0.4)
+
         roi_img, roi_bbox = detect_roi(img)
         if roi_bbox:
             roi_area = roi_bbox[2] * roi_bbox[3]
@@ -263,7 +279,7 @@ def extract_weight_from_image(pil_img):
         if custom_result:
             try:
                 weight = float(custom_result)
-                if 0.01 <= weight <= 500:
+                if 0.001 <= weight <= 1000:
                     logging.info(f"Custom OCR result: {custom_result}, Confidence: 95.0%")
                     return custom_result, 95.0
                 else:
@@ -273,11 +289,8 @@ def extract_weight_from_image(pil_img):
 
         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")
         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, 
+        final_roi = cv2.adaptiveThreshold(preprocessed_roi, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                           cv2.THRESH_BINARY_INV, block_size, 8)
         save_debug_image(final_roi, "11_fallback_adaptive_thresh")
 
@@ -312,9 +325,9 @@ def extract_weight_from_image(pil_img):
                         weight /= 1000  # Convert grams to kilograms
                     elif unit == 'lb':
                         weight *= 0.453592  # Convert pounds to kilograms
-                    range_score = 1.5 if 0.01 <= weight <= 500 else 0.8
+                    range_score = 1.5 if 0.001 <= weight <= 1000 else 0.8
                     digit_count = len(text.replace('.', ''))
-                    digit_score = 1.3 if 2 <= digit_count <= 6 else 0.9
+                    digit_score = 1.3 if 2 <= digit_count <= 7 else 0.9
                     score = conf * range_score * digit_score
                     if roi_bbox:
                         (x_roi, y_roi, w_roi, h_roi) = roi_bbox
@@ -331,6 +344,7 @@ def extract_weight_from_image(pil_img):
                         logging.info(f"Candidate EasyOCR weight: '{text}', Unit: {unit or 'none'}, 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.")
@@ -347,12 +361,12 @@ def extract_weight_from_image(pil_img):
 
         try:
             final_weight = float(best_weight)
-            if final_weight < 0.01 or final_weight > 500:
+            if final_weight < 0.001 or final_weight > 1000:
                 best_conf *= 0.7
         except ValueError:
             pass
 
-        logging.info(f"Final detected weight: {best_weight}, Unit: {unit or 'none'}, Confidence: {round(best_conf * 100, 2)}%")
+        logging.info(f"Final detected weight: {best_weight} kg, Confidence: {round(best_conf * 100, 2)}%")
         return best_weight, round(best_conf * 100, 2)
 
     except Exception as e: