Restructure project sources.

src/fmri/Simulator.cpp

@@ -0,0 +1,233 @@
+#include <cassert>
+#include <iostream>
+#include <optional>
+#include <vector>
+
+#include <caffe/caffe.hpp>
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+
+#include "Simulator.hpp"
+#include "Range.hpp"
+
+using namespace caffe;
+using namespace std;
+using namespace fmri;
+
+struct Simulator::Impl
+{
+    caffe::Net<DType> net;
+    cv::Size input_geometry;
+    optional<cv::Mat> means;
+    unsigned int num_channels;
+    map<string, LayerInfo> layerInfo_;
+
+    Impl(const string& model_file, const string& weights_file, const string& means_file);
+
+    vector<cv::Mat> getWrappedInputLayer();
+    cv::Mat preprocess(cv::Mat original) const;
+    vector<LayerData> simulate(const string &input_file);
+    const map<string, LayerInfo>& layerInfo() const;
+
+    void computeLayerInfo();
+
+    void loadMeans(const string &means_file);
+
+    void ensureNoInPlaceLayers();
+};
+
+// Create simple forwarding functions.
+Simulator::Simulator(const string& model_file, const string& weights_file, const string& means_file) :
+	pImpl(new Impl(model_file, weights_file, means_file))
+{
+}
+
+vector<LayerData> Simulator::simulate(const string& image_file)
+{
+    return pImpl->simulate(image_file);
+}
+
+Simulator::Impl::Impl(const string& model_file, const string& weights_file, const string& means_file) :
+	net(model_file, TEST)
+{
+	net.CopyTrainedLayersFrom(weights_file);
+    ensureNoInPlaceLayers();
+
+	auto input_layer = net.input_blobs()[0];
+	input_geometry = cv::Size(input_layer->width(), input_layer->height());
+	num_channels = input_layer->channels();
+
+	input_layer->Reshape(1, num_channels,
+			input_geometry.height, input_geometry.width);
+	/* Forward dimension change to all layers. */
+	net.Reshape();
+
+    if (!means_file.empty()) {
+        loadMeans(means_file);
+    }
+
+    computeLayerInfo();
+}
+
+void Simulator::Impl::loadMeans(const string &means_file)
+{// Read in the means file
+    BlobProto proto;
+    ReadProtoFromBinaryFileOrDie(means_file, &proto);
+
+    Blob<DType> mean_blob;
+    mean_blob.FromProto(proto);
+
+    CHECK_EQ(mean_blob.channels(), num_channels) << "Number of channels should match!" << endl;
+
+    vector<cv::Mat> channels;
+    float *data = mean_blob.mutable_cpu_data();
+    for (auto i : Range(num_channels)) {
+        (void)i;// Suppress unused warning
+        channels.emplace_back(mean_blob.height(), mean_blob.width(), CV_32FC1, data);
+        data += mean_blob.height() * mean_blob.width();
+    }
+
+    cv::Mat mean;
+    merge(channels, mean);
+
+    this->means = cv::Mat(input_geometry, mean.type(), cv::mean(mean));
+}
+
+vector<LayerData> Simulator::Impl::simulate(const string& image_file)
+{
+	cv::Mat im = cv::imread(image_file, -1);
+
+    assert(!im.empty());
+
+    auto input = preprocess(im);
+    auto channels = getWrappedInputLayer();
+
+    cv::split(input, channels);
+
+    net.Forward();
+
+	vector<LayerData> result;
+
+	const auto& names = net.layer_names();
+	const auto& results = net.top_vecs();
+
+	for (auto i : Range(names.size())) {
+		CHECK_EQ(results[i].size(), 1) << "Multiple outputs per layer are not supported!" << endl;
+		const auto blob = results[i][0];
+
+		result.emplace_back(names[i], blob->shape(), blob->cpu_data());
+	}
+
+    return result;
+}
+
+vector<cv::Mat> Simulator::Impl::getWrappedInputLayer()
+{
+    vector<cv::Mat> channels;
+    auto input_layer = net.input_blobs()[0];
+
+    const int width = input_geometry.width;
+    const int height = input_geometry.height;
+
+    DType* input_data = input_layer->mutable_cpu_data();
+    for (auto i : Range(num_channels)) {
+        (void)i;// Suppress unused warning
+        channels.emplace_back(height, width, CV_32FC1, input_data);
+        input_data += width * height;
+    }
+
+    return channels;
+}
+
+static cv::Mat fix_channels(const int num_channels, cv::Mat original) {
+    cv::Mat converted;
+
+    if (num_channels == 1 && original.channels() == 3) {
+        cv::cvtColor(original, converted, cv::COLOR_BGR2GRAY);
+    } else if (num_channels == 1 && original.channels() == 4) {
+        cv::cvtColor(original, converted, cv::COLOR_BGRA2GRAY);
+    } else if (num_channels == 3 && original.channels() == 1) {
+        cv::cvtColor(original, converted, cv::COLOR_GRAY2BGR);
+    } else if (num_channels == 3 && original.channels() == 4) {
+        cv::cvtColor(original, converted, cv::COLOR_BGRA2BGR);
+    } else {
+		CHECK(num_channels == original.channels()) << "Cannot convert between channel types. ";
+        return original;
+    }
+
+    return converted;
+}
+
+static cv::Mat resize(const cv::Size& targetSize, cv::Mat original)
+{
+    if (targetSize != original.size()) {
+        cv::Mat resized;
+        cv::resize(original, resized, targetSize);
+
+        return resized;
+    }
+
+    return original;
+}
+
+cv::Mat Simulator::Impl::preprocess(cv::Mat original) const
+{
+    auto converted = fix_channels(num_channels, std::move(original));
+
+    auto resized = resize(input_geometry, converted);
+
+    cv::Mat sample_float;
+    resized.convertTo(sample_float, num_channels == 3 ? CV_32FC3 : CV_32FC1);
+
+    if (!means) {
+        return sample_float;
+    }
+
+    cv::Mat normalized;
+    cv::subtract(sample_float, *means, normalized);
+
+    return normalized;
+
+}
+
+const map<string, LayerInfo> &Simulator::Impl::layerInfo() const
+{
+    return layerInfo_;
+}
+
+void Simulator::Impl::computeLayerInfo()
+{
+    const auto& names = net.layer_names();
+    const auto& layers = net.layers();
+
+    CHECK_EQ(names.size(), layers.size()) << "Size mismatch";
+
+    for (auto i : Range(names.size())) {
+        auto& layer = layers[i];
+        LayerInfo layerInfo(names[i], layer->type(), layer->blobs());
+        CHECK_NE(layerInfo.type(), LayerInfo::Type::Split) << "Split layers are not supported!";
+        layerInfo_.emplace(names[i], std::move(layerInfo));
+    }
+}
+
+void Simulator::Impl::ensureNoInPlaceLayers()
+{
+    auto blobList = net.top_vecs();
+    typeof(blobList) uniqueVecs;
+    unique_copy(blobList.begin(), blobList.end(), back_inserter(uniqueVecs));
+
+    LOG_IF(ERROR, blobList.size() != uniqueVecs.size())
+        << "Network file contains in-place layers, layer-state will not be accurate\n"
+        << "If accurate results are desired, see the deinplace script in tools." << endl;
+}
+
+Simulator::~Simulator()
+{
+    // Empty but defined constructor.
+}
+
+const map<string, LayerInfo> & Simulator::layerInfo() const
+{
+    return pImpl->layerInfo();
+}