Spaces:
Sleeping
Sleeping
Update ocr_engine.py
Browse files- ocr_engine.py +48 -170
ocr_engine.py
CHANGED
@@ -32,30 +32,30 @@ def estimate_brightness(img):
|
|
32 |
return np.mean(gray)
|
33 |
|
34 |
def preprocess_image(img):
|
35 |
-
"""Preprocess image with
|
36 |
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
37 |
brightness = estimate_brightness(img)
|
38 |
|
39 |
-
# Apply
|
40 |
-
clahe_clip =
|
41 |
clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=(8, 8))
|
42 |
enhanced = clahe.apply(gray)
|
43 |
save_debug_image(enhanced, "01_preprocess_clahe")
|
44 |
|
45 |
-
#
|
46 |
-
blurred = cv2.GaussianBlur(enhanced, (
|
47 |
save_debug_image(blurred, "02_preprocess_blur")
|
48 |
|
49 |
-
#
|
50 |
-
block_size = max(
|
51 |
thresh = cv2.adaptiveThreshold(
|
52 |
blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
53 |
-
cv2.THRESH_BINARY_INV, block_size,
|
54 |
)
|
55 |
|
56 |
-
#
|
57 |
-
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (
|
58 |
-
thresh = cv2.morphologyEx(thresh, cv2.
|
59 |
save_debug_image(thresh, "03_preprocess_morph")
|
60 |
return thresh, enhanced
|
61 |
|
@@ -63,12 +63,12 @@ def correct_rotation(img):
|
|
63 |
"""Correct image rotation using edge detection."""
|
64 |
try:
|
65 |
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
66 |
-
edges = cv2.Canny(gray,
|
67 |
-
lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=
|
68 |
if lines is not None:
|
69 |
angles = [np.arctan2(line[0][3] - line[0][1], line[0][2] - line[0][0]) * 180 / np.pi for line in lines]
|
70 |
angle = np.median(angles)
|
71 |
-
if abs(angle) > 0.
|
72 |
h, w = img.shape[:2]
|
73 |
center = (w // 2, h // 2)
|
74 |
M = cv2.getRotationMatrix2D(center, angle, 1.0)
|
@@ -81,21 +81,21 @@ def correct_rotation(img):
|
|
81 |
return img
|
82 |
|
83 |
def detect_roi(img):
|
84 |
-
"""Detect region of interest with
|
85 |
try:
|
86 |
save_debug_image(img, "04_original")
|
87 |
thresh, enhanced = preprocess_image(img)
|
88 |
brightness_map = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
89 |
-
block_sizes = [max(
|
90 |
valid_contours = []
|
91 |
img_area = img.shape[0] * img.shape[1]
|
92 |
|
93 |
for block_size in block_sizes:
|
94 |
temp_thresh = cv2.adaptiveThreshold(
|
95 |
enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
96 |
-
cv2.THRESH_BINARY_INV, block_size,
|
97 |
)
|
98 |
-
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (
|
99 |
temp_thresh = cv2.morphologyEx(temp_thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
|
100 |
save_debug_image(temp_thresh, f"05_roi_threshold_block{block_size}")
|
101 |
contours, _ = cv2.findContours(temp_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
@@ -105,15 +105,15 @@ def detect_roi(img):
|
|
105 |
x, y, w, h = cv2.boundingRect(c)
|
106 |
roi_brightness = np.mean(brightness_map[y:y+h, x:x+w])
|
107 |
aspect_ratio = w / h
|
108 |
-
if (
|
109 |
-
0.
|
110 |
valid_contours.append((c, area * roi_brightness))
|
111 |
logging.debug(f"Contour (block {block_size}): Area={area}, Aspect={aspect_ratio:.2f}, Brightness={roi_brightness:.2f}")
|
112 |
|
113 |
if valid_contours:
|
114 |
contour, _ = max(valid_contours, key=lambda x: x[1])
|
115 |
x, y, w, h = cv2.boundingRect(contour)
|
116 |
-
padding = max(5, min(
|
117 |
x, y = max(0, x - padding), max(0, y - padding)
|
118 |
w, h = min(w + 2 * padding, img.shape[1] - x), min(h + 2 * padding, img.shape[0] - y)
|
119 |
roi_img = img[y:y+h, x:x+w]
|
@@ -130,183 +130,61 @@ def detect_roi(img):
|
|
130 |
return img, None
|
131 |
|
132 |
def detect_digit_template(digit_img, brightness):
|
133 |
-
"""Digit recognition with
|
134 |
try:
|
135 |
h, w = digit_img.shape
|
136 |
if h < 5 or w < 2:
|
137 |
logging.debug("Digit image too small for template matching.")
|
138 |
return None
|
139 |
|
140 |
-
# Expanded digit templates for seven-segment display variations
|
141 |
digit_templates = {
|
142 |
-
'0': [
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
153 |
-
],
|
154 |
-
'1': [
|
155 |
-
np.array([[0, 0, 1, 0, 0],
|
156 |
-
[0, 0, 1, 0, 0],
|
157 |
-
[0, 0, 1, 0, 0],
|
158 |
-
[0, 0, 1, 0, 0],
|
159 |
-
[0, 0, 1, 0, 0]], dtype=np.float32),
|
160 |
-
np.array([[0, 1, 0],
|
161 |
-
[0, 1, 0],
|
162 |
-
[0, 1, 0],
|
163 |
-
[0, 1, 0],
|
164 |
-
[0, 1, 0]], dtype=np.float32)
|
165 |
-
],
|
166 |
-
'2': [
|
167 |
-
np.array([[1, 1, 1, 1, 1],
|
168 |
-
[0, 0, 0, 1, 1],
|
169 |
-
[1, 1, 1, 1, 1],
|
170 |
-
[1, 1, 0, 0, 0],
|
171 |
-
[1, 1, 1, 1, 1]], dtype=np.float32),
|
172 |
-
np.array([[1, 1, 1, 1],
|
173 |
-
[0, 0, 1, 1],
|
174 |
-
[1, 1, 1, 1],
|
175 |
-
[1, 1, 0, 0],
|
176 |
-
[1, 1, 1, 1]], dtype=np.float32)
|
177 |
-
],
|
178 |
-
'3': [
|
179 |
-
np.array([[1, 1, 1, 1, 1],
|
180 |
-
[0, 0, 0, 1, 1],
|
181 |
-
[1, 1, 1, 1, 1],
|
182 |
-
[0, 0, 0, 1, 1],
|
183 |
-
[1, 1, 1, 1, 1]], dtype=np.float32),
|
184 |
-
np.array([[1, 1, 1, 1],
|
185 |
-
[0, 0, 1, 1],
|
186 |
-
[1, 1, 1, 1],
|
187 |
-
[0, 0, 1, 1],
|
188 |
-
[1, 1, 1, 1]], dtype=np.float32)
|
189 |
-
],
|
190 |
-
'4': [
|
191 |
-
np.array([[1, 1, 0, 0, 1],
|
192 |
-
[1, 1, 0, 0, 1],
|
193 |
-
[1, 1, 1, 1, 1],
|
194 |
-
[0, 0, 0, 0, 1],
|
195 |
-
[0, 0, 0, 0, 1]], dtype=np.float32),
|
196 |
-
np.array([[1, 0, 0, 1],
|
197 |
-
[1, 0, 0, 1],
|
198 |
-
[1, 1, 1, 1],
|
199 |
-
[0, 0, 0, 1],
|
200 |
-
[0, 0, 0, 1]], dtype=np.float32)
|
201 |
-
],
|
202 |
-
'5': [
|
203 |
-
np.array([[1, 1, 1, 1, 1],
|
204 |
-
[1, 1, 0, 0, 0],
|
205 |
-
[1, 1, 1, 1, 1],
|
206 |
-
[0, 0, 0, 1, 1],
|
207 |
-
[1, 1, 1, 1, 1]], dtype=np.float32),
|
208 |
-
np.array([[1, 1, 1, 1],
|
209 |
-
[1, 1, 0, 0],
|
210 |
-
[1, 1, 1, 1],
|
211 |
-
[0, 0, 1, 1],
|
212 |
-
[1, 1, 1, 1]], dtype=np.float32)
|
213 |
-
],
|
214 |
-
'6': [
|
215 |
-
np.array([[1, 1, 1, 1, 1],
|
216 |
-
[1, 1, 0, 0, 0],
|
217 |
-
[1, 1, 1, 1, 1],
|
218 |
-
[1, 0, 0, 1, 1],
|
219 |
-
[1, 1, 1, 1, 1]], dtype=np.float32),
|
220 |
-
np.array([[1, 1, 1, 1],
|
221 |
-
[1, 1, 0, 0],
|
222 |
-
[1, 1, 1, 1],
|
223 |
-
[1, 0, 1, 1],
|
224 |
-
[1, 1, 1, 1]], dtype=np.float32)
|
225 |
-
],
|
226 |
-
'7': [
|
227 |
-
np.array([[1, 1, 1, 1, 1],
|
228 |
-
[0, 0, 0, 0, 1],
|
229 |
-
[0, 0, 0, 0, 1],
|
230 |
-
[0, 0, 0, 0, 1],
|
231 |
-
[0, 0, 0, 0, 1]], dtype=np.float32),
|
232 |
-
np.array([[1, 1, 1, 1],
|
233 |
-
[0, 0, 0, 1],
|
234 |
-
[0, 0, 0, 1],
|
235 |
-
[0, 0, 0, 1],
|
236 |
-
[0, 0, 0, 1]], dtype=np.float32)
|
237 |
-
],
|
238 |
-
'8': [
|
239 |
-
np.array([[1, 1, 1, 1, 1],
|
240 |
-
[1, 0, 0, 0, 1],
|
241 |
-
[1, 1, 1, 1, 1],
|
242 |
-
[1, 0, 0, 0, 1],
|
243 |
-
[1, 1, 1, 1, 1]], dtype=np.float32),
|
244 |
-
np.array([[1, 1, 1, 1],
|
245 |
-
[1, 0, 0, 1],
|
246 |
-
[1, 1, 1, 1],
|
247 |
-
[1, 0, 0, 1],
|
248 |
-
[1, 1, 1, 1]], dtype=np.float32)
|
249 |
-
],
|
250 |
-
'9': [
|
251 |
-
np.array([[1, 1, 1, 1, 1],
|
252 |
-
[1, 0, 0, 0, 1],
|
253 |
-
[1, 1, 1, 1, 1],
|
254 |
-
[0, 0, 0, 1, 1],
|
255 |
-
[1, 1, 1, 1, 1]], dtype=np.float32),
|
256 |
-
np.array([[1, 1, 1, 1],
|
257 |
-
[1, 0, 0, 1],
|
258 |
-
[1, 1, 1, 1],
|
259 |
-
[0, 0, 1, 1],
|
260 |
-
[1, 1, 1, 1]], dtype=np.float32)
|
261 |
-
],
|
262 |
-
'.': [
|
263 |
-
np.array([[0, 0, 0],
|
264 |
-
[0, 1, 0],
|
265 |
-
[0, 0, 0]], dtype=np.float32),
|
266 |
-
np.array([[0, 0],
|
267 |
-
[1, 0],
|
268 |
-
[0, 0]], dtype=np.float32)
|
269 |
-
]
|
270 |
}
|
271 |
|
272 |
-
# Try multiple sizes for digit image
|
273 |
sizes = [(5, 5), (4, 4), (3, 3)] if h > w else [(3, 3), (2, 2)]
|
274 |
best_match, best_score = None, -1
|
275 |
for size in sizes:
|
276 |
digit_img_resized = cv2.resize(digit_img, size, interpolation=cv2.INTER_AREA)
|
277 |
-
digit_img_resized = (digit_img_resized >
|
278 |
|
279 |
for digit, templates in digit_templates.items():
|
280 |
for template in templates:
|
281 |
-
if digit == '.' and size[0] > 3:
|
282 |
-
continue
|
283 |
-
if digit != '.' and size[0] <= 3:
|
284 |
-
continue
|
285 |
if template.shape[0] != size[0] or template.shape[1] != size[1]:
|
286 |
continue
|
287 |
result = cv2.matchTemplate(digit_img_resized, template, cv2.TM_CCOEFF_NORMED)
|
288 |
_, max_val, _, _ = cv2.minMaxLoc(result)
|
289 |
-
if max_val > 0.
|
290 |
best_score = max_val
|
291 |
best_match = digit
|
292 |
logging.debug(f"Template match: {best_match}, Score: {best_score:.2f}")
|
293 |
-
return best_match if best_score > 0.
|
294 |
except Exception as e:
|
295 |
logging.error(f"Template digit detection failed: {str(e)}")
|
296 |
return None
|
297 |
|
298 |
def perform_ocr(img, roi_bbox):
|
299 |
-
"""Perform OCR with Tesseract and
|
300 |
try:
|
301 |
thresh, enhanced = preprocess_image(img)
|
302 |
brightness = estimate_brightness(img)
|
303 |
pil_img = Image.fromarray(enhanced)
|
304 |
save_debug_image(pil_img, "07_ocr_input")
|
305 |
|
306 |
-
#
|
307 |
configs = [
|
308 |
r'--oem 3 --psm 7 -c tessedit_char_whitelist=0123456789.', # Single line
|
309 |
-
r'--oem 3 --psm 6 -c tessedit_char_whitelist=0123456789.'
|
|
|
310 |
]
|
311 |
for config in configs:
|
312 |
text = pytesseract.image_to_string(pil_img, config=config)
|
@@ -321,13 +199,13 @@ def perform_ocr(img, roi_bbox):
|
|
321 |
logging.info(f"Validated Tesseract text: {text}, Confidence: {confidence:.2f}%")
|
322 |
return text, confidence
|
323 |
|
324 |
-
#
|
325 |
logging.info("Tesseract failed, using template-based detection.")
|
326 |
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
327 |
digits_info = []
|
328 |
for c in contours:
|
329 |
x, y, w, h = cv2.boundingRect(c)
|
330 |
-
if w >
|
331 |
digits_info.append((x, x+w, y, y+h))
|
332 |
|
333 |
if digits_info:
|
@@ -344,7 +222,7 @@ def perform_ocr(img, roi_bbox):
|
|
344 |
digit = detect_digit_template(digit_crop, brightness)
|
345 |
if digit:
|
346 |
recognized_text += digit
|
347 |
-
elif x_min - prev_x_max <
|
348 |
recognized_text += '.'
|
349 |
prev_x_max = x_max
|
350 |
|
@@ -365,19 +243,19 @@ def perform_ocr(img, roi_bbox):
|
|
365 |
return None, 0.0
|
366 |
|
367 |
def extract_weight_from_image(pil_img):
|
368 |
-
"""Extract weight from any digital scale image."""
|
369 |
try:
|
370 |
img = np.array(pil_img)
|
371 |
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
372 |
save_debug_image(img, "00_input_image")
|
373 |
img = correct_rotation(img)
|
374 |
brightness = estimate_brightness(img)
|
375 |
-
conf_threshold = 0.
|
376 |
|
377 |
# Try ROI-based detection
|
378 |
roi_img, roi_bbox = detect_roi(img)
|
379 |
if roi_bbox:
|
380 |
-
conf_threshold *= 1.
|
381 |
|
382 |
result, confidence = perform_ocr(roi_img, roi_bbox)
|
383 |
if result and confidence >= conf_threshold * 100:
|
@@ -390,10 +268,10 @@ def extract_weight_from_image(pil_img):
|
|
390 |
except ValueError:
|
391 |
logging.warning(f"Invalid weight format: {result}")
|
392 |
|
393 |
-
# Full image fallback
|
394 |
logging.info("Primary OCR failed, using full image fallback.")
|
395 |
result, confidence = perform_ocr(img, None)
|
396 |
-
if result and confidence >= conf_threshold * 0.
|
397 |
try:
|
398 |
weight = float(result)
|
399 |
if 0.001 <= weight <= 5000:
|
|
|
32 |
return np.mean(gray)
|
33 |
|
34 |
def preprocess_image(img):
|
35 |
+
"""Preprocess image with enhanced contrast and adaptive thresholding."""
|
36 |
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
37 |
brightness = estimate_brightness(img)
|
38 |
|
39 |
+
# Apply CLAHE with dynamic clip limit
|
40 |
+
clahe_clip = 10.0 if brightness < 80 else 5.0
|
41 |
clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=(8, 8))
|
42 |
enhanced = clahe.apply(gray)
|
43 |
save_debug_image(enhanced, "01_preprocess_clahe")
|
44 |
|
45 |
+
# Stronger blur to reduce noise
|
46 |
+
blurred = cv2.GaussianBlur(enhanced, (7, 7), 1.0)
|
47 |
save_debug_image(blurred, "02_preprocess_blur")
|
48 |
|
49 |
+
# Adaptive thresholding with larger block size
|
50 |
+
block_size = max(11, min(41, int(img.shape[0] / 15) * 2 + 1))
|
51 |
thresh = cv2.adaptiveThreshold(
|
52 |
blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
53 |
+
cv2.THRESH_BINARY_INV, block_size, 5
|
54 |
)
|
55 |
|
56 |
+
# Morphological operations for better digit separation
|
57 |
+
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
|
58 |
+
thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
|
59 |
save_debug_image(thresh, "03_preprocess_morph")
|
60 |
return thresh, enhanced
|
61 |
|
|
|
63 |
"""Correct image rotation using edge detection."""
|
64 |
try:
|
65 |
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
66 |
+
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
|
67 |
+
lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=20, minLineLength=10, maxLineGap=5)
|
68 |
if lines is not None:
|
69 |
angles = [np.arctan2(line[0][3] - line[0][1], line[0][2] - line[0][0]) * 180 / np.pi for line in lines]
|
70 |
angle = np.median(angles)
|
71 |
+
if abs(angle) > 0.5:
|
72 |
h, w = img.shape[:2]
|
73 |
center = (w // 2, h // 2)
|
74 |
M = cv2.getRotationMatrix2D(center, angle, 1.0)
|
|
|
81 |
return img
|
82 |
|
83 |
def detect_roi(img):
|
84 |
+
"""Detect region of interest with relaxed contour analysis."""
|
85 |
try:
|
86 |
save_debug_image(img, "04_original")
|
87 |
thresh, enhanced = preprocess_image(img)
|
88 |
brightness_map = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
89 |
+
block_sizes = [max(11, min(41, int(img.shape[0] / s) * 2 + 1)) for s in [5, 10, 15]]
|
90 |
valid_contours = []
|
91 |
img_area = img.shape[0] * img.shape[1]
|
92 |
|
93 |
for block_size in block_sizes:
|
94 |
temp_thresh = cv2.adaptiveThreshold(
|
95 |
enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
96 |
+
cv2.THRESH_BINARY_INV, block_size, 5
|
97 |
)
|
98 |
+
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
|
99 |
temp_thresh = cv2.morphologyEx(temp_thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
|
100 |
save_debug_image(temp_thresh, f"05_roi_threshold_block{block_size}")
|
101 |
contours, _ = cv2.findContours(temp_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
105 |
x, y, w, h = cv2.boundingRect(c)
|
106 |
roi_brightness = np.mean(brightness_map[y:y+h, x:x+w])
|
107 |
aspect_ratio = w / h
|
108 |
+
if (30 < area < (img_area * 0.98) and
|
109 |
+
0.02 <= aspect_ratio <= 25.0 and w > 15 and h > 5 and roi_brightness > 10):
|
110 |
valid_contours.append((c, area * roi_brightness))
|
111 |
logging.debug(f"Contour (block {block_size}): Area={area}, Aspect={aspect_ratio:.2f}, Brightness={roi_brightness:.2f}")
|
112 |
|
113 |
if valid_contours:
|
114 |
contour, _ = max(valid_contours, key=lambda x: x[1])
|
115 |
x, y, w, h = cv2.boundingRect(contour)
|
116 |
+
padding = max(5, min(25, int(min(w, h) * 0.5)))
|
117 |
x, y = max(0, x - padding), max(0, y - padding)
|
118 |
w, h = min(w + 2 * padding, img.shape[1] - x), min(h + 2 * padding, img.shape[0] - y)
|
119 |
roi_img = img[y:y+h, x:x+w]
|
|
|
130 |
return img, None
|
131 |
|
132 |
def detect_digit_template(digit_img, brightness):
|
133 |
+
"""Digit recognition with adjusted template matching."""
|
134 |
try:
|
135 |
h, w = digit_img.shape
|
136 |
if h < 5 or w < 2:
|
137 |
logging.debug("Digit image too small for template matching.")
|
138 |
return None
|
139 |
|
|
|
140 |
digit_templates = {
|
141 |
+
'0': [np.array([[1, 1, 1, 1, 1], [1, 0, 0, 0, 1], [1, 0, 0, 0, 1], [1, 0, 0, 0, 1], [1, 1, 1, 1, 1]], dtype=np.float32)],
|
142 |
+
'1': [np.array([[0, 0, 1, 0, 0], [0, 0, 1, 0, 0], [0, 0, 1, 0, 0], [0, 0, 1, 0, 0], [0, 0, 1, 0, 0]], dtype=np.float32)],
|
143 |
+
'2': [np.array([[1, 1, 1, 1, 1], [0, 0, 0, 1, 1], [1, 1, 1, 1, 1], [1, 1, 0, 0, 0], [1, 1, 1, 1, 1]], dtype=np.float32)],
|
144 |
+
'3': [np.array([[1, 1, 1, 1, 1], [0, 0, 0, 1, 1], [1, 1, 1, 1, 1], [0, 0, 0, 1, 1], [1, 1, 1, 1, 1]], dtype=np.float32)],
|
145 |
+
'4': [np.array([[1, 1, 0, 0, 1], [1, 1, 0, 0, 1], [1, 1, 1, 1, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1]], dtype=np.float32)],
|
146 |
+
'5': [np.array([[1, 1, 1, 1, 1], [1, 1, 0, 0, 0], [1, 1, 1, 1, 1], [0, 0, 0, 1, 1], [1, 1, 1, 1, 1]], dtype=np.float32)],
|
147 |
+
'6': [np.array([[1, 1, 1, 1, 1], [1, 1, 0, 0, 0], [1, 1, 1, 1, 1], [1, 0, 0, 1, 1], [1, 1, 1, 1, 1]], dtype=np.float32)],
|
148 |
+
'7': [np.array([[1, 1, 1, 1, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1]], dtype=np.float32)],
|
149 |
+
'8': [np.array([[1, 1, 1, 1, 1], [1, 0, 0, 0, 1], [1, 1, 1, 1, 1], [1, 0, 0, 0, 1], [1, 1, 1, 1, 1]], dtype=np.float32)],
|
150 |
+
'9': [np.array([[1, 1, 1, 1, 1], [1, 0, 0, 0, 1], [1, 1, 1, 1, 1], [0, 0, 0, 1, 1], [1, 1, 1, 1, 1]], dtype=np.float32)],
|
151 |
+
'.': [np.array([[0, 0, 0], [0, 1, 0], [0, 0, 0]], dtype=np.float32)]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
152 |
}
|
153 |
|
|
|
154 |
sizes = [(5, 5), (4, 4), (3, 3)] if h > w else [(3, 3), (2, 2)]
|
155 |
best_match, best_score = None, -1
|
156 |
for size in sizes:
|
157 |
digit_img_resized = cv2.resize(digit_img, size, interpolation=cv2.INTER_AREA)
|
158 |
+
digit_img_resized = (digit_img_resized > 90).astype(np.float32) # Adjusted binarization threshold
|
159 |
|
160 |
for digit, templates in digit_templates.items():
|
161 |
for template in templates:
|
|
|
|
|
|
|
|
|
162 |
if template.shape[0] != size[0] or template.shape[1] != size[1]:
|
163 |
continue
|
164 |
result = cv2.matchTemplate(digit_img_resized, template, cv2.TM_CCOEFF_NORMED)
|
165 |
_, max_val, _, _ = cv2.minMaxLoc(result)
|
166 |
+
if max_val > 0.50 and max_val > best_score: # Lowered threshold
|
167 |
best_score = max_val
|
168 |
best_match = digit
|
169 |
logging.debug(f"Template match: {best_match}, Score: {best_score:.2f}")
|
170 |
+
return best_match if best_score > 0.50 else None
|
171 |
except Exception as e:
|
172 |
logging.error(f"Template digit detection failed: {str(e)}")
|
173 |
return None
|
174 |
|
175 |
def perform_ocr(img, roi_bbox):
|
176 |
+
"""Perform OCR with enhanced Tesseract and template fallback."""
|
177 |
try:
|
178 |
thresh, enhanced = preprocess_image(img)
|
179 |
brightness = estimate_brightness(img)
|
180 |
pil_img = Image.fromarray(enhanced)
|
181 |
save_debug_image(pil_img, "07_ocr_input")
|
182 |
|
183 |
+
# Enhanced Tesseract configurations
|
184 |
configs = [
|
185 |
r'--oem 3 --psm 7 -c tessedit_char_whitelist=0123456789.', # Single line
|
186 |
+
r'--oem 3 --psm 6 -c tessedit_char_whitelist=0123456789.', # Block of text
|
187 |
+
r'--oem 3 --psm 10 -c tessedit_char_whitelist=0123456789.' # Single character
|
188 |
]
|
189 |
for config in configs:
|
190 |
text = pytesseract.image_to_string(pil_img, config=config)
|
|
|
199 |
logging.info(f"Validated Tesseract text: {text}, Confidence: {confidence:.2f}%")
|
200 |
return text, confidence
|
201 |
|
202 |
+
# Enhanced template-based detection
|
203 |
logging.info("Tesseract failed, using template-based detection.")
|
204 |
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
205 |
digits_info = []
|
206 |
for c in contours:
|
207 |
x, y, w, h = cv2.boundingRect(c)
|
208 |
+
if w > 3 and h > 4 and 0.02 <= w/h <= 5.0:
|
209 |
digits_info.append((x, x+w, y, y+h))
|
210 |
|
211 |
if digits_info:
|
|
|
222 |
digit = detect_digit_template(digit_crop, brightness)
|
223 |
if digit:
|
224 |
recognized_text += digit
|
225 |
+
elif x_min - prev_x_max < 15 and prev_x_max != -float('inf'):
|
226 |
recognized_text += '.'
|
227 |
prev_x_max = x_max
|
228 |
|
|
|
243 |
return None, 0.0
|
244 |
|
245 |
def extract_weight_from_image(pil_img):
|
246 |
+
"""Extract weight from any digital scale image with adjusted thresholds."""
|
247 |
try:
|
248 |
img = np.array(pil_img)
|
249 |
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
250 |
save_debug_image(img, "00_input_image")
|
251 |
img = correct_rotation(img)
|
252 |
brightness = estimate_brightness(img)
|
253 |
+
conf_threshold = 0.60 if brightness > 70 else 0.40 # Lowered threshold
|
254 |
|
255 |
# Try ROI-based detection
|
256 |
roi_img, roi_bbox = detect_roi(img)
|
257 |
if roi_bbox:
|
258 |
+
conf_threshold *= 1.2 if (roi_bbox[2] * roi_bbox[3]) > (img.shape[0] * img.shape[1] * 0.03) else 1.0
|
259 |
|
260 |
result, confidence = perform_ocr(roi_img, roi_bbox)
|
261 |
if result and confidence >= conf_threshold * 100:
|
|
|
268 |
except ValueError:
|
269 |
logging.warning(f"Invalid weight format: {result}")
|
270 |
|
271 |
+
# Full image fallback with relaxed threshold
|
272 |
logging.info("Primary OCR failed, using full image fallback.")
|
273 |
result, confidence = perform_ocr(img, None)
|
274 |
+
if result and confidence >= conf_threshold * 0.80 * 100:
|
275 |
try:
|
276 |
weight = float(result)
|
277 |
if 0.001 <= weight <= 5000:
|