diff --git a/openbr/plugins/classification/caffe.cpp b/openbr/plugins/classification/caffe.cpp index 73d8aa1..87ed05a 100644 --- a/openbr/plugins/classification/caffe.cpp +++ b/openbr/plugins/classification/caffe.cpp @@ -146,10 +146,16 @@ BR_REGISTER(Transform, CaffeFVTransform) * \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_property QString prefix Prefix of the classifier output. + * \br_property bool saveDist Save the full distribution. */ class CaffeClassifierTransform : public CaffeBaseTransform { Q_OBJECT + Q_PROPERTY(QString prefix READ get_prefix WRITE set_prefix RESET reset_prefix STORED false) + Q_PROPERTY(bool saveDist READ get_saveDist WRITE set_saveDist RESET reset_saveDist STORED false) + BR_PROPERTY(QString, prefix, "") + BR_PROPERTY(bool, saveDist, false) void project(const Template &src, Template &dst) const { @@ -159,21 +165,25 @@ class CaffeClassifierTransform : public CaffeBaseTransform dst = src; QList labels; QList confidences; + QList distributions; foreach (const Mat &m, caffeOutput) { - double maxVal; int maxLoc; - minMaxIdx(m, NULL, &maxVal, NULL, &maxLoc); + double maxVal; int maxLoc[2]; + minMaxIdx(m, NULL, &maxVal, NULL, maxLoc); - labels.append(maxLoc); + labels.append(maxLoc[1]); confidences.append(maxVal); + if (saveDist) distributions.append(OpenCVUtils::matrixToString(m)); } if (labels.size() == 1) { - dst.file.set("Label", labels[0]); - dst.file.set("Confidence", confidences[0]); + dst.file.set(prefix + "Label", labels[0]); + dst.file.set(prefix + "Confidence", confidences[0]); + if (saveDist) dst.file.set(prefix + "Distribution", distributions[0]); } else { - dst.file.setList("Labels", labels); - dst.file.setList("Confidences", confidences); + dst.file.setList(prefix + "Labels", labels); + dst.file.setList(prefix + "Confidences", confidences); + if (saveDist) dst.file.setList(prefix + "Distributions", distributions); } } };