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