Spaces:
Running
Running
File size: 2,188 Bytes
1dbcf19 8348064 1dbcf19 17d218a 65ef5f8 fd59642 1dbcf19 65ef5f8 25cb585 8348064 546e454 25cb585 d07e9f7 8348064 fd59642 909b3d7 fd59642 909b3d7 fd59642 909b3d7 fd59642 1dbcf19 d07e9f7 1dbcf19 25cb585 83f2f91 1dbcf19 0e76d8f 1dbcf19 909b3d7 d07e9f7 546e454 909b3d7 546e454 909b3d7 5607fce fd59642 546e454 1dbcf19 0e76d8f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image, ImageFilter
import torch
import re
# Load TrOCR model and processor
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
def clean_ocr_text(text):
text = text.replace(",", ".").replace("s", "5").replace("o", "0").replace("O", "0")
return re.sub(r"[^\d.kg]", "", text.lower())
def restore_decimal(text):
if re.fullmatch(r"\d{5}", text):
return f"{text[:2]}.{text[2:]}"
elif re.fullmatch(r"\d{4}", text):
return f"{text[:2]}.{text[2:]}"
return text
def extract_unit_from_text(raw_text):
raw_text = raw_text.lower()
if "kg" in raw_text or "kgs" in raw_text or "k9" in raw_text:
return "kg"
elif re.search(r'\dg', raw_text) or "gram" in raw_text:
return "g"
return None
def extract_weight(image):
try:
image = image.resize((image.width * 2, image.height * 2), Image.BICUBIC)
image = image.filter(ImageFilter.SHARPEN)
pixel_values = processor(images=image, return_tensors="pt").pixel_values
generated_ids = model.generate(pixel_values)
raw_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
cleaned = clean_ocr_text(raw_text)
match = re.search(r"(\d{1,3}\.\d{1,3})\s*(kg|g)?", cleaned)
if match:
weight = match.group(1)
unit = match.group(2) or extract_unit_from_text(raw_text)
return f"{weight} {unit or ''}".strip(), raw_text
match = re.search(r"\d{4,5}", cleaned)
if match:
decimal_fixed = restore_decimal(match.group())
unit = extract_unit_from_text(raw_text)
return f"{decimal_fixed} {unit or ''}".strip(), raw_text
match = re.search(r"\d+", cleaned)
if match:
unit = extract_unit_from_text(raw_text)
return f"{match.group()} {unit or ''}".strip(), raw_text
return "Error: No valid weight found", raw_text
except Exception as e:
return f"Error: {str(e)}", ""
|