{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8fa7cf38",
   "metadata": {},
   "source": [
    "# How to create dynamic masks with DALL·E and Segment Anything\n",
    "\n",
    "Segment Anything is a model from Meta that can be used to select portions of images. Combined with DALL·E's ability to inpaint specified portions of images, you can use Segment Anything to easily select any part of an image you'd like to alter.\n",
    "\n",
    "In this notebook, we'll use these tools to become fashion designers and dynamically replace our digital models' outfits with tailored, original creations. The notebook follows this flow:\n",
    "- **Setup:** Initialise your libraries and any location directories.\n",
    "- **Generate original image:** Make an original image that we'll create dynamic masks from.\n",
    "- **Generate mask:** Use Segment Anything to create a dynamic mask.\n",
    "- **Create new image:** Generate a new image with the masked area inpainted with a fresh prompt."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b0d2d8c3",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "To get started we'll need to follow the [instructions](https://github.com/facebookresearch/segment-anything) for using the Segment Anything (SAM) model open-sourced by Meta. As of May 2023, the key steps are:\n",
    "- Install [Pytorch](https://pytorch.org/get-started/locally/) (version 1.7+).\n",
    "- Install the library using ```pip install git+https://github.com/facebookresearch/segment-anything.git```.\n",
    "- Install dependencies using ```pip install opencv-python pycocotools matplotlib onnxruntime onnx```.\n",
    "- Download a [model checkpoint](https://github.com/facebookresearch/segment-anything#model-checkpoints) to use (default size is 2.4 GB).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eefae92",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch torchvision torchaudio\n",
    "!pip install git+https://github.com/facebookresearch/segment-anything.git\n",
    "!pip install opencv-python pycocotools matplotlib onnxruntime onnx\n",
    "!pip install requests\n",
    "!pip install openai\n",
    "!pip install numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fc8b8ed1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "from matplotlib import rcParams\n",
    "import numpy as np\n",
    "import openai\n",
    "import os\n",
    "from PIL import Image\n",
    "import requests\n",
    "from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor\n",
    "import torch\n",
    "\n",
    "# Set directories for generation images and edit images\n",
    "base_image_dir = os.path.join(\"images\", \"01_generations\")\n",
    "mask_dir = os.path.join(\"images\", \"02_masks\")\n",
    "edit_image_dir = os.path.join(\"images\", \"03_edits\")\n",
    "\n",
    "# Point to your downloaded SAM model\n",
    "sam_model_filepath = \"./sam_vit_h_4b8939.pth\"\n",
    "\n",
    "# Initiate SAM model\n",
    "sam = sam_model_registry[\"default\"](checkpoint=sam_model_filepath)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "45d3570d",
   "metadata": {},
   "source": [
    "## Generate original image\n",
    "\n",
    "First we'll create an original image which we'll generate masks from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "11e75cf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_dalle_images(response, filename, image_dir):\n",
    "    # save the images\n",
    "    urls = [datum[\"url\"] for datum in response[\"data\"]]  # extract URLs\n",
    "    images = [requests.get(url).content for url in urls]  # download images\n",
    "    image_names = [f\"{filename}_{i + 1}.png\" for i in range(len(images))]  # create names\n",
    "    filepaths = [os.path.join(image_dir, name) for name in image_names]  # create filepaths\n",
    "    for image, filepath in zip(images, filepaths):  # loop through the variations\n",
    "        with open(filepath, \"wb\") as image_file:  # open the file\n",
    "            image_file.write(image)  # write the image to the file\n",
    "\n",
    "    return filepaths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "21a50249",
   "metadata": {},
   "outputs": [],
   "source": [
    "dalle_prompt = '''\n",
    "Full length, zoomed out photo of our premium Lederhosen-inspired jumpsuit. \n",
    "Showcase the intricate hand-stitched details and high-quality leather, while highlighting the perfect blend of Austrian heritage and modern fashion. \n",
    "This piece appeals to a sophisticated, trendsetting audience who appreciates cultural fusion and innovative design.\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "12e0a11a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate your images\n",
    "generation_response = openai.Image.create(\n",
    "    prompt=dalle_prompt,\n",
    "    n=3,\n",
    "    size=\"1024x1024\",\n",
    "    response_format=\"url\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a53df44b",
   "metadata": {},
   "outputs": [],
   "source": [
    "filepaths = process_dalle_images(generation_response, \"generation\", base_image_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d1d78b52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "images/01_generations/generation_1.png\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 908544 characters]",
      "text/plain": [
       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=1024x1024>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "images/01_generations/generation_2.png\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 1490452 characters]",
      "text/plain": [
       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=1024x1024>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "images/01_generations/generation_3.png\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 1485284 characters]",
      "text/plain": [
       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=1024x1024>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# print the new generations\n",
    "for filepath in filepaths:\n",
    "    print(filepath)\n",
    "    display(Image.open(filepath))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "267beab0",
   "metadata": {},
   "source": [
    "## Generate Mask\n",
    "\n",
    "Next we'll load up one of our images and generate masks.\n",
    "\n",
    "For this demonstration we're picking a UX where we \"click\" on a point on the image to generate masks from. However, there are [example notebooks](https://github.com/facebookresearch/segment-anything/blob/main/notebooks/automatic_mask_generator_example.ipynb) provided by Meta which show how to generate every possible mask for an image, draw a box, and some other useful approaches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b213f887",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'images/01_generations/generation_2.png'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Pick one of your generated images\n",
    "chosen_image = \"images/01_generations/generation_2.png\"\n",
    "chosen_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b05e1b03",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to display mask using matplotlib\n",
    "def show_mask(mask, ax):\n",
    "    color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "\n",
    "\n",
    "# Function to display where we've \"clicked\"\n",
    "def show_points(coords, labels, ax, marker_size=375):\n",
    "    pos_points = coords[labels == 1]\n",
    "    neg_points = coords[labels == 0]\n",
    "    ax.scatter(\n",
    "        pos_points[:, 0],\n",
    "        pos_points[:, 1],\n",
    "        color=\"green\",\n",
    "        marker=\"*\",\n",
    "        s=marker_size,\n",
    "        edgecolor=\"white\",\n",
    "        linewidth=1.25,\n",
    "    )\n",
    "    ax.scatter(\n",
    "        neg_points[:, 0],\n",
    "        neg_points[:, 1],\n",
    "        color=\"red\",\n",
    "        marker=\"*\",\n",
    "        s=marker_size,\n",
    "        edgecolor=\"white\",\n",
    "        linewidth=1.25,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a1fc1dbd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 834196 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Load chosen image using opencv\n",
    "image = cv2.imread(chosen_image)\n",
    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "# Display our chosen image\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(image)\n",
    "plt.axis(\"on\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ff9ed4d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 834664 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Set the pixel coordinates for our \"click\" to assign masks\n",
    "input_point = np.array([[525, 325]])\n",
    "input_label = np.array([1])\n",
    "\n",
    "# Display the point we've clicked on\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(image)\n",
    "show_points(input_point, input_label, plt.gca())\n",
    "plt.axis(\"on\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "2cce9fa5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3, 1024, 1024)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initiate predictor with Segment Anything model\n",
    "predictor = SamPredictor(sam)\n",
    "predictor.set_image(image)\n",
    "\n",
    "# Use the predictor to gather masks for the point we clicked\n",
    "masks, scores, logits = predictor.predict(\n",
    "    point_coords=input_point,\n",
    "    point_labels=input_label,\n",
    "    multimask_output=True,\n",
    ")\n",
    "\n",
    "# Check the shape - should be three masks of the same dimensions as our image\n",
    "masks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "809df45a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 822808 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 801656 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 756616 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Display the possible masks we can select along with their confidence\n",
    "for i, (mask, score) in enumerate(zip(masks, scores)):\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(image)\n",
    "    show_mask(mask, plt.gca())\n",
    "    show_points(input_point, input_label, plt.gca())\n",
    "    plt.title(f\"Mask {i+1}, Score: {score:.3f}\", fontsize=18)\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "42b97534",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose which mask you'd like to use\n",
    "chosen_mask = masks[1]\n",
    "\n",
    "# We'll now reverse the mask so that it is clear and everything else is white\n",
    "chosen_mask = chosen_mask.astype(\"uint8\")\n",
    "chosen_mask[chosen_mask != 0] = 255\n",
    "chosen_mask[chosen_mask == 0] = 1\n",
    "chosen_mask[chosen_mask == 255] = 0\n",
    "chosen_mask[chosen_mask == 1] = 255\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "818a6300",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10412 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=1024x1024>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create a base blank mask\n",
    "width = 1024\n",
    "height = 1024\n",
    "mask = Image.new(\"RGBA\", (width, height), (0, 0, 0, 1))  # create an opaque image mask\n",
    "\n",
    "# Convert mask back to pixels to add our mask replacing the third dimension\n",
    "pix = np.array(mask)\n",
    "pix[:, :, 3] = chosen_mask\n",
    "\n",
    "# Convert pixels back to an RGBA image and display\n",
    "new_mask = Image.fromarray(pix, \"RGBA\")\n",
    "new_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "24c1ec1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll save this mask for re-use for our edit\n",
    "new_mask.save(os.path.join(mask_dir, \"new_mask.png\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70babb01",
   "metadata": {},
   "source": [
    "## Create new image\n",
    "\n",
    "Now we'll combine our original image with the mask and the Edit endpoint for DALLE to inpaint the transparent area according to a new prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "ac812c94",
   "metadata": {},
   "outputs": [],
   "source": [
    "# edit an image\n",
    "\n",
    "# call the OpenAI API\n",
    "edit_response = openai.Image.create_edit(\n",
    "    image=open(chosen_image, \"rb\"),  # from the generation section\n",
    "    mask=open(os.path.join(mask_dir, \"new_mask.png\"), \"rb\"),  # from right above\n",
    "    prompt=\"Brilliant leather Lederhosen with a formal look, detailed, intricate, photorealistic\",  # provide a prompt to fill the space\n",
    "    n=3,\n",
    "    size=\"1024x1024\",\n",
    "    response_format=\"url\",\n",
    ")\n",
    "\n",
    "edit_filepaths = process_dalle_images(edit_response, \"edits\", edit_image_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ea775c1b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x16d8a63b0>"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 123324 characters]",
      "text/plain": [
       "<Figure size 1100x800 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Display your beautiful creations!\n",
    "%matplotlib inline\n",
    "\n",
    "# figure size in inches optional\n",
    "rcParams[\"figure.figsize\"] = 11 ,8\n",
    "\n",
    "# read images\n",
    "img_A = mpimg.imread(edit_filepaths[0])\n",
    "img_B = mpimg.imread(edit_filepaths[1])\n",
    "img_C = mpimg.imread(edit_filepaths[2])\n",
    "\n",
    "# display images\n",
    "fig, ax = plt.subplots(1,3)\n",
    "[a.axis(\"off\") for a in ax]\n",
    "ax[0].imshow(img_A)\n",
    "ax[1].imshow(img_B)\n",
    "ax[2].imshow(img_C)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f516c7d9",
   "metadata": {},
   "source": [
    "Beautiful!\n",
    "\n",
    "Now you too can easily create dynamic masks to extend your images - enjoy the APIs, and please share what you build!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "image_env",
   "language": "python",
   "name": "image_env"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
