Image watermarker

colors.py

class LightColors:
  ONE_COLOR=(254, 254, 254)
  ZERO_COLOR=(253, 253, 253)
  BG_COLOR=(255, 255, 255)

class DarkColors:
  ONE_COLOR=(2, 2, 2)
  ZERO_COLOR=(1, 1, 1)
  BG_COLOR=(0, 0, 0)

read_watermark.py

from PIL import Image
from colors import LightColors, DarkColors

def decode_watermark(input_path: str = "./out.png"):
  """
    Decodes a UID from an image by reversing a watermarking technique where
    the UID is encoded into the first row of pixels, and each pixel is a color
    that corresponds to a bit in the UID.
  """

  # Load the image and convert it to RGB, then get the pixel map
  img = Image.open(input_path).convert("RGB")
  pixels = img.load()

  theme = "day" if pixels[0, 1] == (255, 255, 255) else "night"
  colors = LightColors if theme == "day" else DarkColors

  result = ""
  finished = False
  idx = 0

  while not finished:
    pixel = pixels[idx, 0]

    # Set THEME
    if theme == "" and pixels[idx, 1] == (255, 255, 255):
      theme = "day"
    else:
      theme = "night"

    if pixel == colors.ONE_COLOR:
      result += "1"
    elif pixel == colors.ZERO_COLOR:
      result += "0"
    elif pixel == colors.BG_COLOR:
      if pixels[idx + 1, 0] == colors.BG_COLOR:
        finished = True
      else:
        result += " "
    else:
      raise Exception("Unknown color")

    idx += 1

  uid = "".join([chr(int(x, 2)) for x in result.split(" ")])

  print(f"Found UID: {uid}")

if __name__ == "__main__":
  import sys
  if len(sys.argv) == 2:
    decode_watermark(sys.argv[1])
  else:
    print("Usage: python read_watermark.py <input_path>")

watermark.py

from PIL import Image
from colors import LightColors, DarkColors

def index_to_color(idx: int, image: Image.Image) -> tuple:
  """Grabs the color at a given index in the palette of a provided image"""
  return tuple(image.getpalette()[idx * 3:idx * 3 + 3])

def encode_watermark(uid: str, input_path: str, output_path: str = "./out.png"):
  """
    Encodes a UID into an image using a watermarking technique.
    The UID is encoded into the first row of pixels, where each pixel is a color
    that corresponds to a bit in the UID.
  """

  # Convert input to a list of bytes (['1100010', '1100010', ...])
  uid_bytes = [format(x, 'b') for x in bytearray(uid, 'utf-8')]

  # Join the bytes with a space ('1100010 1100010 ...')
  uid_bytes_space = " ".join(uid_bytes)

  # Load the image, then extract the palette & find the index of the last color in the palette
  img = Image.open(input_path)
  og_palette = img.getpalette()
  color_offset = len(og_palette) // 3

  # Get the pixel map
  pixels = img.load()

  # Check if the first pixel is white or black, then set the theme accordingly
  theme = "day" if index_to_color(pixels[0, 0], img) == (255,255,255) else "night"
  print(f"Recognized theme: {theme}")

  # Set the colors based on the theme
  colors = LightColors if theme == "day" else DarkColors

  # Extend the original palette with the colors we need
  img.putpalette([
    *og_palette,
    *colors.ZERO_COLOR,
    *colors.ONE_COLOR,
    *colors.BG_COLOR # TODO: This adds 1kb extra to the final image...
  ])

  # In the first row of pixels, set each pixel to a color that corresponds
  # to the bit at the same index in the uid_bytes_space string
  if len(uid_bytes_space) <= img.size[0]:
    for i, bit in enumerate(uid_bytes_space):
      if bit == " ":
        # "Pause" before new byte - set the pixel to the background color
        pixels[i, 0] = color_offset + 2 # BG_COLOR
        continue

      if bit == "1":
        pixels[i, 0] = color_offset + 1 # ONE_COLOR
      else:
        pixels[i, 0] = color_offset + 0 # ZERO_COLOR

  img.save(output_path)


if __name__ == "__main__":
  import sys
  if len(sys.argv) == 3:
    encode_watermark(sys.argv[1], sys.argv[2])
  elif len(sys.argv) == 4:
    encode_watermark(sys.argv[1], sys.argv[2], sys.argv[3])
  else:
    print("Usage: python watermark.py <uid> <input_path> [output_path]")