Peter M. Groen / openbr

Commit 8eead12bdccf4b338f8a0ee9d4c20285d7baec8c

Authored by Charles Otto
1 parent c6ac2ae2

Preliminary lmdb gallery plugin

Inline Side-by-side
Showing 3 changed files with 227 additions and 0 deletions
openbr/plugins/cmake/lmdbGallery.cmake 0 → 100644
  View file @8eead12
  1 +set(BR_WITH_CAFFE OFF CACHE BOOL "Build with Caffe")
  2 +
  3 +if (${BR_WITH_CAFFE})
  4 + find_package(Caffe)
  5 + set(BR_THIRDPARTY_LIBS ${BR_THIRDPARTY_LIBS} ${Caffe_LIBRARIES})
  6 + include_directories(${Caffe_INCLUDE_DIRS})
  7 +else()
  8 + set(BR_EXCLUDED_PLUGINS ${BR_EXCLUDED_PLUGINS} plugins/gallery/lmdbGallery.cpp)
  9 +endif()
  10 +
... ...
openbr/plugins/gallery/lmdbGallery.cpp 0 → 100644
  View file @8eead12
  1 +#include <openbr/openbr_plugin.h>
  2 +#include "openbr/plugins/openbr_internal.h"
  3 +#include <QFutureSynchronizer>
  4 +#include <QtConcurrentRun>
  5 +#include <QMutexLocker>
  6 +#include <QWaitCondition>
  7 +
  8 +#include "caffe/util/db.hpp"
  9 +#include "caffe/util/io.hpp"
  10 +
  11 +using namespace br;
  12 +
  13 +class lmdbGallery : public Gallery
  14 +{
  15 + Q_OBJECT
  16 + TemplateList readBlock(bool *done)
  17 + {
  18 + *done = false;
  19 + if (!initialized) {
  20 + db = QSharedPointer<caffe::db::DB>(caffe::db::GetDB("lmdb"));
  21 + db->Open(file.name.toStdString(),caffe::db::READ);
  22 + cursor = QSharedPointer<caffe::db::Cursor>(db->NewCursor());
  23 + initialized = true;
  24 + }
  25 +
  26 + caffe::Datum datum;
  27 + datum.ParseFromString(cursor->value());
  28 +
  29 + cv::Mat img;
  30 + if (datum.encoded()) {
  31 + img = caffe::DecodeDatumToCVMatNative(datum);
  32 + }
  33 + else {
  34 + // create output image of appropriate size
  35 + img.create(datum.height(), datum.width(), CV_MAKETYPE(CV_8U, datum.channels()));
  36 + // copy matrix data from datum.
  37 + for (int h = 0; h < datum.height(); ++h) {
  38 + uchar* ptr = img.ptr<uchar>(h);
  39 + int img_index = 0;
  40 + for (int w = 0; w < datum.width(); ++w) {
  41 + for (int c = 0; c < datum.channels(); ++c) {
  42 + int datum_index = (c * datum.height() + h) * datum.width() + w;
  43 + ptr[img_index++] = (unsigned char) datum.data()[datum_index];
  44 + }
  45 + }
  46 + }
  47 + }
  48 +
  49 + // We acquired the image data, now decode filename from db key
  50 + QString baseKey = cursor->key().c_str();
  51 +
  52 + int idx = baseKey.indexOf("_");
  53 + if (idx != -1)
  54 + baseKey = baseKey.right(baseKey.size() - idx - 1);
  55 +
  56 + TemplateList output;
  57 + output.append(Template(img));
  58 + output.last().file.name = baseKey;
  59 + output.last().file.set("Label", datum.label());
  60 +
  61 + cursor->Next();
  62 +
  63 + if (!cursor->valid())
  64 + *done = true;
  65 +
  66 + return output;
  67 + }
  68 +
  69 + bool initialized;
  70 + QSharedPointer<caffe::db::DB> db;
  71 + QSharedPointer<caffe::db::Cursor> cursor;
  72 +
  73 + QFutureSynchronizer<void> aThread;
  74 + QMutex dataLock;
  75 + QWaitCondition dataWait;
  76 +
  77 + bool should_end;
  78 + TemplateList data;
  79 +
  80 +
  81 + QHash<QString,int> observedLabels;
  82 +
  83 + static void commitLoop(lmdbGallery * base)
  84 + {
  85 + QSharedPointer<caffe::db::Transaction> txn;
  86 +
  87 + int total_count = 0;
  88 +
  89 + // Acquire the lock
  90 + QMutexLocker lock(&base->dataLock);
  91 +
  92 + while (true)
  93 + {
  94 + // wait for data, or end signal
  95 + while(base->data.isEmpty() && !base->should_end)
  96 + base->dataWait.wait(&base->dataLock);
  97 +
  98 + // If should_end, but there is still data, we need another commit
  99 + // round
  100 + if (base->should_end && base->data.isEmpty())
  101 + break;
  102 +
  103 + txn = QSharedPointer<caffe::db::Transaction>(base->db->NewTransaction());
  104 +
  105 + TemplateList working = base->data;
  106 + base->data.clear();
  107 +
  108 + // no longer blocking dataLock
  109 + lock.unlock();
  110 +
  111 + foreach(const Template & t, working) {
  112 + qDebug() << "Converting t to datum";
  113 + // add current image to transaction
  114 + caffe::Datum datum;
  115 + caffe::CVMatToDatum(t.m(), &datum);
  116 +
  117 + qDebug() << "Did conversion";
  118 + QVariant base_label = t.file.value("Label");
  119 + QString label_str = base_label.toString();
  120 +
  121 +
  122 + if (!base->observedLabels.contains(label_str) )
  123 + base->observedLabels[label_str] = base->observedLabels.size();
  124 +
  125 + datum.set_label(base->observedLabels[label_str]);
  126 + qDebug() << "Did labels";
  127 + std::string out;
  128 + datum.SerializeToString(&out);
  129 + qDebug() << "Serialized";
  130 +
  131 + char key_cstr[256];
  132 + int len = snprintf(key_cstr, 256, "%08d_%s", total_count, qPrintable(t.file.name));
  133 + txn->Put(std::string(key_cstr, len), out);
  134 + qDebug() << "added to txn";
  135 + total_count++;
  136 +
  137 + }
  138 + qDebug() << "committing!";
  139 + txn->Commit();
  140 + qDebug() << "Finished commit";
  141 + lock.relock();
  142 + }
  143 +
  144 + }
  145 +
  146 + void write(const Template &t)
  147 + {
  148 + if (!initialized) {
  149 + db = QSharedPointer<caffe::db::DB> (caffe::db::GetDB("lmdb"));
  150 + db->Open(file.name.toStdString(), caffe::db::NEW);
  151 + observedLabels.clear();
  152 + initialized = true;
  153 + should_end = false;
  154 + // fire thread
  155 + aThread.clearFutures();
  156 + aThread.addFuture(QtConcurrent::run(lmdbGallery::commitLoop, this));
  157 + }
  158 +
  159 + QMutexLocker lock(&dataLock);
  160 + data.append(t);
  161 + dataWait.wakeAll();
  162 +
  163 + }
  164 +
  165 + ~lmdbGallery()
  166 + {
  167 + if (initialized) {
  168 + QMutexLocker lock(&dataLock);
  169 + should_end = true;
  170 + dataWait.wakeAll();
  171 + lock.unlock();
  172 +
  173 + aThread.waitForFinished();
  174 + }
  175 + }
  176 +
  177 +
  178 + void init()
  179 + {
  180 + initialized = false;
  181 + should_end = false;
  182 + }
  183 +};
  184 +
  185 +BR_REGISTER(Gallery, lmdbGallery)
  186 +
  187 +
  188 +#include "lmdbGallery.moc"
  189 +
