push sheeet

pkgs/development/python-modules/invisible-watermark/default.nix

@@ -0,0 +1,105 @@
+{
+  lib,
+  stdenv,
+  buildPythonPackage,
+  pythonOlder,
+  fetchFromGitHub,
+  setuptools,
+  opencv-python,
+  torch,
+  onnx,
+  onnxruntime,
+  pillow,
+  pywavelets,
+  numpy,
+  callPackage,
+  withOnnx ? false, # Enables the rivaGan en- and decoding method
+}:
+
+buildPythonPackage rec {
+  pname = "invisible-watermark";
+  version = "0.2.0";
+  pyproject = true;
+  disabled = pythonOlder "3.6";
+
+  src = fetchFromGitHub {
+    owner = "ShieldMnt";
+    repo = "invisible-watermark";
+    # nixpkgs-update: no auto update
+    rev = "e58e451cff7e092457cd915e445b1a20b64a7c8f"; # No git tag, see https://github.com/ShieldMnt/invisible-watermark/issues/22
+    hash = "sha256-6SjVpKFtiiLLU7tZ3hBQr0KT/YEQyywJj0e21/dJRzk=";
+  };
+
+  build-system = [ setuptools ];
+
+  dependencies = [
+    opencv-python
+    torch
+    pillow
+    pywavelets
+    numpy
+  ]
+  ++ lib.optionals withOnnx [
+    onnx
+    onnxruntime
+  ];
+
+  postPatch = ''
+    substituteInPlace imwatermark/rivaGan.py --replace \
+      'You can install it with pip: `pip install onnxruntime`.' \
+      'You can install it with an override: `python3Packages.invisible-watermark.override { withOnnx = true; };`.'
+  '';
+
+  passthru.tests =
+    let
+      image = "${src}/test_vectors/original.jpg";
+      methods = [
+        "dwtDct"
+        "dwtDctSvd"
+        "rivaGan"
+      ];
+      testCases = builtins.concatMap (method: [
+        {
+          method = method;
+          withOnnx = true;
+        }
+        {
+          method = method;
+          withOnnx = false;
+        }
+      ]) methods;
+      createTest =
+        { method, withOnnx }:
+        let
+          testName = "${if withOnnx then "withOnnx" else "withoutOnnx"}-${method}";
+          # This test fails in the sandbox on aarch64-linux, see https://github.com/microsoft/onnxruntime/issues/10038
+          skipTest =
+            stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64 && withOnnx && method == "rivaGan";
+        in
+        lib.optionalAttrs (!skipTest) {
+          "${testName}" = callPackage ./tests/cli.nix {
+            inherit
+              image
+              method
+              testName
+              withOnnx
+              ;
+          };
+        };
+      allTests = map createTest testCases;
+    in
+    (lib.attrsets.mergeAttrsList allTests)
+    // {
+      python = callPackage ./tests/python { inherit image; };
+    };
+
+  pythonImportsCheck = [ "imwatermark" ];
+
+  meta = {
+    description = "Library for creating and decoding invisible image watermarks";
+    mainProgram = "invisible-watermark";
+    homepage = "https://github.com/ShieldMnt/invisible-watermark";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ Luflosi ];
+  };
+}

pkgs/development/python-modules/invisible-watermark/tests/cli.nix

