Caffe Image Classification C++

predict.cpp

#include <cuda_runtime.h>
 
#include <cstring>
#include <cstdlib>
#include <vector>
 
#include <string>
#include <iostream>
#include <stdio.h>
#include "caffe/caffe.hpp"
#include "caffe/util/io.hpp"
#include "caffe/blob.hpp"
 
using namespace caffe;
using namespace std;
 
int main(int argc, char** argv) {
 
    if (argc < 3 || argc > 5) {
        LOG(ERROR) << "test_net model [CPU/GPU] [Device ID]";
        return 1;
    }
    Caffe::set_phase(Caffe::TEST);
    
    //Setting CPU or GPU
    if (argc >= 4 && strcmp(argv[3], "GPU") == 0) {
        Caffe::set_mode(Caffe::GPU);
        int device_id = 0;
        if (argc == 5) {
            device_id = atoi(argv[4]);
        }
        Caffe::SetDevice(device_id);
        LOG(ERROR) << "Using GPU #" << device_id;
    } else {
        LOG(ERROR) << "Using CPU";
        Caffe::set_mode(Caffe::CPU);
    }
 
    //get the net
    Net<float> caffe_test_net(argv[1]);
    //get trained net
    caffe_test_net.CopyTrainedLayersFrom(argv[2]);

    // Run ForwardPrefilled 
    float loss;
    const vector<Blob<float>*>& result =  caffe_test_net.ForwardPrefilled(&loss);
    
    // Now result will contain the argmax results.
    const float* argmaxs = result[0]->cpu_data();
    for (int i = 0; i < result[0]->num(); ++i) {
        LOG(INFO) << " Image: "<< i << " class:" << argmaxs[i];
    }
  
    return 0;
}