!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()));

}