boc/quixzoom-capture-pipeline/video-to-observation/yolo-detector.js

/**
 * QUIXZOOM YOLO Object Detector
 * 
 * Wrapper för YOLOv8/Ultralytics för objektdetektering i video.
 * Fallback till mock om YOLO inte är installerat.
 */

const { execSync } = require('child_process');
const fs = require('fs');
const path = require('path');

class YOLODetector {
  constructor(config = {}) {
    this.config = {
      model: config.model || 'yolov8n.pt',
      confidence: config.confidence || 0.5,
      device: config.device || 'cpu',
      ...config,
    };
    
    this.available = this.checkAvailability();
    this.classMap = this.buildClassMap();
  }

  /**
   * Kolla om YOLO är tillgängligt
   */
  checkAvailability() {
    try {
      execSync('python3 -c "import ultralytics"', { stdio: 'pipe' });
      console.log('[YOLO] Ultralytics available');
      return true;
    } catch {
      console.log('[YOLO] Ultralytics not available, using mock');
      return false;
    }
  }

  /**
   * Mappa YOLO-klasser till QUIXZOOM-objekttyper
   */
  buildClassMap() {
    return {
      // COCO-klasser → QUIXZOOM-typer
      'person': null, // Ignorera
      'bicycle': null,
      'car': null,
      'motorcycle': null,
      'traffic light': 'traffic_sign',
      'fire hydrant': 'utility_box',
      'stop sign': 'traffic_sign',
      'parking meter': 'utility_box',
      'bench': 'bench',
      'bird': null,
      'cat': null,
      'dog': null,
      'backpack': null,
      'umbrella': null,
      'handbag': null,
      'tie': null,
      'suitcase': null,
      'frisbee': null,
      'skis': null,
      'snowboard': null,
      'sports ball': null,
      'kite': null,
      'baseball bat': null,
      'baseball glove': null,
      'skateboard': null,
      'surfboard': null,
      'tennis racket': null,
      'bottle': null,
      'wine glass': null,
      'cup': null,
      'fork': null,
      'knife': null,
      'spoon': null,
      'bowl': null,
      'banana': null,
      'apple': null,
      'sandwich': null,
      'orange': null,
      'broccoli': null,
      'carrot': null,
      'hot dog': null,
      'pizza': null,
      'donut': null,
      'cake': null,
      'chair': null,
      'couch': null,
      'potted plant': 'tree',
      'bed': null,
      'dining table': null,
      'toilet': null,
      'tv': null,
      'laptop': null,
      'mouse': null,
      'remote': null,
      'keyboard': null,
      'cell phone': null,
      'microwave': null,
      'oven': null,
      'toaster': null,
      'sink': null,
      'refrigerator': null,
      'book': null,
      'clock': null,
      'vase': null,
      'scissors': null,
      'teddy bear': null,
      'hair drier': null,
      'toothbrush': null,
    };
  }

  /**
   * Detektera objekt i bild
   */
  async detect(imagePath) {
    if (!this.available) {
      return this.mockDetect(imagePath);
    }

    try {
      // Kör YOLO via Python
      const script = `
from ultralytics import YOLO
import json
import sys

model = YOLO('${this.config.model}')
results = model('${imagePath}', verbose=False)

detections = []
for r in results:
    for box in r.boxes:
        cls = int(box.cls)
        conf = float(box.conf)
        name = model.names[cls]
        xyxy = [float(x) for x in box.xyxy[0]]
        detections.append({
            'class': name,
            'confidence': conf,
            'bbox': xyxy
        })

print(json.dumps(detections))
`;
      
      const output = execSync(`python3 -c "${script}"`, {
        encoding: 'utf-8',
        timeout: 30000,
        stdio: ['pipe', 'pipe', 'pipe'],
      });
      
      const detections = JSON.parse(output.trim());
      return this.mapDetections(detections, imagePath);
      
    } catch (error) {
      console.warn(`[YOLO] Detection failed: ${error.message}`);
      return this.mockDetect(imagePath);
    }
  }

  /**
   * Mappa YOLO-detektioner till QUIXZOOM-format
   */
  mapDetections(detections, imagePath) {
    const objects = [];
    
    for (const det of detections) {
      const quixType = this.classMap[det.class];
      if (!quixType) continue; // Ignorera ointressanta klasser
      
      if (det.confidence < this.config.confidence) continue;
      
      // Beräkna relativ position i bilden
      const bbox = det.bbox;
      const centerX = (bbox[0] + bbox[2]) / 2;
      const centerY = (bbox[1] + bbox[3]) / 2;
      const imgWidth = 1920; // Antaget
      const imgHeight = 1080;
      
      const relX = centerX / imgWidth;
      const relY = centerY / imgHeight;
      
      objects.push({
        type: quixType,
        confidence: det.confidence,
        bbox: bbox,
        position: { x: relX, y: relY },
        size: {
          width: bbox[2] - bbox[0],
          height: bbox[3] - bbox[1],
        },
      });
    }
    
    return objects;
  }

  /**
   * Mock-detektion när YOLO inte är tillgängligt
   */
  mockDetect(imagePath) {
    // Använd bildens metadata för att generera realistiska mock-objekt
    const frameNum = parseInt(path.basename(imagePath).match(/\d+/)?.[0] || '0');
    const objects = [];
    
    // Bangkok-specifika objekt med varierande confidence
    const templates = [
      { type: 'street_lamp', prob: 0.6, confRange: [0.65, 0.95] },
      { type: 'traffic_sign', prob: 0.3, confRange: [0.55, 0.85] },
      { type: 'tree', prob: 0.5, confRange: [0.60, 0.90] },
      { type: 'utility_box', prob: 0.2, confRange: [0.50, 0.80] },
      { type: 'manhole', prob: 0.25, confRange: [0.45, 0.75] },
    ];
    
    // Generera 0-3 objekt per frame
    const numObjects = Math.floor(Math.random() * 3);
    
    for (let i = 0; i < numObjects; i++) {
      const template = templates[Math.floor(Math.random() * templates.length)];
      
      if (Math.random() < template.prob) {
        const confidence = template.confRange[0] + Math.random() * (template.confRange[1] - template.confRange[0]);
        
        objects.push({
          type: template.type,
          confidence: confidence,
          bbox: [100 + i * 200, 200, 300 + i * 200, 500],
          position: { x: 0.3 + i * 0.2, y: 0.5 },
          size: { width: 200, height: 300 },
        });
      }
    }
    
    return objects;
  }

  /**
   * Batch-detektion på flera bilder
   */
  async detectBatch(imagePaths) {
    const results = [];
    
    for (const path of imagePaths) {
      const detections = await this.detect(path);
      results.push({
        image: path,
        objects: detections,
      });
    }
    
    return results;
  }
}

module.exports = YOLODetector;

// Demo
if (require.main === module) {
  const detector = new YOLODetector();
  console.log('YOLO Detector ready');
  console.log('Available:', detector.available);
}