Merge pull request #260784 from Luflosi/update/python3Packages.invisible-watermark

{ lib

, stdenv

, buildPythonPackage

, pythonOlder

, fetchFromGitHub

, pillow

, pywavelets

, numpy

, callPackage

, withOnnx ? false # Enables the rivaGan en- and decoding method

}:

buildPythonPackage rec {

  pname = "invisible-watermark";

  version = "0.1.5";

  version = "0.2.0";

  format = "setuptools";

  disabled = pythonOlder "3.6";

  src = fetchFromGitHub {

    owner = "ShieldMnt";

    repo = "invisible-watermark";

    rev = version;

    hash = "sha256-NGDPEETuM7rYbo8kXYoRWLJWpa/lWLKEvaaiDzSWYZ4=";

    rev = "e58e451cff7e092457cd915e445b1a20b64a7c8f"; # No git tag, see https://github.com/ShieldMnt/invisible-watermark/issues/22

    hash = "sha256-6SjVpKFtiiLLU7tZ3hBQr0KT/YEQyywJj0e21/dJRzk=";

  };

  propagatedBuildInputs = [

    opencv4

    torch

    onnx

    onnxruntime

    pillow

    pywavelets

    numpy

  ] ++ lib.optionals withOnnx [

    onnx

    onnxruntime

  ];

  postPatch = ''

    substituteInPlace setup.py \

      --replace 'opencv-python>=4.1.0.25' 'opencv'

    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.isLinux && stdenv.isAarch64 && withOnnx && method == "rivaGan";

    in lib.optionalAttrs (!skipTest) {

      "${testName}" = callPackage ./tests/cli.nix { inherit image method testName withOnnx; };

    };

    allTests = builtins.map createTest testCases;

  in (lib.attrsets.mergeAttrsList allTests) // {

    python = callPackage ./tests/python { inherit image; };

  };

  pythonImportsCheck = [ "imwatermark" ];

  meta = with lib; {

{ image

, method

, python3Packages

, 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' = python3Packages.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"

''

# -*- 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)

{ image

, invisible-watermark

, opencv4

, python3

, runCommand

, stdenvNoCC

}:

# This test checks if the python code shown in the README works correctly

let

  message = "fnörd1";

  method = "dwtDct";

  pythonWithPackages = python3.withPackages (pp: with pp; [ 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"

''

# -*- 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)