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