Sanjayraju30 commited on
Commit
9d7955b
·
verified ·
1 Parent(s): 81907cd

Create ocr_engine.py

Browse files
Files changed (1) hide show
  1. ocr_engine.py +55 -0
ocr_engine.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from mmocr.utils.ocr import MMOCR
2
+ import numpy as np
3
+ import cv2
4
+ import re
5
+ from PIL import Image
6
+
7
+ ocr = MMOCR(det='DBPANet', recog='SAR', device='cpu')
8
+
9
+ def extract_weight_from_image(pil_img):
10
+ try:
11
+ img = np.array(pil_img.convert("RGB"))[:, :, ::-1]
12
+ result = ocr.readtext(img, print_result=False, output=None)[0]['result']
13
+
14
+ raw_texts = []
15
+ weight_candidates = []
16
+ fallback_weight = None
17
+ fallback_conf = 0.0
18
+
19
+ for text, conf in result:
20
+ original = text
21
+ cleaned = text.lower().strip()
22
+ cleaned = cleaned.replace(",", ".")
23
+ cleaned = cleaned.replace("o", "0").replace("O", "0")
24
+ cleaned = cleaned.replace("s", "5").replace("S", "5")
25
+ cleaned = cleaned.replace("g", "9").replace("G", "6")
26
+ cleaned = cleaned.replace("kg", "").replace("kgs", "")
27
+ cleaned = re.sub(r"[^0-9\.]", "", cleaned)
28
+
29
+ raw_texts.append(f"{original} → {cleaned} (conf: {round(conf, 2)})")
30
+
31
+ if cleaned and cleaned.replace(".", "").isdigit() and not fallback_weight:
32
+ fallback_weight = cleaned
33
+ fallback_conf = conf
34
+
35
+ if cleaned.count(".") <= 1 and re.fullmatch(r"\d{2,4}(\.\d{1,3})?", cleaned):
36
+ weight_candidates.append((cleaned, conf))
37
+
38
+ if weight_candidates:
39
+ best_weight, best_conf = sorted(weight_candidates, key=lambda x: -x[1])[0]
40
+ elif fallback_weight:
41
+ best_weight, best_conf = fallback_weight, fallback_conf
42
+ else:
43
+ return "Not detected", 0.0, "\n".join(raw_texts)
44
+
45
+ if "." in best_weight:
46
+ int_part, dec_part = best_weight.split(".")
47
+ int_part = int_part.lstrip("0") or "0"
48
+ best_weight = f"{int_part}.{dec_part}"
49
+ else:
50
+ best_weight = best_weight.lstrip("0") or "0"
51
+
52
+ return best_weight, round(best_conf * 100, 2), "\n".join(raw_texts)
53
+
54
+ except Exception as e:
55
+ return f"Error: {str(e)}", 0.0, "OCR failed"