pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6
7 # build-system
8 pybind11,
9 setuptools,
10
11 # dependencies
12 einops,
13 numpy,
14 matplotlib,
15 pandas,
16 pytorch-msssim,
17 scipy,
18 torch,
19 torch-geometric,
20 torchvision,
21
22 # optional-dependencies
23 ipywidgets,
24 jupyter,
25
26 # tests
27 plotly,
28 pytestCheckHook,
29}:
30
31buildPythonPackage rec {
32 pname = "compressai";
33 version = "1.2.6";
34 pyproject = true;
35
36 src = fetchFromGitHub {
37 owner = "InterDigitalInc";
38 repo = "CompressAI";
39 tag = "v${version}";
40 hash = "sha256-xvzhhLn0iBzq3h1nro8/83QWEQe9K4zRa3RSZk+hy3Y=";
41 fetchSubmodules = true;
42 };
43
44 build-system = [
45 pybind11
46 setuptools
47 ];
48
49 dependencies = [
50 einops
51 numpy
52 matplotlib
53 pandas
54 pytorch-msssim
55 scipy
56 torch
57 torch-geometric
58 torchvision
59 ];
60
61 optional-dependencies = {
62 tutorials = [
63 ipywidgets
64 jupyter
65 ];
66 };
67
68 pythonImportsCheck = [
69 "compressai"
70 "compressai._CXX"
71 ];
72
73 preCheck = ''
74 # We have to delete the source because otherwise it is used intead the installed package.
75 rm -rf compressai
76
77 export HOME=$(mktemp -d)
78 '';
79
80 nativeCheckInputs = [
81 plotly
82 pytestCheckHook
83 ];
84
85 disabledTests = [
86 # Those tests require internet access to download some weights
87 "test_image_codec"
88 "test_update"
89 "test_eval_model_pretrained"
90 "test_cheng2020_anchor"
91 "test_pretrained"
92
93 # Flaky (AssertionError: assert 0.08889999999999998 < 0.064445)
94 "test_find_close"
95 ];
96
97 disabledTestPaths = lib.optionals stdenv.hostPlatform.isDarwin [
98 # Cause pytest to hang on Darwin after the tests are done
99 "tests/test_eval_model.py"
100 "tests/test_train.py"
101
102 # fails in sandbox as it tries to launch a web browser (which fails due to missing `osascript`)
103 "tests/test_plot.py::test_plot[plotly-ms-ssim-rgb]"
104 ];
105
106 meta = {
107 description = "PyTorch library and evaluation platform for end-to-end compression research";
108 homepage = "https://github.com/InterDigitalInc/CompressAI";
109 license = lib.licenses.bsd3Clear;
110 maintainers = with lib.maintainers; [ GaetanLepage ];
111 };
112}