Sanjayraju30 commited on
Commit
0d59618
·
verified ·
1 Parent(s): 0590b95

Create ocrr_engine.py

Browse files
Files changed (1) hide show
  1. ocrr_engine.py +85 -0
ocrr_engine.py ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pytesseract
2
+ import numpy as np
3
+ import cv2
4
+ import re
5
+ from PIL import Image
6
+ import logging
7
+ import sys
8
+
9
+ # Set up logging
10
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[logging.StreamHandler(sys.stdout)])
11
+
12
+ def preprocess_image(img):
13
+ """Preprocess image for robust OCR, optimized for various weight display formats."""
14
+ try:
15
+ # Convert PIL to OpenCV format
16
+ img = np.array(img)
17
+ img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
18
+
19
+ # Convert to grayscale
20
+ gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
21
+
22
+ # Enhance contrast for diverse lighting conditions
23
+ clahe = cv2.createCLAHE(clipLimit=4.0, tileGridSize=(8, 8))
24
+ enhanced = clahe.apply(gray)
25
+
26
+ # Apply adaptive thresholding with flexible block size
27
+ block_size = max(11, min(31, int(img.shape[0] / 20) * 2 + 1))
28
+ thresh = cv2.adaptiveThreshold(
29
+ enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, block_size, 3
30
+ )
31
+
32
+ # Sharpen to enhance edges
33
+ kernel = np.array([[-1, -1, -1],
34
+ [-1, 9, -1],
35
+ [-1, -1, -1]])
36
+ sharpened = cv2.filter2D(thresh, -1, kernel)
37
+
38
+ # Denoise for noisy images
39
+ denoised = cv2.fastNlMeansDenoising(sharpened, h=10)
40
+
41
+ return denoised
42
+ except Exception as e:
43
+ logging.error(f"Preprocessing failed: {str(e)}")
44
+ return gray # Fallback to grayscale
45
+
46
+ def extract_weight_from_image(pil_img):
47
+ """Extract weight and unit from a digital scale image, supporting various formats."""
48
+ try:
49
+ # Preprocess image
50
+ thresh = preprocess_image(pil_img)
51
+
52
+ # Convert to PIL for Tesseract
53
+ pil_img = Image.fromarray(cv2.cvtColor(thresh, cv2.COLOR_GRAY2RGB))
54
+
55
+ # Try Tesseract with optimized config for numbers and units
56
+ config = r'--oem 3 --psm 7 -c tessedit_char_whitelist=0123456789.,kgKg'
57
+ text = pytesseract.image_to_string(pil_img, config=config)
58
+ logging.info(f"Tesseract raw output: {text}")
59
+
60
+ # Clean and validate text
61
+ text = text.strip().lower()
62
+ text = re.sub(r'\s+', '', text) # Remove extra spaces
63
+
64
+ # Extract weight and unit
65
+ match = re.search(r'(\d*\.?\d+)([kgkg]?)', text)
66
+ if match:
67
+ weight_str = match.group(1)
68
+ unit = match.group(2) if match.group(2) else "g" # Default to grams if no unit
69
+ weight_str = weight_str.replace(',', '.') # Handle decimal formats (e.g., 68,0)
70
+ if re.fullmatch(r'^\d*\.?\d+$', weight_str):
71
+ weight_str = weight_str.lstrip('0') or '0'
72
+ confidence = 95.0 if len(weight_str.replace('.', '')) >= 3 else 90.0
73
+ try:
74
+ weight = float(weight_str)
75
+ if 0.001 <= weight <= 5000:
76
+ logging.info(f"Detected weight: {weight} {unit}, Confidence: {confidence:.2f}%")
77
+ return weight_str, confidence, unit
78
+ except ValueError:
79
+ logging.warning(f"Invalid weight format: {weight_str}")
80
+
81
+ logging.info("No valid weight detected.")
82
+ return "Not detected", 0.0, ""
83
+ except Exception as e:
84
+ logging.error(f"Weight extraction failed: {str(e)}")
85
+ return "Not detected", 0.0, ""