Edge AI Computing

# Edge AI Gauge Reader - TensorFlow Lite
import tensorflow as tf
import numpy as np
import cv2

class EdgeGaugeReader:
    def __init__(self, model_path):
        # Load TensorFlow Lite model
        self.interpreter = tf.lite.Interpreter(model_path=model_path)
        self.interpreter.allocate_tensors()
        
        # Get input and output details
        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()
        
        # Input shape for preprocessing
        self.input_shape = self.input_details[0]['shape']
        self.input_dtype = self.input_details[0]['dtype']
    
    def preprocess_image(self, image):
        """Preprocess image for model input"""
        # Resize to model input size
        height, width = self.input_shape[1:3]
        image = cv2.resize(image, (width, height))
        
        # Normalize pixel values
        if self.input_dtype == np.uint8:
            # Quantized model expects uint8
            image = image.astype(np.uint8)
        else:
            # Float model expects normalized values
            image = image.astype(np.float32) / 255.0
        
        # Add batch dimension
        image = np.expand_dims(image, axis=0)
        return image
    
    def predict(self, image):
        """Run inference on preprocessed image"""
        # Preprocess input
        input_data = self.preprocess_image(image)
        
        # Set input tensor
        self.interpreter.set_tensor(
            self.input_details[0]['index'], 
            input_data
        )
        
        # Run inference
        self.interpreter.invoke()
        
        # Get outputs
        detection_boxes = self.interpreter.get_tensor(
            self.output_details[0]['index']
        )[0]
        
        detection_classes = self.interpreter.get_tensor(
            self.output_details[1]['index']
        )[0]
        
        detection_scores = self.interpreter.get_tensor(
            self.output_details[2]['index']
        )[0]
        
        ocr_text = self.interpreter.get_tensor(
            self.output_details[3]['index']
        )[0]
        
        return {
            'boxes': detection_boxes,
            'classes': detection_classes,
            'scores': detection_scores,
            'text': ocr_text
        }
    
    def read_gauge(self, image, confidence_threshold=0.5):
        """Extract gauge reading from image"""
        results = self.predict(image)
        
        # Filter by confidence
        valid_detections = results['scores'] > confidence_threshold
        
        if not np.any(valid_detections):
            return None
        
        # Get best detection
        best_idx = np.argmax(results['scores'])
        
        if results['scores'][best_idx] < confidence_threshold:
            return None
        
        # Extract reading from OCR text
        reading_text = results['text'][best_idx].decode('utf-8')
        
        try:
            # Parse numeric value
            reading_value = float(reading_text.strip())
            
            return {
                'value': reading_value,
                'text': reading_text,
                'confidence': float(results['scores'][best_idx]),
                'bbox': results['boxes'][best_idx].tolist()
            }
        except ValueError:
            return None

# Usage example
def main():
    # Initialize edge AI model
    model_path = "gauge_reader_quantized.tflite"
    gauge_reader = EdgeGaugeReader(model_path)
    
    # Initialize camera
    cap = cv2.VideoCapture(0)
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # Read gauge value
        result = gauge_reader.read_gauge(frame)
        
        if result:
            # Display result
            value = result['value']
            confidence = result['confidence']
            
            cv2.putText(frame, f"Reading: {value:.2f}", 
                       (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                       1, (0, 255, 0), 2)
            
            cv2.putText(frame, f"Confidence: {confidence:.2f}", 
                       (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 
                       1, (0, 255, 0), 2)
            
            # Draw bounding box
            bbox = result['bbox']
            h, w = frame.shape[:2]
            x1, y1, x2, y2 = [int(coord * dim) for coord, dim 
                             in zip(bbox, [w, h, w, h])]
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        # Display frame
        cv2.imshow('Edge AI Gauge Reader', frame)
        
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()