AI on Edge Devices

#!/usr/bin/env python3
"""
Real-time object detection using TensorFlow Lite on Raspberry Pi
"""

import cv2
import numpy as np
from tflite_runtime.interpreter import Interpreter
from picamera2 import Picamera2
import time

class ObjectDetector:
    def __init__(self, model_path, label_path, threshold=0.5):
        """Initialize TensorFlow Lite interpreter and load labels"""
        self.threshold = threshold
        self.interpreter = Interpreter(model_path=model_path)
        self.interpreter.allocate_tensors()
        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()
        with open(label_path, 'r') as f:
            self.labels = [line.strip() for line in f.readlines()]
        self.input_shape = self.input_details[0]['shape']
        self.height = self.input_shape[1]
        self.width = self.input_shape[2]

    def detect_objects(self, frame):
        """Run object detection on a frame"""
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame_resized = cv2.resize(frame_rgb, (self.width, self.height))
        input_data = np.expand_dims(frame_resized, axis=0)
        self.interpreter.set_tensor(self.input_details[0]['index'], input_data)
        self.interpreter.invoke()
        boxes = self.interpreter.get_tensor(self.output_details[0]['index'])[0]
        classes = self.interpreter.get_tensor(self.output_details[1]['index'])[0]
        scores = self.interpreter.get_tensor(self.output_details[2]['index'])[0]
        return boxes, classes, scores

    def draw_detections(self, frame, boxes, classes, scores):
        """Draw bounding boxes and labels on frame"""
        h, w = frame.shape[:2]
        for i in range(len(scores)):
            if scores[i] > self.threshold:
                ymin, xmin, ymax, xmax = boxes[i]
                left = int(xmin * w)
                top = int(ymin * h)
                right = int(xmax * w)
                bottom = int(ymax * h)
                label = self.labels[int(classes[i])]
                score = scores[i]
                cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)
                cv2.rectangle(frame, (left, top - 20), (left + 100, top), (0, 255, 0), -1)
                cv2.putText(frame, f"{label}: {score:.2f}", (left, top - 5), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
        return frame

def main():
    MODEL_PATH = "models/detect.tflite"
    LABEL_PATH = "models/coco_labels.txt"
    detector = ObjectDetector(MODEL_PATH, LABEL_PATH, threshold=0.5)
    picam2 = Picamera2()
    config = picam2.create_preview_configuration(
        main={"size": (640, 480), "format": "RGB888"}
    )
    picam2.configure(config)
    picam2.start()
    frame_count = 0
    start_time = time.time()
    try:
        while True:
            frame = picam2.capture_array()
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            boxes, classes, scores = detector.detect_objects(frame)
            frame = detector.draw_detections(frame, boxes, classes, scores)
            frame_count += 1
            if frame_count % 30 == 0:
                fps = frame_count / (time.time() - start_time)
                print(f"FPS: {fps:.1f}, Objects: {sum(scores > 0.5)}")
            cv2.imshow("Object Detection", frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    finally:
        picam2.stop()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

{
  "name": "edge-ai-server",
  "version": "1.0.0",
  "main": "server.js",
  "scripts": {
    "start": "node server.js"
  },
  "dependencies": {
    "express": "^4.18.2",
    "cors": "^2.8.5",
    "node-fetch": "^3.3.2",
    "socket.io": "^4.7.2"
  }
}

const express = require('express');
const cors = require('cors');
const http = require('http');
const { Server } = require('socket.io');
const tflite = require('tflite-node');

const app = express();
app.use(cors());
app.use(express.json());

const server = http.createServer(app);
const io = new Server(server, {
  cors: { origin: "*" }
});

const port = process.env.PORT || 3000;
const MODEL_PATH = process.env.MODEL_PATH || "/home/pi/edgeai/models/detect.tflite";
const LABEL_PATH = process.env.LABEL_PATH || "/home/pi/edgeai/models/coco_labels.txt";

let interpreter = null;
let labels = [];
let isModelLoaded = false;
let lastFPS = 0;

// Health endpoint
app.get('/health', (req, res) => {
  res.json({
    status: 'ok',
    modelLoaded: isModelLoaded,
    fps: lastFPS || 0,
    uptime: process.uptime()
  });
});

// Detection API endpoint
app.post('/api/detect', async (req, res) => {
  if (!isModelLoaded) {
    return res.status(503).json({ error: 'Model not loaded' });
  }
  try {
    const { image } = req.body;
    const results = await runDetection(image);
    res.json(results);
  } catch (error) {
    console.error('Detection error:', error);
    res.status(500).json({ error: error.message });
  }
});

// WebSocket for real-time streaming
io.on('connection', (socket) => {
  console.log('New client connected');
  socket.emit('modelStatus', { loaded: isModelLoaded });
  
  socket.on('disconnect', () => console.log('Client disconnected'));
});

server.listen(port, () => {
  console.log(`Edge AI Server running on port ${port}`);
  loadModel();
});

async function loadModel() {
  try {
    interpreter = await tflite.loadSession(MODEL_PATH);
    const labelResponse = await fetch(LABEL_PATH);
    const labelText = await labelResponse.text();
    labels = labelText.trim().split('\n');
    isModelLoaded = true;
    console.log('Model loaded successfully');
    io.emit('modelStatus', { loaded: true });
  } catch (error) {
    console.error('Failed to load model:', error);
  }
}

async function runDetection(imageData) {
  const buffer = Buffer.from(imageData.replace(/^data:image\/[a-z]+;base64,/, ''), 'base64');
  const results = await interpreter.run(buffer);
  return results.map(r => ({
    label: labels[r.classId] || 'unknown',
    confidence: r.confidence.toFixed(2),
    boundingBox: r.boundingBox
  }));
}