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