Final net output and calculation of argmax in C++

[Ankit Dhall]

Hello,

I have been trying to get the output of the network using C++. It contains a 40 channel probability on a 300x300 image. I want to find the argmax of these per-pixel probabilities and get the result as a 1x300x300 image. I could only find few vague code snippets online and have been unable to get the desired result. 

I also tried the WrapInputLayer code which will split the channels but am getting a segmentation fault.

here is my code:

int main(int argc, char** argv) {

  //Setting CPU or GPU

    Caffe::set_mode(Caffe::GPU);

    int device_id = 0;

    Caffe::SetDevice(device_id);

  //get the net

  Net<float> caffe_test_net("/home/ubuntu/caffe_parsenet/models/upnet/rgb/model_definition.prototxt", caffe::TEST);

  //get trained net

  caffe_test_net.CopyTrainedLayersFrom("/home/ubuntu/caffe_parsenet/models/upnet/rgb/pretrained_model.caffemodel");

  //get datum

  Datum datum;

  if (!ReadImageToDatum("./b1-09517_Clipped.jpg", 1, 300, 300, &datum)) {

    LOG(ERROR) << "Error during file reading";

  }

  //get the blob

  Blob<float>* blob = new Blob<float>(1, datum.channels(), datum.height(), datum.width());

  //get the blobproto

  BlobProto blob_proto;

  blob_proto.set_num(1);

  blob_proto.set_channels(datum.channels());

  blob_proto.set_height(datum.height());

  blob_proto.set_width(datum.width());

  const int data_size = datum.channels() * datum.height() * datum.width();

  int size_in_datum = std::max<int>(datum.data().size(),

                                    datum.float_data_size());

  for (int i = 0; i < size_in_datum; ++i) {

    blob_proto.add_data(0.);

  }

 const string& data = datum.data();

  if (data.size() != 0) {

    for (int i = 0; i < size_in_datum; ++i) {

      blob_proto.set_data(i, blob_proto.data(i) + (uint8_t)data[i]);

    }

  }

  //set data into blob
blob->FromProto(blob_proto);

  //fill the vector

  vector<Blob<float>*> bottom;

  bottom.push_back(blob);

  float type = 0.0;

////////////////////////////////////////////////////////////////

const vector<Blob<float>*>& result = caffe_test_net.Forward(bottom, &type);  // forward pass

std::vector<cv::Mat>* input_channels;

Blob<float>* input_layer = caffe_test_net.output_blobs()[0];

  int width = input_layer->width();

  int height = input_layer->height();

  float* input_data = input_layer->mutable_cpu_data();

  for (int i = 0; i < input_layer->channels(); ++i) {

    cv::Mat channel(height, width, CV_32FC1, input_data);

    input_channels->push_back(channel);

    input_data += width * height;

  }
}

Could someone suggest a better way to get the argmax 1x300x300 image? Or help fix the segmentation fault?

Regards,

Ankit