diff --git a/openbr/plugins/classification/caffe.cpp b/openbr/plugins/classification/caffe.cpp index 3f393ce..e35577d 100644 --- a/openbr/plugins/classification/caffe.cpp +++ b/openbr/plugins/classification/caffe.cpp @@ -16,11 +16,46 @@ 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 +{ +public: + CaffeNet() : Net("", caffe::TEST) {} + CaffeNet(const QString &model, caffe::Phase phase) : Net(model.toStdString(), phase) {} +}; + +class CaffeResourceMaker : public ResourceMaker +{ + 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 + * 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 @@ -39,23 +74,22 @@ class CaffeFVTransform : public UntrainableTransform BR_PROPERTY(QString, weights, "") BR_PROPERTY(int, gpuDevice, -1) - QSharedPointer > net; + Resource caffeResource; void init() { - if (gpuDevice >= 0) { - Caffe::SetDevice(gpuDevice); - Caffe::set_mode(Caffe::GPU); - } else { - Caffe::set_mode(Caffe::CPU); - } + caffeResource.setResourceMaker(new CaffeResourceMaker(model, weights, gpuDevice)); + } - net.reset(new Net(model.toStdString(), caffe::TEST)); - net->CopyTrainedLayersFrom(weights.toStdString()); + bool timeVarying() const + { + return gpuDevice < 0 ? false : true; } void project(const Template &src, Template &dst) const { + CaffeNet *net = caffeResource.acquire(); + MemoryDataLayer *data_layer = static_cast *>(net->layers()[0].get()); if (src.size() != data_layer->batch_size()) @@ -70,6 +104,8 @@ class CaffeFVTransform : public UntrainableTransform 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); } };