File size: 1,663 Bytes
1dbcf19
8348064
1dbcf19
17d218a
65ef5f8
8348064
1dbcf19
 
65ef5f8
25cb585
83f2f91
8348064
 
 
83f2f91
25cb585
 
8348064
 
 
 
 
 
1dbcf19
 
8348064
1dbcf19
 
25cb585
 
83f2f91
1dbcf19
8348064
 
1dbcf19
83f2f91
5607fce
8348064
 
5607fce
8348064
5607fce
 
1dbcf19
 
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
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image, ImageFilter
import torch
import re

# Load TrOCR model
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")

def clean_ocr_text(text):
    print("[RAW OCR]", text)
    # Fix common misreads
    text = text.replace(",", ".").replace("s", "5").replace("o", "0").replace("O", "0")
    text = re.sub(r"[^\d\.kg]", "", text.lower())  # Keep digits, dot, kg
    print("[CLEANED OCR]", text)
    return text

def preprocess_image(image):
    # Enlarge + sharpen for better OCR
    image = image.resize((image.width * 2, image.height * 2), Image.BICUBIC)
    image = image.filter(ImageFilter.SHARPEN)
    return image

def extract_weight(image):
    try:
        image = preprocess_image(image)
        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)

        # Try matching decimal weight with unit
        match = re.search(r'(\d{1,3}\.\d{1,3})\s*(kg|g)', cleaned)
        if match:
            return f"{match.group(1)} {match.group(2)}"

        # Fallback: match only decimal number
        fallback = re.search(r'(\d{1,3}\.\d{1,3})', cleaned)
        if fallback:
            return f"{fallback.group(1)} g"  # Default to grams

        return f"No valid weight found | OCR: {cleaned}"
    except Exception as e:
        return f"Error: {str(e)}"