misc update

diff --git a/examples/test_net.cpp b/examples/test_net.cpp
index 30f87b7..f1af817 100644 (file)
--- a/examples/test_net.cpp
+++ b/examples/test_net.cpp
@@ -21,7 +21,7 @@ int main(int argc, char** argv) {
     return 0;
   }
 
-  cudaSetDevice(1);
+  cudaSetDevice(0);
   Caffe::set_phase(Caffe::TEST);
 
   if (argc == 5 && strcmp(argv[4], "GPU") == 0) {

diff --git a/python/caffe/imagenet/wrapper.py b/python/caffe/imagenet/wrapper.py
index 82402f0..ec1d409 100644 (file)
--- a/python/caffe/imagenet/wrapper.py
+++ b/python/caffe/imagenet/wrapper.py
@@ -57,7 +57,10 @@ def prepare_image(filename, center_only = False):
   img = io.imread(filename)
   if img.ndim == 2:
     img = np.tile(img[:, :, np.newaxis], (1, 1, 3))
-  img_reshape = transform.resize(img, (IMAGE_DIM,IMAGE_DIM)) * 255
+  elif img.shape[2] == 4:
+    img = img[:,:,:3]
+  # Resize and convert to BGR
+  img_reshape = (transform.resize(img, (IMAGE_DIM,IMAGE_DIM)) * 255)[:, :, ::-1]
   # subtract main
   img_reshape -= IMAGENET_MEAN
   return oversample(img_reshape, center_only)
@@ -101,6 +104,7 @@ def main(argv):
     print 'Use gpu.'
     net.caffenet.set_mode_gpu()
   files = glob.glob(os.path.join(FLAGS.root, "*." + FLAGS.ext))
+  files.sort()
   print 'A total of %d files' % len(files)
   output = np.empty((len(files), 1000), dtype=np.float32)
   start = time.time()
@@ -109,8 +113,8 @@ def main(argv):
     if i % 1000 == 0:
       print 'Processed %d files, elapsed %.2f s' % (i, time.time() - start)
   # Finally, write the results
-  np.save(FLAG.output, output)
-  print 'Done. Saved to %s.' % FLAGS.output 
+  np.save(FLAGS.output, output)
+  print 'Done. Saved to %s.' % FLAGS.output
 
 
 if __name__ == "__main__":

diff --git a/python/caffe/pycaffe.cpp b/python/caffe/pycaffe.cpp
index 21e3a52..32bba0d 100644 (file)
--- a/python/caffe/pycaffe.cpp
+++ b/python/caffe/pycaffe.cpp
@@ -1,4 +1,8 @@
 // Copyright Yangqing Jia 2013
+// pycaffe provides a wrapper of the caffe::Net class as well as some
+// caffe::Caffe functions so that one could easily call it from Python.
+// Note that for python, we will simply use float as the data type.
+
 #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
 
 #include <boost/python.hpp>
@@ -7,17 +11,20 @@
 
 
 using namespace caffe;
-using namespace boost::python;
+using boost::python::extract;
+using boost::python::len;
+using boost::python::list;
+using boost::python::object;
+
 
-// For python, we will simply use float.
 // A simple wrapper over CaffeNet that runs the forward process.
 struct CaffeNet
 {
   CaffeNet(string param_file, string pretrained_param_file,
-      boost::python::list bottom) {
+      list bottom) {
     vector<int> bottom_vec;
-    for (int i = 0; i < boost::python::len(bottom); ++i) {
-      bottom_vec.push_back(boost::python::extract<int>(bottom[i]));
+    for (int i = 0; i < len(bottom); ++i) {
+      bottom_vec.push_back(extract<int>(bottom[i]));
     }
     net_.reset(new Net<float>(param_file, bottom_vec));
     net_->CopyTrainedLayersFrom(pretrained_param_file);
@@ -25,21 +32,30 @@ struct CaffeNet
 
   virtual ~CaffeNet() {}
 
-  void Forward(boost::python::list bottom, boost::python::list top) {
+  inline void check_array_against_blob(
+      PyArrayObject* arr, Blob<float>* blob) {
+    CHECK(PyArray_FLAGS(arr) & NPY_ARRAY_C_CONTIGUOUS);
+    CHECK_EQ(PyArray_NDIM(arr), 4);
+    CHECK_EQ(PyArray_ITEMSIZE(arr), 4);
+    npy_intp* dims = PyArray_DIMS(arr);
+    CHECK_EQ(dims[0], blob->num());
+    CHECK_EQ(dims[1], blob->channels());
+    CHECK_EQ(dims[2], blob->height());
+    CHECK_EQ(dims[3], blob->width());
+  }
+
+  // The actual forward function. It takes in a python list of numpy arrays as
+  // input and a python list of numpy arrays as output. The input and output
+  // should all have correct shapes, are single-precisionabcdnt- and c contiguous.
+  void Forward(list bottom, list top) {
     vector<Blob<float>*>& input_blobs = net_->input_blobs();
-    CHECK_EQ(boost::python::len(bottom), input_blobs.size());
-    CHECK_EQ(boost::python::len(top), net_->num_outputs());
+    CHECK_EQ(len(bottom), input_blobs.size());
+    CHECK_EQ(len(top), net_->num_outputs());
     // First, copy the input
     for (int i = 0; i < input_blobs.size(); ++i) {
-      boost::python::object elem = bottom[i];
+      object elem = bottom[i];
       PyArrayObject* arr = reinterpret_cast<PyArrayObject*>(elem.ptr());
-      CHECK_EQ(PyArray_NDIM(arr), 4);
-      CHECK_EQ(PyArray_ITEMSIZE(arr), 4);
-      npy_intp* dims = PyArray_DIMS(arr);
-      CHECK_EQ(dims[0], input_blobs[i]->num());
-      CHECK_EQ(dims[1], input_blobs[i]->channels());
-      CHECK_EQ(dims[2], input_blobs[i]->height());
-      CHECK_EQ(dims[3], input_blobs[i]->width());
+      check_array_against_blob(arr, input_blobs[i]);
       switch (Caffe::mode()) {
       case Caffe::CPU:
         memcpy(input_blobs[i]->mutable_cpu_data(), PyArray_DATA(arr),
@@ -57,23 +73,16 @@ struct CaffeNet
     const vector<Blob<float>*>& output_blobs = net_->ForwardPrefilled();
     //LOG(INFO) << "End";
     for (int i = 0; i < output_blobs.size(); ++i) {
-      boost::python::object elem = top[i];
+      object elem = top[i];
       PyArrayObject* arr = reinterpret_cast<PyArrayObject*>(elem.ptr());
-      CHECK(PyArray_FLAGS(arr) & NPY_ARRAY_C_CONTIGUOUS);
-      CHECK_EQ(PyArray_NDIM(arr), 4);
-      CHECK_EQ(PyArray_ITEMSIZE(arr), 4);
-      npy_intp* dims = PyArray_DIMS(arr);
-      CHECK_EQ(dims[0], output_blobs[i]->num());
-      CHECK_EQ(dims[1], output_blobs[i]->channels());
-      CHECK_EQ(dims[2], output_blobs[i]->height());
-      CHECK_EQ(dims[3], output_blobs[i]->width());
+      check_array_against_blob(arr, output_blobs[i]);
       switch (Caffe::mode()) {
       case Caffe::CPU:
         memcpy(PyArray_DATA(arr), output_blobs[i]->cpu_data(),
             sizeof(float) * output_blobs[i]->count());
         break;
       case Caffe::GPU:
-        cudaMemcpy(PyArray_DATA(arr), output_blobs[i]->gpu_data(), 
+        cudaMemcpy(PyArray_DATA(arr), output_blobs[i]->gpu_data(),
             sizeof(float) * output_blobs[i]->count(), cudaMemcpyDeviceToHost);
         break;
       default:
@@ -81,20 +90,24 @@ struct CaffeNet
       }  // switch (Caffe::mode())
     }
   }
-  
+
+  // The caffe::Caffe utility functions.
   void set_mode_cpu() { Caffe::set_mode(Caffe::CPU); }
   void set_mode_gpu() { Caffe::set_mode(Caffe::GPU); }
-  
   void set_phase_train() { Caffe::set_phase(Caffe::TRAIN); }
   void set_phase_test() { Caffe::set_phase(Caffe::TEST); }
-
   void set_device(int device_id) { Caffe::SetDevice(device_id); }
+
+  // The pointer to the internal caffe::Net instant.
        shared_ptr<Net<float> > net_;
 };
 
+
+// The boost python module definition.
 BOOST_PYTHON_MODULE(pycaffe)
 {
-  class_<CaffeNet>("CaffeNet", init<string, string, boost::python::list>())
+  boost::python::class_<CaffeNet>(
+      "CaffeNet", boost::python::init<string, string, list>())
       .def("Forward", &CaffeNet::Forward)
       .def("set_mode_cpu", &CaffeNet::set_mode_cpu)
       .def("set_mode_gpu", &CaffeNet::set_mode_gpu)
@@ -102,4 +115,4 @@ BOOST_PYTHON_MODULE(pycaffe)
       .def("set_phase_test", &CaffeNet::set_phase_test)
       .def("set_device", &CaffeNet::set_device)
   ;
-}
\ No newline at end of file
+}