Sanjayraju30 commited on
Commit
91ebd46
·
verified ·
1 Parent(s): e2a19e4

Update ocr_engine.py

Browse files
Files changed (1) hide show
  1. ocr_engine.py +5 -3
ocr_engine.py CHANGED
@@ -2,9 +2,9 @@ from transformers import TrOCRProcessor, VisionEncoderDecoderModel
2
  from PIL import Image
3
  import re
4
 
5
- # Load smaller/faster model
6
- processor = TrOCRProcessor.from_pretrained("microsoft/trocr-small-printed")
7
- model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-small-printed")
8
 
9
  def extract_weight(image: Image.Image) -> str:
10
  image = image.convert("RGB")
@@ -12,10 +12,12 @@ def extract_weight(image: Image.Image) -> str:
12
  generated_ids = model.generate(pixel_values, max_length=32)
13
  full_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
14
 
 
15
  cleaned = full_text.lower().replace(" ", "")
16
  match = re.search(r"(\d+(\.\d+)?)", cleaned)
17
  weight = match.group(1) if match else None
18
 
 
19
  if any(u in cleaned for u in ["kg", "kgs", "kilogram", "kilo"]):
20
  unit = "kg"
21
  elif any(u in cleaned for u in ["g", "gram", "grams"]):
 
2
  from PIL import Image
3
  import re
4
 
5
+ # Load lightweight, fast, public model
6
+ processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
7
+ model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
8
 
9
  def extract_weight(image: Image.Image) -> str:
10
  image = image.convert("RGB")
 
12
  generated_ids = model.generate(pixel_values, max_length=32)
13
  full_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
14
 
15
+ # Extract number
16
  cleaned = full_text.lower().replace(" ", "")
17
  match = re.search(r"(\d+(\.\d+)?)", cleaned)
18
  weight = match.group(1) if match else None
19
 
20
+ # Detect unit
21
  if any(u in cleaned for u in ["kg", "kgs", "kilogram", "kilo"]):
22
  unit = "kg"
23
  elif any(u in cleaned for u in ["g", "gram", "grams"]):