Update ocr_engine.py

ocr_engine.py

@@ -1,25 +1,28 @@
+from transformers import DonutProcessor, VisionEncoderDecoderModel
 from PIL import Image
-from transformers import AutoProcessor, VisionEncoderDecoderModel
+import torch
 import re
 
-# Load model fine-tuned for 7-segment displays
-processor = AutoProcessor.from_pretrained("roboflow/ocr-7segment")
-model = VisionEncoderDecoderModel.from_pretrained("roboflow/ocr-7segment")
+# Load model
+processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
 
 def extract_weight(image: Image.Image) -> str:
     image = image.convert("RGB")
-    pixel_values = processor(images=image, return_tensors="pt").pixel_values
-    generated_ids = model.generate(pixel_values)
-    full_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    pixel_values = processor(image, return_tensors="pt").pixel_values
 
-    print("OCR Text:", full_text)  # optional debug
+    # Generate output
+    outputs = model.generate(pixel_values, max_length=512)
+    decoded = processor.batch_decode(outputs, skip_special_tokens=True)[0]
 
-    # Extract number (weight)
-    match = re.search(r"(\d+(\.\d+)?)", full_text)
+    print("OCR Output:", decoded)  # Optional for debug
+
+    # Extract number
+    match = re.search(r"(\d+(\.\d+)?)", decoded)
     weight = match.group(1) if match else None
 
     # Detect unit
-    text_lower = full_text.lower().replace(" ", "")
+    text_lower = decoded.lower().replace(" ", "")
     if any(u in text_lower for u in ["kg", "kgs", "kilogram", "kilo"]):
         unit = "kg"
     elif any(u in text_lower for u in ["g", "gram", "grams"]):