pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 pythonOlder,
6 fetchFromGitHub,
7 opencv4,
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 opencv4
34 torch
35 pillow
36 pywavelets
37 numpy
38 ]
39 ++ lib.optionals withOnnx [
40 onnx
41 onnxruntime
42 ];
43
44 postPatch = ''
45 substituteInPlace setup.py \
46 --replace 'opencv-python>=4.1.0.25' 'opencv'
47 substituteInPlace imwatermark/rivaGan.py --replace \
48 'You can install it with pip: `pip install onnxruntime`.' \
49 'You can install it with an override: `python3Packages.invisible-watermark.override { withOnnx = true; };`.'
50 '';
51
52 passthru.tests =
53 let
54 image = "${src}/test_vectors/original.jpg";
55 methods = [
56 "dwtDct"
57 "dwtDctSvd"
58 "rivaGan"
59 ];
60 testCases = builtins.concatMap (method: [
61 {
62 method = method;
63 withOnnx = true;
64 }
65 {
66 method = method;
67 withOnnx = false;
68 }
69 ]) methods;
70 createTest =
71 { method, withOnnx }:
72 let
73 testName = "${if withOnnx then "withOnnx" else "withoutOnnx"}-${method}";
74 # This test fails in the sandbox on aarch64-linux, see https://github.com/microsoft/onnxruntime/issues/10038
75 skipTest = stdenv.isLinux && stdenv.isAarch64 && withOnnx && method == "rivaGan";
76 in
77 lib.optionalAttrs (!skipTest) {
78 "${testName}" = callPackage ./tests/cli.nix {
79 inherit
80 image
81 method
82 testName
83 withOnnx
84 ;
85 };
86 };
87 allTests = builtins.map createTest testCases;
88 in
89 (lib.attrsets.mergeAttrsList allTests)
90 // {
91 python = callPackage ./tests/python { inherit image; };
92 };
93
94 pythonImportsCheck = [ "imwatermark" ];
95
96 meta = with lib; {
97 description = "A library for creating and decoding invisible image watermarks";
98 mainProgram = "invisible-watermark";
99 homepage = "https://github.com/ShieldMnt/invisible-watermark";
100 license = licenses.mit;
101 maintainers = with maintainers; [ Luflosi ];
102 };
103}