images/kd-tree.c

   1 #undef LOCAL_DEBUG
   2
   3 #include "sherlock/sherlock.h"
   4 #include "lib/heap.h"
   5 #include "images/images.h"
   6 #include "images/image-sig.h"
   7 #include "images/kd-tree.h"
   8
   9 #include <alloca.h>
  10
  11 #define SQR(x) ((x)*(x))
  12 #define IMAGE_SEARCH_CMP(x,y) (is->buf[x].dist < is->buf[y].dist)
  13
  14 void
  15 image_search_init(struct image_search *is, struct image_tree *tree, struct image_vector *query, uns max_dist)
  16 {
  17   // FIXME: empty tree
  18   is->tree = tree;
  19   is->nodes = tree->nodes;
  20   is->leaves = tree->leaves;
  21   is->query = *query;
  22   is->max_dist = max_dist;
  23   is->size = 0x1000;
  24   is->buf = xmalloc((is->size + 1) * sizeof(struct image_search_item));
  25   is->heap = xmalloc((is->size + 1) * sizeof(u32));
  26   is->visited = is->count = 1;
  27   is->heap[1] = 1;
  28   struct image_search_item *item = is->buf + 1;
  29   item->index = 1;
  30   item->bbox = tree->bbox;
  31   item->dist = 0;
  32   for (uns i = 0; i < IMAGE_VEC_K; i++)
  33     {
  34       if (query->f[i] < item->bbox.vec[0].f[i])
  35         item->dist += SQR(item->bbox.vec[0].f[i] - query->f[i]);
  36       else if (query->f[i] > item->bbox.vec[1].f[i])
  37         item->dist += SQR(query->f[i] - item->bbox.vec[0].f[i]);
  38       else
  39         {
  40           item->dist = 0;
  41           break;
  42         }
  43     }
  44 }
  45
  46 void
  47 image_search_done(struct image_search *is)
  48 {
  49   xfree(is->buf);
  50   xfree(is->heap);
  51 }
  52
  53 static void
  54 image_search_grow_slow(struct image_search *is)
  55 {
  56   is->size *= 2;
  57   is->buf = xrealloc(is->buf, (is->size + 1) * sizeof(struct image_search_item));
  58   is->heap = xrealloc(is->heap, (is->size + 1) * sizeof(u32));
  59 }
  60
  61 static inline struct image_search_item *
  62 image_search_grow(struct image_search *is)
  63 {
  64   if (is->count == is->visited)
  65   {
  66     if (is->count == is->size)
  67       image_search_grow_slow(is);
  68     is->visited++;
  69     is->heap[is->visited] = is->visited;
  70   }
  71   return is->buf + is->heap[++is->count];
  72 }
  73
  74 static inline uns
  75 image_search_leaf_dist(struct image_search *is, struct image_bbox *bbox, struct image_leaf *leaf)
  76 {
  77   uns dist = 0;
  78   uns flags = leaf->flags;
  79   for (uns i = 0; i < IMAGE_VEC_K; i++)
  80     {
  81       uns bits = IMAGE_LEAF_BITS(i);
  82       uns mask = (1 << bits) - 1;
  83       uns value = flags & mask;
  84       flags >>= bits;
  85       int dif = bbox->vec[0].f[i] + (bbox->vec[1].f[i] - bbox->vec[0].f[i]) * value / ((1 << bits) - 1) - is->query.f[i];
  86       dist += dif * dif;
  87     }
  88   return dist;
  89 }
  90
  91 int
  92 image_search_next(struct image_search *is, oid_t *oid, uns *dist)
  93 {
  94   while (likely(is->count))
  95     {
  96       struct image_search_item *item = is->buf + is->heap[1];
  97       DBG("Main loop... dist=%d count=%d visited=%d size=%d index=0x%08x bbox=[(%s),(%s)]",
  98           item->dist, is->count, is->visited, is->size, item->index,
  99           stk_print_image_vector(&item->bbox.vec[0]), stk_print_image_vector(&item->bbox.vec[1]));
 100       if (unlikely(item->dist > is->max_dist))
 101         {
 102           DBG("Maximum distance reached");
 103           return 0;
 104         }
 105
 106       /* Expand leaf */
 107       if (item->index & IMAGE_SEARCH_ITEM_TYPE)
 108         {
 109           *oid = item->index & ~IMAGE_SEARCH_ITEM_TYPE;
 110           *dist = item->dist;
 111           DBG("Found item %d at distance %d", *oid, *dist);
 112           HEAP_DELMIN(u32, is->heap, is->count, IMAGE_SEARCH_CMP, HEAP_SWAP);
 113           return 1;
 114         }
 115
 116       /* Expand node with leaves */
 117       else if (is->nodes[item->index].val & IMAGE_NODE_LEAF)
 118         {
 119           DBG("Expanding node to list of leaves");
 120           struct image_leaf *leaf = is->leaves + (is->nodes[item->index].val & ~IMAGE_NODE_LEAF);
 121           item->dist = image_search_leaf_dist(is, &item->bbox, leaf);
 122           item->index = IMAGE_SEARCH_ITEM_TYPE | leaf->oid;
 123           HEAP_INCREASE(u32, is->heap, is->count, IMAGE_SEARCH_CMP, HEAP_SWAP, 1);
 124           while (!((leaf++)->flags & IMAGE_LEAF_LAST))
 125             {
 126               struct image_search_item *nitem = image_search_grow(is);
 127               nitem->dist = image_search_leaf_dist(is, &item->bbox, leaf);
 128               nitem->index = IMAGE_SEARCH_ITEM_TYPE | leaf->oid;
 129               HEAP_INSERT(u32, is->heap, is->count, IMAGE_SEARCH_CMP, HEAP_SWAP);
 130             }
 131         }
 132
 133       /* Expand internal node */
 134       else
 135         {
 136           DBG("Expanding internal node");
 137           struct image_search_item *nitem = image_search_grow(is);
 138           uns dim = is->nodes[item->index].val & IMAGE_NODE_DIM;
 139           uns pivot = is->nodes[item->index].val >> 8;
 140           item->index *= 2;
 141           nitem->bbox = item->bbox;
 142           nitem->dist = item->dist;
 143           uns query = is->query.f[dim];
 144           int dif = query - pivot;
 145           if (dif > 0)
 146             {
 147               nitem->index = item->index++;
 148               item->bbox.vec[0].f[dim] = pivot;
 149               nitem->bbox.vec[1].f[dim] = pivot;
 150               if (query > item->bbox.vec[1].f[dim])
 151                 nitem->dist -= SQR(query - item->bbox.vec[1].f[dim]);
 152             }
 153           else
 154             {
 155               nitem->index = item->index + 1;
 156               item->bbox.vec[1].f[dim] = pivot;
 157               nitem->bbox.vec[0].f[dim] = pivot;
 158               if (query < item->bbox.vec[0].f[dim])
 159                 nitem->dist -= SQR(item->bbox.vec[0].f[dim] - query);
 160             }
 161           nitem->dist += SQR(dif);
 162           HEAP_INSERT(u32, is->heap, is->count, IMAGE_SEARCH_CMP, HEAP_SWAP);
 163         }
 164     }
 165   DBG("Heap is empty");
 166   return 0;
 167 }
 168