Spaces:
Running
Running
| import cv2 | |
| import numpy as np | |
| from PIL import Image, PngImagePlugin | |
| def png_encode(im_name, extra): | |
| """Encode watermark using PNG metadata""" | |
| try: | |
| im = Image.open(im_name) | |
| info = PngImagePlugin.PngInfo() | |
| info.add_text("TXT", extra) | |
| im.save("test.png", pnginfo=info) | |
| return "test.png", "Watermark added successfully" | |
| except Exception as e: | |
| return im_name, f"Error adding watermark: {str(e)}" | |
| def to_bin(data): | |
| """Convert data to binary format as string""" | |
| if isinstance(data, str): | |
| return ''.join(format(x, '08b') for x in data.encode('utf-8')) | |
| elif isinstance(data, bytes): | |
| return ''.join(format(x, '08b') for x 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): | |
| """Decode watermark from image""" | |
| try: | |
| # First try PNG metadata method | |
| try: | |
| im = Image.open(image_name) | |
| if "TXT" in im.info: | |
| return im.info["TXT"] | |
| except: | |
| pass | |
| # Steganography method | |
| image = cv2.imread(image_name) | |
| if image is None: | |
| raise ValueError("Could not read image file") | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| binary_data = "" | |
| # Extract binary data from image | |
| 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] | |
| # Convert binary string to bytes | |
| bytes_data = b'' | |
| for i in range(0, len(binary_data), 8): | |
| byte = binary_data[i:i+8] | |
| if len(byte) == 8: | |
| try: | |
| byte_val = int(byte, 2) | |
| bytes_data += bytes([byte_val]) | |
| except ValueError: | |
| continue | |
| # Try to find the end marker in raw bytes | |
| try: | |
| end_marker = b'=====' | |
| if end_marker in bytes_data: | |
| bytes_data = bytes_data.split(end_marker)[0] | |
| # Try to find the delimiter in raw bytes | |
| delimiter = b'#####' | |
| if delimiter in bytes_data: | |
| bytes_data = bytes_data.split(delimiter)[0] | |
| # Decode the bytes to string | |
| decoded_text = bytes_data.decode('utf-8', errors='ignore') | |
| return decoded_text.strip() | |
| except Exception as e: | |
| print(f"Decoding error: {e}") | |
| # If regular decoding fails, try character by character | |
| result = "" | |
| current_bytes = bytearray() | |
| for byte in bytes_data: | |
| current_bytes.append(byte) | |
| try: | |
| char = current_bytes.decode('utf-8') | |
| result += char | |
| current_bytes = bytearray() | |
| except UnicodeDecodeError: | |
| continue | |
| if result: | |
| return result.strip() | |
| return "Error: Could not decode watermark" | |
| except Exception as e: | |
| return f"Error detecting watermark: {str(e)}" | |
| def encode(image_name, secret_data, txt=None): | |
| """Encode watermark using steganography""" | |
| try: | |
| image = cv2.imread(image_name) | |
| if image is None: | |
| raise ValueError("Could not read image file") | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| # Calculate maximum bytes that can be encoded | |
| n_bytes = image.shape[0] * image.shape[1] * 3 // 8 | |
| # Prepare the data with delimiters | |
| secret_data = str(secret_data) | |
| complete_data = secret_data + "#####" + "=====" | |
| # Convert to binary | |
| binary_secret_data = to_bin(complete_data) | |
| # Check if the data can fit in the image | |
| if len(binary_secret_data) > n_bytes * 8: | |
| return image_name, "Watermark is too large for Image Size" | |
| data_index = 0 | |
| binary_len = len(binary_secret_data) | |
| # Embed the data | |
| for i in range(image.shape[0]): | |
| for j in range(image.shape[1]): | |
| if data_index < binary_len: | |
| pixel = image[i, j] | |
| for color_channel in range(3): | |
| if data_index < binary_len: | |
| # Get the binary value of the pixel | |
| binary_value = format(pixel[color_channel], '08b') | |
| # Replace the least significant bit | |
| binary_value = binary_value[:-1] + binary_secret_data[data_index] | |
| # Update the pixel value | |
| image[i, j, color_channel] = int(binary_value, 2) | |
| data_index += 1 | |
| else: | |
| break | |
| # Save the result | |
| output_path = "watermarked_image.png" | |
| cv2.imwrite(output_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR)) | |
| return output_path, "Watermark embedded successfully" | |
| except Exception as e: | |
| return image_name, f"Error encoding watermark: {str(e)}" |