Spaces:
Running
Running
File size: 5,533 Bytes
ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c ee68036 cc6391c |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
import cv2
import numpy as np
import easyocr
import re
from typing import Tuple, List, Optional
from PIL import Image, ImageDraw
import pytz
from datetime import datetime
class WeightDetector:
def __init__(self):
"""Initialize with English and optimized settings"""
self.reader = easyocr.Reader(
['en'],
gpu=True,
model_storage_directory='model',
download_enabled=True
)
self.ist = pytz.timezone('Asia/Kolkata')
def get_current_ist(self) -> str:
"""Get current time in Indian Standard Time"""
return datetime.now(self.ist).strftime('%Y-%m-%d %H:%M:%S %Z')
def preprocess_image(self, image: np.ndarray) -> np.ndarray:
"""Enhanced image preprocessing for digital displays"""
# Convert to grayscale
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# Contrast enhancement
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
contrast_enhanced = clahe.apply(gray)
# Thresholding for digital displays
_, thresh = cv2.threshold(contrast_enhanced, 0, 255,
cv2.THRESH_BINARY + cv2.THRESH_OTSU)
# Noise reduction
denoised = cv2.medianBlur(thresh, 3)
return denoised
def extract_weight_value(self, text: str) -> Optional[float]:
"""Improved weight extraction with better pattern matching"""
# Clean the text
text = text.replace(' ', '').replace(',', '.').lower()
# Patterns for digital scale displays
patterns = [
r'(\d+\.\d+)[gkl]', # 12.34g or 12.34kg
r'(\d+)[gkl]', # 123g or 123kg
r'(\d+\.\d+)', # Just numbers with decimal
r'(\d+)' # Just whole numbers
]
for pattern in patterns:
match = re.search(pattern, text)
if match:
try:
value = float(match.group(1))
# Assume grams if no unit specified
if 'k' in text:
return value * 1000 # Convert kg to g
return value
except ValueError:
continue
return None
def detect_weight(self, image_path: str) -> Tuple[Optional[float], str, Image.Image]:
"""Enhanced weight detection with better error handling"""
try:
# Read image
img = Image.open(image_path).convert("RGB")
img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
# Preprocess
processed = self.preprocess_image(img_cv)
# OCR with configuration optimized for digital displays
results = self.reader.readtext(
processed,
paragraph=False,
detail=1,
allowlist='0123456789.gkGlL',
width_ths=2.0,
text_threshold=0.7
)
# Process all potential weight values
detected_weights = []
for (bbox, text, prob) in results:
weight = self.extract_weight_value(text)
if weight is not None and prob > 0.4: # Minimum confidence
detected_weights.append({
'weight': weight,
'text': text,
'probability': prob,
'bbox': bbox
})
# Prepare output
draw = ImageDraw.Draw(img)
current_time = self.get_current_ist()
if detected_weights:
# Sort by probability and area (larger text is more likely the weight)
detected_weights.sort(
key=lambda x: (
x['probability'],
(x['bbox'][2][0] - x['bbox'][0][0]) * # Width
(x['bbox'][2][1] - x['bbox'][0][1]) # Height
),
reverse=True
)
best_match = detected_weights[0]
# Draw all detections
for item in detected_weights:
bbox = item['bbox']
polygon = [(int(x), int(y)) for [x, y] in bbox]
color = "green" if item == best_match else "red"
draw.polygon(polygon, outline=color, width=2)
label = f"{item['weight']}g (p={item['probability']:.2f})"
draw.text((polygon[0][0], polygon[0][1] - 15), label, fill=color)
# Add timestamp to image
draw.text((10, 10), f"Captured at: {current_time}", fill="blue")
return best_match['weight'], current_time, img
# No weight detected
draw.text((10, 10), f"Captured at: {current_time}", fill="blue")
return None, current_time, img
except Exception as e:
current_time = self.get_current_ist()
error_img = Image.new("RGB", (300, 100), color="white")
draw = ImageDraw.Draw(error_img)
draw.text((10, 10), f"Error: {str(e)}", fill="red")
draw.text((10, 30), f"Time: {current_time}", fill="blue")
return None, current_time, error_img |