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