From 96b3b695ab272f9ece4f7984d224a6038df53bad Mon Sep 17 00:00:00 2001
From: Bert Peters <bert@bertptrs.nl>
Date: Mon, 9 Oct 2017 14:42:30 +0200
Subject: [PATCH] Improve caffe detection function.

---
 CMakeLists.txt                |  4 +--
 cmake/modules/FindCaffe.cmake | 50 +++++++++++++++++++++++++++--------
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 578dddc..80a3451 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -30,8 +30,8 @@ target_link_libraries(fmri ${Boost_LIBRARIES})
 
 # Require Caffe
 find_package(Caffe REQUIRED)
-include_directories(${Caffe_INCLUDE_DIR})
-target_link_libraries(fmri ${Caffe_LIBS})
+include_directories(${Caffe_INCLUDE_DIRS})
+target_link_libraries(fmri ${Caffe_LIBRARIES})
 
 # Require glog
 find_package(Glog REQUIRED)
diff --git a/cmake/modules/FindCaffe.cmake b/cmake/modules/FindCaffe.cmake
index 21fa95c..e267356 100644
--- a/cmake/modules/FindCaffe.cmake
+++ b/cmake/modules/FindCaffe.cmake
@@ -1,14 +1,42 @@
-# Caffe package for CNN Triplet training
-unset(Caffe_FOUND)
+#
+# Try to find the Caffe library and include path.
+# Once done this will define
+#
+# CAFFE_FOUND
+# Caffe_INCLUDE_DIRS
+# Caffe_LIBRARIES
+#
+# A custom location can be specified for the library using Caffe_DIR
+#
 
-find_path(Caffe_INCLUDE_DIR NAMES caffe/caffe.hpp caffe/common.hpp caffe/net.hpp caffe/proto/caffe.pb.h caffe/util/io.hpp
-  HINTS
-  /usr/local/include)
+set (Caffe_DIR "" CACHE STRING
+  "Custom location of the root directory of a Caffe installation")
 
-find_library(Caffe_LIBS NAMES caffe
-  HINTS
-  /usr/local/lib)
+find_path(Caffe_INCLUDE_DIR
+  NAMES caffe.hpp
+  PATH_SUFFIXES include/caffe caffe
+  PATHS ${Caffe_DIR}
+  DOC "The directory where caffe.hpp resides")
 
-if(Caffe_LIBS AND Caffe_INCLUDE_DIR)
-    set(Caffe_FOUND 1)
-endif()
+find_library(Caffe_LIBRARY
+  NAMES caffe
+  PATH_SUFFIXES build/lib
+  PATHS ${Caffe_DIR}
+  DOC "The Caffe library")
+
+
+# handle the QUIETLY and REQUIRED arguments and set CAFFE_FOUND to TRUE if
+# all listed variables are TRUE
+include(FindPackageHandleStandardArgs)
+FIND_PACKAGE_HANDLE_STANDARD_ARGS(Caffe
+  "Could NOT find Caffe. Declare Caffe_DIR (either using -DCaffe_DIR=<path> flag or via CMakeGUI/ccmake) to point to root directory of library."
+  Caffe_LIBRARY Caffe_INCLUDE_DIR)
+
+if (CAFFE_FOUND)
+  set(Caffe_LIBRARIES ${Caffe_LIBRARY})
+  set(Caffe_INCLUDE_DIRS ${Caffe_INCLUDE_DIR})
+endif(CAFFE_FOUND)
+
+mark_as_advanced(
+  Caffe_INCLUDE_DIR
+  Caffe_LIBRARY)