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