weight_detector.py

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()
            }