Demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9899c070",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "from IPython.display import clear_output\n",
    "import torch.nn.functional as F\n",
    "import PIL.Image\n",
    "from io import BytesIO\n",
    "import IPython.display\n",
    "import numpy as np\n",
    "import skimage.io\n",
    "import image_transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "593c3dbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_dir = \"sample_images\"\n",
    "target_image_dir = \"sample_target_images\"\n",
    "out_dir = \"out_dir\"\n",
    "encoder_model_path = \"encoder_model.pth\"\n",
    "decoder_model_path = \"decoder_model.pth\"\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e48f8d46",
   "metadata": {},
   "source": [
    "### Initialize Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1229030a",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_model = torch.jit.load(\"encoder_model.pth\")\n",
    "decoder_model = torch.jit.load(\"decoder_model.pth\")\n",
    "encoder_model.eval().to(device)\n",
    "decoder_model.eval().to(device)\n",
    "clear_output()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4671df24",
   "metadata": {},
   "source": [
    "### Utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "320efc44",
   "metadata": {},
   "outputs": [],
   "source": [
    "def showarray(a, fmt='png'):\n",
    "    \"\"\"\n",
    "    takes a numpy array (0 to 1) of size h, w, 3\n",
    "    \"\"\"\n",
    "    a = np.uint8(a*255.)\n",
    "    f = BytesIO()\n",
    "    PIL.Image.fromarray(a).save(f, fmt)\n",
    "    IPython.display.display(IPython.display.Image(data=f.getvalue()))\n",
    "    \n",
    "def text_to_bits(text, encoding='utf-8', errors='surrogatepass'):\n",
    "    bits = bin(int.from_bytes(text.encode(encoding, errors), 'big'))[2:]\n",
    "    return bits.zfill(8 * ((len(bits) + 7) // 8))\n",
    "\n",
    "def text_from_bits(bits, encoding='utf-8', errors='surrogatepass'):\n",
    "    n = int(bits, 2)\n",
    "    return n.to_bytes((n.bit_length() + 7) // 8, 'big').decode(encoding, errors) or '\\0'\n",
    "\n",
    "def load_images(image_filepaths, img_size=256):\n",
    "    image_batch_np = []\n",
    "    for file_path in image_filepaths:\n",
    "        image_from_file = skimage.io.imread(file_path)/255.0\n",
    "        image_from_file = image_from_file[:, :, :3]\n",
    "        image_batch_np.append(image_from_file)\n",
    "    image_batch_np = np.stack(image_batch_np, axis=0)\n",
    "    image_batch = torch.from_numpy(image_batch_np).float()\n",
    "    image_batch = image_batch.permute(0, 3, 1, 2)\n",
    "\n",
    "    h, w = image_batch.shape[2:]\n",
    "    if h > w:\n",
    "        image_batch = image_batch[:, :, int((h-w)/2):int((h+w)/2), :]\n",
    "    elif w > h:\n",
    "        image_batch = image_batch[:, :, :, int((w-h)/2):int((w+h)/2)]\n",
    "    image_batch = F.interpolate(image_batch, size=(img_size, img_size), mode='bilinear', align_corners=True)\n",
    "\n",
    "    return image_batch\n",
    "\n",
    "def save_images(image_batch, out_dir, prefix=\"\"):\n",
    "    if not os.path.exists(out_dir):\n",
    "        os.makedirs(out_dir)\n",
    "    \n",
    "    image_paths = []\n",
    "    for img_idx in range(image_batch.shape[0]):\n",
    "        image_np = image_batch[img_idx].permute(1, 2, 0).cpu().numpy()\n",
    "        image_np = np.uint8(image_np*255.)\n",
    "        file_path = os.path.join(out_dir, \"{}_{}.png\".format(prefix, img_idx))\n",
    "        PIL.Image.fromarray(image_np).save(file_path)\n",
    "        image_paths.append(file_path)\n",
    "    \n",
    "    return image_paths\n",
    "\n",
    "def find_image_paths(image_dir):\n",
    "    image_paths = []\n",
    "    for img_file in os.listdir(image_dir):\n",
    "        if img_file.endswith(\".png\") or img_file.endswith(\".jpg\"):\n",
    "            image_paths.append(os.path.join(image_dir, img_file))\n",
    "    image_paths.sort()\n",
    "    return image_paths\n",
    "\n",
    "def decode_images(image_paths, secret_numpy, decoder_model):\n",
    "    image_batch = load_images(image_paths)\n",
    "    with torch.no_grad():\n",
    "        image_batch = image_batch.to(device)\n",
    "        decoded_secrets, _ = decoder_model(image_batch)\n",
    "\n",
    "    predicted_secrets = (F.sigmoid(decoded_secrets) > 0.5).long()\n",
    "    secrets = torch.from_numpy(secret_numpy).repeat(predicted_secrets.shape[0], 1).to(device) \n",
    "    secret_accuracy = (predicted_secrets == secrets).float().mean().item()\n",
    "\n",
    "    decoding_results = []\n",
    "    for img_idx, image_path in enumerate(image_paths):\n",
    "        image_predicted_secret = predicted_secrets[img_idx].cpu().numpy().tolist()\n",
    "        image_predicted_secret_bits = \"\".join([str(b) for b in image_predicted_secret][:secrete_num_bits])\n",
    "        try:\n",
    "            image_predicted_secret_text = text_from_bits(image_predicted_secret_bits)\n",
    "        except:\n",
    "            image_predicted_secret_text = \"could not decode\"\n",
    "        \n",
    "        decoding_results.append({\n",
    "            \"image_path\": image_path,\n",
    "            \"image_predicted_secret_text\": image_predicted_secret_text,\n",
    "            \"bit_accuracy\": (predicted_secrets[img_idx] == secrets[img_idx]).float().mean().item()\n",
    "        })\n",
    "    \n",
    "    return secret_accuracy, decoding_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "259ef177",
   "metadata": {},
   "source": [
    "## Sign images with a secret and visualie the encoded images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7634d560",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "secret_text = \"sample\"\n",
    "secret_size = 128\n",
    "secret_bits = text_to_bits(secret_text)\n",
    "secrete_num_bits = len(secret_bits)\n",
    "\n",
    "assert secrete_num_bits <= secret_size\n",
    "\n",
    "secret_bits = secret_bits + \"\".join([\"0\"]*(secret_size-secrete_num_bits))\n",
    "secret_numpy = np.array([[ int(c) for c in  secret_bits ]])\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "original_image_paths = find_image_paths(img_dir)\n",
    "original_image_paths = original_image_paths[:3]\n",
    "images = load_images(original_image_paths)\n",
    "\n",
    "images = images.to(device)\n",
    "secrets = torch.from_numpy(secret_numpy).repeat(images.shape[0], 1).to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    encoded_images, secret_images = encoder_model(images, secrets)\n",
    "    signed_image_dir = os.path.join(out_dir, \"signed_images\")\n",
    "    encoded_image_paths = save_images(encoded_images, signed_image_dir)\n",
    "\n",
    "\n",
    "for sidx in range(len(encoded_image_paths)):\n",
    "    original_image_numpy = images[sidx].permute(1, 2, 0).cpu().numpy() \n",
    "    encoded_image_numpy = encoded_images[sidx].permute(1, 2, 0).cpu().numpy()\n",
    "    residual = (encoded_image_numpy - original_image_numpy)\n",
    "    rmin, rmax = np.min(residual), np.max(residual)\n",
    "    residual_scaled = (residual-rmin)/(rmax - rmin)\n",
    "    original_encoded_image = np.concatenate( (original_image_numpy, encoded_image_numpy, residual_scaled), axis=1)\n",
    "    print(\"Original Image,\", \"Signed Image,\", \"Perturbation (Scaled for Visualization)\")\n",
    "    showarray(original_encoded_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bebb16bc",
   "metadata": {},
   "source": [
    "## Apply benign transformations on signed images and decode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ce63ae6",
   "metadata": {},
   "outputs": [],
   "source": [
    "benign_dir = os.path.join(out_dir, \"benign_transformed_images\")\n",
    "benign_tranform_list, benign_file_paths = image_transforms.apply_benign_transforms(encoded_image_paths, benign_dir)\n",
    "\n",
    "for key in benign_tranform_list:\n",
    "    secret_accuracy, decoding_results = decode_images(benign_file_paths[key], secret_numpy, decoder_model)\n",
    "    for row in decoding_results[:1]:\n",
    "        IPython.display.display(IPython.display.Image(row['image_path']))\n",
    "        print(\"Transform : {}\".format(key))\n",
    "        print(\"Predicted secret: {}\".format(row['image_predicted_secret_text']))\n",
    "        print(\"Bit accuracy: {}\".format(row['bit_accuracy']))\n",
    "        print(\"Image path: {}\".format(row['image_path']))\n",
    "    print (\"-----------------------------------------------------\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0d35928",
   "metadata": {},
   "source": [
    "### Apply malicious (face-swap) transform on signed images and decode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53f58281",
   "metadata": {},
   "outputs": [],
   "source": [
    "mal_dir = os.path.join(out_dir, \"mal_transformed_images\")\n",
    "target_image_paths = find_image_paths(target_image_dir)\n",
    "mal_tranform_list, mal_file_paths = image_transforms.apply_malicious_transforms(encoded_image_paths[1:], target_image_paths[1:2], mal_dir)\n",
    "\n",
    "for key in mal_tranform_list:\n",
    "    secret_accuracy, decoding_results = decode_images(mal_file_paths[key], secret_numpy, decoder_model)\n",
    "    for row in decoding_results:\n",
    "        print(\"Transform : {}\".format(key))\n",
    "        print(\"Predicted secret: {}\".format(row['image_predicted_secret_text']))\n",
    "        print(\"Bit accuracy: {}\".format(row['bit_accuracy']))\n",
    "        IPython.display.display(IPython.display.Image(row['image_path']))\n",
    "        print (\"-----------------------------------------------------\\n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05a467e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "alternate_deepfakes_dir = \"alternate_deepfakes\"\n",
    "alternate_deepfakes_dir_filepaths = find_image_paths(alternate_deepfakes_dir)\n",
    "secret_accuracy, decoding_results = decode_images(alternate_deepfakes_dir_filepaths, secret_numpy, decoder_model)\n",
    "for row in decoding_results:\n",
    "    print(\"Transform : {}\".format(key))\n",
    "    print(\"Predicted secret: {}\".format(row['image_predicted_secret_text']))\n",
    "    print(\"Bit accuracy: {}\".format(row['bit_accuracy']))\n",
    "    IPython.display.display(IPython.display.Image(row['image_path']))\n",
    "    print (\"-\"*100)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "envDF",
   "language": "python",
   "name": "envdf"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}