images/image-idx.c

   1 #define LOCAL_DEBUG
   2
   3 #include "sherlock/sherlock.h"
   4 #include "lib/mempool.h"
   5 #include "lib/conf.h"
   6 #include "lib/fastbuf.h"
   7 #include "lib/chartype.h"
   8 #include "sherlock/object.h"
   9 #include "lib/url.h"
  10 #include "lib/unicode.h"
  11 #include "sherlock/lizard-fb.h"
  12 #include "sherlock/tagged-text.h"
  13 #include "charset/charconv.h"
  14 #include "charset/unicat.h"
  15 #include "charset/fb-charconv.h"
  16 #include "indexer/indexer.h"
  17 #include "indexer/lexicon.h"
  18 #include "indexer/params.h"
  19 #include "utils/dumpconfig.h"
  20 #include "lang/lang.h"
  21 #include "lib/base224.h"
  22 #include "lib/bbuf.h"
  23
  24 #include "images/images.h"
  25 #include "images/image-obj.h"
  26 #include "images/image-sig.h"
  27 #include "images/image-dup.h"
  28
  29 #include <stdlib.h>
  30 #include <fcntl.h>
  31 #include <string.h>
  32
  33 /* This should happen in gatherer or scanner */
  34 static void
  35 generate_signatures(uns limit)
  36 {
  37   struct fastbuf *fb_cards = index_bopen("cards", O_RDONLY);
  38   struct fastbuf *fb_card_attrs = index_bopen("card-attrs", O_RDONLY);
  39   struct fastbuf *fb_signatures = index_bopen("image-sig", O_CREAT | O_WRONLY | O_TRUNC);
  40   struct card_attr ca;
  41   struct image_signature sig;
  42   struct mempool *pool = mp_new(1 << 16);
  43   struct buck2obj_buf *bob = buck2obj_alloc();
  44   uns count = 0;
  45
  46   log(L_INFO, "Generating image signatures");
  47   bputl(fb_signatures, 0);
  48   imo_decompress_thumbnails_init();
  49
  50   for (oid_t oid = 0; bread(fb_card_attrs, &ca, sizeof(ca)); oid++)
  51     if ((uns)((ca.type_flags >> 4) - 8) < 4)
  52       {
  53         bsetpos(fb_cards, (sh_off_t)ca.card << CARD_POS_SHIFT);
  54         uns buck_len = bgetl(fb_cards) - (LIZARD_COMPRESS_HEADER - 1);
  55         uns buck_type = bgetc(fb_cards) + BUCKET_TYPE_PLAIN;
  56         mp_flush(pool);
  57         struct odes *obj = obj_read_bucket(bob, pool, buck_type, buck_len, fb_cards, NULL);
  58         struct oattr *attr;
  59         if (!obj)
  60           die("Failed to read card");
  61         if (attr = obj_find_attr(obj, 'N'))
  62           {
  63             byte *url = obj_find_aval(obj_find_attr(obj, 'U' + OBJ_ATTR_SON)->son, 'U');
  64             DBG("Reading oid=%d url=%s", oid, url);
  65             struct image_obj imo;
  66             imo_init(&imo, pool, obj);
  67             if (imo_decompress_thumbnail(&imo))
  68               {
  69                 if (compute_image_signature(&imo.thumb, &sig))
  70                   {
  71                     bwrite(fb_signatures, &oid, sizeof(oid));
  72                     bwrite(fb_signatures, &sig.vec, sizeof(struct image_vector));
  73                     bputc(fb_signatures, sig.len);
  74                     if (sig.len)
  75                       bwrite(fb_signatures, sig.reg, sig.len * sizeof(struct image_region));
  76                     count++;
  77                     if (count >= limit)
  78                       break;
  79                   }
  80                 else
  81                   DBG("Cannot create signature");
  82               }
  83             else
  84               DBG("Cannot decompress thumbnail");
  85           }
  86       }
  87   brewind(fb_signatures);
  88   bputl(fb_signatures, count);
  89   DBG("%d signatures written", count);
  90
  91   imo_decompress_thumbnails_done();
  92   buck2obj_free(bob);
  93   mp_delete(pool);
  94   bclose(fb_cards);
  95   bclose(fb_card_attrs);
  96   bclose(fb_signatures);
  97 }
  98
  99 struct signature_record {
 100   oid_t oid;
 101   struct image_vector vec;
 102 };
 103
 104 #define ASORT_PREFIX(x) build_search_tree_##x
 105 #define ASORT_KEY_TYPE struct signature_record *
 106 #define ASORT_ELT(i) rec[i]
 107 #define ASORT_LT(x,y) x->vec.f[dim] < y->vec.f[dim]
 108 #define ASORT_EXTRA_ARGS , uns dim, struct signature_record **rec
 109 #include "lib/arraysort.h"
 110
 111 #if 0
 112 #define DBG_KD(x...) DBG(x)
 113 #else
 114 #define DBG_KD(x...) do{}while(0)
 115 #endif
 116
 117 static void
 118 build_search_tree(void)
 119 {
 120   log(L_INFO, "Building KD-tree");
 121
 122   struct fastbuf *fb_signatures = index_bopen("image-sig", O_RDONLY);
 123   struct image_tree tree;
 124   tree.count = bgetl(fb_signatures);
 125   ASSERT(tree.count < 0x80000000);
 126   if (!tree.count)
 127     {
 128       /* FIXME */
 129       bclose(fb_signatures);
 130       die("There are no signatures");
 131     }
 132   else
 133     {
 134       DBG("Reading %d signatures", tree.count);
 135       struct signature_record *records = xmalloc(tree.count * sizeof(struct signature_record));
 136       struct signature_record **precords = xmalloc(tree.count * sizeof(void *));
 137       for (uns i = 0; i < tree.count; i++)
 138         {
 139           bread(fb_signatures, &records[i].oid, sizeof(oid_t));
 140           bread(fb_signatures, &records[i].vec, sizeof(struct image_vector));
 141           uns len = bgetc(fb_signatures);
 142           bskip(fb_signatures, len * sizeof(struct image_region));
 143           precords[i] = records + i;
 144           if (likely(i))
 145             for (uns j = 0; j < IMAGE_VEC_K; j++)
 146               {
 147                 tree.bbox.vec[0].f[j] = MIN(tree.bbox.vec[0].f[j], records[i].vec.f[j]);
 148                 tree.bbox.vec[1].f[j] = MAX(tree.bbox.vec[1].f[j], records[i].vec.f[j]);
 149               }
 150           else
 151             tree.bbox.vec[0] = tree.bbox.vec[1] = records[0].vec;
 152         }
 153       bclose(fb_signatures);
 154
 155       for (tree.depth = 1; (uns)(2 << tree.depth) < tree.count; tree.depth++);
 156       DBG("depth=%d nodes=%d bbox=[(%s), (%s)]", tree.depth, 1 << tree.depth,
 157           stk_print_image_vector(tree.bbox.vec + 0), stk_print_image_vector(tree.bbox.vec + 1));
 158       uns leaves_index = 1 << (tree.depth - 1);
 159       tree.nodes = xmalloc_zero((1 << tree.depth) * sizeof(struct image_node));
 160       tree.leaves = xmalloc_zero(tree.count * sizeof(struct image_leaf));
 161
 162       /* Initialize recursion */
 163       struct stk {
 164         struct image_bbox bbox;
 165         uns index, count;
 166         struct signature_record **start;
 167       } stk_top[32], *stk = stk_top + 1;
 168       stk->index = 1;
 169       stk->start = precords;
 170       stk->count = tree.count;
 171       stk->bbox.vec[0] = tree.bbox.vec[0];
 172       for (uns i = 0; i < IMAGE_VEC_K; i++)
 173         stk->bbox.vec[1].f[i] = tree.bbox.vec[1].f[i] - tree.bbox.vec[0].f[i];
 174       uns entry_index = 0;
 175
 176       /* Main loop */
 177       while (stk != stk_top)
 178         {
 179           DBG_KD("Main loop... depth=%d index=%d count=%d, start=%d, min=%s dif=%s",
 180               stk - stk_top, stk->index, stk->count, stk->start - precords,
 181               stk_print_image_vector(stk->bbox.vec + 0), stk_print_image_vector(stk->bbox.vec + 1));
 182           ASSERT(stk->count);
 183
 184           /* Create leaf node */
 185           if (stk->index >= leaves_index || stk->count < 2)
 186             {
 187               tree.nodes[stk->index].val = IMAGE_NODE_LEAF | entry_index;
 188               for (; stk->count--; stk->start++)
 189                 {
 190                   struct image_leaf *leaf = &tree.leaves[entry_index++];
 191                   struct signature_record *record = *stk->start;
 192                   leaf->oid = record->oid;
 193                   leaf->flags = 0;
 194                   for (uns i = IMAGE_VEC_K; i--; )
 195                     {
 196                       uns bits = IMAGE_LEAF_BITS(i);
 197                       leaf->flags <<= bits;
 198                       if (stk->bbox.vec[1].f[i])
 199                         {
 200                           uns value =
 201                             (record->vec.f[i] - stk->bbox.vec[0].f[i]) *
 202                             ((1 << bits) - 1) / stk->bbox.vec[1].f[i];
 203                           ASSERT(value < (uns)(1 << bits));
 204                           leaf->flags |= value;
 205                         }
 206                     }
 207                   if (!stk->count)
 208                     leaf->flags |= IMAGE_LEAF_LAST;
 209                   DBG_KD("Creating leaf node; oid=%d vec=(%s) flags=0x%08x",
 210                       leaf->oid, stk_print_image_vector(&record->vec), leaf->flags);
 211                 }
 212               stk--;
 213             }
 214
 215           /* Create internal node */
 216           else
 217             {
 218               /* Select dimension to splis */
 219               uns dim = 0;
 220               for (uns i = 1; i < IMAGE_VEC_K; i++)
 221                 if (stk->bbox.vec[1].f[i] > stk->bbox.vec[1].f[dim])
 222                   dim = i;
 223
 224               /* Sort... FIXME: we only need the median */
 225               build_search_tree_sort(stk->count, dim, stk->start);
 226
 227               /* Split in the middle */
 228               uns index = stk->index;
 229               stk[1].index = stk[0].index * 2;
 230               stk[0].index = stk[1].index + 1;
 231               stk[1].count = stk[0].count >> 1;
 232               stk[0].count -= stk[1].count;
 233               stk[1].start = stk[0].start;
 234               stk[0].start += stk[1].count;
 235
 236               /* Choose split value */
 237               uns lval = stk->start[-1]->vec.f[dim];
 238               uns rval = stk->start[0]->vec.f[dim];
 239               uns pivot = stk->bbox.vec[0].f[dim] + (stk->bbox.vec[1].f[dim] >> 1);
 240               if (pivot <= lval)
 241                 pivot = lval;
 242               else if (pivot >= rval)
 243                 pivot = rval;
 244
 245               DBG_KD("Created internal node; dim=%d pivot=%d", dim, pivot);
 246
 247               /* Split the box */
 248               stk[1].bbox = stk[0].bbox;
 249               stk[1].bbox.vec[1].f[dim] = pivot - stk[0].bbox.vec[0].f[dim];
 250               stk[0].bbox.vec[0].f[dim] += stk[1].bbox.vec[1].f[dim];
 251               stk[0].bbox.vec[1].f[dim] -= stk[1].bbox.vec[1].f[dim];
 252
 253               /* Fill the node structure */
 254               tree.nodes[index].val = dim + (pivot << 8);
 255               stk++;
 256             }
 257         }
 258
 259       DBG("Tree constructed, saving...");
 260
 261       struct fastbuf *fb_tree = index_bopen("image-tree", O_CREAT | O_WRONLY | O_TRUNC);
 262       bputl(fb_tree, tree.count);
 263       bputl(fb_tree, tree.depth);
 264       bwrite(fb_tree, &tree.bbox, sizeof(struct image_bbox));
 265       bwrite(fb_tree, tree.nodes + 1, ((1 << tree.depth) - 1) * sizeof(struct image_node));
 266       bwrite(fb_tree, tree.leaves, tree.count * sizeof(struct image_leaf));
 267       bclose(fb_tree);
 268
 269       xfree(tree.leaves);
 270       xfree(tree.nodes);
 271       xfree(precords);
 272       xfree(records);
 273     }
 274 }
 275
 276
 277 static char *shortopts = CF_SHORT_OPTS "";
 278 static struct option longopts[] =
 279 {
 280   CF_LONG_OPTS
 281   { NULL, 0, 0, 0 }
 282 };
 283
 284 static char *help = "\
 285 Usage: image-indexer [<options>]\n\
 286 \n\
 287 Options:\n" CF_USAGE;
 288
 289 static void NONRET
 290 usage(byte *msg)
 291 {
 292   if (msg)
 293   {
 294     fputs(msg, stderr);
 295     fputc('\n', stderr);
 296   }
 297   fputs(help, stderr);
 298   exit(1);
 299 }
 300
 301
 302 int
 303 main(int argc UNUSED, char **argv)
 304 {
 305   int opt;
 306
 307   log_init(argv[0]);
 308   while ((opt = cf_getopt(argc, argv, shortopts, longopts, NULL)) >= 0)
 309     switch (opt)
 310     {
 311       default:
 312       usage("Invalid option");
 313     }
 314   if (optind != argc)
 315     usage("Invalid usage");
 316
 317   generate_signatures(~0U);
 318   build_search_tree();
 319
 320   return 0;
 321 }