Update weight_detector.py

weight_detector.py

@@ -2,55 +2,60 @@ 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
+from skimage import filters
 
 class WeightDetector:
     def __init__(self):
-        """Initialize with English and optimized settings"""
+        """OCR optimized for 7-segment displays"""
         self.reader = easyocr.Reader(
             ['en'],
             gpu=True,
             model_storage_directory='model',
-            download_enabled=True
+            download_enabled=True,
+            recog_network='english_g2'  # Better for digital displays
         )
         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
+        """Get current IST time"""
+        return datetime.now(self.ist).strftime('%Y-%m-%d %H:%M:%S IST')
+
+    def is_blurry(self, image: np.ndarray, threshold=100) -> bool:
+        """Check if image is blurry using Laplacian variance"""
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        variance = cv2.Laplacian(gray, cv2.CV_64F).var()
+        return variance < threshold
+
+    def preprocess_7segment(self, image: np.ndarray) -> np.ndarray:
+        """Optimized preprocessing for 7-segment displays"""
         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)
+        # Adaptive thresholding for digital displays
+        thresh = cv2.adaptiveThreshold(
+            gray, 255,
+            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
+            cv2.THRESH_BINARY_INV, 11, 2
+        )
         
-        # Noise reduction
-        denoised = cv2.medianBlur(thresh, 3)
+        # Remove small noise
+        kernel = np.ones((2, 2), np.uint8)
+        cleaned = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
         
-        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()
+        return cleaned
+
+    def extract_weight(self, text: str) -> Optional[float]:
+        """Extract weight value (handles decimals, units like g/kg)"""
+        text = text.replace(" ", "").replace(",", ".")
         
-        # Patterns for digital scale displays
+        # Patterns for digital scales (e.g., "0.000g", "12.34 kg")
         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
+            r'(\d+\.\d+)\s*[gGkK]',  # 12.34g or 12.34kg
+            r'(\d+)\s*[gGkK]',        # 123g or 123kg
+            r'(\d+\.\d+)',            # Decimal only
+            r'(\d+)'                   # Whole number
         ]
         
         for pattern in patterns:
@@ -58,85 +63,93 @@ class WeightDetector:
             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
+                    if 'k' in text.lower():  # Convert kg to g
+                        return value * 1000
                     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"""
+
+    def detect_weight(self, image_path: str) -> dict:
+        """Detect weight from image with error checks"""
         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)
+            # Check for blur
+            if self.is_blurry(img_cv):
+                return {
+                    "weight": None,
+                    "message": "⚠️ Image is blurry! Ensure clear focus.",
+                    "image": img,
+                    "time": self.get_current_ist()
+                }
             
-            # OCR with configuration optimized for digital displays
+            # Preprocess for 7-segment digits
+            processed = self.preprocess_7segment(img_cv)
+            
+            # OCR with 7-segment optimization
             results = self.reader.readtext(
                 processed,
+                allowlist='0123456789.gkGKlL',
                 paragraph=False,
-                detail=1,
-                allowlist='0123456789.gkGlL',
-                width_ths=2.0,
-                text_threshold=0.7
+                text_threshold=0.7,
+                width_ths=1.5
             )
             
-            # Process all potential weight values
+            # Extract and validate weights
             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
+                weight = self.extract_weight(text)
+                if weight and prob > 0.5:  # Minimum confidence
                     detected_weights.append({
-                        'weight': weight,
-                        'text': text,
-                        'probability': prob,
-                        'bbox': bbox
+                        "weight": weight,
+                        "text": text,
+                        "probability": prob,
+                        "bbox": bbox
                     })
             
-            # Prepare output
-            draw = ImageDraw.Draw(img)
-            current_time = self.get_current_ist()
-            
+            # Select best match (highest confidence + largest area)
             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
-                    ),
+                    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
+                # Draw annotations
+                draw = ImageDraw.Draw(img)
                 for item in detected_weights:
-                    bbox = item['bbox']
+                    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})"
+                    label = f"{item['weight']}g (Conf: {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")
+                # Add timestamp
+                draw.text((10, 10), f"Time: {self.get_current_ist()}", fill="blue")
                 
-                return best_match['weight'], current_time, img
+                return {
+                    "weight": best_match["weight"],
+                    "message": f"✅ Detected: {best_match['weight']}g (Conf: {best_match['probability']:.2f})",
+                    "image": img,
+                    "time": self.get_current_ist()
+                }
             
-            # No weight detected
-            draw.text((10, 10), f"Captured at: {current_time}", fill="blue")
-            return None, current_time, img
+            return {
+                "weight": None,
+                "message": "❌ No weight detected. Ensure clear 7-segment digits.",
+                "image": img,
+                "time": self.get_current_ist()
+            }
             
         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
+            return {
+                "weight": None,
+                "message": f"⚠️ Error: {str(e)}",
+                "image": None,
+                "time": self.get_current_ist()
+            }