Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| from PIL import Image, PngImagePlugin | |
| import os | |
| import uuid | |
| def png_encode(im_name, extra): | |
| try: | |
| im = Image.open(im_name) | |
| info = PngImagePlugin.PngInfo() | |
| info.add_text("TXT", extra) | |
| unique_id = str(uuid.uuid4())[:8] | |
| filename, ext = os.path.splitext(im_name) | |
| new_filename = f"{filename}_{unique_id}{ext}" | |
| im.save(new_filename, pnginfo=info) | |
| width, height = im.size | |
| rect_width, rect_height = 200, 50 | |
| x = width - rect_width - 10 | |
| y = height - rect_height - 10 | |
| png_encode_coords = (x, y, rect_width, rect_height) | |
| return new_filename, "", png_encode_coords | |
| except Exception as e: | |
| return None, str(e), None | |
| def to_bin(data): | |
| if isinstance(data, str): | |
| return ''.join([format(ord(i), "08b") for i in data]) | |
| elif isinstance(data, bytes): | |
| return ''.join([format(i, "08b") for i in data]) | |
| elif isinstance(data, np.ndarray): | |
| return [format(i, "08b") for i in data] | |
| elif isinstance(data, int) or isinstance(data, np.uint8): | |
| return format(data, "08b") | |
| else: | |
| raise TypeError("Type not supported.") | |
| def decode(image_name, txt=None): | |
| try: | |
| BGRimage = cv2.imread(f"{image_name}") | |
| image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB) | |
| binary_data = "" | |
| for row in image: | |
| for pixel in row: | |
| r, g, b = to_bin(pixel) | |
| binary_data += r[-1] | |
| binary_data += g[-1] | |
| binary_data += b[-1] | |
| all_bytes = [binary_data[i: i + 8] for i in range(0, len(binary_data), 8)] | |
| decoded_data = "" | |
| for byte in all_bytes: | |
| decoded_data += chr(int(byte, 2)) | |
| if decoded_data[-5:] == "=====": | |
| break | |
| this = decoded_data[:-5].split("#####", 1)[0] | |
| return this | |
| except Exception as e: | |
| return str(e) | |
| def encode(image_name, secret_data, txt=None): | |
| msg = "" | |
| BGRimage = cv2.imread(f"{image_name}") | |
| image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB) | |
| n_bytes = image.shape[0] * image.shape[1] * 3 // 8 | |
| print("[*] Maximum bytes to encode:", n_bytes) | |
| secret_data1 = secret_data | |
| while True: | |
| if len(secret_data1) + 5 < (n_bytes): | |
| secret_data1 = f'{secret_data1}#####' | |
| elif len(secret_data1) + 5 >= (n_bytes): | |
| break | |
| secret_data = secret_data1 | |
| if len(secret_data) > n_bytes: | |
| msg = "Watermark is too large for Image Size" | |
| return image_name, msg | |
| secret_data += "=====" | |
| data_index = 0 | |
| binary_secret_data = to_bin(secret_data) | |
| data_len = len(binary_secret_data) | |
| for row in image: | |
| for pixel in row: | |
| r, g, b = to_bin(pixel) | |
| if data_index < data_len: | |
| pixel[0] = int(r[:-1] + binary_secret_data[data_index], 2) | |
| data_index += 1 | |
| if data_index < data_len: | |
| pixel[1] = int(g[:-1] + binary_secret_data[data_index], 2) | |
| data_index += 1 | |
| if data_index < data_len: | |
| pixel[2] = int(b[:-1] + binary_secret_data[data_index], 2) | |
| data_index += 1 | |
| if data_index >= data_len: | |
| break | |
| return image, msg |