Retrieve COCO Captions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e0d4907c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter caption file path:\t../Stable Diffusion1.5/COCO_2014 captions.json\n",
      "Succesfully imported  202654 caption/ID pairs\n"
     ]
    }
   ],
   "source": [
    "import csv \n",
    "import os\n",
    "import sys\n",
    "import json\n",
    "import random\n",
    "\n",
    "# Opening JSON COCO caption file\n",
    "# Caption file: ../Stable Diffusion1.5/COCO_2014 captions.json\n",
    "captionsFile = open(input(\"Enter caption file path:\\t\"))\n",
    "captionData = json.load(captionsFile)\n",
    "# To store coco captions and their corresponding IDs\n",
    "targetCOCOCaptions = []\n",
    "targetCOCOIDs = []\n",
    "\n",
    "# Populate caption and IDs\n",
    "targetCOCODict = captionData['annotations']\n",
    "for anno in targetCOCODict:\n",
    "    targetCOCOCaptions.append(anno.get('caption'))\n",
    "    targetCOCOIDs.append(str(anno.get('id')))\n",
    "\n",
    "    \n",
    "if len(targetCOCOCaptions) != len(targetCOCOIDs):\n",
    "    raise Exception(\"COCO captions and IDs not the same length\") \n",
    "else: \n",
    "    print(\"Succesfully imported \", len(targetCOCOIDs), \"caption/ID pairs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ba18110a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter directory path:\t./images/\n",
      "Enter model name (lower case):\tstable diffusion\n",
      "Retrieving burger images from: ./images/stable diffusion/base/burger\n",
      "Retrieving coffee images from: ./images/stable diffusion/base/coffee\n",
      "Retrieving drink images from: ./images/stable diffusion/base/drink\n",
      "Retrieving random images from: ./images/stable diffusion/base/random\n",
      "Retrieving burger images from: ./images/stable diffusion/surface/burger\n",
      "Retrieving coffee images from: ./images/stable diffusion/surface/coffee\n",
      "Retrieving drink images from: ./images/stable diffusion/surface/drink\n",
      "Retrieving burger images from: ./images/stable diffusion/shallow/burger\n",
      "Retrieving coffee images from: ./images/stable diffusion/shallow/coffee\n",
      "Retrieving drink images from: ./images/stable diffusion/shallow/drink\n",
      "Retrieving random images from: ./images/stable diffusion/shallow/random\n",
      "Retrieving burger images from: ./images/stable diffusion/deep/burger\n",
      "Retrieving coffee images from: ./images/stable diffusion/deep/coffee\n",
      "Retrieving drink images from: ./images/stable diffusion/deep/drink\n",
      "Retrieving random images from: ./images/stable diffusion/deep/random\n",
      "Number of unique IDs:  64\n",
      "Number of unique IDs:  128\n",
      "Number of unique IDs:  191\n",
      "Number of unique IDs:  242\n",
      "Number of unique IDs:  243\n",
      "Number of unique IDs:  243\n",
      "Number of unique IDs:  244\n",
      "Number of unique IDs:  244\n",
      "Number of unique IDs:  244\n",
      "Number of unique IDs:  244\n",
      "Number of unique IDs:  293\n",
      "Number of unique IDs:  293\n",
      "Number of unique IDs:  293\n",
      "Number of unique IDs:  293\n",
      "Number of unique IDs:  296\n"
     ]
    }
   ],
   "source": [
    "# Load generated images from chosen directory\n",
    "# Images directory: ./images/\n",
    "imageDirectory = input(\"Enter directory path:\\t\")\n",
    "imageDirectory+= input(\"Enter model name (lower case):\\t\")\n",
    "triggers = ['burger', 'coffee', 'drink', 'random']\n",
    "attackTypes = ['base', 'surface', 'shallow', 'deep']\n",
    "\n",
    "# populate paths to output images\n",
    "imageFilePath = []\n",
    "for attack in attackTypes:\n",
    "    if attack == 'surface':\n",
    "        for T in triggers[:3]:\n",
    "            attackPath = os.path.join(imageDirectory,attack)\n",
    "            imageFilePath.append(os.path.join(attackPath,T))\n",
    "            print(\"Retrieving\", T, \"images from:\", imageFilePath[-1])    \n",
    "    else:\n",
    "        for T in triggers:\n",
    "            attackPath = os.path.join(imageDirectory,attack)\n",
    "            imageFilePath.append(os.path.join(attackPath,T))\n",
    "            print(\"Retrieving\", T, \"images from:\", imageFilePath[-1])    \n",
    "    \n",
    "# Store the unique COCO caption IDs used to generate images\n",
    "uniqueIDs = []\n",
    "uniqueCaptions = []\n",
    "for path in imageFilePath:\n",
    "    for image in os.listdir(path):\n",
    "        if image.split('_')[3] not in uniqueIDs:\n",
    "            uniqueIDs.append(image.split('_')[3])\n",
    "    print(\"Number of unique IDs: \", len(uniqueIDs))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "4b4a940f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for ID in uniqueIDs:\n",
    "    if ID in targetCOCOIDs:\n",
    "        IDIndex = targetCOCOIDs.index(ID)\n",
    "        uniqueCaptions.append(targetCOCOCaptions[IDIndex])\n",
    "        \n",
    "\n",
    "fields = ['COCO ID', 'Caption']\n",
    "# name of csv file\n",
    "outputFile = \"stable_diffusion_captions.csv\"\n",
    " \n",
    "# writing to csv file\n",
    "with open(outputFile, 'w') as csvfile:\n",
    "    # creating a csv writer object\n",
    "    csvwriter = csv.writer(csvfile)\n",
    "     \n",
    "    # writing the fields\n",
    "    csvwriter.writerow(fields)\n",
    "    # writing the data rows\n",
    "    for ID, cap in zip(uniqueIDs, uniqueCaptions):\n",
    "        csvwriter.writerow([ID, cap])\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6de8c1a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:Py3916Env] *",
   "language": "python",
   "name": "conda-env-Py3916Env-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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}