Peter M. Groen / openbr

Commit 2849f3b9fa162aa797d7d5a717014538d8d37a8f

Authored by bhklein
1 parent 902b8bdf

add prefix to output, optionally save full distribution

Inline Side-by-side
Showing 1 changed file with 17 additions and 7 deletions
openbr/plugins/classification/caffe.cpp
  View file @2849f3b
... ... @@ -146,10 +146,16 @@ BR_REGISTER(Transform, CaffeFVTransform)
146 146 * \br_property QString model path to prototxt model file
147 147 * \br_property QString weights path to caffemodel file
148 148 * \br_property int gpuDevice ID of GPU to use. gpuDevice < 0 runs on the CPU only.
  149 + * \br_property QString prefix Prefix of the classifier output.
  150 + * \br_property bool saveDist Save the full distribution.
149 151 */
150 152 class CaffeClassifierTransform : public CaffeBaseTransform
151 153 {
152 154 Q_OBJECT
  155 + Q_PROPERTY(QString prefix READ get_prefix WRITE set_prefix RESET reset_prefix STORED false)
  156 + Q_PROPERTY(bool saveDist READ get_saveDist WRITE set_saveDist RESET reset_saveDist STORED false)
  157 + BR_PROPERTY(QString, prefix, "")
  158 + BR_PROPERTY(bool, saveDist, false)
153 159  
154 160 void project(const Template &src, Template &dst) const
155 161 {
... ... @@ -159,21 +165,25 @@ class CaffeClassifierTransform : public CaffeBaseTransform
159 165 dst = src;
160 166  
161 167 QList<int> labels; QList<float> confidences;
  168 + QList<QString> distributions;
162 169  
163 170 foreach (const Mat &m, caffeOutput) {
164   - double maxVal; int maxLoc;
165   - minMaxIdx(m, NULL, &maxVal, NULL, &maxLoc);
  171 + double maxVal; int maxLoc[2];
  172 + minMaxIdx(m, NULL, &maxVal, NULL, maxLoc);
166 173  
167   - labels.append(maxLoc);
  174 + labels.append(maxLoc[1]);
168 175 confidences.append(maxVal);
  176 + if (saveDist) distributions.append(OpenCVUtils::matrixToString(m));
169 177 }
170 178  
171 179 if (labels.size() == 1) {
172   - dst.file.set("Label", labels[0]);
173   - dst.file.set("Confidence", confidences[0]);
  180 + dst.file.set(prefix + "Label", labels[0]);
  181 + dst.file.set(prefix + "Confidence", confidences[0]);
  182 + if (saveDist) dst.file.set(prefix + "Distribution", distributions[0]);
174 183 } else {
175   - dst.file.setList<int>("Labels", labels);
176   - dst.file.setList<float>("Confidences", confidences);
  184 + dst.file.setList<int>(prefix + "Labels", labels);
  185 + dst.file.setList<float>(prefix + "Confidences", confidences);
  186 + if (saveDist) dst.file.setList<QString>(prefix + "Distributions", distributions);
177 187 }
178 188 }
179 189 };
... ...