Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| from PIL import Image, PngImagePlugin | |
| import os | |
| import uuid | |
| 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 encode(image_path, secret_data): | |
| try: | |
| # Read the image | |
| image = cv2.imread(image_path) | |
| if image is None: | |
| return None, "Failed to read image" | |
| # Calculate maximum bytes for encoding | |
| n_bytes = image.shape[0] * image.shape[1] * 3 // 8 | |
| # Check if the secret data is too large | |
| secret_data_with_delimiter = f'{secret_data}#####' | |
| if len(secret_data_with_delimiter) + 5 >= n_bytes: | |
| return None, "Watermark is too large for Image Size" | |
| # Add ending delimiter | |
| secret_data_with_delimiter += "=====" | |
| # Convert data to binary | |
| binary_secret_data = to_bin(secret_data_with_delimiter) | |
| data_len = len(binary_secret_data) | |
| data_index = 0 | |
| # Encode the data into the image | |
| for i in range(image.shape[0]): | |
| for j in range(image.shape[1]): | |
| pixel = image[i, j] | |
| for color_channel in range(3): | |
| if data_index < data_len: | |
| # Get the binary value of the pixel | |
| binary_value = to_bin(pixel[color_channel]) | |
| # Replace the least significant bit | |
| new_binary = binary_value[:-1] + binary_secret_data[data_index] | |
| # Update the pixel with the new value | |
| pixel[color_channel] = int(new_binary, 2) | |
| data_index += 1 | |
| image[i, j] = pixel | |
| if data_index >= data_len: | |
| break | |
| if data_index >= data_len: | |
| break | |
| return image, None | |
| except Exception as e: | |
| return None, f"Error during encoding: {str(e)}" | |
| def decode(image_path): | |
| try: | |
| # Read the image | |
| image = cv2.imread(image_path) | |
| if image is None: | |
| return "Failed to read image" | |
| binary_data = "" | |
| for row in image: | |
| for pixel in row: | |
| for color_channel in pixel: | |
| binary_data += to_bin(color_channel)[-1] | |
| # Convert binary to text | |
| 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 | |
| # Remove delimiters | |
| return decoded_data[:-5].split("#####")[0] | |
| except Exception as e: | |
| return f"Error during decoding: {str(e)}" | |
| 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, "PNG", pnginfo=info) | |
| width, height = im.size | |
| rect_width, rect_height = 200, 50 | |
| x = width - rect_width - 10 | |
| y = height - rect_height - 10 | |
| global png_encode_coords | |
| png_encode_coords = (x, y, rect_width, rect_height) | |
| return new_filename, None | |
| except Exception as e: | |
| return None, str(e) | |