boc/quixzoom-capture-pipeline/weak-supervision/pipeline.js

/**
 * QUIXZOOM Weak Supervision Pipeline
 * 
 * Kombinerar flera modeller för att ge första förslag:
 * - YOLOv8 (COCO)
 * - Grounding DINO (text-prompted detection)
 * - SAM2 (segmentation)
 * - OCR (text detection)
 * - Depth estimation (3D position)
 * 
 * Output: Förslag med confidence som människa kan korrigera
 */

const YOLODetector = require('../video-to-observation/yolo-detector');
const fs = require('fs');
const path = require('path');

class WeakSupervisionPipeline {
  constructor(config = {}) {
    this.config = {
      confidenceThreshold: config.confidenceThreshold || 0.5,
      iouThreshold: config.iouThreshold || 0.5,
      ...config,
    };
    
    this.yolo = new YOLODetector();
    this.models = {
      yolo: { weight: 0.3, available: this.yolo.available },
      groundingDINO: { weight: 0.3, available: false }, // Mock
      sam2: { weight: 0.2, available: false }, // Mock
      ocr: { weight: 0.1, available: false }, // Mock
      depth: { weight: 0.1, available: false }, // Mock
    };
  }

  /**
   * Hierarkiska etiketter för QUIXZOOM
   */
  static HIERARCHY = {
    infrastructure: {
      road: {
        asphalt: { id: 'road_asphalt', parent: 'road' },
        crack: { id: 'road_crack', parent: 'road' },
        pothole: { id: 'road_pothole', parent: 'road' },
        lane_marking: { id: 'road_lane_marking', parent: 'road' },
      },
      lighting: {
        street_lamp: { id: 'lighting_street_lamp', parent: 'lighting' },
        traffic_light: { id: 'lighting_traffic_light', parent: 'lighting' },
        flood_light: { id: 'lighting_flood_light', parent: 'lighting' },
      },
      utility: {
        electrical_cabinet: { id: 'utility_electrical_cabinet', parent: 'utility' },
        manhole: { id: 'utility_manhole', parent: 'utility' },
        drain: { id: 'utility_drain', parent: 'utility' },
        hydrant: { id: 'utility_hydrant', parent: 'utility' },
      },
      signage: {
        stop: { id: 'signage_stop', parent: 'signage' },
        speed: { id: 'signage_speed', parent: 'signage' },
        direction: { id: 'signage_direction', parent: 'signage' },
        warning: { id: 'signage_warning', parent: 'signage' },
      },
      vegetation: {
        tree: { id: 'vegetation_tree', parent: 'vegetation' },
        bush: { id: 'vegetation_bush', parent: 'vegetation' },
        grass: { id: 'vegetation_grass', parent: 'vegetation' },
      },
      furniture: {
        bench: { id: 'furniture_bench', parent: 'furniture' },
        trash_can: { id: 'furniture_trash_can', parent: 'furniture' },
        bike_rack: { id: 'furniture_bike_rack', parent: 'furniture' },
      },
    },
  };

  /**
   * Kör alla modeller på en bild
   */
  async detect(imagePath) {
    const results = {
      image: imagePath,
      proposals: [],
      modelResults: {},
    };

    // 1. YOLO
    if (this.models.yolo.available) {
      const yoloResults = await this.yolo.detect(imagePath);
      results.modelResults.yolo = yoloResults;
      this.addProposals(results.proposals, yoloResults, 'yolo', this.models.yolo.weight);
    }

    // 2. Grounding DINO (mock)
    const groundingResults = await this.mockGroundingDINO(imagePath);
    results.modelResults.groundingDINO = groundingResults;
    this.addProposals(results.proposals, groundingResults, 'groundingDINO', this.models.groundingDINO.weight);

    // 3. SAM2 (mock)
    const samResults = await this.mockSAM2(imagePath);
    results.modelResults.sam2 = samResults;
    this.addProposals(results.proposals, samResults, 'sam2', this.models.sam2.weight);

    // 4. Konsolidera förslag
    const consolidated = this.consolidateProposals(results.proposals);
    results.consolidated = consolidated;

    return results;
  }

  /**
   * Lägg till förslag från modell
   */
  addProposals(proposals, detections, modelName, weight) {
    for (const det of detections) {
      proposals.push({
        type: det.type,
        bbox: det.bbox,
        confidence: det.confidence * weight,
        model: modelName,
        position: det.position,
        size: det.size,
      });
    }
  }

  /**
   * Konsolidera förslag från flera modeller
   */
  consolidateProposals(proposals) {
    // Gruppera efter typ
    const byType = {};
    for (const prop of proposals) {
      if (!byType[prop.type]) byType[prop.type] = [];
      byType[prop.type].push(prop);
    }

    const consolidated = [];

    for (const [type, typeProposals] of Object.entries(byType)) {
      // Sortera efter confidence
      typeProposals.sort((a, b) => b.confidence - a.confidence);

      // Non-maximum suppression
      const selected = [];
      for (const prop of typeProposals) {
        let overlap = false;
        for (const sel of selected) {
          const iou = this.calculateIoU(prop.bbox, sel.bbox);
          if (iou > this.config.iouThreshold) {
            overlap = true;
            // Öka confidence om flera modeller håller med
            sel.confidence = Math.min(1, sel.confidence + prop.confidence * 0.3);
            sel.supportingModels = sel.supportingModels || [sel.model];
            sel.supportingModels.push(prop.model);
            break;
          }
        }
        if (!overlap) {
          selected.push({ ...prop, supportingModels: [prop.model] });
        }
      }

      consolidated.push(...selected);
    }

    // Sortera efter confidence
    consolidated.sort((a, b) => b.confidence - a.confidence);

    return consolidated;
  }

