RenAzum commited on
Commit
f6cb6df
·
1 Parent(s): 027e587

Add application

Browse files
Files changed (4) hide show
  1. .gitignore +1 -0
  2. app.py +207 -0
  3. requirements.txt +4 -0
  4. utils.py +120 -0
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ venv
app.py ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from PIL import Image, ImageDraw, PngImagePlugin
3
+ import io
4
+ import os
5
+ import cv2
6
+ import numpy as np
7
+ import uuid
8
+
9
+
10
+ # Helper functions
11
+ def to_bin(data):
12
+ if isinstance(data, str):
13
+ return ''.join([format(ord(i), "08b") for i in data])
14
+ elif isinstance(data, bytes):
15
+ return ''.join([format(i, "08b") for i in data])
16
+ elif isinstance(data, np.ndarray):
17
+ return [format(i, "08b") for i in data]
18
+ elif isinstance(data, int) or isinstance(data, np.uint8):
19
+ return format(data, "08b")
20
+ else:
21
+ raise TypeError("Type not supported.")
22
+
23
+
24
+ def encode(image_path, secret_data):
25
+ try:
26
+ # Read the image
27
+ image = cv2.imread(image_path)
28
+ if image is None:
29
+ return None, "Failed to read image"
30
+
31
+ # Calculate maximum bytes for encoding
32
+ n_bytes = (image.shape[0] * image.shape[1] * 3) // 8
33
+
34
+ # Prepare secret data
35
+ secret_data_with_delimiter = f'{secret_data}#####'
36
+ if len(secret_data_with_delimiter) >= n_bytes:
37
+ return None, "Watermark is too large for Image Size"
38
+
39
+ secret_data_with_delimiter += "====="
40
+ binary_secret_data = to_bin(secret_data_with_delimiter)
41
+ data_len = len(binary_secret_data)
42
+
43
+ # Create a copy of the image
44
+ watermarked_image = image.copy()
45
+ data_index = 0
46
+
47
+ # Encode the data
48
+ for i in range(watermarked_image.shape[0]):
49
+ for j in range(watermarked_image.shape[1]):
50
+ for k in range(3): # RGB channels
51
+ if data_index < data_len:
52
+ # Get the current pixel value
53
+ pixel = watermarked_image[i, j, k]
54
+ # Convert to binary and modify the least significant bit
55
+ binary_pixel = format(pixel, '08b')
56
+ new_pixel = binary_pixel[:-1] + binary_secret_data[data_index]
57
+ # Update the pixel value
58
+ watermarked_image[i, j, k] = int(new_pixel, 2)
59
+ data_index += 1
60
+ else:
61
+ break
62
+ if data_index >= data_len:
63
+ break
64
+ if data_index >= data_len:
65
+ break
66
+
67
+ return watermarked_image, None
68
+ except Exception as e:
69
+ return None, f"Error during encoding: {str(e)}"
70
+
71
+
72
+ def decode(image_path):
73
+ try:
74
+ # Read the image
75
+ image = cv2.imread(image_path)
76
+ if image is None:
77
+ return "Failed to read image"
78
+
79
+ binary_data = ""
80
+ # Extract the least significant bits
81
+ for i in range(image.shape[0]):
82
+ for j in range(image.shape[1]):
83
+ for k in range(3):
84
+ pixel = format(image[i, j, k], '08b')
85
+ binary_data += pixel[-1]
86
+
87
+ # Convert binary to text
88
+ decoded_data = ""
89
+ for i in range(0, len(binary_data), 8):
90
+ byte = binary_data[i:i+8]
91
+ if len(byte) == 8:
92
+ decoded_data += chr(int(byte, 2))
93
+ # Check for ending delimiter
94
+ if decoded_data[-5:] == "=====":
95
+ break
96
+
97
+ # Remove delimiters and return the message
98
+ if "#####" in decoded_data:
99
+ return decoded_data.split("#####")[0]
100
+ return "No watermark found"
101
+
102
+ except Exception as e:
103
+ return f"Error during decoding: {str(e)}"
104
+
105
+
106
+ def png_encode(im_name, extra):
107
+ try:
108
+ im = Image.open(im_name)
109
+ info = PngImagePlugin.PngInfo()
110
+ info.add_text("TXT", extra)
111
+
112
+ unique_id = str(uuid.uuid4())[:8]
113
+ filename, ext = os.path.splitext(im_name)
114
+ new_filename = f"{filename}_{unique_id}{ext}"
115
+ im.save(new_filename, "PNG", pnginfo=info)
116
+
117
+ width, height = im.size
118
+ rect_width, rect_height = 200, 50
119
+ x = width - rect_width - 10
120
+ y = height - rect_height - 10
121
+ global png_encode_coords
122
+ png_encode_coords = (x, y, rect_width, rect_height)
123
+
124
+ return new_filename, None
125
+
126
+ except Exception as e:
127
+ return None, str(e)
128
+
129
+
130
+ def highlight_watermark(image, coords=(0, 0, 100, 50), color="red", width=5):
131
+ try:
132
+ if isinstance(image, np.ndarray):
133
+ image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
134
+ draw = ImageDraw.Draw(image)
135
+ x, y, w, h = coords
136
+ draw.rectangle((max(0, x - 5), max(0, y - 5), x + w + 5, y + h + 5), outline=color, width=width)
137
+ return image
138
+ except Exception as e:
139
+ print(f"Error highlighting: {e}")
140
+ return image
141
+
142
+
143
+ # Gradio function handlers
144
+ def choose_encode(inp_im, inp_mark, cho):
145
+ try:
146
+ if not inp_im:
147
+ return None, "Please upload an image.", None
148
+ if not inp_mark:
149
+ return None, "Please enter watermark text.", None
150
+
151
+ if cho == "stegan":
152
+ watermarked_image, error_msg = encode(inp_im, inp_mark)
153
+ if error_msg:
154
+ return None, error_msg, None
155
+
156
+ output_path = os.path.splitext(inp_im)[0] + "_watermarked.png"
157
+ cv2.imwrite(output_path, watermarked_image)
158
+ return output_path, "Steganography watermark added successfully.", output_path
159
+
160
+ elif cho == "pnginfo":
161
+ output_path, error_msg = png_encode(inp_im, inp_mark)
162
+ if error_msg:
163
+ return None, error_msg, None
164
+ return output_path, "PNG metadata watermark added successfully.", output_path
165
+ except Exception as e:
166
+ return None, f"An unexpected error occurred: {e}", None
167
+
168
+
169
+ def detect_watermark(det_im):
170
+ if not det_im:
171
+ return "Please upload an image."
172
+
173
+ detected_text = decode(det_im)
174
+ if not detected_text or detected_text.startswith("Error"):
175
+ return "No watermark detected or not encoded with this tool."
176
+
177
+ return detected_text
178
+
179
+
180
+ # Gradio app
181
+ with gr.Blocks() as app:
182
+ with gr.Tab("Add Watermark"):
183
+ cho = gr.Radio(choices=["stegan", "pnginfo"], value="stegan", label="Watermark Method")
184
+ with gr.Row():
185
+ with gr.Column():
186
+ inp_im = gr.Image(label="Input Image", type="filepath")
187
+ inp_mark = gr.Textbox(label="Watermark Text")
188
+ mark_btn = gr.Button("Add Watermark")
189
+ msg_box = gr.Textbox(label="System Message")
190
+ file_output = gr.File(label="Download Watermarked Image")
191
+ with gr.Column():
192
+ out_im = gr.Image(label="Watermarked Image")
193
+
194
+ mark_btn.click(choose_encode, inputs=[inp_im, inp_mark, cho], outputs=[out_im, msg_box, file_output])
195
+
196
+ with gr.Tab("Detect Watermark"):
197
+ with gr.Row():
198
+ with gr.Column():
199
+ det_im = gr.Image(label="Watermarked Image", type="filepath")
200
+ det_btn = gr.Button("Detect Watermark")
201
+ with gr.Column():
202
+ det_msg = gr.Textbox(label="Detected Watermark", lines=6)
203
+
204
+ det_btn.click(detect_watermark, inputs=[det_im], outputs=[det_msg])
205
+
206
+ if __name__ == "__main__":
207
+ app.launch()
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ gradio==3.50.2
2
+ pillow==10.0.0
3
+ numpy==1.24.3
4
+ opencv-python-headless==4.8.0.76
utils.py ADDED
@@ -0,0 +1,120 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ from PIL import Image, PngImagePlugin
4
+ import os
5
+ import uuid
6
+
7
+ def to_bin(data):
8
+ if isinstance(data, str):
9
+ return ''.join([format(ord(i), "08b") for i in data])
10
+ elif isinstance(data, bytes):
11
+ return ''.join([format(i, "08b") for i in data])
12
+ elif isinstance(data, np.ndarray):
13
+ return [format(i, "08b") for i in data]
14
+ elif isinstance(data, int) or isinstance(data, np.uint8):
15
+ return format(data, "08b")
16
+ else:
17
+ raise TypeError("Type not supported.")
18
+
19
+ def encode(image_path, secret_data):
20
+ try:
21
+ # Read the image
22
+ image = cv2.imread(image_path)
23
+ if image is None:
24
+ return None, "Failed to read image"
25
+
26
+ # Calculate maximum bytes for encoding
27
+ n_bytes = (image.shape[0] * image.shape[1] * 3) // 8
28
+
29
+ # Prepare secret data
30
+ secret_data_with_delimiter = f'{secret_data}#####'
31
+ if len(secret_data_with_delimiter) + 5 >= n_bytes:
32
+ return None, "Watermark is too large for Image Size"
33
+
34
+ secret_data_with_delimiter += "====="
35
+ binary_secret_data = to_bin(secret_data_with_delimiter)
36
+ data_len = len(binary_secret_data)
37
+
38
+ # Create a copy of the image
39
+ watermarked_image = image.copy()
40
+ data_index = 0
41
+
42
+ # Encode the data
43
+ for i in range(watermarked_image.shape[0]):
44
+ for j in range(watermarked_image.shape[1]):
45
+ for k in range(3): # RGB channels
46
+ if data_index < data_len:
47
+ # Get the current pixel value
48
+ pixel = watermarked_image[i, j, k]
49
+ # Convert to binary and modify the least significant bit
50
+ binary_pixel = format(pixel, '08b')
51
+ new_pixel = binary_pixel[:-1] + binary_secret_data[data_index]
52
+ # Update the pixel value
53
+ watermarked_image[i, j, k] = int(new_pixel, 2)
54
+ data_index += 1
55
+ else:
56
+ break
57
+ if data_index >= data_len:
58
+ break
59
+ if data_index >= data_len:
60
+ break
61
+
62
+ return watermarked_image, None
63
+ except Exception as e:
64
+ return None, f"Error during encoding: {str(e)}"
65
+
66
+ def decode(image_path):
67
+ try:
68
+ # Read the image
69
+ image = cv2.imread(image_path)
70
+ if image is None:
71
+ return "Failed to read image"
72
+
73
+ binary_data = ""
74
+ # Extract the least significant bits
75
+ for i in range(image.shape[0]):
76
+ for j in range(image.shape[1]):
77
+ for k in range(3):
78
+ pixel = format(image[i, j, k], '08b')
79
+ binary_data += pixel[-1]
80
+
81
+ # Convert binary to text
82
+ decoded_data = ""
83
+ for i in range(0, len(binary_data), 8):
84
+ byte = binary_data[i:i+8]
85
+ if len(byte) == 8:
86
+ decoded_data += chr(int(byte, 2))
87
+ # Check for ending delimiter
88
+ if decoded_data[-5:] == "=====":
89
+ break
90
+
91
+ # Remove delimiters and return the message
92
+ if "#####" in decoded_data:
93
+ return decoded_data.split("#####")[0]
94
+ return "No watermark found"
95
+
96
+ except Exception as e:
97
+ return f"Error during decoding: {str(e)}"
98
+
99
+ def png_encode(im_name, extra):
100
+ try:
101
+ im = Image.open(im_name)
102
+ info = PngImagePlugin.PngInfo()
103
+ info.add_text("TXT", extra)
104
+
105
+ unique_id = str(uuid.uuid4())[:8]
106
+ filename, ext = os.path.splitext(im_name)
107
+ new_filename = f"{filename}_{unique_id}{ext}"
108
+ im.save(new_filename, "PNG", pnginfo=info)
109
+
110
+ width, height = im.size
111
+ rect_width, rect_height = 200, 50
112
+ x = width - rect_width - 10
113
+ y = height - rect_height - 10
114
+ global png_encode_coords
115
+ png_encode_coords = (x, y, rect_width, rect_height)
116
+
117
+ return new_filename, None
118
+
119
+ except Exception as e:
120
+ return None, str(e)