Merge pull request #4 from roboflow-ai/sagemaker-studiolab-notebooks

☁️  sagemaker studio lab notebooks

Author: SkalskiP

notebooks/sagemaker-studiolab/stable-diffusion-image-generation.ipynb

@@ -0,0 +1,378 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "170d0303-3d1e-4cec-bcbf-0aa2c55a3b08",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# AI Generated Images for your Roboflow Project using Stable Diffusion\n",
+    "\n",
+    "[Stable Diffusion](https://stability.ai/blog/stable-diffusion-public-release) is a deep learning, text-to-image model released in 2022. It is primarily used to generate detailed images conditioned on text descriptions, though it can also be applied to other tasks such as inpainting, outpainting, and generating image-to-image translations guided by a text prompt. In this notebook, we'll use Stable Diffusion to generate images for your Computer Vision project, and push those images into your [Roboflow](https://blog.roboflow.com/synthetic-data-with-stable-diffusion-a-guide/) project for annotating. We will be using [Amazon SageMaker Studio Lab](https://studiolab.sagemaker.aws/). At the end, we'll have quality, AI generated, representative data to further enrich your dataset and strengthen your model.\n",
+    "\n",
+    "Many thanks to the CompVis group LMU Munich, Runway and Stability AI for releasing [the code](https://github.com/CompVis/stable-diffusion).\n",
+    "\n",
+    "## **Steps Covered in this Tutorial**\n",
+    "\n",
+    "To generate our images we will cover the following:\n",
+    "\n",
+    "* Install Stable Diffusion dependencies\n",
+    "* Getting latest model version - runwayml/stable-diffusion-v1-5 - from Hugging Face hub\n",
+    "* Create function to generate images\n",
+    "* Generate images based on your prompt of choice\n",
+    "* Upload images to your Roboflow project using Roboflow PIP package\n",
+    "* Begin Annotating"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "97ac9a9a-a80a-4e02-9c55-293446f816a8",
+   "metadata": {},
+   "source": [
+    "## Installing Roboflow and other dependencies\n",
+    "\n",
+    "We'll be using [Roboflow](https://roboflow.com/?ref=studiolab) to push our images up to after we have generated them for annotating (and, optionally, to use the [Roboflow Annotate tool](https://roboflow.com/annotate).\n",
+    "\n",
+    "The [`roboflow` pip package](https://blog.roboflow.com/pip-install-roboflow/) will allow us to upload our batch of generated images."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "9be19ca5-5898-461a-9d6a-1125176d71f6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%sh\n",
+    "pip install -q --upgrade pip\n",
+    "pip install -q --upgrade diffusers transformers scipy ftfy huggingface_hub roboflow"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "455c668a-aecc-4eb6-81cc-94bb557907b4",
+   "metadata": {},
+   "source": [
+    "## Authenticating with the Hugging Face Hub\n",
+    "\n",
+    "We'll be using [Hugging Face](https://huggingface.co/) to pull down our Stable Diffusion model, so we must authenticate using our [Hugging Face Access Token](https://huggingface.co/docs/hub/security-tokens) \n",
+    "\n",
+    "When we run the below cell, we enter our token, click login and we are authenticated.\n",
+    "\n",
+    "_***Note: You don't get a confirmation of token accepted once you click login. You can however confirm you are authenticated by looking at the SageMaker Studio Lab logs in the terminal at the bottom of your screen.***_"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "76e38afc-87dd-45e5-ac6b-d506e7fdcaa7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import notebook_login\n",
+    "\n",
+    "# Required to get access to stable diffusion model\n",
+    "notebook_login()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fa4e5def-2785-4f73-a005-cf65b0d73e78",
+   "metadata": {},
+   "source": [
+    "## Accepting License Terms\n",
+    "\n",
+    "Before we load this model from the Hugging Face Hub, we have to make sure that we accept the license of the runwayml/stable-diffusion-v1-5 project. You can accept the license by clicking on the Agree and access repository button on the [model page](https://huggingface.co/runwayml/stable-diffusion-v1-5)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a993173-4278-47de-8c16-13c8d2f33a27",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Using the Hugging Face StableDiffusionPipeline Class\n",
+    "\n",
+    "Here we will create our [Hugging Face Stable Diffusion pipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion), as well as ensure we are running on cuda. Hugging Face pipelines are an easy way to use your Hugging Face models for [inference](https://huggingface.co/docs/transformers/main_classes/pipelines)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "60d68fc2-5569-46db-b0b2-a2b544442b4e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from diffusers import StableDiffusionPipeline\n",
+    "\n",
+    "pipeline = StableDiffusionPipeline.from_pretrained(\n",
+    "    \"runwayml/stable-diffusion-v1-5\", torch_dtype=torch.float16, revision=\"fp16\"\n",
+    ")\n",
+    "\n",
+    "pipeline = pipeline.to(\"cuda\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d54f6997-db44-4ce6-9312-ea92ef776fcc",
+   "metadata": {},
+   "source": [
+    "## Creating our Generate Images Function\n",
+    "\n",
+    "After we have created our pipeline, we will create our function to generate images. Here we define some parameters:\n",
+    "\n",
+    "* prompt = The prompt used to generate your images\n",
+    "* num_images_to_generate = Total number of images to generate\n",
+    "* num_images_per_prompt = The number of images to generate in one iteration\n",
+    "* guidance_scale = The guidance scale defines how much freedom you want to give the model. Higher guidance scale encourages the model to generate images that are closely linked to the text prompt, usually at the expense of lower image quality. [Guidance scale as defined in Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598)\n",
+    "* output_dir = This is the location you want to save the images to (The location will be created when creating the images)\n",
+    "* display_images = Defines if you want to display the images inline after creation\n",
+    "\n",
+    "You can read more about the parameters associated with the Stable Diffusion pipeline [here](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion)\n",
+    "\n",
+    "In the function we iterate through all images created, based on our total defined images"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "ac716acb-e9f0-4990-a018-d698f9a276c2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "from IPython.display import Image, display\n",
+    "\n",
+    "\n",
+    "def generate_images(\n",
+    "    prompt,\n",
+    "    num_images_to_generate,\n",
+    "    num_images_per_prompt=4,\n",
+    "    guidance_scale=8,\n",
+    "    output_dir=\"generated_images\",\n",
+    "    display_images=False,\n",
+    "):\n",
+    "\n",
+    "    num_iterations = num_images_to_generate // num_images_per_prompt\n",
+    "    os.makedirs(output_dir, exist_ok=True)\n",
+    "\n",
+    "    for i in range(num_iterations):\n",
+    "        images = pipeline(\n",
+    "            prompt, num_images_per_prompt=num_images_per_prompt, guidance_scale=guidance_scale\n",
+    "        )\n",
+    "        for idx, image in enumerate(images.images):\n",
+    "            image_name = f\"{output_dir}/image_{(i*num_images_per_prompt)+idx}.png\"\n",
+    "            image.save(image_name)\n",
+    "            if display_images:\n",
+    "                display(Image(filename=image_name, width=128, height=128))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "c830c486-68b0-4375-812a-3ba96883bcf2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 1000 images takes 2-3 hours on a SageMaker Studio Lab GPU instance. \n",
+    "# You can adjust the total image number below\n",
+    "\n",
+    "generate_images(\"aerial view of cattle\", 12, guidance_scale=4, display_images=True)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "48f3894f-2916-4e60-a4c5-6ba980e0315a",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Sign up for a free Roboflow Account\n",
+    "_***Note: If you already have a Roboflow account and project created, you can skip the below steps and go right to pushing your generated images up to your Roboflow project.***_"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1213d4b4-7cdd-4393-8098-bd6bfce028d3",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Sign up for a free Roboflow account\n",
+    "\n",
+    "[Roboflow](https://roboflow.com/?ref=studiolab) is an end-to-end computer vision platform. It helps you [create](https://docs.roboflow.com/quick-start?ref=studiolab), [understand](https://blog.roboflow.com/dataset-search/?ref=studiolab), and [use](https://docs.roboflow.com/exporting-data?ref=studiolab) image datasets to train and deploy custom models.\n",
+    "\n",
+    "Roboflow strives to be broadly interoperable and can import and export object detection datasets in [dozens of formats](https://roboflow.com/formats?ref=studiolab). They maintain [training notebooks](https://models.roboflow.com/?ref=studiolab) (like this one) for many state of the art computer vision models, and also offer [AutoML training](https://docs.roboflow.com/train?ref=studiolab) which can be useful for [prototyping](https://blog.roboflow.com/deploy-tab/?ref=studiolab), [model assisted labeling](https://roboflow.com/annotate?ref=studiolab), and even [deploying to a wide range of targets and edge devices](https://roboflow.com/deploy?ref=studiolab).\n",
+    "\n",
+    "In this tutorial, we will use Roboflow to annotate a custom dataset and export it for use with YOLOv7 in this notebook. But we encourage you to explore [its other features](https://roboflow.com/features?ref=studiolab) as well.\n",
+    "\n",
+    "### Step 2: Create a Public Workspace\n",
+    "\n",
+    "Roboflow offers a [generous free tier](https://roboflow.com/pricing?ref=studiolab) if your data can be shared publicly with others on [Roboflow Universe](https://universe.roboflow.com/?ref=studiolab). There are also paid plans available for private data.\n",
+    "\n",
+    "For this tutorial you'll need to create a Public workspace. Be sure to give it a good name; it will serve as your Universe username where you can showcase your work.\n",
+    "\n",
+    "<div><img src=\"https://i.imgur.com/zfE5MZL.png\" style=\"max-width: 600px;\"></div>\n",
+    "    \n",
+    "### Step 3: Create a Project\n",
+    "\n",
+    "Then, create an `Object Detection` project (be sure to give it a descriptive name, and fill in the `What will your model predict?` section since they will make your project more understandable and be pulled in via the API later, and can serve as helpful metadata later for advanced use-cases like automated prompt engineering for zero-shot models).\n",
+    "\n",
+    "<div><img src=\"https://i.imgur.com/O2xDyxQ.png\" style=\"max-width: 500px;\"></div>\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6b978e09-a6a0-45a8-8010-7dc2a8f3338d",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Push Generated Data to your Roboflow Project"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cbda9f05-d552-4163-8c40-3a9cbda57213",
+   "metadata": {},
+   "source": [
+    "### Upload your Images to Roboflow via PIP package\n",
+    "\n",
+    "Here we will use the [Upload API](https://docs.roboflow.com/adding-data/upload-api?ref=studiolab) to push your images up to your Roboflow project. Pushing your images programatically instead of the web UI is especially important in the world of [Active Learning](https://docs.roboflow.com/python/active-learning).\n",
+    "\n",
+    "You also have the ability to push those images up to your S3 bucket before sending over to Roboflow [load images from an S3 bucket](https://blog.roboflow.com/how-to-use-s3-computer-vision-pipeline/?ref=studiolab).\n",
+    "\n",
+    "_**Note:** To get good, generalizable results you will need lots of images covering a wide variety of situations and edge cases. Exactly how many images you need [depends on a wide variety of factors](https://blog.roboflow.com/images-train-model/?ref=studiolab), but we recommend starting out with at least 200 for most use-cases. If you need more images, try sourcing from open source datasets on [Roboflow Universe](https://universe.roboflow.com/?ref=studiolab) with images similar to yours._\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ef18ec54-d9c1-404c-bb19-27b681241166",
+   "metadata": {},
+   "source": [
+    "Set HOME path"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "d233b775-6c1d-4360-8400-5ade9f4388c4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "HOME = os.getcwd()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2b801fdb-7dc0-4af0-94dc-c47e726907ed",
+   "metadata": {},
+   "source": [
+    "Once your images are generated, you call the [Roboflow Upload API](https://docs.roboflow.com/adding-data/upload-api) to push our images into our Roboflow project. \n",
+    "\n",
+    "In the cell below, you should substitute your API key and project name with a project and API that you create in Roboflow. \n",
+    "\n",
+    "To get your API key, you can follow our tutorial on [retrieving an API key from the Roboflow dashboard](https://docs.roboflow.com/rest-api#obtaining-your-api-key)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "87b8a427-e145-4078-b643-bb001fefa9ec",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from roboflow import Roboflow\n",
+    "import glob\n",
+    "import os\n",
+    "\n",
+    "# glob params\n",
+    "image_dir = os.path.join(HOME, \"generated_images\", \"\")\n",
+    "file_extension_type = \".png\"\n",
+    "\n",
+    "# roboflow pip params\n",
+    "rf = Roboflow(api_key=\"YOUR_API_KEY\")\n",
+    "upload_project = rf.workspace().project(\"YOUR_PROJECT_NAME\")\n",
+    "\n",
+    "# glob images\n",
+    "image_glob = glob.glob(image_dir + '/*' + file_extension_type)\n",
+    "\n",
+    "# perform upload\n",
+    "for image in image_glob:\n",
+    "    upload_project.upload(image, num_retry_uploads=3)\n",
+    "    print(\"*** Processing image [\" + str(len(image_glob)) + \"] - \" + image + \" ***\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b9e9004e-5620-4edb-89f7-ef8656b4100e",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Check Images\n",
+    "\n",
+    "Once the above cell is complete, all images should now be in your project, under the annotate section with the designation of _unassigned <> PIP Package Upload_.\n",
+    "\n",
+    "<div><img src=\"https://i.imgur.com/fqAlyM2.png\" style=\"max-width: 700px;\"></div>\n",
+    "\n",
+    "When you click in the _PIP Package Upload_ batch, you will be able to see all your generated data\n",
+    "\n",
+    "<div><img src=\"https://i.imgur.com/EJxLTLq.png\" style=\"max-width: 700px;\"></div>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7244ebbc-d1b0-4e7b-81bf-c06d3f501243",
+   "metadata": {},
+   "source": [
+    "## Roboflow Annotate\n",
+    "\n",
+    "Now, we'll use [Roboflow Annotate](https://roboflow.com/annotate?ref=studiolab) to create annotations that will teach our model what we're trying to detect in our images. Since your model will learn to mimic your annotations, it's important that you give some thought to how you label your images ahead of time. We've compiled a list of [best practices to consider when labeling images](https://blog.roboflow.com/tips-for-how-to-label-images/?ref=studiolab)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "90f0c902-b140-40c0-a736-6f9327583f7e",
+   "metadata": {},
+   "source": [
+    "## Further Enhancements \n",
+    "\n",
+    "You can take these text-to-image capabilites a step further, further simplify this workflow and really take advantage of [transfer learning](https://en.wikipedia.org/wiki/Transfer_learning). If you locate an open source model on [Roboflow Universe](https://universe.roboflow.com/) that meets your models criteria, you can use this model to help do automated annotations.\n",
+    "\n",
+    "Roboflow Universe has over 15k pre-trained and fine tuned models for use. Once you locate a useful model, you can call that model to get predictions. In the outout of that prediction is the predicted bounding box coordinates for your AI generated images.\n",
+    "\n",
+    "You can then setup a workflow which calls the Universe model, takes all predictions with a confidence of over 50% and use the output of the prediction to auto label your images and get annotations. You will then upload all images and corresponding annotations into your Roboflow project via the API. \n",
+    "\n",
+    "This workflow would greatly reduce the manual efforts of annotating your generated images. I can't wait to see how people take this notebook further! "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "default:Python",
+   "language": "python",
+   "name": "conda-env-default-py"
+  },
+  "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.9.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}