@@ -0,0 +1,67 @@
+{
+  image,
+  method,
+  invisible-watermark,
+  runCommand,
+  testName,
+  withOnnx,
+}:
+
+# This file runs one test case.
+# There are six test cases in total. method can have three possible values and
+# withOnnx two possible values. 3 * 2 = 6.
+#
+# The case where the method is rivaGan and invisible-watermark is built
+# without onnx is expected to fail and this case is handled accordingly.
+#
+# The test works by first encoding a message into a test image,
+# then decoding the message from the image again and checking
+# if the message was decoded correctly.
+
+let
+  message =
+    if method == "rivaGan" then
+      "asdf" # rivaGan only supports 32 bits
+    else
+      "fnörd1";
+  length = (builtins.stringLength message) * 8;
+  invisible-watermark' = invisible-watermark.override { inherit withOnnx; };
+  expected-exit-code = if method == "rivaGan" && !withOnnx then "1" else "0";
+in
+runCommand "invisible-watermark-test-${testName}" { nativeBuildInputs = [ invisible-watermark' ]; }
+  ''
+    set +e
+    invisible-watermark \
+      --verbose \
+      --action encode \
+      --type bytes \
+      --method '${method}' \
+      --watermark '${message}' \
+      --output output.png \
+      '${image}'
+    exit_code="$?"
+    set -euf -o pipefail
+    if [ "$exit_code" != '${expected-exit-code}' ]; then
+      echo "Exit code of invisible-watermark was $exit_code while ${expected-exit-code} was expected."
+      exit 1
+    fi
+    if [ '${expected-exit-code}' == '1' ]; then
+      echo 'invisible-watermark failed as expected.'
+      touch "$out"
+      exit 0
+    fi
+    decoded_message="$(invisible-watermark \
+                        --action decode \
+                        --type bytes \
+                        --method '${method}' \
+                        --length '${toString length}' \
+                        output.png \
+                      )"
+
+    if [ '${message}' != "$decoded_message" ]; then
+      echo "invisible-watermark did not decode the watermark correctly."
+      echo "The original message was ${message} but the decoded message was $decoded_message."
+      exit 1
+    fi
+    touch "$out"
+  ''

pkgs/development/python-modules/invisible-watermark/tests/python/decode.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+
+import os
+import cv2
+from imwatermark import WatermarkDecoder
+
+input_file = os.environ['image']
+output_file_path = os.environ['out']
+num_bits = int(os.environ['num_bits'])
+method = os.environ['method']
+
+bgr = cv2.imread(input_file)
+
+decoder = WatermarkDecoder('bytes', num_bits)
+watermark = decoder.decode(bgr, method)
+message = watermark.decode('utf-8')
+
+with open(output_file_path, 'w') as f:
+    f.write(message)

pkgs/development/python-modules/invisible-watermark/tests/python/default.nix

@@ -0,0 +1,46 @@
+{
+  image,
+  invisible-watermark,
+  opencv4,
+  python,
+  runCommand,
+  stdenvNoCC,
+}:
+
+# This test checks if the python code shown in the README works correctly
+
+let
+  message = "fnörd1";
+  method = "dwtDct";
+
+  pythonWithPackages = python.withPackages (_: [
+    invisible-watermark
+    opencv4
+  ]);
+  pythonInterpreter = pythonWithPackages.interpreter;
+
+  encode = stdenvNoCC.mkDerivation {
+    name = "encode";
+    realBuilder = pythonInterpreter;
+    args = [ ./encode.py ];
+    inherit image message method;
+  };
+
+  decode = stdenvNoCC.mkDerivation {
+    name = "decode";
+    realBuilder = pythonInterpreter;
+    args = [ ./decode.py ];
+    inherit method;
+    image = "${encode}/test_wm.png";
+    num_bits = (builtins.stringLength message) * 8;
+  };
+in
+runCommand "invisible-watermark-test-python" { } ''
+  decoded_message="$(cat '${decode}')"
+  if [ '${message}' != "$decoded_message" ]; then
+    echo "invisible-watermark did not decode the watermark correctly."
+    echo "The original message was ${message} but the decoded message was $decoded_message."
+    exit 1
+  fi
+  touch "$out"
+''

pkgs/development/python-modules/invisible-watermark/tests/python/encode.py

@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+
+import os
+import cv2
+from imwatermark import WatermarkEncoder
+
+input_file_path = os.environ['image']
+output_dir = os.environ['out']
+message = os.environ['message']
+method = os.environ['method']
+
+os.mkdir(output_dir)
+
+bgr = cv2.imread(input_file_path)
+
+encoder = WatermarkEncoder()
+encoder.set_watermark('bytes', message.encode('utf-8'))
+bgr_encoded = encoder.encode(bgr, method)
+
+output_file = os.path.join(output_dir, 'test_wm.png')
+cv2.imwrite(output_file, bgr_encoded)