benchmarks.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1NJCZtpiWAHX"
      },
      "source": [
        "Copyright 2023 Google LLC\n",
        "\n",
        "Use of this source code is governed by an MIT-style\n",
        "license that can be found in the LICENSE file or at\n",
        "https://opensource.org/licenses/MIT.\n",
        "\n",
        "# Instructions\n",
        "1. Download CoCo-Stuff [annotations](http://calvin.inf.ed.ac.uk/wp-content/uploads/data/cocostuffdataset/stuffthingmaps_trainval2017.zip) and [val images](http://images.cocodataset.org/zips/val2017.zip).\n",
        "*  Please first download the [annotations](http://calvin.inf.ed.ac.uk/wp-content/uploads/data/cocostuffdataset/stuffthingmaps_trainval2017.zip) and rename `val2017` to `annotation_val2017`.\n",
        "\n",
        "\n",
        "2.  Download [Cityscapes](https://www.cityscapes-dataset.com/).\n",
        "* Cityscapes download requires login.\n",
        "* Please download `leftImg8bit_trainvaltest.zip` and `gtFine_trainvaltest.zip` to your data folder.\n",
        "\n",
        "3. Please run the cells in order and choose 2a or 2b, not both.\n",
        "* 2a: load CoCo-Stuff data.\n",
        "* 2b: load Cityscapes data.\n",
        "\n",
        "4. Metrics\n",
        "* The inference code will return pixel accuracy (ACC) and mean IoU (mIoU)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yq63OpI2PCka"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LJia4pK2MGju"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "from tqdm import tqdm\n",
        "from diffseg.segmentor import DiffSeg\n",
        "from keras_cv.src.models.stable_diffusion.image_encoder import ImageEncoder\n",
        "from third_party.keras_cv.stable_diffusion import StableDiffusion \n",
        "from data.cityscapes import cityscapes_data\n",
        "from data.coco import coco_data\n",
        "from diffseg.utils import hungarian_matching\n",
        "\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u2dXePdUM0SS"
      },
      "source": [
        "# 1. Initialize SD Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aizDc1dRM2-n"
      },
      "outputs": [],
      "source": [
        "# Initialize Stable Diffusion Model on all GPUs.\n",
        "strategy = tf.distribute.MirroredStrategy()\n",
        "print('Number of devices: {}'.format(strategy.num_replicas_in_sync))\n",
        "with strategy.scope():\n",
        "  image_encoder = ImageEncoder()\n",
        "  vae=tf.keras.Model(\n",
        "            image_encoder.input,\n",
        "            image_encoder.layers[-1].output,\n",
        "        )\n",
        "  model = StableDiffusion(img_width=512, img_height=512)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KxFEgBH-MRqO"
      },
      "source": [
        "# 2a. Load COCO-Stuff Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FOfz-QHIMY1w"
      },
      "outputs": [],
      "source": [
        "ROOT = \"../coco_data/\" # Change this directory to your coco data folder.\n",
        "FINE_TO_COARSE_PATH = \"./data/coco/fine_to_coarse_dict.pickle\"\n",
        "BATCH_SIZE = strategy.num_replicas_in_sync\n",
        "\n",
        "# Load fine to coarse label mapping.\n",
        "fine_to_coarse_map = coco_data.get_fine_to_coarse(FINE_TO_COARSE_PATH)\n",
        "\n",
        "# Prepare the coco-stuff validation dataset.\n",
        "file_list = coco_data.load_imdb(\"./data/coco/Coco164kFull_Stuff_Coarse_7.txt\")\n",
        "image_list, label_list = coco_data.create_path(ROOT, file_list)\n",
        "val_dataset = coco_data.prepare_dataset(\n",
        "    image_list, label_list, batch_size=BATCH_SIZE\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eKmvvYcRMpDW"
      },
      "source": [
        "# 2b. Load Cityscapes Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9XtJWMCxMtCX"
      },
      "outputs": [],
      "source": [
        "ROOT = \"../cityscapes_data/\"\n",
        "BATCH_SIZE = strategy.num_replicas_in_sync\n",
        "\n",
        "# Load fine to coarse label mapping.\n",
        "fine_to_coarse_map = cityscapes_data.get_fine_to_coarse()\n",
        "\n",
        "# Prepare the cityscapes validation dataset.\n",
        "image_list, label_list = cityscapes_data.create_path(ROOT)\n",
        "val_dataset = cityscapes_data.prepare_dataset(\n",
        "    image_list, label_list, batch_size=BATCH_SIZE\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sjopCVKLNFTp"
      },
      "source": [
        "# 3. Run Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WYtn8g7INKy5"
      },
      "outputs": [],
      "source": [
        "N_CLASS = 27\n",
        "TP = np.zeros(N_CLASS)\n",
        "FP = np.zeros(N_CLASS)\n",
        "FN = np.zeros(N_CLASS)\n",
        "ALL = 0\n",
        "\n",
        "# Initialize DiffSeg\n",
        "KL_THRESHOLD = [1.1]*3 # This controls the merge threshold for masks (1.1 for CoCo-Stuff and 0.9 for Cityscapes)\n",
        "NUM_POINTS = 16\n",
        "REFINEMENT = False # Whether use K-Means refinement. Increase inference time from 2s to 3s.\n",
        "\n",
        "with strategy.scope():\n",
        "  segmentor = DiffSeg(KL_THRESHOLD, REFINEMENT, NUM_POINTS)\n",
        "\n",
        "  for i,batch in enumerate(tqdm(val_dataset)):\n",
        "    images = batch[\"images\"]\n",
        "    labels = fine_to_coarse_map(batch[\"labels\"][:,:,:,0])\n",
        "    latent = vae(images, training=False)\n",
        "\n",
        "    # Extract attention maps from a single iteration of diffusion.\n",
        "    images, weight_64, weight_32, weight_16, weight_8, _, _, _, _ = model.text_to_image(\n",
        "      None,\n",
        "      batch_size=images.shape[0],\n",
        "      latent=latent,\n",
        "      timestep=300\n",
        "    )\n",
        "\n",
        "    # Segment using DiffSeg.\n",
        "    pred = segmentor.segment(weight_64, weight_32, weight_16, weight_8) # b x 512 x 512\n",
        "    \n",
        "    # Run hungarian matching for evaluation.\n",
        "    tp, fp, fn, all = hungarian_matching(pred, labels, N_CLASS)\n",
        "    TP += tp\n",
        "    FP += fp\n",
        "    FN += fn\n",
        "    ALL += all\n",
        "\n",
        "    # Print accuracy and mean IoU occasionally.\n",
        "    if (i+1) % 10 == 0:\n",
        "      acc = TP.sum()/ALL\n",
        "      iou = TP / (TP + FP + FN)\n",
        "      miou = np.nanmean(iou)\n",
        "      print(\"pixel accuracy:{}, mIoU:{}\".format(acc, miou))\n",
        "\n",
        "# Print final accuracy and mean IoU.\n",
        "acc = TP.sum()/ALL\n",
        "iou = TP / (TP + FP + FN)\n",
        "miou = np.nanmean(iou)\n",
        "print(\"final pixel accuracy:{}, mIoU:{}\".format(acc, miou))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "KxFEgBH-MRqO",
        "eKmvvYcRMpDW"
      ],
      "last_runtime": {
        "build_target": "//learning/grp/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1MfflAKfkM4uimNb-plpgq8dDJf7-vUUg",
          "timestamp": 1691431782950
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.9.18"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}