Sanjayraju30 commited on
Commit
5217dbe
·
verified ·
1 Parent(s): 4aec32e

Update ocr_engine.py

Browse files
Files changed (1) hide show
  1. ocr_engine.py +12 -20
ocr_engine.py CHANGED
@@ -1,38 +1,30 @@
1
- from transformers import TrOCRProcessor, VisionEncoderDecoderModel
2
  from PIL import Image
 
3
  import re
4
 
5
- # Load model + processor
6
- processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
7
- model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1")
8
 
9
  def extract_weight(image: Image.Image) -> str:
10
  image = image.convert("RGB")
11
  pixel_values = processor(images=image, return_tensors="pt").pixel_values
12
- generated_ids = model.generate(pixel_values, max_length=20)
13
  full_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
14
 
15
- # For debugging (optional):
16
- print("OCR Output:", full_text)
17
 
18
  # Extract number (weight)
19
  match = re.search(r"(\d+(\.\d+)?)", full_text)
20
- if match:
21
- weight = match.group(1)
22
- else:
23
- return "No valid weight detected"
24
 
25
- # Detect unit — smarter match
26
  text_lower = full_text.lower().replace(" ", "")
27
- if any(unit in text_lower for unit in ["kg", "kgs", "kilogram", "kilo", "k.g"]):
28
  unit = "kg"
29
- elif any(unit in text_lower for unit in ["g", "gram", "grams"]):
30
  unit = "grams"
31
  else:
32
- # Smart fallback: use value
33
- if float(weight) >= 5:
34
- unit = "kg"
35
- else:
36
- unit = "grams"
37
 
38
- return f"{weight} {unit}"
 
 
1
  from PIL import Image
2
+ from transformers import AutoProcessor, VisionEncoderDecoderModel
3
  import re
4
 
5
+ # Load model fine-tuned for 7-segment displays
6
+ processor = AutoProcessor.from_pretrained("roboflow/ocr-7segment")
7
+ model = VisionEncoderDecoderModel.from_pretrained("roboflow/ocr-7segment")
8
 
9
  def extract_weight(image: Image.Image) -> str:
10
  image = image.convert("RGB")
11
  pixel_values = processor(images=image, return_tensors="pt").pixel_values
12
+ generated_ids = model.generate(pixel_values)
13
  full_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
14
 
15
+ print("OCR Text:", full_text) # optional debug
 
16
 
17
  # Extract number (weight)
18
  match = re.search(r"(\d+(\.\d+)?)", full_text)
19
+ weight = match.group(1) if match else None
 
 
 
20
 
21
+ # Detect unit
22
  text_lower = full_text.lower().replace(" ", "")
23
+ if any(u in text_lower for u in ["kg", "kgs", "kilogram", "kilo"]):
24
  unit = "kg"
25
+ elif any(u in text_lower for u in ["g", "gram", "grams"]):
26
  unit = "grams"
27
  else:
28
+ unit = "kg" if weight and float(weight) >= 5 else "grams"
 
 
 
 
29
 
30
+ return f"{weight} {unit}" if weight else "No valid weight detected"