Spaces:
Sleeping
Sleeping
File size: 4,467 Bytes
23f56a9 04d1279 23f56a9 04d1279 23f56a9 04d1279 864e9e3 23f56a9 04d1279 23f56a9 04d1279 23f56a9 04d1279 23f56a9 04d1279 29136da 04d1279 29136da 04d1279 83fb57e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import gradio as gr
from PIL import Image, ImageDraw
import utils
import io
import os
import cv2
def highlight_watermark(image, coords=(0, 0, 100, 50), color="red", width=5):
"""Highlights watermark area. Adapts rectangle size based on coordinates."""
try:
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 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":
image_pil = Image.open(inp_im)
image_cv2, msg = utils.encode(inp_im, inp_mark)
if msg:
return None, msg, None
image_pil = Image.fromarray(cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB))
highlighted_image = highlight_watermark(image_pil.copy(), (0, 0, image_pil.width, image_pil.height), color="blue")
img_byte_arr = io.BytesIO()
highlighted_image.save(img_byte_arr, format='PNG')
img_byte_arr.seek(0)
download_name = os.path.splitext(os.path.basename(inp_im))[0] + "_watermarked.png"
return highlighted_image, "Steganography watermark added successfully.", (download_name, img_byte_arr.read())
elif cho == "pnginfo":
out_im_path, out_msg = utils.png_encode(inp_im, inp_mark)
if out_msg: # Handle potential errors from png_encode
return None, out_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)
img_byte_arr = io.BytesIO()
highlighted_image.save(img_byte_arr, format='PNG')
img_byte_arr.seek(0)
download_name = os.path.splitext(os.path.basename(inp_im))[0] + "_watermarked.png"
return highlighted_image, "PNG metadata watermark added successfully.", (download_name, img_byte_arr.read())
except Exception as e:
return None, f"Error processing PNG: {e}", None
except FileNotFoundError:
return None, "Image file not found.", None
except ValueError as e:
return None, f"Invalid input: {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."
det_msg = utils.decode(det_im)
if det_msg is None or det_msg == "":
return "No watermark detected or not encoded with this tool."
return det_msg
with gr.Blocks() as app:
with gr.Tab("Add Watermark"):
cho = gr.Radio(choices=["stegan", "pnginfo"], value="stegan")
with gr.Row():
with gr.Column():
inp_im = gr.Image(label="Input Image", type="filepath")
inp_mark = gr.Textbox(label="Watermark")
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()
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")
with gr.Column():
det_msg = gr.Textbox(label="Detected Watermark", lines=6, max_lines=50)
mark_btn.click(choose_encode, [inp_im, inp_mark, cho], [out_im, msg_box, file_output])
det_btn.click(utils.decode, [det_im], det_msg)
app.launch() |