|
|
|
|
|
""" |
|
invisible-watermark is the utility to add hidden watermark and extract watermark from image |
|
|
|
embed watermark: ./invisible-watermark -v -a encode -t bytes -m dwtDct -w 'hello' -o ./test_vectors/wm.png ./test_vectors/original.jpg |
|
|
|
decode watermark: ./invisible-watermark -v -a decode -t bytes -m dwtDct -l 40 ./test_vectors/wm.png |
|
""" |
|
|
|
import argparse |
|
import os |
|
import shutil |
|
import sys |
|
import time |
|
import uuid |
|
import cv2 |
|
import base64 |
|
import time |
|
import pprint |
|
|
|
from imwatermark import WatermarkEncoder, WatermarkDecoder |
|
|
|
pp = pprint.PrettyPrinter(indent=2) |
|
|
|
|
|
class CustomArgParseFormatter( |
|
argparse.ArgumentDefaultsHelpFormatter, |
|
argparse.RawDescriptionHelpFormatter): |
|
pass |
|
|
|
|
|
def main(): |
|
description = __doc__.format() |
|
parser = argparse.ArgumentParser(description=description, |
|
formatter_class=CustomArgParseFormatter) |
|
|
|
parser.add_argument('input', |
|
help='The path of input') |
|
parser.add_argument('-a', '--action', required=True, help='encode|decode') |
|
parser.add_argument('-t', '--type', default='bits', help='bytes|b16|bits|uuid|ipv4') |
|
parser.add_argument('-m', '--method', default='maxDct', help='dwtDct|dwtDctSvd|rivaGan') |
|
parser.add_argument('-w', '--watermark', default='', help='embedded string') |
|
parser.add_argument('-l', '--length', default=0, type=int, |
|
help='watermark bits length, required for bytes|b16|bits watermark') |
|
parser.add_argument('-o', '--output', required=False, |
|
help='The path of output') |
|
parser.add_argument('-v', '--verbose', required=False, default=False, |
|
action='store_true', help='print info') |
|
args = parser.parse_args() |
|
action = args.action |
|
inputFile = args.input |
|
|
|
if action == 'encode': |
|
if not args.output: |
|
sys.stderr.write('output is required. exiting...\n') |
|
sys.exit(1) |
|
if not args.watermark: |
|
sys.stderr.write('watermark is required. exiting...\n') |
|
sys.exit(1) |
|
|
|
encoder = WatermarkEncoder() |
|
wmType = args.type |
|
if args.method == 'rivaGan': |
|
WatermarkEncoder.loadModel() |
|
bgr = cv2.imread(args.input) |
|
if args.type == 'bytes': |
|
wm = args.watermark.encode('utf-8') |
|
elif args.type == 'b16': |
|
wm = args.watermark.upper().encode('utf-8') |
|
elif args.type == 'bits': |
|
wm = [int(c) % 2 for c in args.watermark] |
|
else: |
|
wm = args.watermark |
|
encoder.set_watermark(wmType, wm) |
|
start = time.time() |
|
bgr_encoded = encoder.encode(bgr, args.method) |
|
if args.verbose: |
|
print('watermark length:', encoder.get_length()) |
|
print('encode time ms:', (time.time()-start) * 1000) |
|
cv2.imwrite(args.output, bgr_encoded) |
|
elif action == 'decode': |
|
if args.type in ['bytes', 'bits', 'b16']: |
|
if args.length <= 0: |
|
sys.stderr.write('length is required for bytes watermark decoding\n') |
|
sys.exit(1) |
|
wmType = args.type |
|
decoder = WatermarkDecoder(wmType, args.length) |
|
else: |
|
decoder = WatermarkDecoder(args.type) |
|
if args.method == 'rivaGan': |
|
WatermarkDecoder.loadModel() |
|
bgr = cv2.imread(args.input) |
|
start = time.time() |
|
wm = decoder.decode(bgr, args.method) |
|
if args.verbose: |
|
print('decode time ms:', (time.time()-start) * 1000) |
|
if args.type in ['bytes', 'b16']: |
|
wm = wm.decode('utf-8') |
|
print(wm) |
|
else: |
|
sys.stderr.write('unknown action\n') |
|
sys.exit(1) |
|
|
|
if __name__ == '__main__': |
|
sys.exit(main()) |
|
|
