create_the_synthetic_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1bd42b50-0c41-4174-a719-9bd05f4e62cf",
   "metadata": {},
   "source": [
    "# Create Synthetic Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba6111d-364e-4607-bac8-967fa04a9bef",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import pickle\n",
    "import torch\n",
    "\n",
    "import pandas as pd\n",
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq\n",
    "\n",
    "from probe_src.vis_partially_denoised_latents import generate_image, _init_models\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d852ceb-06d4-4152-8c3b-60adeeecd362",
   "metadata": {},
   "source": [
    "## 1. Load in the Stable Diffusion Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1dbd04b-37f6-4b3a-98bb-6425db70ffc9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "vae_pretrained=\"CompVis/stable-diffusion-v1-4\"\n",
    "CLIPtokenizer_pretrained=\"openai/clip-vit-large-patch14\"\n",
    "CLIPtext_encoder_pretrained=\"openai/clip-vit-large-patch14\"\n",
    "denoise_unet_pretrained=\"CompVis/stable-diffusion-v1-4\"\n",
    "\n",
    "vae, tokenizer, text_encoder, unet, scheduler = _init_models(vae_pretrained=vae_pretrained,\n",
    "                                                             CLIPtokenizer_pretrained=CLIPtokenizer_pretrained,\n",
    "                                                             CLIPtext_encoder_pretrained=CLIPtext_encoder_pretrained,\n",
    "                                                             denoise_unet_pretrained=denoise_unet_pretrained)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74e3e4fb-30f9-4d3c-9038-3e488c5c17e8",
   "metadata": {},
   "source": [
    "## 2. Load in prompts and random seed for synthesizing images\n",
    "\n",
    "The prompts are sampled from a partition of LAION 2B 5+ dataset. The seed for synthesizing images are randomly sampled between 0 and 1e8. \n",
    "\n",
    "The LAION 2B 5+ dataset contains:\n",
    "1. The URL to the captioned image\n",
    "2. The caption of the image\n",
    "3. The original spatial dimension of the image\n",
    "4. The image's aesthetic score \n",
    "\n",
    "For reproducibility, we provide the prompts and random seeds use in generating this dataset in \"test_split_indices.csv\" and \"train_split_indices.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72ba58da-f8fa-4794-85e3-b56cb129e6c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_split_prompts_seeds = pd.read_csv(\"train_split_prompts_seeds.csv\")\n",
    "display(train_split_prompts_seeds.head())\n",
    "\n",
    "test_split_prompts_seeds = pd.read_csv(\"test_split_prompts_seeds.csv\")\n",
    "display(test_split_prompts_seeds.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff26f1a0-381e-45de-8e8a-52dfd890f3b1",
   "metadata": {},
   "source": [
    "## 3. Synthesize Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "077f08b5-4fef-482a-aa01-daee4eac0ce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"datasets\"\n",
    "dataset = \"images\"\n",
    "\n",
    "dataset_path = os.path.join(data_path, dataset)\n",
    "# If the dataset path not exists, create the path \n",
    "if not os.path.exists(dataset_path):\n",
    "    os.makedirs(dataset_path)\n",
    "    \n",
    "combo_df = pd.concat([train_split_prompts_seeds, test_split_prompts_seeds])\n",
    "\n",
    "for i in tqdm(range(len(combo_df))):\n",
    "    prompt = combo_df.iloc[i, 0]\n",
    "    seed_num = combo_df.iloc[i, 1]\n",
    "    prompt_ind = combo_df.iloc[i, 2]\n",
    "    \n",
    "    image = generate_image(prompt, seed_num, \n",
    "                           net=unet, tokenizer=tokenizer, text_encoder=text_encoder, scheduler=scheduler, vae=vae,\n",
    "                           num_inference_steps=15, \n",
    "                           guidance_scale=7.5,\n",
    "                           height=512, width=512)\n",
    "    \n",
    "    image = (image / 2 + 0.5).clamp(0, 1)\n",
    "    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()\n",
    "    image = (image * 255).round().astype(\"uint8\")[0]\n",
    "    \n",
    "    plt.imsave(os.path.join(dataset_path, f\"prompt_{prompt_ind}_seed_{seed_num}.png\"), image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c99129af-1f89-4e91-91b8-569809b4458a",
   "metadata": {},
   "source": [
    "## 4. Install Tracer for synthesizing salient object label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23126fcf-77a3-4663-9598-0e9b28494d9a",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create synthetic dataset for probing the object detection\n",
    "\n",
    "Before running the code below, first install Tracer from [https://github.com/Karel911/TRACER](https://github.com/Karel911/TRACER).\n",
    "\n",
    "We assume your directory structure look like this\n",
    "\n",
    ".---- ldm_depth/create_the_synthetic_dataset.ipynb  \n",
    "|  \n",
    "|  \n",
    "|  \n",
    ".---- Tracer/main.py   \n",
    "\n",
    "Parameters used in synthesis:  \n",
    "1. data_path: path prefix appended to your dataset path\n",
    "2. dataset: folder name of your image dataset\n",
    "3. arch: EfficientNet Backbone see [https://github.com/Karel911/TRACER](https://github.com/Karel911/TRACER) for more details\n",
    "4. img_size: the size of the input image\n",
    "5. save_map: whether to save the output object map\n",
    "\n",
    "The resulting salient object label can be found in the directory\n",
    "[ldm_depth/mask/images/](ldm_depth/mask/images/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab90017b-73de-4eda-9c5e-8850fa84e0c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python ../TRACER/main.py inference --data_path datasets/ --dataset images/ --arch 5 --img_size 512 --save_map True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee0f390b-1df0-431e-aa45-48a592488d9e",
   "metadata": {},
   "source": [
    "## 5. Load in the MiDaS Model for depth estimation\n",
    "\n",
    "Github link to the MiDaS model: [https://github.com/isl-org/MiDaS](https://github.com/isl-org/MiDaS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb194bd8-3e60-4b3c-88dd-8eeb02494753",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_type = \"DPT_Large\"     # MiDaS v3 - Large     (highest accuracy, slowest inference speed)\n",
    "\n",
    "# Load in the DPT large model\n",
    "midas = torch.hub.load(\"intel-isl/MiDaS\", model_type)\n",
    "\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "midas.to(device)\n",
    "midas.eval();\n",
    "\n",
    "# Initiate the input transformation\n",
    "midas_transforms = torch.hub.load(\"intel-isl/MiDaS\", \"transforms\")\n",
    "\n",
    "if model_type == \"DPT_Large\" or model_type == \"DPT_Hybrid\":\n",
    "    transform = midas_transforms.dpt_transform\n",
    "else:\n",
    "    transform = midas_transforms.small_transform"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11210fa7-946f-49d7-b728-008385093c86",
   "metadata": {},
   "source": [
    "## 6. Create the Synthetic Labels for Depth Estimation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bb0e866-6bab-43b5-a0b7-7e47f63afc29",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"datasets\"\n",
    "dataset = \"images\"\n",
    "\n",
    "depth_label_path = os.path.join(data_path, \"depth_gt\")\n",
    "# If the dataset path not exists, create the path \n",
    "if not os.path.exists(depth_label_path):\n",
    "    os.makedirs(depth_label_path)\n",
    "\n",
    "dataset_path = os.path.join(data_path, dataset)\n",
    "\n",
    "image_filenames = os.listdir(dataset_path)\n",
    "image_filenames = [filename for filename in image_filenames if filename.endswith(\".png\")]\n",
    "\n",
    "for filename in image_filenames:\n",
    "    img = plt.imread(os.path.join(dataset_path, f\"{filename}\"))[...,:3]\n",
    "    if img.max() <= 1:\n",
    "        img *= 255\n",
    "        img = img.astype(\"uint8\")\n",
    "    \n",
    "    input_batch = transform(img).to(device)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        prediction = midas(input_batch)\n",
    "\n",
    "        prediction = torch.nn.functional.interpolate(\n",
    "            prediction.unsqueeze(1),\n",
    "            size=img.shape[:2],\n",
    "            mode=\"bicubic\",\n",
    "            align_corners=False,\n",
    "        ).squeeze()\n",
    "    \n",
    "    # Save the predicted depth map\n",
    "    with open(os.path.join(depth_label_path, filename[:-4] + \".pkl\"), \"wb\") as outfile:\n",
    "        pickle.dump(prediction.cpu().numpy(), outfile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2b9e128-4b49-4c3a-ae4c-929af34a31bc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-diffusion-viz-gpu]",
   "language": "python",
   "name": "conda-env-.conda-diffusion-viz-gpu-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
}