edit_propagation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3306ccce-4b17-41a9-831d-add6cccddc0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import gc\n",
    "import imageio\n",
    "from PIL import Image\n",
    "from torchvision.transforms import PILToTensor\n",
    "import os\n",
    "import json\n",
    "from PIL import Image, ImageDraw\n",
    "import torch.nn.functional as F\n",
    "import cv2\n",
    "import glob\n",
    "from torchvision.transforms import PILToTensor\n",
    "from src.models.dift_sd import SDFeaturizer4Eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "081cd585-9d9d-4ffe-8c9b-6c6360d2e4ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_grid(h, w, device, normalize=False, homogeneous=False):\n",
    "    if normalize:\n",
    "        lin_y = torch.linspace(-1., 1., steps=h, device=device)\n",
    "        lin_x = torch.linspace(-1., 1., steps=w, device=device)\n",
    "    else:\n",
    "        lin_y = torch.arange(0, h, device=device)\n",
    "        lin_x = torch.arange(0, w, device=device)\n",
    "    grid_y, grid_x = torch.meshgrid((lin_y, lin_x))\n",
    "    grid = torch.stack((grid_x, grid_y), -1)\n",
    "    if homogeneous:\n",
    "        grid = torch.cat([grid, torch.ones_like(grid[..., :1])], dim=-1)\n",
    "    return grid  # [h, w, 2 or 3]\n",
    "\n",
    "\n",
    "def normalize_coords(coords, h, w, no_shift=False):\n",
    "    assert coords.shape[-1] == 2\n",
    "    if no_shift:\n",
    "        return coords / torch.tensor([w-1., h-1.], device=coords.device) * 2\n",
    "    else:\n",
    "        return coords / torch.tensor([w-1., h-1.], device=coords.device) * 2 - 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a13b459-4698-4a9c-803f-d7ba8adb6962",
   "metadata": {},
   "outputs": [],
   "source": [
    "cat = 'cat'\n",
    "dift = SDFeaturizer4Eval(cat_list=['cat'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0606e9dd-9e51-49ec-bf37-1f2bc9f78a84",
   "metadata": {},
   "outputs": [],
   "source": [
    "src_img = Image.open('./assets/guitar_cat.jpg').convert('RGB')\n",
    "trg_img = Image.open('./assets/painting_cat.jpg').convert('RGB')\n",
    "sticker = imageio.imread('./assets/cartoon.png')\n",
    "sticker_color, sticker_mask = sticker[..., :3], sticker[..., 3]\n",
    "\n",
    "assert np.array(src_img).shape[:2] == sticker.shape[:2]\n",
    "h_src, w_src = sticker.shape[:2]\n",
    "h_trg, w_trg = np.array(trg_img).shape[:2]\n",
    "\n",
    "sd_feat_src = dift.forward(src_img, cat)\n",
    "sd_feat_trg = dift.forward(trg_img, cat)\n",
    "\n",
    "sd_feat_src = F.normalize(sd_feat_src.squeeze(), p=2, dim=0)\n",
    "sd_feat_trg = F.normalize(sd_feat_trg.squeeze(), p=2, dim=0)\n",
    "feat_dim = sd_feat_src.shape[0]\n",
    "\n",
    "grid_src = gen_grid(h_src, w_src, device='cuda')\n",
    "grid_trg = gen_grid(h_trg, w_trg, device='cuda')\n",
    "\n",
    "coord_src = grid_src[sticker_mask > 0]\n",
    "coord_src = coord_src[torch.randperm(len(coord_src))][:1000]\n",
    "coord_src_normed = normalize_coords(coord_src, h_src, w_src)\n",
    "grid_trg_normed = normalize_coords(grid_trg, h_trg, w_trg)\n",
    "\n",
    "feat_src = F.grid_sample(sd_feat_src[None], coord_src_normed[None, None], align_corners=True).squeeze().T\n",
    "feat_trg = F.grid_sample(sd_feat_trg[None], grid_trg_normed[None], align_corners=True).squeeze()\n",
    "feat_trg_flattened = feat_trg.permute(1, 2, 0).reshape(-1, feat_dim)\n",
    "\n",
    "distances = torch.cdist(feat_src, feat_trg_flattened)\n",
    "_, indices = torch.min(distances, dim=1)\n",
    "\n",
    "src_pts = coord_src.reshape(-1, 2).cpu().numpy()\n",
    "trg_pts = grid_trg.reshape(-1, 2)[indices].cpu().numpy()\n",
    "\n",
    "M, mask = cv2.findHomography(src_pts, trg_pts, cv2.RANSAC, 5.0)\n",
    "sticker_out = cv2.warpPerspective(sticker, M, (w_trg, h_trg))\n",
    "\n",
    "sticker_out_alpha = sticker_out[..., 3:] / 255\n",
    "sticker_alpha = sticker[..., 3:] / 255\n",
    "\n",
    "trg_img_with_sticker = sticker_out_alpha * sticker_out[..., :3] + (1 - sticker_out_alpha) * trg_img\n",
    "src_img_with_sticker = sticker_alpha * sticker[..., :3] + (1 - sticker_alpha) * src_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88723600-c18f-4eb1-aec7-feb4112e2610",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 2, figsize=(10, 10))\n",
    "\n",
    "axs[0, 0].imshow(src_img)\n",
    "axs[0, 0].set_title(\"Source Image\")\n",
    "axs[0, 0].axis('off')\n",
    "\n",
    "axs[0, 1].imshow(src_img_with_sticker.astype(np.uint8))\n",
    "axs[0, 1].set_title(\"Source Image with Edits\")\n",
    "axs[0, 1].axis('off')\n",
    "\n",
    "axs[1, 0].imshow(trg_img)\n",
    "axs[1, 0].set_title(\"Target Image\")\n",
    "axs[1, 0].axis('off')\n",
    "\n",
    "axs[1, 1].imshow(trg_img_with_sticker.astype(np.uint8))\n",
    "axs[1, 1].set_title(\"Target Image with Propagated Edits\")\n",
    "axs[1, 1].axis('off')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}