augmentation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bcd3a43c-b69c-42ce-924a-56c109882878",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import json\n",
    "from imgaug import augmenters as iaa\n",
    "import imgaug as ia\n",
    "import skimage.io as io\n",
    "import imageio\n",
    "from pycocotools.coco import COCO\n",
    "import pycocotools.mask as cocomask\n",
    "from itertools import groupby"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "69ecec9a-f274-4503-bebf-1fba8c8a7503",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = ['train', 'valid', 'test', 'difficult']\n",
    "root = os.getcwd()\n",
    "\n",
    "def binary_mask_to_rle(binary_mask):\n",
    "    rle = {'counts': [], 'size': list(binary_mask.shape)}\n",
    "    counts = rle.get('counts')\n",
    "    for i, (value, elements) in enumerate(groupby(binary_mask.ravel(order='F'))):\n",
    "        if i == 0 and value == 1:\n",
    "            counts.append(0)\n",
    "        counts.append(len(list(elements)))\n",
    "    return rle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "19c10579-d453-4328-8a09-f9c65ca37d2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.01s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Sveta\\anaconda3\\envs\\CudaLab\\lib\\site-packages\\imgaug\\imgaug.py:184: DeprecationWarning: Function `SegmentationMapOnImage()` is deprecated. Use `SegmentationMapsOnImage` instead. (Note the plural 'Maps' instead of old 'Map'.).\n",
      "  warn_deprecated(msg, stacklevel=3)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.00s)\n",
      "creating index...\n",
      "index created!\n",
      "loading annotations into memory...\n",
      "Done (t=0.00s)\n",
      "creating index...\n",
      "index created!\n",
      "loading annotations into memory...\n",
      "Done (t=0.00s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "#can add different augmentation\n",
    "rot = iaa.Flipud(1.0)\n",
    "\n",
    "#old folders\n",
    "old_folder_images = os.path.join(root, 'images_split')\n",
    "old_folder_annotations = os.path.join(root, 'annotations_split')\n",
    "\n",
    "#create new folders\n",
    "new_folder_images = os.path.join(root, 'images')\n",
    "new_folder_annotations = os.path.join(root, 'annotations')\n",
    "os.makedirs(os.path.join(root, 'images'), exist_ok=True)\n",
    "os.makedirs(os.path.join(root, 'annotations'), exist_ok=True)\n",
    "\n",
    "#iterate through datasets\n",
    "for dataset in datasets:\n",
    "    #create new folders\n",
    "    new_path_images = os.path.join(new_folder_images, dataset)\n",
    "    os.makedirs(new_path_images, exist_ok=True)\n",
    "    \n",
    "    #annotations path\n",
    "    old_path_annotations = os.path.join(old_folder_annotations, dataset + '.json')\n",
    "    new_path_annotations = os.path.join(new_folder_annotations, dataset + '.json')\n",
    "    \n",
    "    #load data\n",
    "    coco = COCO(old_path_annotations)\n",
    "    with open(old_path_annotations) as f:\n",
    "        data = json.load(f)\n",
    "        cat_ids = coco.getCatIds()\n",
    "        anns_ids = coco.getAnnIds(catIds=cat_ids, iscrowd=None)\n",
    "        anns = coco.loadAnns(anns_ids)\n",
    "    \n",
    "    #iterate through annotations\n",
    "    for item in anns:\n",
    "        #read image\n",
    "        image_name = coco.loadImgs(item['image_id'])[0]['file_name']\n",
    "        image_path = os.path.join(old_folder_images, dataset, image_name)\n",
    "        image = io.imread(image_path)\n",
    "        \n",
    "        #read mask\n",
    "        mask = coco.annToMask(item)\n",
    "        segmap = ia.SegmentationMapOnImage(mask, shape=image.shape)\n",
    "        \n",
    "        #read bounding box\n",
    "        bb = ia.augmentables.bbs.BoundingBoxesOnImage([ia.augmentables.bbs.BoundingBox(item['bbox'][0], item['bbox'][1], \n",
    "                                                             item['bbox'][0] + item['bbox'][2],\n",
    "                                                             item['bbox'][1] + item['bbox'][3])], shape=image.shape)\n",
    "        \n",
    "        #apply augmentation\n",
    "        image_aug = rot.augment_image(image)\n",
    "        segmap_aug = rot.augment_segmentation_maps(segmap)\n",
    "        bb_aug = rot.augment_bounding_boxes(bb)\n",
    "        \n",
    "        #save augmentation\n",
    "        io.imsave(os.path.join(new_path_images, image_name), image_aug)\n",
    "        item['segmentation'] = binary_mask_to_rle(segmap_aug.get_arr())\n",
    "        item['bbox'] = [bb_aug[0].x1, bb_aug[0].y1, bb_aug[0].width, bb_aug[0].height]\n",
    "    \n",
    "    #save json\n",
    "    data[\"annotations\"] = anns\n",
    "    with open(new_path_annotations, 'w') as f:\n",
    "        json.dump(data, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "224473d1-a4b6-4557-aa38-af3284ac166f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}