Merge pull request #389 from biometrics/caffe

Caffe

Merge pull request #389 from biometrics/caffe
Caffe
JordanCheney
2 parents f17baf81 b5c16f94
Showing 5 changed files with 351 additions and 0 deletions
openbr/plugins/classification/caffe.cpp
openbr/plugins/cmake/caffe.cmake
openbr/plugins/gallery/lmdb.cpp
openbr/plugins/imgproc/pad.cpp
openbr/plugins/imgproc/roi.cpp
+#include <openbr/plugins/openbr_internal.h>
+#include <openbr/core/opencvutils.h>
+#include <openbr/core/qtutils.h>
+
+#include <opencv2/imgproc/imgproc.hpp>
+#include <caffe/caffe.hpp>
+
+using caffe::Caffe;
+using caffe::Net;
+using caffe::MemoryDataLayer;
+using caffe::Blob;
+using caffe::shared_ptr;
+
+using namespace cv;
+
+namespace br
+{
+
+// Net doesn't expose a default constructor which is expected by the default resource allocator.
+// To get around that we make this custom stub class which has a default constructor that passes
+// empty values to the Net constructor.
+class CaffeNet : public Net<float>
+{
+public:
+    CaffeNet() : Net<float>("", caffe::TEST) {}
+    CaffeNet(const QString &model, caffe::Phase phase) : Net<float>(model.toStdString(), phase) {}
+};
+
+class CaffeResourceMaker : public ResourceMaker<CaffeNet>
+{
+    QString model;
+    QString weights;
+    int gpuDevice;
+
+public:
+    CaffeResourceMaker(const QString &model, const QString &weights, int gpuDevice) : model(model), weights(weights), gpuDevice(gpuDevice) {}
+
+private:
+    CaffeNet *make() const
+    {
+        if (gpuDevice >= 0) {
+            Caffe::SetDevice(gpuDevice);
+            Caffe::set_mode(Caffe::GPU);
+        } else {
+            Caffe::set_mode(Caffe::CPU);
+        }
+
+        CaffeNet *net = new CaffeNet(model, caffe::TEST);
+        net->CopyTrainedLayersFrom(weights.toStdString());
+        return net;
+    }
+};
+
+/*!
+ * \brief A transform that wraps the Caffe deep learning library. This transform expects the input to a given Caffe model to be a MemoryDataLayer.
+ * The output of the Caffe network is treated as a feature vector and is stored in dst. Batch processing is possible. For a given batch size set in
+ * the memory data layer, src is expected to have an equal number of mats. Dst will always have the same size (number of mats) as src and the ordering
+ * will be preserved, so dst[1] is the output of src[1] after it passes through the neural net.
+ * \author Jordan Cheney \cite jcheney
+ * \br_property QString model path to prototxt model file
+ * \br_property QString weights path to caffemodel file
+ * \br_property int gpuDevice ID of GPU to use. gpuDevice < 0 runs on the CPU only.
+ * \br_link Caffe Integration Tutorial ../tutorials.md#caffe
+ * \br_link Caffe website http://caffe.berkeleyvision.org
+ */
+class CaffeFVTransform : public UntrainableTransform
+{
+    Q_OBJECT
+
+    Q_PROPERTY(QString model READ get_model WRITE set_model RESET reset_model STORED false)
+    Q_PROPERTY(QString weights READ get_weights WRITE set_weights RESET reset_weights STORED false)
+    Q_PROPERTY(int gpuDevice READ get_gpuDevice WRITE set_gpuDevice RESET reset_gpuDevice STORED false)
+    BR_PROPERTY(QString, model, "")
+    BR_PROPERTY(QString, weights, "")
+    BR_PROPERTY(int, gpuDevice, -1)
+
+    Resource<CaffeNet> caffeResource;
+
+    void init()
+    {
+        caffeResource.setResourceMaker(new CaffeResourceMaker(model, weights, gpuDevice));
+    }
+
+    bool timeVarying() const
+    {
+        return gpuDevice < 0 ? false : true;
+    }
+
+    void project(const Template &src, Template &dst) const
+    {
+        CaffeNet *net = caffeResource.acquire();
+
+        MemoryDataLayer<float> *data_layer = static_cast<MemoryDataLayer<float> *>(net->layers()[0].get());
+
+        if (src.size() != data_layer->batch_size())
+            qFatal("src should have %d (batch size) mats. It has %d mats.", data_layer->batch_size(), src.size());
+
+        dst.file = src.file;
+
+        data_layer->AddMatVector(src.toVector().toStdVector(), std::vector<int>(src.size(), 0));
+
+        Blob<float> *output = net->ForwardPrefilled()[1]; // index 0 is the labels from the data layer (in this case the 0 array we passed in above).
+                                                          // index 1 is the ouput of the final layer, which is what we want
+        int dim_features = output->count() / data_layer->batch_size();
+        for (int n = 0; n < data_layer->batch_size(); n++)
+            dst += Mat(1, dim_features, CV_32FC1, output->mutable_cpu_data() + output->offset(n));
+
+        caffeResource.release(net);
+    }
+};
+
+BR_REGISTER(Transform, CaffeFVTransform)
+
+} // namespace br
+
+#include "classification/caffe.moc"
+set(BR_WITH_CAFFE OFF CACHE BOOL "Build with Caffe")
+
+if(${BR_WITH_CAFFE})
+  find_package(Caffe)
+  include_directories(${Caffe_INCLUDE_DIRS})
+  add_definitions(${Caffe_DEFINITIONS})
+  set(BR_THIRDPARTY_LIBS ${BR_THIRDPARTY_LIBS} ${Caffe_LIBRARIES})
+else()
+  set(BR_EXCLUDED_PLUGINS ${BR_EXCLUDED_PLUGINS} plugins/classification/caffe.cpp)
+  set(BR_EXCLUDED_PLUGINS ${BR_EXCLUDED_PLUGINS} plugins/gallery/lmdb.cpp)
+endif()
+#include <openbr/openbr_plugin.h>
+#include "openbr/plugins/openbr_internal.h"
+#include <openbr/core/qtutils.h>
+
+#include <QFutureSynchronizer>
+#include <QtConcurrentRun>
+#include <QMutexLocker>
+#include <QWaitCondition>
+
+#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::DB>(caffe::db::GetDB("lmdb"));
+            db->Open(file.name.toStdString(),caffe::db::READ);
+            cursor = QSharedPointer<caffe::db::Cursor>(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<uchar>(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<caffe::db::DB> db;
+    QSharedPointer<caffe::db::Cursor> cursor;
+
+    QFutureSynchronizer<void> aThread;
+    QMutex dataLock;
+    QWaitCondition dataWait;
+
+    bool should_end;
+    TemplateList data;
+
+    QHash<QString, int> observedLabels;
+
+    static void commitLoop(lmdbGallery * base)
+    {
+        QSharedPointer<caffe::db::Transaction> 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<caffe::db::Transaction>(base->db->NewTransaction());
+
+            TemplateList working = base->data;
+            base->data.clear();
+
+            // no longer blocking dataLock
+            lock.unlock();
+
+            foreach(const Template &t, working) {
+                // add current image to transaction
+                caffe::Datum datum;
+                caffe::CVMatToDatum(t.m(), &datum);
+
+                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]);
+
+                std::string out;
+                datum.SerializeToString(&out);
+
+                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);
+
+                total_count++;
+            }
+
+            txn->Commit();
+            lock.relock();
+        }
+    }
+
+    void write(const Template &t)
+    {
+        if (!initialized) {
+            db = QSharedPointer<caffe::db::DB> (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 "gallery/lmdb.moc"
+
+#include <openbr/plugins/openbr_internal.h>
+
+using namespace cv;
+
+namespace br
+{
+
+class PadTransform : public UntrainableTransform
+{
+    Q_OBJECT
+
+    Q_PROPERTY(int padSize READ get_padSize WRITE set_padSize RESET reset_padSize STORED false)
+    Q_PROPERTY(int padValue READ get_padValue WRITE set_padValue RESET reset_padValue STORED false)
+    BR_PROPERTY(int, padSize, 0)
+    BR_PROPERTY(int, padValue, 0)
+
+    void project(const Template &src, Template &dst) const
+    {
+        dst.file = src.file;
+
+        foreach (const Mat &m, src) {
+            Mat padded = padValue * Mat::ones(m.rows + 2*padSize, m.cols + 2*padSize, m.type());
+            padded(Rect(padSize, padSize, padded.cols - padSize, padded.rows - padSize)) = m;
+            dst += padded;
+        }
+    }
+};
+
+BR_REGISTER(Transform, PadTransform)
+
+} // namespace br
+
+#include "imgproc/pad.moc"
@@ -26,12 +26,16 @@ namespace br
  * \ingroup transforms
  * \brief Crops the rectangular regions of interest.
  * \author Josh Klontz \cite jklontz
+ * \br_property QString propName Optional property name for a rectangle in metadata. If no propName is given the transform will use rects stored in the file.rects field or build a rectangle using "X", "Y", "Width", and "Height" fields if they exist.
+ * \br_property bool copyOnCrop If true make a clone of each crop before appending the crop to dst. This guarantees that the crops will be continuous in memory, which is an occasionally useful property. Default is false.
  */
 class ROITransform : public UntrainableTransform
 {
     Q_OBJECT
     Q_PROPERTY(QString propName READ get_propName WRITE set_propName RESET reset_propName STORED false)
+    Q_PROPERTY(bool copyOnCrop READ get_copyOnCrop WRITE set_copyOnCrop RESET reset_copyOnCrop STORED false)
     BR_PROPERTY(QString, propName, "")
+    BR_PROPERTY(bool, copyOnCrop, false)
     void project(const Template &src, Template &dst) const
     {
@@ -52,6 +56,10 @@ class ROITransform : public UntrainableTransform
                 qWarning("No rects present in file.");
         }
         dst.file.clearRects();
+
+        if (copyOnCrop)
+            for (int i = 0; i < dst.size(); i++)
+                dst.replace(i, dst[i].clone());
     }
 };