Update ocr_engine.py

ocr_engine.py

@@ -1,25 +1,24 @@
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import cv2
+import pytesseract
 from PIL import Image
-import torch
 import re
 
-# Load processor and model once
-processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+def extract_weight(img_path):
+    img = cv2.imread(img_path)
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    
+    # OCR
+    text = pytesseract.image_to_string(gray, config='--psm 6')
 
-def extract_weight(image):
-    try:
-        # OCR inference
-        pixel_values = processor(images=image, return_tensors="pt").pixel_values
-        generated_ids = model.generate(pixel_values)
-        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
-        print("OCR Output:", text)
+    # Clean and lower text
+    text = text.lower().replace('\n', ' ').strip()
 
-        # Extract valid float or integer from OCR result
-        match = re.search(r'\d{2,5}(\.\d{1,2})?', text)  # e.g., 65 or 5325.0
-        if match:
-            return match.group() + " kg"
-        else:
-            return "No valid weight found"
-    except Exception as e:
-        return f"Error: {str(e)}"
+    # Regex to find pattern like 52.25 g or 75.8 kg
+    match = re.search(r'(\d+\.\d+|\d+)\s*(kg|g)', text)
+
+    if match:
+        number = match.group(1)
+        unit = match.group(2)
+        return f"{number} {unit}"
+    else:
+        return "Weight not detected"