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