File size: 2,269 Bytes
a29f826
da9f292
5d670ae
da9f292
a29f826
da9f292
a29f826
 
da9f292
363a646
65ed4c1
a29f826
 
da9f292
a29f826
 
3ca006e
ddf8948
ee1d691
ddf8948
 
 
a29f826
 
 
ddf8948
 
 
61b752b
ddf8948
61b752b
 
 
ddf8948
 
477d4fe
ddf8948
 
61b752b
ddf8948
 
 
 
acddb2f
103f82b
ddf8948
 
 
 
 
61b752b
ddf8948
 
 
 
 
 
 
103f82b
ddf8948
8fe1b94
65ed4c1
2132698
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
58
59
60
61
62
63
64
from mmocr.apis import MMOCRInferencer
import numpy as np
import cv2
import re
from PIL import Image

# Initialize MMOCR
ocr = MMOCRInferencer(det='DBNet', recog='SAR', device='cpu')  # or 'cuda' if GPU available

def extract_weight_from_image(pil_img):
    try:
        # Convert PIL to OpenCV image (BGR)
        img = np.array(pil_img.convert("RGB"))[:, :, ::-1]

        # Run MMOCR inference
        result = ocr(img)

        raw_texts = []
        weight_candidates = []
        fallback_weight = None
        fallback_conf = 0.0

        for item in result['predictions'][0]:
            text = item['text']
            conf = item.get('score', 0.8)  # Fallback confidence
            original = text
            cleaned = text.lower().strip()

            # Fix common OCR misreads
            cleaned = cleaned.replace(",", ".")
            cleaned = cleaned.replace("o", "0").replace("O", "0")
            cleaned = cleaned.replace("s", "5").replace("S", "5")
            cleaned = cleaned.replace("g", "9").replace("G", "6")
            cleaned = cleaned.replace("kg", "").replace("kgs", "")
            cleaned = re.sub(r"[^0-9\.]", "", cleaned)

            raw_texts.append(f"{original}{cleaned} (conf: {round(conf, 2)})")

            if cleaned and cleaned.replace(".", "").isdigit() and not fallback_weight:
                fallback_weight = cleaned
                fallback_conf = conf

            if cleaned.count(".") <= 1 and re.fullmatch(r"\d{2,4}(\.\d{1,3})?", cleaned):
                weight_candidates.append((cleaned, conf))

        if weight_candidates:
            best_weight, best_conf = sorted(weight_candidates, key=lambda x: -x[1])[0]
        elif fallback_weight:
            best_weight, best_conf = fallback_weight, fallback_conf
        else:
            return "Not detected", 0.0, "\n".join(raw_texts)

        if "." in best_weight:
            int_part, dec_part = best_weight.split(".")
            int_part = int_part.lstrip("0") or "0"
            best_weight = f"{int_part}.{dec_part}"
        else:
            best_weight = best_weight.lstrip("0") or "0"

        return best_weight, round(best_conf * 100, 2), "\n".join(raw_texts)

    except Exception as e:
        return f"Error: {str(e)}", 0.0, "OCR failed"