object_diffusion_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jermwatt/yolo-diffusion/blob/main/object_diffusion_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kVbKjgHovhAm"
      },
      "source": [
        "# Object detection + stable diffusion in python"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XvZ-6yqoLT9L"
      },
      "source": [
        "This notebook allows for experimentation with object detection + segmentation --> diffusion.\n",
        "\n",
        "To use it first:\n",
        "\n",
        "- select a gpu runtime: make sure you have chosen a GPU runtime from the options above - choose A100 if available.\n",
        "- activate the first few cells to\n",
        "  - perform a machine check: make sure your runtime is employing a GPU with sufficient memory\n",
        "  - download associated files and installs\n",
        "- start experimenting!\n",
        "\n",
        "### Some example experiments\n",
        "\n",
        "Here are the sort of experiments you can perform in this notebook.  On the left below we have a picture with two default yolo-recognizable objects: a `pizza` and a `cart` (there are around 70 pre-trained objects you can play with by default).  Using the functionality here you can replace either one independently.\n",
        "\n",
        "So you can replace the `pizza` as shown below.\n",
        "\n",
        "<img src='https://github.com/jermwatt/yolo-diffusion/blob/main/test_data/cat_cash.png?raw=true'>\n",
        "\n",
        "Or just as easily the `cat` as shown here.\n",
        "\n",
        "<img src='https://github.com/jermwatt/yolo-diffusion/blob/main/test_data/squirrel_pizza.png?raw=true'>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "riciSGvmLT9L"
      },
      "source": [
        "## 1.  machine setup and installs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NqvrZr526EDJ"
      },
      "source": [
        "### 1.1. pre-launch machine check\n",
        "\n",
        "Activate the code block below to double check that you are using a GPU runtime for the experiments.  You don't have to use one if you can't, but the experiments will run significantly faster (in particular the diffusion steps) if you are using one."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XZEkR12-eG-Z"
      },
      "outputs": [],
      "source": [
        "# check for GPU runtime\n",
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Not connected to a GPU')\n",
        "else:\n",
        "  print(gpu_info)\n",
        "\n",
        "# check for memory\n",
        "from psutil import virtual_memory\n",
        "ram_gb = virtual_memory().total / 1e9\n",
        "print('Your runtime has {:.1f} gigabytes of available RAM\\n'.format(ram_gb))\n",
        "\n",
        "if ram_gb < 20:\n",
        "  print('Not using a high-RAM runtime')\n",
        "else:\n",
        "  print('You are using a high-RAM runtime!')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YGc-wRzC6ArW"
      },
      "source": [
        "### 1.2. installs\n",
        "\n",
        "Activate the cell below to install all necessary packages to your collab enviroment for these experiments."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-iHfZS44dDgY"
      },
      "outputs": [],
      "source": [
        "# install reqiured libraries\n",
        "!pip install \"ultralytics==8.0.111\" \"transformers==4.29.2\" \"timm==0.9.2\" \"diffusers==0.16.1\" \"safetensors==0.3.1\" \"accelerate==0.19.0\""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Activate the cell below to pull all helper functions required to run this demo."
      ],
      "metadata": {
        "id": "wtFjIXPcLWwr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# pull some images from the web\n",
        "import urllib.request\n",
        "\n",
        "def download_file(url, output_path):\n",
        "    urllib.request.urlretrieve(url, output_path)\n",
        "\n",
        "# pull segmentation module\n",
        "url = \"https://raw.githubusercontent.com/jermwatt/yolo-diffusion/main/segmenter.py\"\n",
        "local_path = \"/content/segmenter.py\"\n",
        "download_file(url,local_path)\n",
        "\n",
        "# pull diffusion module\n",
        "url = \"https://raw.githubusercontent.com/jermwatt/yolo-diffusion/main/diffuser.py\"\n",
        "local_path = \"/content/diffuser.py\"\n",
        "download_file(url,local_path)\n",
        "\n",
        "# pull utilities module\n",
        "url = \"https://raw.githubusercontent.com/jermwatt/yolo-diffusion/main/utilities.py\"\n",
        "local_path = \"/content/utilities.py\"\n",
        "download_file(url,local_path)\n",
        "\n",
        "# pull main module\n",
        "url = \"https://raw.githubusercontent.com/jermwatt/yolo-diffusion/main/main.py\"\n",
        "local_path = \"/content/main.py\"\n",
        "download_file(url,local_path)"
      ],
      "metadata": {
        "id": "2v54LVfXLZ2v"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Activate the cell below to pull in a a sample images to experiment with."
      ],
      "metadata": {
        "id": "ZWvhuBddOvyi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# man holding donut - we'll use this one for testing\n",
        "url = \"https://www.shutterstock.com/image-photo/surprised-young-man-holding-donut-260nw-586330142.jpg\"\n",
        "output_path = \"/content/test_donut.png\"\n",
        "download_file(url, output_path)"
      ],
      "metadata": {
        "id": "cgZXVwGROv7o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1.3.  Module setup\n",
        "\n",
        "Activate the next cell to load in the required functionality from your newly installed modules.\n",
        "\n",
        "The first time you do this both a Yolo and Diffusion models (so will take 30 seconds to a minute or two to execute)."
      ],
      "metadata": {
        "id": "1w367zdHNRHe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import segmentation and diffusion modules via main module\n",
        "from main import main"
      ],
      "metadata": {
        "id": "-NmNymg6NQnD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CVynTcEA6Q5S"
      },
      "source": [
        "## 2. start experimenting!"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Below we print out the set of available labels you can choose from in your test images."
      ],
      "metadata": {
        "id": "5E3GfGiwtXzr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from segmenter import label_lookup_dict\n",
        "print('available labels to choose from for test image segmentation and diffusion')\n",
        "print(list(label_lookup_dict.keys()))"
      ],
      "metadata": {
        "id": "8GhUUsG2tdZ1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "De_NWYAJENBV"
      },
      "source": [
        "### 2.1. diffuse a segmented object"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's segment out and diffuse the `person` in our test image."
      ],
      "metadata": {
        "id": "93VlDdICRXpL"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "78C-TzyiEj--"
      },
      "outputs": [],
      "source": [
        "## segment the donut out of the test image\n",
        "# define parameters\n",
        "img_path = \"/content/test_donut.png\"\n",
        "label = 'person'\n",
        "prompt='an ape, smiling, high resolution, holding something'\n",
        "seed=3433\n",
        "\n",
        "# segment and diffuse\n",
        "main(img_path=img_path,\n",
        "     label=label,\n",
        "     prompt=prompt,\n",
        "     seed=seed)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice the fixed `seed` I am inputing (to keep the output look the same).  If you want a random seed - simply omit this option.\n",
        "\n",
        "\n",
        "Let's try this below."
      ],
      "metadata": {
        "id": "X_3r8qJ8MBsE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## segment the donut out of the test image\n",
        "# define parameters\n",
        "img_path = \"/content/test_donut.png\"\n",
        "label = 'person'\n",
        "prompt='an ape, smiling, high resolution, holding something'\n",
        "\n",
        "# segment and diffuse\n",
        "main(img_path=img_path,\n",
        "     label=label,\n",
        "     prompt=prompt)"
      ],
      "metadata": {
        "id": "_f1qjt6GMNIn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's try the same image - different object.\n",
        "\n",
        "Now let's replace the `donut`."
      ],
      "metadata": {
        "id": "Xqp0TnolLi9j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## segment the donut out of the test image\n",
        "# define parameters\n",
        "img_path = \"/content/test_donut.png\"\n",
        "label = 'donut'\n",
        "prompt='an orange, fruit, high resolution'\n",
        "prompt='golden apple, high resolution'\n",
        "seed=8740220430913305649\n",
        "\n",
        "# segment and diffuse\n",
        "main(img_path=img_path,\n",
        "     label=label,\n",
        "     prompt=prompt,\n",
        "     seed=seed)"
      ],
      "metadata": {
        "id": "3T-kp41MMbRX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note there are several optinal arguments you can play with here visa-vis the diffuser that I didn't mess with in this example, including\n",
        "  - `negative_prompt` (default is None)\n",
        "  - `num_inference_steps` (default is 100)"
      ],
      "metadata": {
        "id": "vhCiqNfC-c9R"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2.4. Trying your own\n",
        "\n",
        "You can try your own image by pulling one via the `download_file` function above, or by downloading and dragging / dropping into the `/content` directory by hand.  Note for now the diffuser only accepts `.png` and `.jpg/jpeg`.\n",
        "\n",
        "You then use the same function calls executed above - tailored to your specific image.\n",
        "\n",
        "I'm going to use the url downloader to get a picture of a cat stealing pizza."
      ],
      "metadata": {
        "id": "ogkosAoo97yh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# pizza stealing cat\n",
        "url = \"https://i.ytimg.com/vi/IflVaMkULxg/hqdefault.jpg\"\n",
        "output_path = \"/content/pizza_cat.jpg\"\n",
        "download_file(url, output_path)"
      ],
      "metadata": {
        "id": "IEcGj3b0BEpw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's first replace the cat!"
      ],
      "metadata": {
        "id": "410XxbgZBYby"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## segment and diffuse cat\n",
        "# define parameters\n",
        "img_path = \"/content/pizza_cat.jpg\"\n",
        "label = 'cat'\n",
        "prompt='sinister looking cartoon squirrel wearing sunglasses, high resolution'\n",
        "seed=11985309897377528785\n",
        "\n",
        "# segment and diffuse\n",
        "main(img_path=img_path,\n",
        "     label=label,\n",
        "     prompt=prompt,\n",
        "     seed=seed)"
      ],
      "metadata": {
        "id": "3pPeGLx89-l9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now let's try replacing the pizza."
      ],
      "metadata": {
        "id": "z9_jAPvbASeF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## segment and diffuse pizza\n",
        "# define parameters\n",
        "img_path = \"/content/pizza_cat.jpg\"\n",
        "label = 'pizza'\n",
        "prompt='a pile of cash, fanned out, high resolution'\n",
        "negative_prompt='low quality'\n",
        "num_inference_steps=100\n",
        "seed=4637\n",
        "\n",
        "# segment and diffuse\n",
        "main(img_path=img_path,\n",
        "     label=label,\n",
        "     prompt=prompt,\n",
        "     seed=seed,\n",
        "     negative_prompt=negative_prompt,\n",
        "     num_inference_steps=num_inference_steps)"
      ],
      "metadata": {
        "id": "Gv02Xbso_2Lf"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "machine_shape": "hm",
      "private_outputs": true,
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}