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