pyrox.dev
1{ lib
2, stdenv
3, buildPythonPackage
4, pythonOlder
5, fetchFromGitHub
6, opencv4
7, torch
8, onnx
9, onnxruntime
10, pillow
11, pywavelets
12, numpy
13, callPackage
14, withOnnx ? false # Enables the rivaGan en- and decoding method
15}:
16
17buildPythonPackage rec {
18 pname = "invisible-watermark";
19 version = "0.2.0";
20 format = "setuptools";
21 disabled = pythonOlder "3.6";
22
23 src = fetchFromGitHub {
24 owner = "ShieldMnt";
25 repo = "invisible-watermark";
26 rev = "e58e451cff7e092457cd915e445b1a20b64a7c8f"; # No git tag, see https://github.com/ShieldMnt/invisible-watermark/issues/22
27 hash = "sha256-6SjVpKFtiiLLU7tZ3hBQr0KT/YEQyywJj0e21/dJRzk=";
28 };
29
30 propagatedBuildInputs = [
31 opencv4
32 torch
33 pillow
34 pywavelets
35 numpy
36 ] ++ lib.optionals withOnnx [
37 onnx
38 onnxruntime
39 ];
40
41 postPatch = ''
42 substituteInPlace setup.py \
43 --replace 'opencv-python>=4.1.0.25' 'opencv'
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 = let
50 image = "${src}/test_vectors/original.jpg";
51 methods = [ "dwtDct" "dwtDctSvd" "rivaGan" ];
52 testCases = builtins.concatMap (method: [
53 { method = method; withOnnx = true; }
54 { method = method; withOnnx = false; }
55 ]) methods;
56 createTest = { method, withOnnx }: let
57 testName = "${if withOnnx then "withOnnx" else "withoutOnnx"}-${method}";
58 # This test fails in the sandbox on aarch64-linux, see https://github.com/microsoft/onnxruntime/issues/10038
59 skipTest = stdenv.isLinux && stdenv.isAarch64 && withOnnx && method == "rivaGan";
60 in lib.optionalAttrs (!skipTest) {
61 "${testName}" = callPackage ./tests/cli.nix { inherit image method testName withOnnx; };
62 };
63 allTests = builtins.map createTest testCases;
64 in (lib.attrsets.mergeAttrsList allTests) // {
65 python = callPackage ./tests/python { inherit image; };
66 };
67
68 pythonImportsCheck = [ "imwatermark" ];
69
70 meta = with lib; {
71 description = "A library for creating and decoding invisible image watermarks";
72 homepage = "https://github.com/ShieldMnt/invisible-watermark";
73 license = licenses.mit;
74 maintainers = with maintainers; [ Luflosi ];
75 };
76}