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