pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 pythonOlder,
6 fetchFromGitHub,
7 setuptools,
8 opencv-python,
9 torch,
10 onnx,
11 onnxruntime,
12 pillow,
13 pywavelets,
14 numpy,
15 callPackage,
16 withOnnx ? false, # Enables the rivaGan en- and decoding method
17}:
18
19buildPythonPackage rec {
20 pname = "invisible-watermark";
21 version = "0.2.0";
22 pyproject = true;
23 disabled = pythonOlder "3.6";
24
25 src = fetchFromGitHub {
26 owner = "ShieldMnt";
27 repo = "invisible-watermark";
28 # nixpkgs-update: no auto update
29 rev = "e58e451cff7e092457cd915e445b1a20b64a7c8f"; # No git tag, see https://github.com/ShieldMnt/invisible-watermark/issues/22
30 hash = "sha256-6SjVpKFtiiLLU7tZ3hBQr0KT/YEQyywJj0e21/dJRzk=";
31 };
32
33 build-system = [ setuptools ];
34
35 dependencies = [
36 opencv-python
37 torch
38 pillow
39 pywavelets
40 numpy
41 ]
42 ++ lib.optionals withOnnx [
43 onnx
44 onnxruntime
45 ];
46
47 postPatch = ''
48 substituteInPlace imwatermark/rivaGan.py --replace \
49 'You can install it with pip: `pip install onnxruntime`.' \
50 'You can install it with an override: `python3Packages.invisible-watermark.override { withOnnx = true; };`.'
51 '';
52
53 passthru.tests =
54 let
55 image = "${src}/test_vectors/original.jpg";
56 methods = [
57 "dwtDct"
58 "dwtDctSvd"
59 "rivaGan"
60 ];
61 testCases = builtins.concatMap (method: [
62 {
63 method = method;
64 withOnnx = true;
65 }
66 {
67 method = method;
68 withOnnx = false;
69 }
70 ]) methods;
71 createTest =
72 { method, withOnnx }:
73 let
74 testName = "${if withOnnx then "withOnnx" else "withoutOnnx"}-${method}";
75 # This test fails in the sandbox on aarch64-linux, see https://github.com/microsoft/onnxruntime/issues/10038
76 skipTest =
77 stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64 && withOnnx && method == "rivaGan";
78 in
79 lib.optionalAttrs (!skipTest) {
80 "${testName}" = callPackage ./tests/cli.nix {
81 inherit
82 image
83 method
84 testName
85 withOnnx
86 ;
87 };
88 };
89 allTests = builtins.map createTest testCases;
90 in
91 (lib.attrsets.mergeAttrsList allTests)
92 // {
93 python = callPackage ./tests/python { inherit image; };
94 };
95
96 pythonImportsCheck = [ "imwatermark" ];
97
98 meta = {
99 description = "Library for creating and decoding invisible image watermarks";
100 mainProgram = "invisible-watermark";
101 homepage = "https://github.com/ShieldMnt/invisible-watermark";
102 license = lib.licenses.mit;
103 maintainers = with lib.maintainers; [ Luflosi ];
104 };
105}