pkgs/development/python-modules/compressai/default.nix

   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
  30 buildPythonPackage 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 }