images/image-search.h

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