Spaces:
Running
Running
File size: 4,336 Bytes
de08da3 |
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
import argparse
from PIL import Image
class Steganography:
BLACK_PIXEL = (0, 0, 0)
def _int_to_bin(self, rgb):
"""Convert an integer tuple to a binary (string) tuple.
:param rgb: An integer tuple like (220, 110, 96)
:return: A string tuple like ("00101010", "11101011", "00010110")
"""
r, g, b = rgb
return f'{r:08b}', f'{g:08b}', f'{b:08b}'
def _bin_to_int(self, rgb):
"""Convert a binary (string) tuple to an integer tuple.
:param rgb: A string tuple like ("00101010", "11101011", "00010110")
:return: Return an int tuple like (220, 110, 96)
"""
r, g, b = rgb
return int(r, 2), int(g, 2), int(b, 2)
def _merge_rgb(self, rgb1, rgb2, digit):
"""Merge two RGB tuples.
:param rgb1: An integer tuple like (220, 110, 96)
:param rgb2: An integer tuple like (240, 95, 105)
:return: An integer tuple with the two RGB values merged.
"""
r1, g1, b1 = self._int_to_bin(rgb1)
r2, g2, b2 = self._int_to_bin(rgb2)
rgb = r1[:digit] + r2[:8-digit], g1[:digit] + g2[:8-digit], b1[:digit] + b2[:8-digit]
return self._bin_to_int(rgb)
def _unmerge_rgb(self, rgb, digit):
"""Unmerge RGB.
:param rgb: An integer tuple like (220, 110, 96)
:return: An integer tuple with the two RGB values merged.
"""
r, g, b = self._int_to_bin(rgb)
# Extract the last 4 bits (corresponding to the hidden image)
# Concatenate 4 zero bits because we are working with 8 bit
new_rgb = r[digit:] + '0'*digit, g[digit:] + '0'*digit, b[digit:] + '0'*digit
return self._bin_to_int(new_rgb)
def merge(self, image1, image2, digit=4):
"""Merge image2 into image1.
:param image1: First image
:param image2: Second image
:return: A new merged image.
"""
# Check the images dimensions
if image2.size[0] > image1.size[0] or image2.size[1] > image1.size[1]:
raise ValueError('Image 2 should be smaller than Image 1!')
# Get the pixel map of the two images
map1 = image1.load()
map2 = image2.load()
new_image = Image.new(image1.mode, image1.size)
new_map = new_image.load()
for i in range(image1.size[0]):
for j in range(image1.size[1]):
is_valid = lambda: i < image2.size[0] and j < image2.size[1]
rgb1 = map1[i ,j]
rgb2 = map2[i, j] if is_valid() else self.BLACK_PIXEL
new_map[i, j] = self._merge_rgb(rgb1, rgb2, digit)
return new_image
def unmerge(self, image, digit=4, binarization=True):
"""Unmerge an image.
:param image: The input image.
:return: The unmerged/extracted image.
"""
pixel_map = image.load()
# Create the new image and load the pixel map
new_image = Image.new(image.mode, image.size)
new_map = new_image.load()
for i in range(image.size[0]):
for j in range(image.size[1]):
r, g, b = self._unmerge_rgb(pixel_map[i, j], digit)
r = 255 if r >= 128 else 0
g = 255 if g >= 128 else 0
b = 255 if b >= 128 else 0
new_map[i, j] = r, g, b
return new_image
def main():
parser = argparse.ArgumentParser(description='Steganography')
subparser = parser.add_subparsers(dest='command')
merge = subparser.add_parser('merge')
merge.add_argument('--image1', required=True, help='Image1 path')
merge.add_argument('--image2', required=True, help='Image2 path')
merge.add_argument('--output', required=True, help='Output path')
unmerge = subparser.add_parser('unmerge')
unmerge.add_argument('--image', required=True, help='Image path')
unmerge.add_argument('--output', required=True, help='Output path')
args = parser.parse_args()
if args.command == 'merge':
image1 = Image.open(args.image1)
image2 = Image.open(args.image2)
Steganography().merge(image1, image2).save(args.output)
elif args.command == 'unmerge':
image = Image.open(args.image)
Steganography().unmerge(image).save(args.output)
if __name__ == '__main__':
main()
|