Spaces:
Runtime error
Runtime error
Update ocr_engine.py
Browse files- ocr_engine.py +19 -17
ocr_engine.py
CHANGED
@@ -9,31 +9,33 @@ def extract_weight_from_image(pil_img):
|
|
9 |
try:
|
10 |
img = np.array(pil_img)
|
11 |
|
12 |
-
# Convert to grayscale and resize
|
13 |
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
14 |
gray = cv2.resize(gray, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)
|
15 |
|
16 |
-
# Histogram equalization
|
17 |
gray = cv2.equalizeHist(gray)
|
18 |
-
|
19 |
-
# Adaptive threshold to enhance text
|
20 |
thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
21 |
cv2.THRESH_BINARY, 11, 2)
|
22 |
-
|
23 |
-
# Invert for LCD-like contrast
|
24 |
thresh = cv2.bitwise_not(thresh)
|
25 |
|
26 |
-
# OCR
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
37 |
|
38 |
except Exception as e:
|
39 |
return f"Error: {str(e)}", 0.0
|
|
|
9 |
try:
|
10 |
img = np.array(pil_img)
|
11 |
|
12 |
+
# Convert to grayscale and resize
|
13 |
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
14 |
gray = cv2.resize(gray, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)
|
15 |
|
16 |
+
# Histogram equalization and adaptive threshold
|
17 |
gray = cv2.equalizeHist(gray)
|
|
|
|
|
18 |
thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
19 |
cv2.THRESH_BINARY, 11, 2)
|
|
|
|
|
20 |
thresh = cv2.bitwise_not(thresh)
|
21 |
|
22 |
+
# OCR with bounding boxes
|
23 |
+
results = reader.readtext(thresh)
|
24 |
+
|
25 |
+
# Filter potential weight values
|
26 |
+
candidates = []
|
27 |
+
for (bbox, text, confidence) in results:
|
28 |
+
# Clean text
|
29 |
+
clean_text = text.replace('kg', '').strip()
|
30 |
+
if re.fullmatch(r"\d{2,4}(\.\d{1,2})?", clean_text):
|
31 |
+
candidates.append((clean_text, confidence))
|
32 |
+
|
33 |
+
if not candidates:
|
34 |
+
return "Not detected", 0.0
|
35 |
+
|
36 |
+
# Choose the highest confidence match
|
37 |
+
best_weight, conf = sorted(candidates, key=lambda x: -x[1])[0]
|
38 |
+
return best_weight, round(conf, 2)
|
39 |
|
40 |
except Exception as e:
|
41 |
return f"Error: {str(e)}", 0.0
|