... ...
share/openbr/cmake/FindLMDB.cmake 0 → 100644
  View file @8eead12
  1 +# Try to find the LMBD libraries and headers
  2 +# LMDB_FOUND - system has LMDB lib
  3 +# LMDB_INCLUDE_DIR - the LMDB include directory
  4 +# LMDB_LIBRARIES - Libraries needed to use LMDB
  5 +
  6 +# FindCWD based on FindGMP by:
  7 +# Copyright (c) 2006, Laurent Montel, <montel@kde.org>
  8 +#
  9 +# Redistribution and use is allowed according to the terms of the BSD license.
  10 +
  11 +# Adapted from FindCWD by:
  12 +# Copyright 2013 Conrad Steenberg <conrad.steenberg@gmail.com>
  13 +# Aug 31, 2013
  14 +
  15 +find_path(LMDB_INCLUDE_DIR NAMES lmdb.h PATHS "$ENV{LMDB_DIR}/include")
  16 +find_library(LMDB_LIBRARIES NAMES lmdb PATHS "$ENV{LMDB_DIR}/lib" )
  17 +
  18 +include(FindPackageHandleStandardArgs)
  19 +find_package_handle_standard_args(LMDB DEFAULT_MSG LMDB_INCLUDE_DIR LMDB_LIBRARIES)
  20 +
  21 +if(LMDB_FOUND)
  22 + message(STATUS "Found lmdb (include: ${LMDB_INCLUDE_DIR}, library: ${LMDB_LIBRARIES})")
  23 + mark_as_advanced(LMDB_INCLUDE_DIR LMDB_LIBRARIES)
  24 +
  25 + caffe_parse_header(${LMDB_INCLUDE_DIR}/lmdb.h
  26 + LMDB_VERSION_LINES MDB_VERSION_MAJOR MDB_VERSION_MINOR MDB_VERSION_PATCH)
  27 + set(LMDB_VERSION "${MDB_VERSION_MAJOR}.${MDB_VERSION_MINOR}.${MDB_VERSION_PATCH}")
  28 +endif()
... ...