Spaces:
Running
Running
Update ocr_engine.py
Browse files- ocr_engine.py +9 -12
ocr_engine.py
CHANGED
@@ -8,7 +8,6 @@ processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
|
|
8 |
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
|
9 |
|
10 |
def clean_ocr_text(text):
|
11 |
-
# Fix common OCR misreads
|
12 |
text = text.replace(",", ".").replace("s", "5").replace("o", "0").replace("O", "0")
|
13 |
return re.sub(r"[^\d.kg]", "", text.lower())
|
14 |
|
@@ -21,41 +20,39 @@ def restore_decimal(text):
|
|
21 |
|
22 |
def extract_unit_from_text(raw_text):
|
23 |
raw_text = raw_text.lower()
|
24 |
-
if "kg" in raw_text:
|
25 |
return "kg"
|
26 |
-
elif "
|
27 |
return "g"
|
28 |
-
return
|
29 |
|
30 |
def extract_weight(image):
|
31 |
try:
|
32 |
-
# Resize & sharpen image
|
33 |
image = image.resize((image.width * 2, image.height * 2), Image.BICUBIC)
|
34 |
image = image.filter(ImageFilter.SHARPEN)
|
35 |
|
36 |
-
# OCR inference
|
37 |
pixel_values = processor(images=image, return_tensors="pt").pixel_values
|
38 |
generated_ids = model.generate(pixel_values)
|
39 |
raw_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
|
40 |
|
41 |
cleaned = clean_ocr_text(raw_text)
|
42 |
|
43 |
-
# Case 1: Match decimal with unit
|
44 |
match = re.search(r"(\d{1,3}\.\d{1,3})\s*(kg|g)?", cleaned)
|
45 |
if match:
|
46 |
-
|
|
|
|
|
47 |
|
48 |
-
# Case 2: Large number fallback like 53255 → 52.255
|
49 |
match = re.search(r"\d{4,5}", cleaned)
|
50 |
if match:
|
51 |
decimal_fixed = restore_decimal(match.group())
|
52 |
unit = extract_unit_from_text(raw_text)
|
53 |
-
return f"{decimal_fixed} {unit}", raw_text
|
54 |
|
55 |
-
# Final fallback: plain number
|
56 |
match = re.search(r"\d+", cleaned)
|
57 |
if match:
|
58 |
-
|
|
|
59 |
|
60 |
return "Error: No valid weight found", raw_text
|
61 |
|
|
|
8 |
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
|
9 |
|
10 |
def clean_ocr_text(text):
|
|
|
11 |
text = text.replace(",", ".").replace("s", "5").replace("o", "0").replace("O", "0")
|
12 |
return re.sub(r"[^\d.kg]", "", text.lower())
|
13 |
|
|
|
20 |
|
21 |
def extract_unit_from_text(raw_text):
|
22 |
raw_text = raw_text.lower()
|
23 |
+
if "kg" in raw_text or "kgs" in raw_text or "k9" in raw_text:
|
24 |
return "kg"
|
25 |
+
elif re.search(r'\dg', raw_text) or "gram" in raw_text:
|
26 |
return "g"
|
27 |
+
return None
|
28 |
|
29 |
def extract_weight(image):
|
30 |
try:
|
|
|
31 |
image = image.resize((image.width * 2, image.height * 2), Image.BICUBIC)
|
32 |
image = image.filter(ImageFilter.SHARPEN)
|
33 |
|
|
|
34 |
pixel_values = processor(images=image, return_tensors="pt").pixel_values
|
35 |
generated_ids = model.generate(pixel_values)
|
36 |
raw_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
|
37 |
|
38 |
cleaned = clean_ocr_text(raw_text)
|
39 |
|
|
|
40 |
match = re.search(r"(\d{1,3}\.\d{1,3})\s*(kg|g)?", cleaned)
|
41 |
if match:
|
42 |
+
weight = match.group(1)
|
43 |
+
unit = match.group(2) or extract_unit_from_text(raw_text)
|
44 |
+
return f"{weight} {unit or ''}".strip(), raw_text
|
45 |
|
|
|
46 |
match = re.search(r"\d{4,5}", cleaned)
|
47 |
if match:
|
48 |
decimal_fixed = restore_decimal(match.group())
|
49 |
unit = extract_unit_from_text(raw_text)
|
50 |
+
return f"{decimal_fixed} {unit or ''}".strip(), raw_text
|
51 |
|
|
|
52 |
match = re.search(r"\d+", cleaned)
|
53 |
if match:
|
54 |
+
unit = extract_unit_from_text(raw_text)
|
55 |
+
return f"{match.group()} {unit or ''}".strip(), raw_text
|
56 |
|
57 |
return "Error: No valid weight found", raw_text
|
58 |
|