weight_detector.py

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