Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from PIL import Image, ImageDraw | |
| import utils | |
| import io | |
| import os | |
| import cv2 | |
| import numpy as np | |
| import tempfile | |
| def highlight_watermark(image, coords=(0, 0, 100, 50), color="red", width=5): | |
| try: | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
| draw = ImageDraw.Draw(image) | |
| x, y, w, h = coords | |
| draw.rectangle((max(0, x - 5), max(0, y - 5), x + w + 5, y + h + 5), outline=color, width=width) | |
| return image | |
| except Exception as e: | |
| print(f"Error highlighting: {e}") | |
| return image | |
| def create_output_file(image, original_filename): | |
| try: | |
| # Create temp directory if it doesn't exist | |
| temp_dir = "temp" | |
| os.makedirs(temp_dir, exist_ok=True) | |
| # Generate output filename | |
| base_name = os.path.splitext(os.path.basename(original_filename))[0] | |
| output_path = os.path.join(temp_dir, f"{base_name}_watermarked.png") | |
| # Save image | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
| image.save(output_path, format='PNG') | |
| return output_path | |
| except Exception as e: | |
| print(f"Error creating output file: {e}") | |
| return None | |
| def choose_encode(inp_im, inp_mark, cho): | |
| try: | |
| if not inp_im: | |
| return None, "Please upload an image.", None | |
| if not inp_mark: | |
| return None, "Please enter watermark text.", None | |
| valid_image_types = ["png", "jpg", "jpeg", "gif", "bmp", "tiff", "webp"] | |
| if not any(inp_im.lower().endswith("." + ext) for ext in valid_image_types): | |
| return None, "Invalid image format. Use PNG or JPG.", None | |
| if cho == "stegan": | |
| watermarked_image, error_msg = utils.encode(inp_im, inp_mark) | |
| if error_msg: | |
| return None, error_msg, None | |
| image_pil = Image.fromarray(cv2.cvtColor(watermarked_image, cv2.COLOR_BGR2RGB)) | |
| highlighted_image = highlight_watermark(image_pil.copy(), (0, 0, image_pil.width, image_pil.height), color="blue") | |
| output_path = create_output_file(highlighted_image, inp_im) | |
| if not output_path: | |
| return None, "Error saving watermarked image.", None | |
| return highlighted_image, "Steganography watermark added successfully.", output_path | |
| elif cho == "pnginfo": | |
| out_im_path, error_msg = utils.png_encode(inp_im, inp_mark) | |
| if error_msg: | |
| return None, error_msg, None | |
| try: | |
| img = Image.open(out_im_path) | |
| coords = utils.png_encode_coords if hasattr(utils, 'png_encode_coords') else (50, 50, 200, 50) | |
| highlighted_image = highlight_watermark(img.copy(), coords) | |
| output_path = create_output_file(highlighted_image, inp_im) | |
| if not output_path: | |
| return None, "Error saving watermarked image.", None | |
| return highlighted_image, "PNG metadata watermark added successfully.", output_path | |
| except Exception as e: | |
| return None, f"Error processing PNG: {e}", None | |
| except Exception as e: | |
| return None, f"An unexpected error occurred: {e}", None | |
| def detect_watermark(det_im): | |
| if not det_im: | |
| return "Please upload an image." | |
| valid_image_types = ["png", "jpg", "jpeg", "gif", "bmp", "tiff", "webp"] | |
| if not any(det_im.lower().endswith("." + ext) for ext in valid_image_types): | |
| return "Invalid image format. Use PNG or JPG." | |
| detected_text = utils.decode(det_im) | |
| if not detected_text or detected_text.startswith("Error"): | |
| return "No watermark detected or not encoded with this tool." | |
| return detected_text | |
| # Create Gradio interface | |
| with gr.Blocks() as app: | |
| with gr.Tab("Add Watermark"): | |
| cho = gr.Radio(choices=["stegan", "pnginfo"], value="stegan", label="Watermark Method") | |
| with gr.Row(): | |
| with gr.Column(): | |
| inp_im = gr.Image(label="Input Image", type="filepath") | |
| inp_mark = gr.Textbox(label="Watermark Text") | |
| mark_btn = gr.Button("Add Watermark") | |
| msg_box = gr.Textbox(label="System Message") | |
| with gr.Column(): | |
| out_im = gr.Image(label="Watermarked Image") | |
| file_output = gr.File(label="Download Watermarked Image") | |
| with gr.Tab("Detect Watermark"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| det_im = gr.Image(label="Watermarked Image", type="filepath") | |
| det_btn = gr.Button("Detect Watermark") | |
| with gr.Column(): | |
| det_msg = gr.Textbox(label="Detected Watermark", lines=6) | |
| mark_btn.click(choose_encode, inputs=[inp_im, inp_mark, cho], outputs=[out_im, msg_box, file_output]) | |
| det_btn.click(detect_watermark, inputs=[det_im], outputs=[det_msg]) | |
| if __name__ == "__main__": | |
| app.launch() |