Rename weight_detector.py to ocr_engine.py

ocr_engine.py

@@ -0,0 +1,11 @@
+import cv2
+import pytesseract
+from PIL import Image
+import numpy as np
+
+def extract_weight(image: Image.Image) -> str:
+    img = np.array(image.convert("RGB"))
+    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    text = pytesseract.image_to_string(gray, config="--psm 7 digits")
+    weight = ''.join(filter(lambda x: x in '0123456789.', text))
+    return weight if weight else "No valid weight detected"

weight_detector.py

@@ -1,155 +0,0 @@
-import cv2
-import numpy as np
-import easyocr
-import re
-from PIL import Image, ImageDraw
-import pytz
-from datetime import datetime
-from skimage import filters
-
-class WeightDetector:
-    def __init__(self):
-        """OCR optimized for 7-segment displays"""
-        self.reader = easyocr.Reader(
-            ['en'],
-            gpu=True,
-            model_storage_directory='model',
-            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 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)
-        
-        # Adaptive thresholding for digital displays
-        thresh = cv2.adaptiveThreshold(
-            gray, 255,
-            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
-            cv2.THRESH_BINARY_INV, 11, 2
-        )
-        
-        # Remove small noise
-        kernel = np.ones((2, 2), np.uint8)
-        cleaned = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
-        
-        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 scales (e.g., "0.000g", "12.34 kg")
-        patterns = [
-            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:
-            match = re.search(pattern, text)
-            if match:
-                try:
-                    value = float(match.group(1))
-                    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) -> dict:
-        """Detect weight from image with error checks"""
-        try:
-            img = Image.open(image_path).convert("RGB")
-            img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-            
-            # 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()
-                }
-            
-            # 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,
-                text_threshold=0.7,
-                width_ths=1.5
-            )
-            
-            # Extract and validate weights
-            detected_weights = []
-            for (bbox, text, prob) in results:
-                weight = self.extract_weight(text)
-                if weight and prob > 0.5:  # Minimum confidence
-                    detected_weights.append({
-                        "weight": weight,
-                        "text": text,
-                        "probability": prob,
-                        "bbox": bbox
-                    })
-            
-            # Select best match (highest confidence + largest area)
-            if detected_weights:
-                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 annotations
-                draw = ImageDraw.Draw(img)
-                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 (Conf: {item['probability']:.2f})"
-                    draw.text((polygon[0][0], polygon[0][1] - 15), label, fill=color)
-                
-                # Add timestamp
-                draw.text((10, 10), f"Time: {self.get_current_ist()}", fill="blue")
-                
-                return {
-                    "weight": best_match["weight"],
-                    "message": f"✅ Detected: {best_match['weight']}g (Conf: {best_match['probability']:.2f})",
-                    "image": img,
-                    "time": self.get_current_ist()
-                }
-            
-            return {
-                "weight": None,
-                "message": "❌ No weight detected. Ensure clear 7-segment digits.",
-                "image": img,
-                "time": self.get_current_ist()
-            }
-            
-        except Exception as e:
-            return {
-                "weight": None,
-                "message": f"⚠️ Error: {str(e)}",
-                "image": None,
-                "time": self.get_current_ist()
-            }