|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# |
| 3 | +# Copyright(C) 2021 wuyaoping |
| 4 | +# |
| 5 | +# DCT algorithm has great a robust but lower capacity. |
| 6 | + |
| 7 | +import numpy as np |
| 8 | +import os.path as osp |
| 9 | +import cv2 |
| 10 | + |
| 11 | +FLAG = '%' |
| 12 | +# Select a part location from the middle frequency |
| 13 | +LOC_MAX = (4, 1) |
| 14 | +LOC_MIN = (3, 2) |
| 15 | +# The difference between MAX and MIN, |
| 16 | +# bigger to improve robust but make picture low quality. |
| 17 | +ALPHA = 1 |
| 18 | + |
| 19 | +# Quantizer table |
| 20 | +TABLE = np.array([ |
| 21 | + [16, 11, 10, 16, 24, 40, 51, 61], |
| 22 | + [12, 12, 14, 19, 26, 58, 60, 55], |
| 23 | + [14, 13, 16, 24, 40, 57, 69, 56], |
| 24 | + [14, 17, 22, 29, 51, 87, 80, 62], |
| 25 | + [18, 22, 37, 56, 68, 109, 103, 77], |
| 26 | + [24, 35, 55, 64, 81, 104, 113, 92], |
| 27 | + [49, 64, 78, 87, 103, 121, 120, 101], |
| 28 | + [72, 92, 95, 98, 112, 100, 103, 99] |
| 29 | +]) |
| 30 | + |
| 31 | + |
| 32 | +def insert(path, txt): |
| 33 | + img = cv2.imread(path, cv2.IMREAD_ANYCOLOR) |
| 34 | + txt = "{}{}{}".format(len(txt), FLAG, txt) |
| 35 | + row, col = img.shape[:2] |
| 36 | + max_bytes = (row // 8) * (col // 8) // 8 |
| 37 | + assert max_bytes >= len( |
| 38 | + txt), "Message overflow the capacity:{}".format(max_bytes) |
| 39 | + img = cv2.cvtColor(img, cv2.COLOR_BGR2YUV) |
| 40 | + # Just use the Y plane to store message, you can use all plane |
| 41 | + y, u, v = cv2.split(img) |
| 42 | + y = y.astype(np.float32) |
| 43 | + blocks = [] |
| 44 | + # Quantize blocks |
| 45 | + for r_idx in range(0, 8 * (row // 8), 8): |
| 46 | + for c_idx in range(0, 8 * (col // 8), 8): |
| 47 | + quantized = cv2.dct(y[r_idx: r_idx+8, c_idx: c_idx+8]) / TABLE |
| 48 | + blocks.append(quantized) |
| 49 | + for idx in range(len(txt)): |
| 50 | + encode(blocks[idx*8: (idx+1)*8], txt[idx]) |
| 51 | + |
| 52 | + idx = 0 |
| 53 | + # Restore Y plane |
| 54 | + for r_idx in range(0, 8 * (row // 8), 8): |
| 55 | + for c_idx in range(0, 8 * (col // 8), 8): |
| 56 | + y[r_idx: r_idx+8, c_idx: c_idx+8] = cv2.idct(blocks[idx] * TABLE) |
| 57 | + idx += 1 |
| 58 | + y = y.astype(np.uint8) |
| 59 | + img = cv2.cvtColor(cv2.merge((y, u, v)), cv2.COLOR_YUV2BGR) |
| 60 | + filename, _ = osp.splitext(path) |
| 61 | + # DCT algorithm can save message even if jpg |
| 62 | + filename += '_dct_embeded' + '.jpg' |
| 63 | + cv2.imwrite(filename, img) |
| 64 | + return filename |
| 65 | + |
| 66 | + |
| 67 | +# Encode a char into the blocks |
| 68 | +def encode(blocks, data): |
| 69 | + data = ord(data) |
| 70 | + for idx in range(len(blocks)): |
| 71 | + bit_val = (data >> idx) & 1 |
| 72 | + max_val = max(blocks[idx][LOC_MAX], blocks[idx][LOC_MIN]) |
| 73 | + min_val = min(blocks[idx][LOC_MAX], blocks[idx][LOC_MIN]) |
| 74 | + if max_val - min_val <= ALPHA: |
| 75 | + max_val = min_val + ALPHA + 1e-3 |
| 76 | + if bit_val == 1: |
| 77 | + blocks[idx][LOC_MAX] = max_val |
| 78 | + blocks[idx][LOC_MIN] = min_val |
| 79 | + else: |
| 80 | + blocks[idx][LOC_MAX] = min_val |
| 81 | + blocks[idx][LOC_MIN] = max_val |
| 82 | + |
| 83 | + |
| 84 | +# Decode a char from the blocks |
| 85 | +def decode(blocks): |
| 86 | + val = 0 |
| 87 | + for idx in range(len(blocks)): |
| 88 | + if blocks[idx][LOC_MAX] > blocks[idx][LOC_MIN]: |
| 89 | + val |= 1 << idx |
| 90 | + return chr(val) |
| 91 | + |
| 92 | + |
| 93 | +def extract(path): |
| 94 | + img = cv2.imread(path, cv2.IMREAD_ANYCOLOR) |
| 95 | + row, col = img.shape[:2] |
| 96 | + max_bytes = (row // 8) * (col // 8) // 8 |
| 97 | + img = cv2.cvtColor(img, cv2.COLOR_BGR2YUV) |
| 98 | + y, u, v = cv2.split(img) |
| 99 | + y = y.astype(np.float32) |
| 100 | + blocks = [] |
| 101 | + for r_idx in range(0, 8 * (row // 8), 8): |
| 102 | + for c_idx in range(0, 8 * (col // 8), 8): |
| 103 | + quantized = cv2.dct(y[r_idx: r_idx+8, c_idx: c_idx+8]) / TABLE |
| 104 | + blocks.append(quantized) |
| 105 | + res = '' |
| 106 | + idx = 0 |
| 107 | + # Extract the length of the message |
| 108 | + while idx < max_bytes: |
| 109 | + ch = decode(blocks[idx*8: (idx+1)*8]) |
| 110 | + idx += 1 |
| 111 | + if ch == FLAG: |
| 112 | + break |
| 113 | + res += ch |
| 114 | + end = int(res) + idx |
| 115 | + assert end <= max_bytes, "Input image isn't correct." |
| 116 | + res = '' |
| 117 | + while idx < end: |
| 118 | + res += decode(blocks[idx*8: (idx+1)*8]) |
| 119 | + idx += 1 |
| 120 | + return res |
| 121 | + |
| 122 | + |
| 123 | +if __name__ == '__main__': |
| 124 | + data = 'A collection of simple python mini projects to enhance your Python skills.' |
| 125 | + res_path = insert('./example.png', data) |
| 126 | + res = extract(res_path) |
| 127 | + print(res) |
0 commit comments