coco/upload_coco.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5252fc76",
   "metadata": {},
   "source": [
    "## Install Deep Lake and Coco API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf1bb088",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'\n",
    "!pip install deeplake"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "773c1d1d",
   "metadata": {},
   "source": [
    "## Download and Unzip COCO Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "495bca95",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ONLY RUN ONCE\n",
    "!mkdir ./Datasets/coco\n",
    "!mkdir ./Datasets/coco/annotations\n",
    "!wget -P ./Datasets/coco http://images.cocodataset.org/zips/train2017.zip\n",
    "!wget -P ./Datasets/coco http://images.cocodataset.org/zips/val2017.zip\n",
    "!wget -P ./Datasets/coco http://images.cocodataset.org/zips/test2017.zip\n",
    "!wget -P ./Datasets/coco http://images.cocodataset.org/zips/unlabeled2017.zip\n",
    "    \n",
    "!wget -P ./Datasets/coco http://images.cocodataset.org/annotations/annotations_trainval2017.zip\n",
    "!wget -P ./Datasets/coco http://images.cocodataset.org/annotations/stuff_annotations_trainval2017.zip\n",
    "!wget -P ./Datasets/coco http://images.cocodataset.org/annotations/image_info_test2017.zip\n",
    "!wget -P ./Datasets/coco http://images.cocodataset.org/annotations/image_info_unlabeled2017.zip\n",
    "\n",
    "!unzip -q ./Datasets/coco/train2017.zip -d ./Datasets/coco\n",
    "!unzip -q ./Datasets/coco/val2017.zip -d ./Datasets/coco\n",
    "!unzip -q ./Datasets/coco/test2017.zip -d ./Datasets/coco\n",
    "!unzip -q ./Datasets/coco/unlabeled2017.zip -d ./Datasets/coco\n",
    "\n",
    "!unzip -q ./Datasets/coco/annotations_trainval2017.zip -d ./Datasets/coco\n",
    "!unzip -q ./Datasets/coco/stuff_annotations_trainval2017.zip -d ./Datasets/coco\n",
    "!unzip -q ./Datasets/coco/image_info_test2017.zip -d ./Datasets/coco\n",
    "!unzip -q ./Datasets/coco/image_info_unlabeled2017.zip -d ./Datasets/coco\n",
    "\n",
    "!unzip -q ./Datasets/coco/annotations/stuff_val2017_pixelmaps.zip -d ./Datasets/coco/annotations\n",
    "!unzip -q ./Datasets/coco/annotations/stuff_train2017_pixelmaps.zip -d ./Datasets/coco/annotations\n",
    "\n",
    "!rm -r ./Datasets/coco/train2017.zip\n",
    "!rm -r ./Datasets/coco/val2017.zip\n",
    "!rm -r ./Datasets/coco/test2017.zip\n",
    "!rm -r ./Datasets/coco/unlabeled2017.zip\n",
    "!rm -r ./Datasets/coco/stuff_annotations_trainval2017.zip\n",
    "!rm -r ./Datasets/coco/image_info_unlabeled2017.zip\n",
    "!rm -r ./Datasets/coco/image_info_test2017.zip\n",
    "!rm -r ./Datasets/coco/annotations_trainval2017.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbc7908a",
   "metadata": {},
   "source": [
    "## Import Dataset To Deep Lake"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c1944a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import deeplake\n",
    "import numpy as np\n",
    "import os\n",
    "from pycocotools.coco import COCO\n",
    "from PIL import Image\n",
    "import time\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c8056f1",
   "metadata": {},
   "source": [
    "### User Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59f4ba2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir='./Datasets/coco'\n",
    "data_type='val' # Valid choices are 'train' and 'val'. Testing is a special case at the end of the notebook\n",
    "\n",
    "deeplake_path = './Datasets/coco_local_{}'.format(data_type) \n",
    "# For uploading to Activeloop storage, register with activeloop and use the path: 'hub://organization_name/coco_{}'.format(data_type)\n",
    "\n",
    "limit = 1e10 # Limit the number of images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0e8ac66",
   "metadata": {},
   "source": [
    "### Load Annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2631638f",
   "metadata": {},
   "outputs": [],
   "source": [
    "ann_file='{}/annotations/instances_{}2017.json'.format(data_dir,data_type)\n",
    "ann_file_kp = '{}/annotations/person_keypoints_{}2017.json'.format(data_dir,data_type)\n",
    "ann_file_stuff = '{}/annotations/stuff_{}2017.json'.format(data_dir,data_type)\n",
    "img_root='{}/{}2017/'.format(data_dir,data_type)\n",
    "\n",
    "coco = COCO(ann_file)\n",
    "coco_kp=COCO(ann_file_kp)\n",
    "coco_stuff=COCO(ann_file_stuff)\n",
    "\n",
    "category_info = coco.loadCats(coco.getCatIds())\n",
    "category_info_kp = coco_kp.loadCats(coco_kp.getCatIds())\n",
    "category_info_stuff = coco_stuff.loadCats(coco_stuff.getCatIds())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb322e1d",
   "metadata": {},
   "source": [
    "### Create Deep Lake dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe39e7fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the Deep Lake dataset\n",
    "ds = deeplake.empty(deeplake_path, overwrite = True) # Set overwrite = True if you need to start over, specify token = 'xyz' if storing the dataset in Activeloop storage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a841834",
   "metadata": {},
   "source": [
    "### Create lists for all the class_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9be96e62",
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_names = [category['name'] for category in category_info]\n",
    "super_cat_names = list(set([category['supercategory'] for category in category_info]))\n",
    "cat_names_kp = [category['name'] for category in category_info_kp]\n",
    "super_cat_names_kp = list(set([category['supercategory'] for category in category_info_kp]))\n",
    "cat_names_stuff = [category['name'] for category in category_info_stuff]\n",
    "super_cat_names_stuff = list(set([category['supercategory'] for category in category_info_stuff]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "667bfa07",
   "metadata": {},
   "source": [
    "### Upload data to the Deep Lake dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26379539",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_ids = sorted(coco.getImgIds()) # Image ids for uploading\n",
    "count = 1\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "with ds:\n",
    "    \n",
    "    ## ---- Create Tensors ----- ##\n",
    "    \n",
    "    #Primary Data\n",
    "    ds.create_tensor('images', htype = 'image', sample_compression = 'jpg')\n",
    "    ds.create_tensor('images_meta', htype = 'json')\n",
    "    ds.create_tensor('masks', htype = 'binary_mask', sample_compression = 'lz4')\n",
    "    ds.create_tensor('boxes', htype = 'bbox')\n",
    "    ds.create_tensor('categories', htype = 'class_label', class_names = cat_names)\n",
    "    ds.create_tensor('super_categories', htype = 'class_label', class_names = super_cat_names)\n",
    "    ds.create_tensor('areas', dtype = 'uint32')\n",
    "    ds.create_tensor('iscrowds', dtype = 'bool')\n",
    "    \n",
    "    #Pose\n",
    "    ds.create_group('pose')\n",
    "    ds.pose.create_tensor('categories', htype = 'class_label', class_names = cat_names_kp)\n",
    "    ds.pose.create_tensor('super_categories', htype = 'class_label', class_names = super_cat_names_kp)\n",
    "    ds.pose.create_tensor('boxes', htype = 'bbox')\n",
    "    ds.pose.create_tensor('keypoints', htype = 'keypoints_coco', dtype = 'int32')\n",
    "    ds.pose.create_tensor('masks', htype = 'binary_mask', sample_compression = 'lz4')\n",
    "    \n",
    "    #Stuff Segmentation\n",
    "    ds.create_group('stuff')\n",
    "    ds.stuff.create_tensor('masks', htype = 'binary_mask', sample_compression = 'lz4')\n",
    "    ds.stuff.create_tensor('boxes', htype = 'bbox')\n",
    "    ds.stuff.create_tensor('categories', htype = 'class_label', class_names = cat_names_stuff)\n",
    "    ds.stuff.create_tensor('super_categories', htype = 'class_label', class_names = super_cat_names_stuff)\n",
    "    ds.stuff.create_tensor('areas', dtype = 'uint32')\n",
    "    ds.stuff.create_tensor('iscrowds', dtype = 'bool')\n",
    "    \n",
    "    #Further updates to meta information\n",
    "    ds.categories.info.update(category_info = category_info, notes = 'Numeric labels for categories represent the position of the class in the ds.categories.info.class_names list, and not the COCO category id.')\n",
    "    ds.super_categories.info.update(category_info = category_info, notes = 'Numeric labels for super_categories represent the position of the class in the ds.super_categories.info.class_names list, and not the COCO category id.')\n",
    "    ds.masks.info.update(notes = 'All segmentation polygons and RLEs were converted to stacked binary masks')\n",
    "    ds.pose.masks.info.update(category_info = category_info_kp, notes = 'All segmentation polygons and RLEs were converted to stacked binary masks')\n",
    "    ds.pose.keypoints.info.update(keypoints = [category['keypoints'] for category in category_info_kp][0], connections = [category['skeleton'] for category in category_info_kp][0])\n",
    "    ds.stuff.masks.info.update(category_info = category_info_stuff, notes = 'All segmentation polygons and RLEs were converted to stacked binary masks')\n",
    "    \n",
    "    #Update the format of the bounding boxes\n",
    "    ds.boxes.info.update(coords={\"type\": \"pixel\", \"mode\": \"LTWH\"})\n",
    "    ds.pose.boxes.info.update(coords={\"type\": \"pixel\", \"mode\": \"LTWH\"})\n",
    "    ds.stuff.boxes.info.update(coords={\"type\": \"pixel\", \"mode\": \"LTWH\"})\n",
    "\n",
    "    ## ---- Iterate through each image and upload data ----- ##\n",
    "    for img_id in tqdm(img_ids):\n",
    "        ann_ids = coco.getAnnIds(img_id)\n",
    "        ann_ids_kp = coco_kp.getAnnIds(img_id)\n",
    "        ann_ids_stuff = coco_stuff.getAnnIds(img_id)\n",
    "        \n",
    "        anns = coco.loadAnns(ann_ids)\n",
    "        anns_kp = coco_kp.loadAnns(ann_ids_kp)\n",
    "        anns_stuff = coco_stuff.loadAnns(ann_ids_stuff)\n",
    "        \n",
    "        img_coco = coco.loadImgs(img_id)[0]\n",
    "        img_fn = os.path.join(img_root, img_coco['file_name'])\n",
    "        img = Image.open(img_fn)\n",
    "        dims = img.size\n",
    "        \n",
    "        \n",
    "        #Iterate through annotations and parse each\n",
    "        \n",
    "        #First Create empty arrays for all annotations\n",
    "        masks = np.zeros((dims[1], dims[0], len(anns)))\n",
    "        boxes = np.zeros((len(anns),4))\n",
    "\n",
    "        #Categories could also be appended as a list of text, but in this example, we append the numeric value of the labels directly\n",
    "        categories = np.zeros((len(anns)))\n",
    "        supercats = np.zeros((len(anns)))\n",
    "\n",
    "        areas = np.zeros((len(anns)))\n",
    "        iscrowds = np.zeros((len(anns)))\n",
    "        supercats = np.zeros((len(anns)))\n",
    "        \n",
    "        #Then populate the arrays with the annotations data\n",
    "        for i, ann in enumerate(anns):\n",
    "            mask = coco.annToMask(ann) #Convert annotation to mask\n",
    "            masks[:,:,i] = mask\n",
    "            boxes[i,:] = ann['bbox']\n",
    "            \n",
    "            # Do a brute force search and make no assumptions between order of relationship of category ids\n",
    "            categories[i] = cat_names.index([category_info[i]['name'] for i in range(len(category_info)) if category_info[i]['id']==ann['category_id']][0])\n",
    "            supercats[i] = super_cat_names.index([category_info[i]['supercategory'] for i in range(len(category_info)) if category_info[i]['id']==ann['category_id']][0])\n",
    "            \n",
    "            areas[i] = ann['area']\n",
    "            iscrowds[i] = ann['iscrowd']            \n",
    "\n",
    "            if 'segmentation' not in ann:\n",
    "                print('--- No segmentation found in annotations. ---')\n",
    "                print('Annotation length: {}'.format(len(anns)))\n",
    "                print('--- image id: {} ---'.format(img_id))        \n",
    "        \n",
    "        #Iterate through keypoints and parse each\n",
    "\n",
    "        #Categories could also be appended as a list of text, but in this example, we append the numeric value of the labels directly\n",
    "        categories_kp = np.zeros((len(anns_kp)))\n",
    "        supercats_kp = np.zeros((len(anns_kp)))\n",
    "\n",
    "        masks_kp = np.zeros((dims[1], dims[0], len(anns_kp)))\n",
    "        boxes_kp = np.zeros((len(anns_kp),4))\n",
    "        keypoints_kp = np.zeros((51,len(anns_kp)))\n",
    "\n",
    "        for j, ann_kp in enumerate(anns_kp):\n",
    "            categories_kp[j] = cat_names_kp.index([category_info_kp[i]['name'] for i in range(len(category_info_kp)) if category_info_kp[i]['id']==ann_kp['category_id']][0])\n",
    "            supercats_kp[j] = super_cat_names_kp.index([category_info_kp[i]['supercategory'] for i in range(len(category_info_kp)) if category_info_kp[i]['id']==ann_kp['category_id']][0])\n",
    "            mask_kp = coco.annToMask(ann_kp) #Convert annotation to mask\n",
    "            masks_kp[:,:,j] = mask_kp\n",
    "            boxes_kp[j,:] = ann_kp['bbox']\n",
    "            keypoints_kp[:,j] = np.array(ann_kp['keypoints'])\n",
    "\n",
    "            \n",
    "        #Iterate through stuff and parse each\n",
    "        \n",
    "        masks_stuff = np.zeros((dims[1], dims[0], len(anns_stuff)))\n",
    "        boxes_stuff = np.zeros((len(anns_stuff),4))\n",
    "        \n",
    "        #Categories could also be appended as a list of text, but in this example, we append the numeric value of the labels directly\n",
    "        categories_stuff = np.zeros((len(anns_stuff)))\n",
    "        supercats_stuff = np.zeros((len(anns_stuff)))\n",
    "\n",
    "        areas_stuff = np.zeros((len(anns_stuff)))\n",
    "        iscrowds_stuff = np.zeros((len(anns_stuff)))\n",
    "        supercats_stuff = np.zeros((len(anns_stuff)))\n",
    "        \n",
    "        for k, ann_stuff in enumerate(anns_stuff):\n",
    "            mask_stuff = coco.annToMask(ann_stuff) #Convert annotation to mask\n",
    "            masks_stuff[:,:,k] = mask_stuff\n",
    "            boxes_stuff[k,:] = ann['bbox']\n",
    "            \n",
    "            # Do a brute force search and make no assumptions between order of relationship of category ids\n",
    "            categories_stuff[k] = cat_names_stuff.index([category_info_stuff[i]['name'] for i in range(len(category_info_stuff)) if category_info_stuff[i]['id']==ann_stuff['category_id']][0])\n",
    "            supercats_stuff[k] = super_cat_names_stuff.index([category_info_stuff[i]['supercategory'] for i in range(len(category_info_stuff)) if category_info_stuff[i]['id']==ann_stuff['category_id']][0])\n",
    "            \n",
    "            areas_stuff[k] = ann_stuff['area']\n",
    "            iscrowds_stuff[k] = ann_stuff['iscrowd']            \n",
    "\n",
    "            if 'segmentation' not in ann_stuff:\n",
    "                print('--- No segmentation found in stuff annotations. ---')\n",
    "                print('Annotation length: {}'.format(len(anns)))\n",
    "                print('--- image id: {} ---'.format(img_id))        \n",
    "            \n",
    "            \n",
    "        #Append data to Deep Lake and verify that Deep Lake can open the image. Only do this after all annotations have been parsed.\n",
    "        try:\n",
    "            ds.append({\n",
    "                'images': deeplake.read(img_fn, verify = True),\n",
    "                'images_meta': img_coco,\n",
    "                'masks': masks.astype('bool'),\n",
    "                'boxes': boxes.astype('float32'),\n",
    "                'categories': categories.astype('uint32'),\n",
    "                'super_categories': supercats.astype('uint32'),\n",
    "                'areas': areas.astype('uint32'),\n",
    "                'iscrowds': iscrowds.astype('bool'),\n",
    "                'pose/categories': categories_kp.astype('uint32'),\n",
    "                'pose/super_categories': supercats_kp.astype('uint32'),\n",
    "                'pose/boxes': boxes_kp.astype('float32'),\n",
    "                'pose/masks': masks_kp.astype('bool'),\n",
    "                'pose/keypoints': keypoints_kp.astype('int32'),\n",
    "                'stuff/masks': masks_stuff.astype('bool'),\n",
    "                'stuff/boxes': boxes_stuff.astype('float32'),\n",
    "                'stuff/categories': categories_stuff.astype('uint32'),\n",
    "                'stuff/super_categories': supercats_stuff.astype('uint32'),\n",
    "                'stuff/areas': areas_stuff.astype('uint32'),\n",
    "                'stuff/iscrowds': iscrowds_stuff.astype('bool')\n",
    "            })\n",
    "\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "\n",
    "        if count>=limit:\n",
    "            break\n",
    "            \n",
    "        count+=1   \n",
    "    \n",
    "end_time = time.time()\n",
    "\n",
    "print('Upload took {} seconds'.format(end_time-start_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d80c6d2",
   "metadata": {},
   "source": [
    "## Special case - COCO Test dataset without annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c04af449",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir='./Datasets/coco'\n",
    "data_type='test'\n",
    "\n",
    "deeplake_path = './Datasets/coco_local_{}'.format(data_type) \n",
    "# For uploading to Activeloop storage, register with activeloop and use the path: 'hub://organization_name/coco_{}'.format(data_type)\n",
    "\n",
    "limit = 1e10 # Limit the number of images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a60a444",
   "metadata": {},
   "outputs": [],
   "source": [
    "ann_file='{}/annotations/image_info_{}2017.json'.format(data_dir,data_type) #There are no actual annotations, just images\n",
    "img_root='{}/{}2017/'.format(data_dir,data_type)\n",
    "\n",
    "coco = COCO(ann_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efc66c76",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Specify dataset path\n",
    "ds = deeplake.empty(deeplake_path, overwrite = True) # Set overwrite = True if you need to start over, specify token = 'xyz' if storing the dataset in Activeloop storage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "257785e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_ids = sorted(coco.getImgIds()) # Image ids for uploading\n",
    "count = 1\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "with ds:\n",
    "    \n",
    "    ## ---- Create Tensors ----- ##\n",
    "    \n",
    "    ds.create_tensor('images', htype = 'image', sample_compression = 'jpg')\n",
    "    ds.create_tensor('images_meta', htype = 'json')\n",
    "    \n",
    "    \n",
    "    ## ---- Iterate through each image and upload data ----- ##\n",
    "        \n",
    "    for img_id in tqdm(img_ids):\n",
    "\n",
    "        img_coco = coco.loadImgs(img_id)[0]\n",
    "        img_fn = os.path.join(img_root, img_coco['file_name'])\n",
    "        img = Image.open(img_fn)\n",
    "        dims = img.size\n",
    "                \n",
    "        #Append data to Deep Lake and verify that Deep Lake can open the image\n",
    "        try:\n",
    "            ds.append({\n",
    "                'images': deeplake.read(img_fn, verify = True),\n",
    "                'images_meta': img_coco\n",
    "            })\n",
    "\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "\n",
    "        if count>=limit:\n",
    "            break\n",
    "            \n",
    "        count+=1   \n",
    "    \n",
    "end_time = time.time()\n",
    "\n",
    "print('Upload took {} seconds'.format(end_time-start_time))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.10 64-bit",
   "language": "python",
   "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.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}