From 8eead12bdccf4b338f8a0ee9d4c20285d7baec8c Mon Sep 17 00:00:00 2001 From: Charles Otto Date: Fri, 26 Jun 2015 17:54:14 -0700 Subject: [PATCH] Preliminary lmdb gallery plugin --- openbr/plugins/cmake/lmdbGallery.cmake | 10 ++++++++++ openbr/plugins/gallery/lmdbGallery.cpp | 189 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ share/openbr/cmake/FindLMDB.cmake | 28 ++++++++++++++++++++++++++++ 3 files changed, 227 insertions(+), 0 deletions(-) create mode 100644 openbr/plugins/cmake/lmdbGallery.cmake create mode 100644 openbr/plugins/gallery/lmdbGallery.cpp create mode 100644 share/openbr/cmake/FindLMDB.cmake diff --git a/openbr/plugins/cmake/lmdbGallery.cmake b/openbr/plugins/cmake/lmdbGallery.cmake new file mode 100644 index 0000000..28dc467 --- /dev/null +++ b/openbr/plugins/cmake/lmdbGallery.cmake @@ -0,0 +1,10 @@ +set(BR_WITH_CAFFE OFF CACHE BOOL "Build with Caffe") + +if (${BR_WITH_CAFFE}) + find_package(Caffe) + set(BR_THIRDPARTY_LIBS ${BR_THIRDPARTY_LIBS} ${Caffe_LIBRARIES}) + include_directories(${Caffe_INCLUDE_DIRS}) +else() + set(BR_EXCLUDED_PLUGINS ${BR_EXCLUDED_PLUGINS} plugins/gallery/lmdbGallery.cpp) +endif() + diff --git a/openbr/plugins/gallery/lmdbGallery.cpp b/openbr/plugins/gallery/lmdbGallery.cpp new file mode 100644 index 0000000..8d41709 --- /dev/null +++ b/openbr/plugins/gallery/lmdbGallery.cpp @@ -0,0 +1,189 @@ +#include +#include "openbr/plugins/openbr_internal.h" +#include +#include +#include +#include + +#include "caffe/util/db.hpp" +#include "caffe/util/io.hpp" + +using namespace br; + +class lmdbGallery : public Gallery +{ + Q_OBJECT + TemplateList readBlock(bool *done) + { + *done = false; + if (!initialized) { + db = QSharedPointer(caffe::db::GetDB("lmdb")); + db->Open(file.name.toStdString(),caffe::db::READ); + cursor = QSharedPointer(db->NewCursor()); + initialized = true; + } + + caffe::Datum datum; + datum.ParseFromString(cursor->value()); + + cv::Mat img; + if (datum.encoded()) { + img = caffe::DecodeDatumToCVMatNative(datum); + } + else { + // create output image of appropriate size + img.create(datum.height(), datum.width(), CV_MAKETYPE(CV_8U, datum.channels())); + // copy matrix data from datum. + for (int h = 0; h < datum.height(); ++h) { + uchar* ptr = img.ptr(h); + int img_index = 0; + for (int w = 0; w < datum.width(); ++w) { + for (int c = 0; c < datum.channels(); ++c) { + int datum_index = (c * datum.height() + h) * datum.width() + w; + ptr[img_index++] = (unsigned char) datum.data()[datum_index]; + } + } + } + } + + // We acquired the image data, now decode filename from db key + QString baseKey = cursor->key().c_str(); + + int idx = baseKey.indexOf("_"); + if (idx != -1) + baseKey = baseKey.right(baseKey.size() - idx - 1); + + TemplateList output; + output.append(Template(img)); + output.last().file.name = baseKey; + output.last().file.set("Label", datum.label()); + + cursor->Next(); + + if (!cursor->valid()) + *done = true; + + return output; + } + + bool initialized; + QSharedPointer db; + QSharedPointer cursor; + + QFutureSynchronizer aThread; + QMutex dataLock; + QWaitCondition dataWait; + + bool should_end; + TemplateList data; + + + QHash observedLabels; + + static void commitLoop(lmdbGallery * base) + { + QSharedPointer txn; + + int total_count = 0; + + // Acquire the lock + QMutexLocker lock(&base->dataLock); + + while (true) + { + // wait for data, or end signal + while(base->data.isEmpty() && !base->should_end) + base->dataWait.wait(&base->dataLock); + + // If should_end, but there is still data, we need another commit + // round + if (base->should_end && base->data.isEmpty()) + break; + + txn = QSharedPointer(base->db->NewTransaction()); + + TemplateList working = base->data; + base->data.clear(); + + // no longer blocking dataLock + lock.unlock(); + + foreach(const Template & t, working) { + qDebug() << "Converting t to datum"; + // add current image to transaction + caffe::Datum datum; + caffe::CVMatToDatum(t.m(), &datum); + + qDebug() << "Did conversion"; + QVariant base_label = t.file.value("Label"); + QString label_str = base_label.toString(); + + + if (!base->observedLabels.contains(label_str) ) + base->observedLabels[label_str] = base->observedLabels.size(); + + datum.set_label(base->observedLabels[label_str]); + qDebug() << "Did labels"; + std::string out; + datum.SerializeToString(&out); + qDebug() << "Serialized"; + + char key_cstr[256]; + int len = snprintf(key_cstr, 256, "%08d_%s", total_count, qPrintable(t.file.name)); + txn->Put(std::string(key_cstr, len), out); + qDebug() << "added to txn"; + total_count++; + + } + qDebug() << "committing!"; + txn->Commit(); + qDebug() << "Finished commit"; + lock.relock(); + } + + } + + void write(const Template &t) + { + if (!initialized) { + db = QSharedPointer (caffe::db::GetDB("lmdb")); + db->Open(file.name.toStdString(), caffe::db::NEW); + observedLabels.clear(); + initialized = true; + should_end = false; + // fire thread + aThread.clearFutures(); + aThread.addFuture(QtConcurrent::run(lmdbGallery::commitLoop, this)); + } + + QMutexLocker lock(&dataLock); + data.append(t); + dataWait.wakeAll(); + + } + + ~lmdbGallery() + { + if (initialized) { + QMutexLocker lock(&dataLock); + should_end = true; + dataWait.wakeAll(); + lock.unlock(); + + aThread.waitForFinished(); + } + } + + + void init() + { + initialized = false; + should_end = false; + } +}; + +BR_REGISTER(Gallery, lmdbGallery) + + +#include "lmdbGallery.moc" + diff --git a/share/openbr/cmake/FindLMDB.cmake b/share/openbr/cmake/FindLMDB.cmake new file mode 100644 index 0000000..8a817fd --- /dev/null +++ b/share/openbr/cmake/FindLMDB.cmake @@ -0,0 +1,28 @@ +# Try to find the LMBD libraries and headers +# LMDB_FOUND - system has LMDB lib +# LMDB_INCLUDE_DIR - the LMDB include directory +# LMDB_LIBRARIES - Libraries needed to use LMDB + +# FindCWD based on FindGMP by: +# Copyright (c) 2006, Laurent Montel, +# +# Redistribution and use is allowed according to the terms of the BSD license. + +# Adapted from FindCWD by: +# Copyright 2013 Conrad Steenberg +# Aug 31, 2013 + +find_path(LMDB_INCLUDE_DIR NAMES lmdb.h PATHS "$ENV{LMDB_DIR}/include") +find_library(LMDB_LIBRARIES NAMES lmdb PATHS "$ENV{LMDB_DIR}/lib" ) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(LMDB DEFAULT_MSG LMDB_INCLUDE_DIR LMDB_LIBRARIES) + +if(LMDB_FOUND) + message(STATUS "Found lmdb (include: ${LMDB_INCLUDE_DIR}, library: ${LMDB_LIBRARIES})") + mark_as_advanced(LMDB_INCLUDE_DIR LMDB_LIBRARIES) + + caffe_parse_header(${LMDB_INCLUDE_DIR}/lmdb.h + LMDB_VERSION_LINES MDB_VERSION_MAJOR MDB_VERSION_MINOR MDB_VERSION_PATCH) + set(LMDB_VERSION "${MDB_VERSION_MAJOR}.${MDB_VERSION_MINOR}.${MDB_VERSION_PATCH}") +endif() -- libgit2 0.21.4