fantos commited on
Commit
2697417
ยท
verified ยท
1 Parent(s): 6528fbf

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +29 -24
utils.py CHANGED
@@ -7,8 +7,7 @@ def png_encode(im_name, extra):
7
  try:
8
  im = Image.open(im_name)
9
  info = PngImagePlugin.PngInfo()
10
- # Encode the text as UTF-8 before adding to metadata
11
- info.add_text("TXT", extra.encode('utf-8').decode('utf-8'))
12
  im.save("test.png", pnginfo=info)
13
  return "test.png", "Watermark added successfully"
14
  except Exception as e:
@@ -17,9 +16,10 @@ def png_encode(im_name, extra):
17
  def to_bin(data):
18
  """Convert data to binary format as string"""
19
  if isinstance(data, str):
20
- return ''.join([format(ord(i), "08b") for i in data])
 
21
  elif isinstance(data, bytes):
22
- return ''.join([format(i, "08b") for i in data])
23
  elif isinstance(data, np.ndarray):
24
  return [format(i, "08b") for i in data]
25
  elif isinstance(data, int) or isinstance(data, np.uint8):
@@ -38,7 +38,7 @@ def decode(image_name, txt=None):
38
  except:
39
  pass
40
 
41
- # If PNG metadata fails, try steganography method
42
  image = cv2.imread(image_name)
43
  if image is None:
44
  raise ValueError("Could not read image file")
@@ -53,23 +53,27 @@ def decode(image_name, txt=None):
53
  binary_data += g[-1]
54
  binary_data += b[-1]
55
 
56
- # Process binary data in chunks of 8 bits
57
- decoded_data = ""
58
  for i in range(0, len(binary_data), 8):
59
  byte = binary_data[i:i+8]
60
- if len(byte) == 8: # Ensure we have a complete byte
61
  try:
62
- decoded_data += chr(int(byte, 2))
63
  except ValueError:
64
  continue
65
 
66
- if decoded_data.endswith("====="):
67
- break
68
-
69
- if "#####" in decoded_data:
70
- decoded_data = decoded_data.split("#####")[0]
71
-
72
- return decoded_data.strip().rstrip("=")
 
 
 
 
73
 
74
  except Exception as e:
75
  return f"Error detecting watermark: {str(e)}"
@@ -86,17 +90,18 @@ def encode(image_name, secret_data, txt=None):
86
  # Calculate maximum bytes that can be encoded
87
  n_bytes = image.shape[0] * image.shape[1] * 3 // 8
88
 
89
- # Prepare the secret data
90
- secret_data = str(secret_data).encode('utf-8').decode('utf-8') # Ensure proper encoding
91
- secret_data_with_delimiter = secret_data + "#####"
 
92
 
93
- if len(secret_data_with_delimiter) + 5 >= n_bytes:
 
 
94
  return image_name, "Watermark is too large for Image Size"
95
 
96
- secret_data_with_delimiter += "====="
97
-
98
- # Convert secret data to binary
99
- binary_secret_data = to_bin(secret_data_with_delimiter)
100
  data_len = len(binary_secret_data)
101
  data_index = 0
102
 
 
7
  try:
8
  im = Image.open(im_name)
9
  info = PngImagePlugin.PngInfo()
10
+ info.add_text("TXT", extra)
 
11
  im.save("test.png", pnginfo=info)
12
  return "test.png", "Watermark added successfully"
13
  except Exception as e:
 
16
  def to_bin(data):
17
  """Convert data to binary format as string"""
18
  if isinstance(data, str):
19
+ # UTF-8๋กœ ์ธ์ฝ”๋”ฉํ•˜์—ฌ ๋ฐ”์ดํŠธ๋กœ ๋ณ€ํ™˜ ํ›„ ๋ฐ”์ด๋„ˆ๋ฆฌ ๋ฌธ์ž์—ด๋กœ ๋ณ€ํ™˜
20
+ return ''.join(format(b, '08b') for b in data.encode('utf-8'))
21
  elif isinstance(data, bytes):
22
+ return ''.join(format(i, "08b") for i in data)
23
  elif isinstance(data, np.ndarray):
24
  return [format(i, "08b") for i in data]
25
  elif isinstance(data, int) or isinstance(data, np.uint8):
 
38
  except:
39
  pass
40
 
41
+ # Steganography method
42
  image = cv2.imread(image_name)
43
  if image is None:
44
  raise ValueError("Could not read image file")
 
53
  binary_data += g[-1]
54
  binary_data += b[-1]
55
 
56
+ # 8๋น„ํŠธ์”ฉ ๋ฌถ์–ด์„œ ๋ฐ”์ดํŠธ ๋ฐฐ์—ด ์ƒ์„ฑ
57
+ bytes_array = bytearray()
58
  for i in range(0, len(binary_data), 8):
59
  byte = binary_data[i:i+8]
60
+ if len(byte) == 8:
61
  try:
62
+ bytes_array.append(int(byte, 2))
63
  except ValueError:
64
  continue
65
 
66
+ # UTF-8๋กœ ๋””์ฝ”๋”ฉ
67
+ try:
68
+ decoded_data = bytes_array.decode('utf-8')
69
+ if "=====" in decoded_data:
70
+ decoded_data = decoded_data.split("=====")[0]
71
+ if "#####" in decoded_data:
72
+ decoded_data = decoded_data.split("#####")[0]
73
+ return decoded_data
74
+ except UnicodeDecodeError:
75
+ # UTF-8 ๋””์ฝ”๋”ฉ์— ์‹คํŒจํ•œ ๊ฒฝ์šฐ
76
+ return "Error: Invalid character encoding detected"
77
 
78
  except Exception as e:
79
  return f"Error detecting watermark: {str(e)}"
 
90
  # Calculate maximum bytes that can be encoded
91
  n_bytes = image.shape[0] * image.shape[1] * 3 // 8
92
 
93
+ # UTF-8๋กœ ์ธ์ฝ”๋”ฉํ•˜์—ฌ ๋ฐ”์ดํŠธ ์ˆ˜ ๊ณ„์‚ฐ
94
+ secret_bytes = secret_data.encode('utf-8')
95
+ delimiter_bytes = "#####".encode('utf-8')
96
+ end_marker_bytes = "=====".encode('utf-8')
97
 
98
+ total_bytes = len(secret_bytes) + len(delimiter_bytes) + len(end_marker_bytes)
99
+
100
+ if total_bytes >= n_bytes:
101
  return image_name, "Watermark is too large for Image Size"
102
 
103
+ # Prepare binary data
104
+ binary_secret_data = to_bin(secret_data) + to_bin("#####") + to_bin("=====")
 
 
105
  data_len = len(binary_secret_data)
106
  data_index = 0
107