caffe: add Python support

Move to CMake in the process.

2 files changed, 41 insertions, 30 deletions

diff --git a/pkgs/applications/science/math/caffe/default.nix b/pkgs/applications/science/math/caffe/default.nix
index 9773c30d6d8d..5c6fe9c573d5 100644
--- a/pkgs/applications/science/math/caffe/default.nix
+++ b/pkgs/applications/science/math/caffe/default.nix
@@ -1,24 +1,26 @@
-{ stdenv
-, openblas
-, boost
-, cudaSupport ? true
-, cudnnSupport ? false
-, cudnn ? null
-, cudatoolkit
+{ stdenv, lib
 , fetchFromGitHub
+, cmake
+, boost
 , google-gflags
 , glog
-, hdf5
+, hdf5-cpp
 , leveldb
 , lmdb
 , opencv
 , protobuf
 , snappy
+, atlas
+, doxygen
+, cudaSupport ? true, cudatoolkit
+, cudnnSupport ? false, cudnn ? null
+, pythonSupport ? false, python ? null, numpy ? null
 }:
 
+assert cudnnSupport -> cudaSupport;
+assert pythonSupport -> (python != null && numpy != null);
 
-let optional = stdenv.lib.optional;
-in stdenv.mkDerivation rec {
+stdenv.mkDerivation rec {
   name = "caffe-${version}";
   version = "1.0-rc5";
 
@@ -29,32 +31,34 @@ in stdenv.mkDerivation rec {
     sha256 = "0lfmmc0n6xvkpygvxclzrvd0zigb4yfc5612anv2ahlxpfi9031c";
   };
 
-  preConfigure = "mv Makefile.config.example Makefile.config";
+  enableParallelBuilding = true;
 
-  makeFlags = [ "BLAS=open"
-                (if !cudaSupport then "CPU_ONLY=1" else "CUDA_DIR=${cudatoolkit}") ]
-              ++ optional cudnnSupport "USE_CUDNN=1";
+  nativeBuildInputs = [ cmake doxygen ];
 
-  # too many issues with tests to run them for now
-  doCheck = false;
-  checkTarget = "runtest";
+  cmakeFlags = [ "-DCUDA_ARCH_NAME=All" ]
+               ++ lib.optional (!cudaSupport) "-DCPU_ONLY=ON"
+               ++ lib.optional (!pythonSupport) "-DBUILD_python=OFF";
 
-  enableParallelBuilding = true;
+  buildInputs = [ boost google-gflags glog protobuf hdf5-cpp lmdb leveldb snappy opencv atlas ]
+                ++ lib.optional cudaSupport cudatoolkit
+                ++ lib.optional cudnnSupport cudnn
+                ++ lib.optionals pythonSupport [ python numpy ];
+
+  propagatedBuildInputs = lib.optional pythonSupport python.pkgs.protobuf;
 
-  buildInputs = [ openblas boost google-gflags glog hdf5 leveldb lmdb opencv
-                  protobuf snappy ]
-                ++ optional cudaSupport cudatoolkit
-                ++ optional cudnnSupport cudnn;
+  outputs = [ "out" "bin" ];
+  # Don't propagate bin.
+  outputBin = "out";
 
-  installPhase = ''
-    mkdir -p $out/{bin,share,lib}
-    for bin in $(find build/tools -executable -type f -name '*.bin');
-    do
-      cp $bin $out/bin/$(basename $bin .bin)
-    done
+  postInstall = ''
+    # Internal static library.
+    rm $out/lib/libproto.a
 
-    cp -r build/examples $out/share
-    cp -r build/lib $out
+    moveToOutput "bin" "$bin"
+  '' + lib.optionalString pythonSupport ''
+    mkdir -p $out/${python.sitePackages}
+    mv $out/python/caffe $out/${python.sitePackages}
+    rm -rf $out/python
   '';
 
   meta = with stdenv.lib; {
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 5ab25a2cbb17..f0a1ed21d3c6 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -1565,6 +1565,13 @@ in {
     };
   };
 
+  caffe = pkgs.caffe.override {
+    python = self.python;
+    boost = self.boost;
+    numpy = self.numpy;
+    pythonSupport = true;
+  };
+
   capstone = buildPythonPackage rec {
     name = "capstone-3.0.4";
     src = pkgs.fetchurl {