!7

// multi-resolution image search
// uses only one prototype image

cmodule ImageRecogSpike {
  switchable S prototypeImageID = #1102884;
  switchable S inputImageID = #1102883;
  switchable bool showResizedProtos;
  switchable int numberOfCandidatesToShow = 3;
  switchable int nTimes = 10;

  // how big do we think the prototype is in the whole image
  // (percentage by height)
  switchable double assumedPrototypeHeightPercentage = 70;
  
  transient L<Int> widths = ll(1, 2, 4, 8, 16, 32, 64, 128);
  transient new L<IBWIntegralImage> integralImages;
  transient BWIntegralImage baseImage;
  transient JTabbedPane tabs = jTabbedPane();
  transient ImageSurface isResult;
  transient BWImage prototypeImage;
  transient new L<OneLevel> levels; // recognizers for each granularity
  
  transient new Best<Rect> allBest;
  transient S status;

  class OneLevel extends SteppableAndBest<Rect> {
    IBWIntegralImage ii; // scaled integral image
    BWImage image;       // scaled image
    BWImage prototype;   // scaled prototype
    float minSimilarity = 0.5f;
    ImageSurface is;
    new Map<Rect, Double> allScores;

    // candidates are top-left corner of rect to try in our coordinates
    L<Pt> candidatesQueue = syncLinkedList();
    new Set<Pt> candidatesTried;
    Iterator<Pt> candidatesStream;
    
    void reset {
      clear(allScores);
      best.clear();
      clear(candidatesQueue);
      clear(candidatesTried);
      makeCandidatesStream();
    }

    *(IBWIntegralImage *ii) {
      image = iBWIntegralImageToBWImage(ii);
      // get assumed height of prototype in scaled-down image
      int ph = max(1, iround(ii.getHeight()*assumedPrototypeHeightPercentage/100.0));
      // resize prototype
      prototype = bwResizeToHeightSmooth(prototypeImage, ph);
      if (showResizedProtos)
        addTab(tabs, "proto " + ii.getWidth(),
          jFullCenterScroll(jPixelatedZoomedImageSurface(4.0, prototype)));
      makeCandidatesStream();
    }
    
    void makeCandidatesStream {
      candidatesStream = mapI rectTopLeftCorner(allSubRectsOfSizeIterator(prototype.getWidth(), prototype.getHeight(),
        imageRect(image)));
    }
    
    int width() { ret ii.getWidth(); }
    int height() { ret ii.getHeight(); }
    
    public bool step() {
      Pt p = nextCandidate();
      if (p != null) ret true with tryCandidate(p);
      false;
    }

    Pt nextCandidate() {
      try object Pt p = popFirst(candidatesQueue);
      ret nextFromIterator(candidatesStream);
    }

    void tryCandidate(Pt p) {
      if (!candidatesTried.add(p)) ret;

      int x = p.x, y = p.y, wp = prototype.getWidth(), hp = prototype.getHeight();

      float maxError = (1f-minSimilarity)*wp*hp;
      float diff = bwImageSectionsSimilarity(image, prototype, x, y, maxError);
      if (diff <= maxError)
        putInBestAndMap(best, allScores, new Rect(x, y, wp, hp), 1-diff/(wp*hp);
    }

    void showBest() {
      overlaySelectionsOnImageSurface(is, nBest(numberOfCandidatesToShow));
    }
    
    L<Rect> nBest(int n) {
      ret topTenKeysByValue(n, allScores);
    }
  }
  
  start-thread {
    dm_watchField numberOfCandidatesToShow(r showBest);
    if (baseImage == null)
      baseImage = BWIntegralImage(loadImage2(inputImageID));
    addTab(tabs, "Result", isResult = jPixelatedZoomedImageSurface(1.0, iBWIntegralImageToBWImage(baseImage)));
    if (prototypeImage == null)
      prototypeImage = loadBWImage(prototypeImageID);
    for (int w : widths) {
      IBWIntegralImage ii = scaledIBWIntegralImage(baseImage, w);
      integralImages.add(ii);
      ImageSurface is = jPixelatedZoomedImageSurface(
        doubleRatio(baseImage.getWidth(), w), iBWIntegralImageToBWImage(ii));
      addTab(tabs, "w=" + w, jFullCenterScroll(is));
      levels.add(setAll(new OneLevel(ii), +is));
    }
    addTab(tabs, "Prototype", jFullCenterScroll(jPixelatedZoomedImageSurface(4, prototypeImage)));

    repeat nTimes {
      recogWithTiming();
    }
    showBest();
  }

  void recogWithTiming {  
    time "Recognition" { 
      recog();
    }
    setField(status := "Recognized in " + lastTiming() + " ms: " + allBest);
  }
  
  void recog {
    reset();
    //fullSearch();
    twoLevels(16, 128, 5);
  }
  
  void fullSearch {
    print("Steps: " + stepAll_roundRobin(levels));
    copyBestFromLevel(last(levels));
  }
  
  void copyBestFromLevel(OneLevel lvl) {
    if (lvl.best.has())
      allBest.put(scaleBetweenWidths(lvl.width(), baseImage.getWidth(), lvl.best!), lvl.best.score);
  }
  
  void twoLevels(int width1, int width2, int candidatesToUse) {
    int spiralSize = width2/width1*2;
    int pixelsPerSpiral = sqr(spiralSize);
    print("Spiral size: " + spiralSize + " / " + pixelsPerSpiral);
    
    int i1 = indexOf(widths, width1), i2 = indexOf(widths, width2);
    if (i1 < 0) fail("Width " + width1 + " not in list: " + widths);
    if (i2 < 0) fail("Width " + width2 + " not in list: " + widths);
    OneLevel lvl1 = levels.get(i1), lvl2 = levels.get(i2);
    
    // Full search in level 1
    stepAll(lvl1);
    
    // Circle around n best candidates found in level 1
    L<Pt> candidates = scaleBetweenLevels(lvl1, lvl2, map topLeftCorner(lvl1.nBest(candidatesToUse)));
    for (Pt p : combineIterators_roundRobin(map(candidates, c ->
      pixelSpiral(c.x, c.y, lvl2.width(), lvl2.height(), pixelsPerSpiral))))
      lvl2.tryCandidate(p);
      
    copyBestFromLevel(lvl2);
  }
  
  Rect scaleBetweenWidths(int w1, int w2, Rect r) {
    ret scaleRect(doubleRatio(w2, w1), r);
  }
  
  Pt scaleBetweenWidths(int w1, int w2, Pt p) {
    ret scalePt(doubleRatio(w2, w1), p);
  }
  
  Pt scaleBetweenLevels(OneLevel lvl1, OneLevel lvl2, Pt p) {
    ret scaleBetweenWidths(lvl1.width(), lvl2.width(), p);
  }
  
  L<Pt> scaleBetweenLevels(OneLevel lvl1, OneLevel lvl2, Iterable<Pt> l) {
    ret map(l, p -> scaleBetweenLevels(lvl1, lvl2, p));
  }

  void showBest {  
    for (OneLevel l : levels) l.showBest();
    setImageSurfaceSelection(isResult, allBest!);
    print("Best: " + allBest);
    print("Candidates tried: " + mapNonNulls(levels, l ->
      empty(l.candidatesTried) ? null : l(l.candidatesTried) + " for w=" + l.width()));
  }
  
  void reset {
    for (OneLevel l : levels) l.reset();
    allBest.clear();
  }

  visual withCenteredButtons(tabs,
    vstackWithSpacing(
      withLabel("# of candidates to show:", dm_fieldSpinner numberOfCandidatesToShow(1, 100)),
      dm_centeredLabel status()));
}