Multi-Gauge Detection System

# Multi-Gauge Detection System
import cv2
import numpy as np
import torch
from ultralytics import YOLO
import concurrent.futures
import json
from datetime import datetime

class MultiGaugeDetector:
    def __init__(self, model_path, config_path):
        # Load YOLO model for gauge detection
        self.yolo_model = YOLO(model_path)
        
        # Load configuration
        with open(config_path, 'r') as f:
            self.config = json.load(f)
        
        # Initialize gauge readers
        self.analog_reader = AnalogGaugeReader()
        self.digital_reader = DigitalGaugeReader()
        
        # Performance tracking
        self.stats = {
            'total_detections': 0,
            'successful_readings': 0,
            'processing_times': []
        }
    
    def detect_gauges(self, image):
        """Detect all gauges in the image"""
        start_time = time.time()
        
        # Run YOLO detection
        results = self.yolo_model(image, conf=0.3, iou=0.5)
        
        # Parse results
        detections = []
        for result in results:
            boxes = result.boxes
            if boxes is not None:
                for box in boxes:
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    confidence = box.conf[0].cpu().numpy()
                    class_id = int(box.cls[0].cpu().numpy())
                    
                    # Extract gauge region
                    gauge_roi = image[int(y1):int(y2), int(x1):int(x2)]
                    
                    detection = {
                        'bbox': [x1, y1, x2, y2],
                        'confidence': float(confidence),
                        'class_id': class_id,
                        'class_name': self.config['classes'][class_id],
                        'roi': gauge_roi,
                        'timestamp': datetime.now().isoformat()
                    }
                    detections.append(detection)
        
        detection_time = time.time() - start_time
        self.stats['processing_times'].append(detection_time)
        
        return detections
    
    def classify_gauge_type(self, detection):
        """Classify gauge as analog or digital"""
        roi = detection['roi']
        class_name = detection['class_name']
        
        # Simple rule-based classification
        # In production, use a CNN classifier
        if 'digital' in class_name.lower() or 'lcd' in class_name.lower():
            return 'digital'
        elif 'analog' in class_name.lower() or 'needle' in class_name.lower():
            return 'analog'
        else:
            # Use image analysis for unknown types
            return self._analyze_gauge_type(roi)
    
    def _analyze_gauge_type(self, roi):
        """Analyze ROI to determine gauge type"""
        # Convert to grayscale
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
        
        # Detect circular shapes (analog gauges)
        circles = cv2.HoughCircles(
            gray, cv2.HOUGH_GRADIENT, 1, 20,
            param1=50, param2=30, minRadius=10, maxRadius=100
        )
        
        if circles is not None:
            return 'analog'
        
        # Detect rectangular displays (digital gauges)
        edges = cv2.Canny(gray, 50, 150)
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        for contour in contours:
            # Check for rectangular shapes
            epsilon = 0.02 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            if len(approx) == 4:  # Rectangle
                return 'digital'
        
        return 'unknown'
    
    def process_batch(self, detections):
        """Process multiple gauges in parallel"""
        analog_batch = []
        digital_batch = []
        
        # Group by type
        for detection in detections:
            gauge_type = self.classify_gauge_type(detection)
            
            if gauge_type == 'analog':
                analog_batch.append(detection)
            elif gauge_type == 'digital':
                digital_batch.append(detection)
        
        results = []
        
        # Process batches in parallel
        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
            futures = []
            
            # Submit analog batch
            if analog_batch:
                futures.append(
                    executor.submit(self._process_analog_batch, analog_batch)
                )
            
            # Submit digital batch
            if digital_batch:
                futures.append(
                    executor.submit(self._process_digital_batch, digital_batch)
                )
            
            # Collect results
            for future in concurrent.futures.as_completed(futures):
                batch_results = future.result()
                results.extend(batch_results)
        
        return results
    
    def _process_analog_batch(self, batch):
        """Process analog gauges"""
        results = []
        
        for detection in batch:
            try:
                reading = self.analog_reader.read_gauge(detection['roi'])
                
                result = {
                    'detection': detection,
                    'reading': reading,
                    'type': 'analog',
                    'status': 'success' if reading else 'failed'
                }
                results.append(result)
                
            except Exception as e:
                result = {
                    'detection': detection,
                    'reading': None,
                    'type': 'analog',
                    'status': 'error',
                    'error': str(e)
                }
                results.append(result)
        
        return results
    
    def _process_digital_batch(self, batch):
        """Process digital gauges"""
        results = []
        
        # Batch OCR processing for efficiency
        rois = [det['roi'] for det in batch]
        readings = self.digital_reader.batch_read(rois)
        
        for detection, reading in zip(batch, readings):
            result = {
                'detection': detection,
                'reading': reading,
                'type': 'digital',
                'status': 'success' if reading else 'failed'
            }
            results.append(result)
        
        return results
    
    def process_image(self, image):
        """Main processing pipeline"""
        start_time = time.time()
        
        # Step 1: Detect all gauges
        detections = self.detect_gauges(image)
        
        if not detections:
            return {
                'status': 'no_gauges_detected',
                'processing_time': time.time() - start_time
            }
        
        # Step 2: Process all detections
        results = self.process_batch(detections)
        
        # Step 3: Apply quality filters
        filtered_results = self._apply_quality_filters(results)
        
        # Step 4: Generate summary
        summary = self._generate_summary(filtered_results)
        
        total_time = time.time() - start_time
        
        return {
            'status': 'success',
            'total_gauges': len(detections),
            'successful_readings': len([r for r in filtered_results if r['status'] == 'success']),
            'results': filtered_results,
            'summary': summary,
            'processing_time': total_time
        }
    
    def _apply_quality_filters(self, results):
        """Apply confidence and quality filters"""
        filtered = []
        
        for result in results:
            # Check detection confidence
            if result['detection']['confidence'] < self.config['min_detection_confidence']:
                result['status'] = 'low_confidence'
                continue
            
            # Check reading quality
            if result['reading'] and 'confidence' in result['reading']:
                if result['reading']['confidence'] < self.config['min_reading_confidence']:
                    result['status'] = 'low_reading_confidence'
                    continue
            
            filtered.append(result)
        
        return filtered
    
    def _generate_summary(self, results):
        """Generate processing summary"""
        total = len(results)
        successful = len([r for r in results if r['status'] == 'success'])
        
        analog_count = len([r for r in results if r['type'] == 'analog'])
        digital_count = len([r for r in results if r['type'] == 'digital'])
        
        return {
            'total_detected': total,
            'successful_readings': successful,
            'success_rate': successful / total if total > 0 else 0,
            'analog_gauges': analog_count,
            'digital_gauges': digital_count,
            'average_confidence': np.mean([
                r['detection']['confidence'] for r in results
            ]) if results else 0
        }

# Usage example
def main():
    # Initialize detector
    detector = MultiGaugeDetector(
        model_path='models/gauge_detector_yolov8.pt',
        config_path='config/gauge_config.json'
    )
    
    # Load test image
    image = cv2.imread('test_images/control_panel.jpg')
    
    # Process image
    results = detector.process_image(image)
    
    # Print results
    print(f"Processing Status: {results['status']}")
    print(f"Total Gauges Detected: {results['total_gauges']}")
    print(f"Successful Readings: {results['successful_readings']}")
    print(f"Processing Time: {results['processing_time']:.3f}s")
    
    # Display detailed results
    for i, result in enumerate(results['results']):
        print(f"\nGauge {i+1}:")
        print(f"  Type: {result['type']}")
        print(f"  Class: {result['detection']['class_name']}")
        print(f"  Confidence: {result['detection']['confidence']:.3f}")
        
        if result['reading']:
            print(f"  Reading: {result['reading']['value']}")
            print(f"  Unit: {result['reading'].get('unit', 'N/A')}")

if __name__ == "__main__":
    main()