  /**
   * Beräkna IoU (Intersection over Union)
   */
  calculateIoU(bbox1, bbox2) {
    const x1 = Math.max(bbox1[0], bbox2[0]);
    const y1 = Math.max(bbox1[1], bbox2[1]);
    const x2 = Math.min(bbox1[2], bbox2[2]);
    const y2 = Math.min(bbox1[3], bbox2[3]);

    const intersection = Math.max(0, x2 - x1) * Math.max(0, y2 - y1);
    const area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1]);
    const area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1]);
    const union = area1 + area2 - intersection;

    return union > 0 ? intersection / union : 0;
  }

  /**
   * Mock: Grounding DINO
   * I verkligheten: text-prompted detection
   */
  async mockGroundingDINO(imagePath) {
    // Simulera detektion med text-prompt
    const frameNum = parseInt(path.basename(imagePath).match(/\d+/)?.[0] || '0');
    const objects = [];

    // Grounding DINO är bra på specifika objekt
    const prompts = [
      { type: 'street_lamp', text: 'street lamp post light', prob: 0.6 },
      { type: 'traffic_sign', text: 'traffic sign road sign', prob: 0.5 },
      { type: 'manhole', text: 'manhole cover sewer', prob: 0.3 },
      { type: 'electrical_cabinet', text: 'electrical box utility cabinet', prob: 0.2 },
    ];

    for (const prompt of prompts) {
      if (Math.random() < prompt.prob) {
        objects.push({
          type: prompt.type,
          confidence: 0.5 + Math.random() * 0.4,
          bbox: [100 + frameNum * 10, 200, 300 + frameNum * 10, 500],
          position: { x: 0.3, y: 0.5 },
          size: { width: 200, height: 300 },
        });
      }
    }

    return objects;
  }

  /**
   * Mock: SAM2 (Segment Anything Model 2)
   * I verkligheten: segmentering av objekt
   */
  async mockSAM2(imagePath) {
    // SAM2 ger exakta segmenteringsmasker
    const frameNum = parseInt(path.basename(imagePath).match(/\d+/)?.[0] || '0');
    const objects = [];

    // SAM2 är bra på att segmentera oregelbundna former
    const segmentable = [
      { type: 'tree', prob: 0.7 },
      { type: 'bush', prob: 0.5 },
      { type: 'pothole', prob: 0.4 },
      { type: 'crack', prob: 0.3 },
    ];

    for (const item of segmentable) {
      if (Math.random() < item.prob) {
        objects.push({
          type: item.type,
          confidence: 0.6 + Math.random() * 0.3,
          bbox: [50 + frameNum * 15, 150, 250 + frameNum * 15, 450],
          position: { x: 0.2, y: 0.4 },
          size: { width: 200, height: 300 },
          mask: true, // SAM2 ger mask
        });
      }
    }

    return objects;
  }

  /**
   * Exportera förslag för human review
   */
  exportForReview(results, outputPath) {
    const reviewData = {
      version: '1.0',
      generatedAt: new Date().toISOString(),
      image: results.image,
      proposals: results.consolidated.map(prop => ({
        id: `prop_${Math.random().toString(36).substr(2, 9)}`,
        type: prop.type,
        typeLabel: this.getTypeLabel(prop.type),
        bbox: prop.bbox,
        confidence: prop.confidence,
        model: prop.model,
        supportingModels: prop.supportingModels || [prop.model],
        needsReview: prop.confidence < 0.7,
        status: 'pending', // pending, approved, rejected, corrected
      })),
    };

    fs.writeFileSync(outputPath, JSON.stringify(reviewData, null, 2));
    return reviewData;
  }

  /**
   * Hämta hierarkisk etikett
   */
  getTypeLabel(type) {
    const labels = {
      street_lamp: 'Gatlykta',
      traffic_sign: 'Trafikskylt',
      tree: 'Träd',
      manhole: 'Brunn',
      electrical_cabinet: 'Elskåp',
      bench: 'Bänk',
      trash_can: 'Papperskorg',
      bike_rack: 'Cykelställ',
      pothole: 'Hål i väg',
      crack: 'Spricka',
      asphalt: 'Asfalt',
      lane_marking: 'Vägmarkering',
      traffic_light: 'Trafikljus',
      flood_light: 'Strålkastare',
      drain: 'Avlopp',
      hydrant: 'Brandpost',
      stop: 'Stoppskylt',
      speed: 'Hastighetsskylt',
      direction: 'Vägvisare',
      warning: 'Varningsskylt',
      bush: 'Buske',
      grass: 'Gräs',
    };
    return labels[type] || type;
  }

  /**
   * Hämta förälder-typ
   */
  getParentType(type) {
    for (const [category, types] of Object.entries(WeakSupervisionPipeline.HIERARCHY.infrastructure)) {
      if (types[type]) return category;
    }
    return null;
  }
}

module.exports = WeakSupervisionPipeline;