Update weight_detector.py

weight_detector.py

@@ -4,64 +4,93 @@ 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 the OCR reader with English language support"""
-        self.reader = easyocr.Reader(['en'])
+        """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 preprocess_image(self, image_path: str) -> np.ndarray:
-        """Preprocess the image for better OCR results"""
-        img = cv2.imread(image_path)
-        if img is None:
-            raise ValueError("Could not read image from path")
-            
+    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(img, cv2.COLOR_BGR2GRAY)
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
         
-        # Apply adaptive thresholding
-        processed = cv2.adaptiveThreshold(
-            gray, 255, 
-            cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-            cv2.THRESH_BINARY, 11, 2
-        )
+        # 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 processed
+        return denoised
     
     def extract_weight_value(self, text: str) -> Optional[float]:
-        """Extract weight value from text using regex patterns"""
-        # Common weight patterns: 12.34g, 56.78 kg, 90.12 lbs, etc.
+        """Improved weight extraction with better pattern matching"""
+        # Clean the text
+        text = text.replace(' ', '').replace(',', '.').lower()
+        
+        # Patterns for digital scale displays
         patterns = [
-            r'(\d+\.\d+)\s*(g|kg|grams|kilograms|lb|lbs|pounds)',
-            r'(\d+)\s*(g|kg|grams|kilograms|lb|lbs|pounds)',
-            r'(\d+\.\d+)',  # Just numbers with decimal
-            r'(\d+)'        # Just whole numbers
+            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, re.IGNORECASE)
+            match = re.search(pattern, text)
             if match:
                 try:
-                    return float(match.group(1))
+                    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], List[dict], Image.Image]:
-        """Detect weight from an image and return value, metadata, and annotated image"""
+    def detect_weight(self, image_path: str) -> Tuple[Optional[float], str, Image.Image]:
+        """Enhanced weight detection with better error handling"""
         try:
-            # Read and preprocess image
+            # Read image
             img = Image.open(image_path).convert("RGB")
             img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
             
-            # Perform OCR
-            results = self.reader.readtext(img_cv)
+            # 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
+            )
             
-            # Find the most likely weight value
+            # Process all potential weight values
             detected_weights = []
             for (bbox, text, prob) in results:
                 weight = self.extract_weight_value(text)
-                if weight is not None:
+                if weight is not None and prob > 0.4:  # Minimum confidence
                     detected_weights.append({
                         'weight': weight,
                         'text': text,
@@ -69,26 +98,45 @@ class WeightDetector:
                         'bbox': bbox
                     })
             
-            # Sort by probability and get the highest
+            # Prepare output
+            draw = ImageDraw.Draw(img)
+            current_time = self.get_current_ist()
+            
             if detected_weights:
-                detected_weights.sort(key=lambda x: x['probability'], reverse=True)
+                # 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 bounding boxes on image
-                draw = ImageDraw.Draw(img)
+                # Draw all detections
                 for item in detected_weights:
                     bbox = item['bbox']
-                    # Convert bbox coordinates to tuple of tuples
                     polygon = [(int(x), int(y)) for [x, y] in bbox]
-                    draw.polygon(polygon, outline="red", width=2)
-                    # Add text label
+                    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] - 10), label, fill="red")
+                    draw.text((polygon[0][0], polygon[0][1] - 15), label, fill=color)
                 
-                return best_match['weight'], detected_weights, img
+                # Add timestamp to image
+                draw.text((10, 10), f"Captured at: {current_time}", fill="blue")
                 
-            return None, [], img
+                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:
-            print(f"Error processing image: {e}")
-            return None, [], Image.new("RGB", (100, 100), color="white")
+            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