File size: 2,057 Bytes
1dbcf19
8348064
1dbcf19
17d218a
65ef5f8
fd59642
1dbcf19
 
65ef5f8
25cb585
fd59642
8348064
fd59642
25cb585
d07e9f7
 
 
 
 
 
8348064
fd59642
 
 
 
 
 
 
 
1dbcf19
 
fd59642
d07e9f7
 
 
fd59642
1dbcf19
 
25cb585
 
83f2f91
1dbcf19
fd59642
0e76d8f
1dbcf19
fd59642
d07e9f7
fd59642
 
 
 
 
5607fce
fd59642
1dbcf19
0e76d8f
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
47
48
49
50
51
52
53
54
55
56
57
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image, ImageFilter
import torch
import re

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

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

def restore_decimal(text):
    if re.fullmatch(r"\d{5}", text):
        return f"{text[:2]}.{text[2:]}"
    elif re.fullmatch(r"\d{4}", text):
        return f"{text[:2]}.{text[2:]}"
    return text

def extract_unit_from_text(raw_text):
    raw_text = raw_text.lower()
    if "kg" in raw_text:
        return "kg"
    elif "g" in raw_text:
        return "g"
    return "g"  # fallback if unit not found

def extract_weight(image):
    try:
        # Enhance image
        image = image.resize((image.width * 2, image.height * 2), Image.BICUBIC)
        image = image.filter(ImageFilter.SHARPEN)

        # OCR inference
        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 direct match (e.g., 52.25 kg or 250.5g)
        match = re.search(r"(\d{1,3}\.\d{1,3})\s*(kg|g)?", cleaned)
        if match:
            return f"{match.group(1)} {match.group(2) or ''}".strip(), raw_text

        # Fallback if no decimal found: convert big number like 53255 to 52.255
        fallback_match = re.search(r"\d{4,5}", cleaned)
        if fallback_match:
            decimal_fixed = restore_decimal(fallback_match.group())
            return decimal_fixed, raw_text

        return "Error: No valid weight found", raw_text
    except Exception as e:
        return f"Error: {str(e)}", ""