tjh.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, pytestCheckHook
5, pythonOlder
6, torch
7, torchvision
8, opencv4
9, yapf
10, packaging
11, pillow
12, addict
13, ninja
14, which
15, onnx
16, onnxruntime
17, scipy
18, pyturbojpeg
19, tifffile
20, lmdb
21, symlinkJoin
22}:
23
24let
25 inherit (torch) cudaCapabilities cudaPackages cudaSupport;
26 inherit (cudaPackages) backendStdenv cudaVersion;
27
28 cuda-common-redist = with cudaPackages; [
29 cuda_cccl # <thrust/*>
30 libcublas # cublas_v2.h
31 libcusolver # cusolverDn.h
32 libcusparse # cusparse.h
33 ];
34
35 cuda-native-redist = symlinkJoin {
36 name = "cuda-native-redist-${cudaVersion}";
37 paths = with cudaPackages; [
38 cuda_cudart # cuda_runtime.h
39 cuda_nvcc
40 ] ++ cuda-common-redist;
41 };
42
43 cuda-redist = symlinkJoin {
44 name = "cuda-redist-${cudaVersion}";
45 paths = cuda-common-redist;
46 };
47
48in
49buildPythonPackage rec {
50 pname = "mmcv";
51 version = "1.7.1";
52 format = "setuptools";
53
54 disabled = pythonOlder "3.6";
55
56 src = fetchFromGitHub {
57 owner = "open-mmlab";
58 repo = pname;
59 rev = "refs/tags/v${version}";
60 hash = "sha256-b4MLBPNRCcPq1osUvqo71PCWVX7lOjAH+dXttd4ZapU";
61 };
62
63 preConfigure = ''
64 export MMCV_WITH_OPS=1
65 '' + lib.optionalString cudaSupport ''
66 export CC=${backendStdenv.cc}/bin/cc
67 export CXX=${backendStdenv.cc}/bin/c++
68 export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}"
69 export FORCE_CUDA=1
70 '';
71
72 postPatch = ''
73 substituteInPlace setup.py --replace "cpu_use = 4" "cpu_use = $NIX_BUILD_CORES"
74 '';
75
76 preCheck = ''
77 # remove the conflicting source directory
78 rm -rf mmcv
79 '';
80
81 # test_cnn test_ops really requires gpus to be useful.
82 # some of the tests take exceedingly long time.
83 # the rest of the tests are disabled due to sandbox env.
84 disabledTests = [
85 "test_cnn"
86 "test_ops"
87 "test_fileclient"
88 "test_load_model_zoo"
89 "test_processing"
90 "test_checkpoint"
91 "test_hub"
92 "test_reader"
93 ];
94
95 nativeBuildInputs = [ ninja which ]
96 ++ lib.optionals cudaSupport [ cuda-native-redist ];
97
98 buildInputs = [ torch ] ++ lib.optionals cudaSupport [ cuda-redist ];
99
100 nativeCheckInputs = [ pytestCheckHook torchvision lmdb onnx onnxruntime scipy pyturbojpeg tifffile ];
101
102 propagatedBuildInputs = [
103 torch
104 opencv4
105 yapf
106 packaging
107 pillow
108 addict
109 ];
110
111 pythonImportsCheck = [
112 "mmcv"
113 ];
114
115 meta = with lib; {
116 description = "A Foundational Library for Computer Vision Research";
117 homepage = "https://github.com/open-mmlab/mmcv";
118 changelog = "https://github.com/open-mmlab/mmcv/releases/tag/v${version}";
119 license = with licenses; [ asl20 ];
120 maintainers = with maintainers; [ rxiao ];
121 };
122}