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 # Fails in python >= 3.13
53 # exec(compile(f.read(), version_file, "exec")) does not populate the locals() namesp
54 # In python 3.13, the locals() dictionary in a function does not automatically update with
55 # changes made by exec().
56 # https://peps.python.org/pep-0558/
57 ''
58 substituteInPlace setup.py \
59 --replace-fail "cpu_use = 4" "cpu_use = $NIX_BUILD_CORES" \
60 --replace-fail "return locals()['__version__']" "return '${version}'"
61 '';
62
63 nativeBuildInputs = [
64 ninja
65 which
66 ];
67
68 buildInputs = [
69 pybind11
70 torch
71 ]
72 ++ lib.optionals cudaSupport (
73 with cudaPackages;
74 [
75 cuda_cudart # cuda_runtime.h
76 cuda_cccl # <thrust/*>
77 libcublas # cublas_v2.h
78 libcusolver # cusolverDn.h
79 libcusparse # cusparse.h
80 ]
81 );
82
83 dependencies = [
84 addict
85 mmengine
86 numpy
87 packaging
88 pillow
89 pyyaml
90 yapf
91
92 # opencv4
93 # torch
94 ];
95
96 env.CUDA_HOME = lib.optionalString cudaSupport (lib.getDev cudaPackages.cuda_nvcc);
97
98 preConfigure = ''
99 export MMCV_WITH_OPS=1
100 ''
101 + lib.optionalString cudaSupport ''
102 export CC=${lib.getExe' backendStdenv.cc "cc"}
103 export CXX=${lib.getExe' backendStdenv.cc "c++"}
104 export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}"
105 export FORCE_CUDA=1
106 '';
107
108 pythonImportsCheck = [ "mmcv" ];
109
110 nativeCheckInputs = [
111 lmdb
112 onnx
113 onnxruntime
114 pytestCheckHook
115 pyturbojpeg
116 scipy
117 tifffile
118 torchvision
119 ];
120
121 # remove the conflicting source directory
122 preCheck = ''
123 rm -rf mmcv
124 '';
125
126 # test_cnn test_ops really requires gpus to be useful.
127 # some of the tests take exceedingly long time.
128 # the rest of the tests are disabled due to sandbox env.
129 disabledTests = [
130 "test_cnn"
131 "test_ops"
132 "test_fileclient"
133 "test_load_model_zoo"
134 "test_processing"
135 "test_checkpoint"
136 "test_hub"
137 "test_reader"
138 ]
139 ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
140 # flaky numerical tests (AssertionError)
141 "test_ycbcr2rgb"
142 "test_ycbcr2bgr"
143 "test_tensor2imgs"
144 ];
145
146 meta = {
147 description = "Foundational Library for Computer Vision Research";
148 homepage = "https://github.com/open-mmlab/mmcv";
149 changelog = "https://github.com/open-mmlab/mmcv/releases/tag/v${version}";
150 license = with lib.licenses; [ asl20 ];
151 maintainers = with lib.maintainers; [ rxiao ];
152 };
153}