envelopecnn.cpp

#include "envelopecnn.h"

using namespace std;
using namespace tensorflow;


EnvelopeCNN::EnvelopeCNN(string modelpath, string graphpath) {
    SessionOptions options = SessionOptions();
    ConfigProto* config = &options.config;
    (*config->mutable_device_count())["GPU"] = 0;
    config->mutable_gpu_options()->set_visible_device_list("");
    config->set_allow_soft_placement(true);

    status = NewSession(options, &session);
    if (status != tensorflow::Status::OK()) {cout << status << endl;}
    data = Tensor(DT_FLOAT, TensorShape({1, 7, 96, 2, 2}));
    status = LoadSavedModel(options, RunOptions(), modelpath, {kSavedModelTagServe}, &model);
    if (status != tensorflow::Status::OK()) {cout << status << endl;}
    status = ReadBinaryProto(Env::Default(), graphpath, &graph_def);
    if (status != tensorflow::Status::OK()) {cout << status << endl;}
    status = session->Create(graph_def);
    if (status != tensorflow::Status::OK()) {cout << status << endl;}
}

void EnvelopeCNN::predict(float**** inp) {
    Tensor data = Tensor(DT_FLOAT, TensorShape({1, 7, 96, 2, 2}));
    float* data_ = dat.flat<float>().data();
    unsigned int count = 0;
    for(unsigned int a = 0; a < 7; a++) {
        for(unsigned int b = 0; b < 96; b++) {
            data_[count] = inp[a][b][0][0];
            data_[count+1] = inp[a][b][0][1];
            data_[count+2] = inp[a][b][1][0];
            data_[count+3] = inp[a][b][0][1];
            count+=4;
        }
    }

    tensor_dict feed_dict = {
            {"input", data},
    };

    std::vector<tensorflow::Tensor> outputs;
    status = session->Run(feed_dict, {"output0/Reshape", "output1/Reshape", "output2/Reshape", "output3/Reshape"}, {}, &outputs);
    if (status != tensorflow::Status::OK()) {cout << "SHIT HAPPENED:\n" << status << endl;}
}