Spaces:
Running
Running
Update ocr_engine.py
Browse files- ocr_engine.py +41 -41
ocr_engine.py
CHANGED
@@ -35,22 +35,22 @@ def preprocess_image(img):
|
|
35 |
"""Preprocess image with aggressive contrast and noise handling."""
|
36 |
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
37 |
brightness = estimate_brightness(img)
|
38 |
-
# Maximum CLAHE for
|
39 |
-
clahe_clip =
|
40 |
-
clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=(
|
41 |
enhanced = clahe.apply(gray)
|
42 |
save_debug_image(enhanced, "01_preprocess_clahe")
|
43 |
-
#
|
44 |
-
blurred = cv2.bilateralFilter(enhanced,
|
45 |
save_debug_image(blurred, "02_preprocess_blur")
|
46 |
-
# Adaptive thresholding with
|
47 |
-
block_size = max(
|
48 |
thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
49 |
-
cv2.THRESH_BINARY_INV, block_size,
|
50 |
-
# Morphological operations for digit segmentation
|
51 |
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
52 |
thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
|
53 |
-
thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=
|
54 |
save_debug_image(thresh, "03_preprocess_morph")
|
55 |
return thresh, enhanced
|
56 |
|
@@ -58,12 +58,12 @@ def correct_rotation(img):
|
|
58 |
"""Correct image rotation using edge detection."""
|
59 |
try:
|
60 |
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
61 |
-
edges = cv2.Canny(gray,
|
62 |
-
lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=
|
63 |
if lines is not None:
|
64 |
angles = [np.arctan2(line[0][3] - line[0][1], line[0][2] - line[0][0]) * 180 / np.pi for line in lines]
|
65 |
angle = np.median(angles)
|
66 |
-
if abs(angle) > 0.
|
67 |
h, w = img.shape[:2]
|
68 |
center = (w // 2, h // 2)
|
69 |
M = cv2.getRotationMatrix2D(center, angle, 1.0)
|
@@ -81,15 +81,15 @@ def detect_roi(img):
|
|
81 |
save_debug_image(img, "04_original")
|
82 |
thresh, enhanced = preprocess_image(img)
|
83 |
brightness_map = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
84 |
-
block_sizes = [max(
|
85 |
valid_contours = []
|
86 |
img_area = img.shape[0] * img.shape[1]
|
87 |
|
88 |
for block_size in block_sizes:
|
89 |
temp_thresh = cv2.adaptiveThreshold(enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
90 |
-
cv2.THRESH_BINARY_INV, block_size,
|
91 |
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
|
92 |
-
temp_thresh = cv2.morphologyEx(temp_thresh, cv2.MORPH_CLOSE, kernel, iterations=
|
93 |
save_debug_image(temp_thresh, f"05_roi_threshold_block{block_size}")
|
94 |
contours, _ = cv2.findContours(temp_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
95 |
|
@@ -98,15 +98,15 @@ def detect_roi(img):
|
|
98 |
x, y, w, h = cv2.boundingRect(c)
|
99 |
roi_brightness = np.mean(brightness_map[y:y+h, x:x+w])
|
100 |
aspect_ratio = w / h
|
101 |
-
if (
|
102 |
-
0.
|
103 |
valid_contours.append((c, area * roi_brightness))
|
104 |
logging.debug(f"Contour (block {block_size}): Area={area}, Aspect={aspect_ratio:.2f}, Brightness={roi_brightness:.2f}")
|
105 |
|
106 |
if valid_contours:
|
107 |
contour, _ = max(valid_contours, key=lambda x: x[1])
|
108 |
x, y, w, h = cv2.boundingRect(contour)
|
109 |
-
padding = max(
|
110 |
x, y = max(0, x - padding), max(0, y - padding)
|
111 |
w, h = min(w + 2 * padding, img.shape[1] - x), min(h + 2 * padding, img.shape[0] - y)
|
112 |
roi_img = img[y:y+h, x:x+w]
|
@@ -123,14 +123,14 @@ def detect_roi(img):
|
|
123 |
return img, None
|
124 |
|
125 |
def detect_digit_template(digit_img, brightness):
|
126 |
-
"""Digit recognition using template matching with
|
127 |
try:
|
128 |
h, w = digit_img.shape
|
129 |
-
if h <
|
130 |
logging.debug("Digit image too small for template matching.")
|
131 |
return None
|
132 |
|
133 |
-
#
|
134 |
digit_templates = {
|
135 |
'0': np.array([[1, 1, 1, 1, 1],
|
136 |
[1, 0, 0, 0, 1],
|
@@ -143,29 +143,29 @@ def detect_digit_template(digit_img, brightness):
|
|
143 |
[0, 0, 1, 0, 0],
|
144 |
[0, 0, 1, 0, 0]]),
|
145 |
'2': np.array([[1, 1, 1, 1, 1],
|
146 |
-
[0, 0, 0,
|
147 |
[1, 1, 1, 1, 1],
|
148 |
-
[1,
|
149 |
[1, 1, 1, 1, 1]]),
|
150 |
'3': np.array([[1, 1, 1, 1, 1],
|
151 |
-
[0, 0, 0,
|
152 |
-
[
|
153 |
-
[0, 0, 0,
|
154 |
[1, 1, 1, 1, 1]]),
|
155 |
-
'4': np.array([[1,
|
156 |
-
[1,
|
157 |
[1, 1, 1, 1, 1],
|
158 |
[0, 0, 0, 0, 1],
|
159 |
[0, 0, 0, 0, 1]]),
|
160 |
'5': np.array([[1, 1, 1, 1, 1],
|
161 |
-
[1,
|
162 |
[1, 1, 1, 1, 1],
|
163 |
-
[0, 0, 0,
|
164 |
[1, 1, 1, 1, 1]]),
|
165 |
'6': np.array([[1, 1, 1, 1, 1],
|
166 |
-
[1,
|
167 |
[1, 1, 1, 1, 1],
|
168 |
-
[1, 0, 0,
|
169 |
[1, 1, 1, 1, 1]]),
|
170 |
'7': np.array([[1, 1, 1, 1, 1],
|
171 |
[0, 0, 0, 0, 1],
|
@@ -180,7 +180,7 @@ def detect_digit_template(digit_img, brightness):
|
|
180 |
'9': np.array([[1, 1, 1, 1, 1],
|
181 |
[1, 0, 0, 0, 1],
|
182 |
[1, 1, 1, 1, 1],
|
183 |
-
[0, 0, 0,
|
184 |
[1, 1, 1, 1, 1]]),
|
185 |
'.': np.array([[0, 0, 0],
|
186 |
[0, 1, 0],
|
@@ -195,11 +195,11 @@ def detect_digit_template(digit_img, brightness):
|
|
195 |
digit_img_resized = cv2.resize(digit_img, (3, 3), interpolation=cv2.INTER_NEAREST)
|
196 |
result = cv2.matchTemplate(digit_img_resized, template, cv2.TM_CCOEFF_NORMED)
|
197 |
_, max_val, _, _ = cv2.minMaxLoc(result)
|
198 |
-
if max_val > 0.
|
199 |
best_score = max_val
|
200 |
best_match = digit
|
201 |
logging.debug(f"Template match: {best_match}, Score: {best_score:.2f}")
|
202 |
-
return best_match if best_score > 0.
|
203 |
except Exception as e:
|
204 |
logging.error(f"Template digit detection failed: {str(e)}")
|
205 |
return None
|
@@ -234,7 +234,7 @@ def perform_ocr(img, roi_bbox):
|
|
234 |
digits_info = []
|
235 |
for c in contours:
|
236 |
x, y, w, h = cv2.boundingRect(c)
|
237 |
-
if w >
|
238 |
digits_info.append((x, x+w, y, y+h))
|
239 |
|
240 |
if digits_info:
|
@@ -251,7 +251,7 @@ def perform_ocr(img, roi_bbox):
|
|
251 |
digit = detect_digit_template(digit_crop, brightness)
|
252 |
if digit:
|
253 |
recognized_text += digit
|
254 |
-
elif x_min - prev_x_max <
|
255 |
recognized_text += '.'
|
256 |
prev_x_max = x_max
|
257 |
|
@@ -279,11 +279,11 @@ def extract_weight_from_image(pil_img):
|
|
279 |
save_debug_image(img, "00_input_image")
|
280 |
img = correct_rotation(img)
|
281 |
brightness = estimate_brightness(img)
|
282 |
-
conf_threshold = 0.
|
283 |
|
284 |
roi_img, roi_bbox = detect_roi(img)
|
285 |
if roi_bbox:
|
286 |
-
conf_threshold *= 1.05 if (roi_bbox[2] * roi_bbox[3]) > (img.shape[0] * img.shape[1] * 0.
|
287 |
|
288 |
result, confidence = perform_ocr(roi_img, roi_bbox)
|
289 |
if result and confidence >= conf_threshold * 100:
|
@@ -298,7 +298,7 @@ def extract_weight_from_image(pil_img):
|
|
298 |
|
299 |
logging.info("Primary OCR failed, using full image fallback.")
|
300 |
result, confidence = perform_ocr(img, None)
|
301 |
-
if result and confidence >= conf_threshold * 0.
|
302 |
try:
|
303 |
weight = float(result)
|
304 |
if 0.01 <= weight <= 1000:
|
|
|
35 |
"""Preprocess image with aggressive contrast and noise handling."""
|
36 |
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
37 |
brightness = estimate_brightness(img)
|
38 |
+
# Maximum CLAHE with adjusted clip for better digit enhancement
|
39 |
+
clahe_clip = 12.0 if brightness < 80 else 8.0
|
40 |
+
clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=(4, 4))
|
41 |
enhanced = clahe.apply(gray)
|
42 |
save_debug_image(enhanced, "01_preprocess_clahe")
|
43 |
+
# Stronger edge-preserving blur
|
44 |
+
blurred = cv2.bilateralFilter(enhanced, 7, 100, 100)
|
45 |
save_debug_image(blurred, "02_preprocess_blur")
|
46 |
+
# Adaptive thresholding with smaller blocks
|
47 |
+
block_size = max(3, min(11, int(img.shape[0] / 40) * 2 + 1))
|
48 |
thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
49 |
+
cv2.THRESH_BINARY_INV, block_size, 2)
|
50 |
+
# Morphological operations for robust digit segmentation
|
51 |
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
52 |
thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
|
53 |
+
thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=6)
|
54 |
save_debug_image(thresh, "03_preprocess_morph")
|
55 |
return thresh, enhanced
|
56 |
|
|
|
58 |
"""Correct image rotation using edge detection."""
|
59 |
try:
|
60 |
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
61 |
+
edges = cv2.Canny(gray, 15, 60, apertureSize=3)
|
62 |
+
lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=20, minLineLength=10, maxLineGap=3)
|
63 |
if lines is not None:
|
64 |
angles = [np.arctan2(line[0][3] - line[0][1], line[0][2] - line[0][0]) * 180 / np.pi for line in lines]
|
65 |
angle = np.median(angles)
|
66 |
+
if abs(angle) > 0.2:
|
67 |
h, w = img.shape[:2]
|
68 |
center = (w // 2, h // 2)
|
69 |
M = cv2.getRotationMatrix2D(center, angle, 1.0)
|
|
|
81 |
save_debug_image(img, "04_original")
|
82 |
thresh, enhanced = preprocess_image(img)
|
83 |
brightness_map = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
84 |
+
block_sizes = [max(3, min(11, int(img.shape[0] / s) * 2 + 1)) for s in [4, 8, 12]]
|
85 |
valid_contours = []
|
86 |
img_area = img.shape[0] * img.shape[1]
|
87 |
|
88 |
for block_size in block_sizes:
|
89 |
temp_thresh = cv2.adaptiveThreshold(enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
90 |
+
cv2.THRESH_BINARY_INV, block_size, 2)
|
91 |
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
|
92 |
+
temp_thresh = cv2.morphologyEx(temp_thresh, cv2.MORPH_CLOSE, kernel, iterations=6)
|
93 |
save_debug_image(temp_thresh, f"05_roi_threshold_block{block_size}")
|
94 |
contours, _ = cv2.findContours(temp_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
95 |
|
|
|
98 |
x, y, w, h = cv2.boundingRect(c)
|
99 |
roi_brightness = np.mean(brightness_map[y:y+h, x:x+w])
|
100 |
aspect_ratio = w / h
|
101 |
+
if (150 < area < (img_area * 0.8) and
|
102 |
+
0.15 <= aspect_ratio <= 12.0 and w > 40 and h > 15 and roi_brightness > 30):
|
103 |
valid_contours.append((c, area * roi_brightness))
|
104 |
logging.debug(f"Contour (block {block_size}): Area={area}, Aspect={aspect_ratio:.2f}, Brightness={roi_brightness:.2f}")
|
105 |
|
106 |
if valid_contours:
|
107 |
contour, _ = max(valid_contours, key=lambda x: x[1])
|
108 |
x, y, w, h = cv2.boundingRect(contour)
|
109 |
+
padding = max(10, min(30, int(min(w, h) * 0.25)))
|
110 |
x, y = max(0, x - padding), max(0, y - padding)
|
111 |
w, h = min(w + 2 * padding, img.shape[1] - x), min(h + 2 * padding, img.shape[0] - y)
|
112 |
roi_img = img[y:y+h, x:x+w]
|
|
|
123 |
return img, None
|
124 |
|
125 |
def detect_digit_template(digit_img, brightness):
|
126 |
+
"""Digit recognition using template matching with adjusted patterns."""
|
127 |
try:
|
128 |
h, w = digit_img.shape
|
129 |
+
if h < 8 or w < 4:
|
130 |
logging.debug("Digit image too small for template matching.")
|
131 |
return None
|
132 |
|
133 |
+
# Adjusted digit templates for seven-segment display
|
134 |
digit_templates = {
|
135 |
'0': np.array([[1, 1, 1, 1, 1],
|
136 |
[1, 0, 0, 0, 1],
|
|
|
143 |
[0, 0, 1, 0, 0],
|
144 |
[0, 0, 1, 0, 0]]),
|
145 |
'2': np.array([[1, 1, 1, 1, 1],
|
146 |
+
[0, 0, 0, 1, 1],
|
147 |
[1, 1, 1, 1, 1],
|
148 |
+
[1, 1, 0, 0, 0],
|
149 |
[1, 1, 1, 1, 1]]),
|
150 |
'3': np.array([[1, 1, 1, 1, 1],
|
151 |
+
[0, 0, 0, 1, 1],
|
152 |
+
[0, 1, 1, 1, 1],
|
153 |
+
[0, 0, 0, 1, 1],
|
154 |
[1, 1, 1, 1, 1]]),
|
155 |
+
'4': np.array([[1, 1, 0, 0, 1],
|
156 |
+
[1, 1, 0, 0, 1],
|
157 |
[1, 1, 1, 1, 1],
|
158 |
[0, 0, 0, 0, 1],
|
159 |
[0, 0, 0, 0, 1]]),
|
160 |
'5': np.array([[1, 1, 1, 1, 1],
|
161 |
+
[1, 1, 0, 0, 0],
|
162 |
[1, 1, 1, 1, 1],
|
163 |
+
[0, 0, 0, 1, 1],
|
164 |
[1, 1, 1, 1, 1]]),
|
165 |
'6': np.array([[1, 1, 1, 1, 1],
|
166 |
+
[1, 1, 0, 0, 0],
|
167 |
[1, 1, 1, 1, 1],
|
168 |
+
[1, 0, 0, 1, 1],
|
169 |
[1, 1, 1, 1, 1]]),
|
170 |
'7': np.array([[1, 1, 1, 1, 1],
|
171 |
[0, 0, 0, 0, 1],
|
|
|
180 |
'9': np.array([[1, 1, 1, 1, 1],
|
181 |
[1, 0, 0, 0, 1],
|
182 |
[1, 1, 1, 1, 1],
|
183 |
+
[0, 0, 0, 1, 1],
|
184 |
[1, 1, 1, 1, 1]]),
|
185 |
'.': np.array([[0, 0, 0],
|
186 |
[0, 1, 0],
|
|
|
195 |
digit_img_resized = cv2.resize(digit_img, (3, 3), interpolation=cv2.INTER_NEAREST)
|
196 |
result = cv2.matchTemplate(digit_img_resized, template, cv2.TM_CCOEFF_NORMED)
|
197 |
_, max_val, _, _ = cv2.minMaxLoc(result)
|
198 |
+
if max_val > 0.65 and max_val > best_score: # Lowered threshold for better match
|
199 |
best_score = max_val
|
200 |
best_match = digit
|
201 |
logging.debug(f"Template match: {best_match}, Score: {best_score:.2f}")
|
202 |
+
return best_match if best_score > 0.65 else None
|
203 |
except Exception as e:
|
204 |
logging.error(f"Template digit detection failed: {str(e)}")
|
205 |
return None
|
|
|
234 |
digits_info = []
|
235 |
for c in contours:
|
236 |
x, y, w, h = cv2.boundingRect(c)
|
237 |
+
if w > 6 and h > 8 and 0.1 <= w/h <= 2.5: # Loosened size and aspect ratio
|
238 |
digits_info.append((x, x+w, y, y+h))
|
239 |
|
240 |
if digits_info:
|
|
|
251 |
digit = detect_digit_template(digit_crop, brightness)
|
252 |
if digit:
|
253 |
recognized_text += digit
|
254 |
+
elif x_min - prev_x_max < 6 and prev_x_max != -float('inf'): # Adjusted decimal gap
|
255 |
recognized_text += '.'
|
256 |
prev_x_max = x_max
|
257 |
|
|
|
279 |
save_debug_image(img, "00_input_image")
|
280 |
img = correct_rotation(img)
|
281 |
brightness = estimate_brightness(img)
|
282 |
+
conf_threshold = 0.75 if brightness > 100 else 0.55 # Lowered threshold
|
283 |
|
284 |
roi_img, roi_bbox = detect_roi(img)
|
285 |
if roi_bbox:
|
286 |
+
conf_threshold *= 1.05 if (roi_bbox[2] * roi_bbox[3]) > (img.shape[0] * img.shape[1] * 0.15) else 1.0
|
287 |
|
288 |
result, confidence = perform_ocr(roi_img, roi_bbox)
|
289 |
if result and confidence >= conf_threshold * 100:
|
|
|
298 |
|
299 |
logging.info("Primary OCR failed, using full image fallback.")
|
300 |
result, confidence = perform_ocr(img, None)
|
301 |
+
if result and confidence >= conf_threshold * 0.8 * 100: # Adjusted fallback threshold
|
302 |
try:
|
303 |
weight = float(result)
|
304 |
if 0.01 <= weight <= 1000:
|