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