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