create_LungSegData.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import cv2\n",
    "from skimage.measure import label\n",
    "from skimage.color import label2rgb\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "#path to cohen segmentation masks and data\n",
    "cohen_ann=\"../data_sources/covid-chestxray-dataset/annotations/lungVAE-masks\"\n",
    "cohen_img=\"../data_sources/covid-chestxray-dataset/images\"\n",
    "\n",
    "#path to XSLOR segmentation training data masks and CXR images\n",
    "XLSor = \"../data_sources/XLSor_data/data/Augmentation\"\n",
    "\n",
    "#path to NIH segmentation test set\n",
    "NIH_img = \"../data_sources/XLSor_data/data/NIH/images\"\n",
    "NIH_ann = \"../data_sources/XLSor_data/data/NIH/masks\"\n",
    "\n",
    "#path to test data folders\n",
    "cohen_test_save_dir = \"../datasets/segmentation/cohen_test/\"\n",
    "NIH_test_save_dir = \"../datasets/segmentation/NIH_test/\"\n",
    "\n",
    "#path to train data folder\n",
    "train_save_dir = \"../datasets/segmentation/train/\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#first lets create the test data\n",
    "\n",
    "#annotation files:\n",
    "ann_files = os.listdir(cohen_ann)\n",
    "img_files = os.listdir(cohen_img)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import ceil, floor\n",
    "def squarify(M,val):\n",
    "    \n",
    "    (a,b)=M.shape\n",
    "    if a>b:\n",
    "        padding=((0,0),(ceil((a-b)/2),floor((a-b)/2)))\n",
    "    else:\n",
    "        padding=((ceil((b-a)/2),floor((b-a)/2)),(0,0))\n",
    "        \n",
    "    out=np.pad(M,padding,mode='constant',constant_values=val)\n",
    "    \n",
    "    if (out.shape[0]!=out.shape[1]):\n",
    "        print('error when squraifying')\n",
    "        exit()\n",
    "   \n",
    "    return np.pad(M,padding,mode='constant',constant_values=val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "pairs={}\n",
    "for img_file in img_files:\n",
    "    for ann_file in ann_files:\n",
    "        key_check = ann_file.split(\"_mask\")[0]\n",
    "        value_check = '.'.join(img_file.split('.')[:-1])\n",
    "       \n",
    "        if key_check==value_check:\n",
    "            if ann_file in pairs.keys():\n",
    "                print(\"key error, already exists\")\n",
    "                '''print(ann_file)\n",
    "                print(pairs[ann_file])\n",
    "                print(img_file)'''\n",
    "            else:\n",
    "                pairs[ann_file]=img_file\n",
    "                '''print(ann_file)\n",
    "                print(img_file)'''\n",
    "            \n",
    "            count+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "206\n",
      "206\n"
     ]
    }
   ],
   "source": [
    "print(len(pairs))\n",
    "print(count)\n",
    "show = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-18-929b55a69716>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     34\u001b[0m     \u001b[0mcv2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimwrite\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg_out_path\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 36\u001b[0;31m     \u001b[0mdisp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlabel2rgb\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mann\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbg_label\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     37\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     38\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/open_CV/lib/python3.7/site-packages/skimage/_shared/utils.py\u001b[0m in \u001b[0;36mfixed_func\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     64\u001b[0m                 \u001b[0;31m# warn that arg_name default value changed:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     65\u001b[0m                 \u001b[0mwarnings\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwarn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwarning_msg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mFutureWarning\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstacklevel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 66\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     67\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     68\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mfixed_func\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/open_CV/lib/python3.7/site-packages/skimage/color/colorlabel.py\u001b[0m in \u001b[0;36mlabel2rgb\u001b[0;34m(label, image, colors, alpha, bg_label, bg_color, image_alpha, kind)\u001b[0m\n\u001b[1;32m    112\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mkind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'overlay'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    113\u001b[0m         return _label2rgb_overlay(label, image, colors, alpha, bg_label,\n\u001b[0;32m--> 114\u001b[0;31m                                   bg_color, image_alpha)\n\u001b[0m\u001b[1;32m    115\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0mkind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'avg'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    116\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_label2rgb_avg\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlabel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mimage\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbg_label\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbg_color\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/open_CV/lib/python3.7/site-packages/skimage/color/colorlabel.py\u001b[0m in \u001b[0;36m_label2rgb_overlay\u001b[0;34m(label, image, colors, alpha, bg_label, bg_color, image_alpha)\u001b[0m\n\u001b[1;32m    189\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mimage\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    190\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 191\u001b[0;31m     \u001b[0mdense_labels\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmapped_labels_flat\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    192\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    193\u001b[0m     \u001b[0mlabel_to_color\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mc\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdense_labels\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcolor_cycle\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "from skimage.measure import label\n",
    "from skimage.color import label2rgb\n",
    "\n",
    "count = 0\n",
    "\n",
    "for ann_file,img_file in pairs.items():\n",
    "    \n",
    "    \n",
    "    \n",
    "    ann = cv2.imread(os.path.join(cohen_ann,ann_file),0)\n",
    "    img = cv2.imread(os.path.join(cohen_img,img_file),0)\n",
    "    \n",
    "    img = squarify(img,0)\n",
    "    ann = squarify(ann,0)\n",
    "    img = cv2.resize(img,(512,512))\n",
    "    ann = (cv2.resize(ann,(512,512))>(255/2)).astype(int)*255\n",
    "    \n",
    "    new_name= \"IMAGE_\"+str(count)\n",
    "    count+=1\n",
    "    \n",
    "    \n",
    "    '''cv2.imwrite(save_dir+new_name+\"_mask.png\",ann)\n",
    "    cv2.imwrite(save_dir+new_name+\".png\",img)'''\n",
    "    \n",
    "    ann_out_path = os.path.join(cohen_test_save_dir,new_name+\"_mask.png\")\n",
    "    img_out_path = os.path.join(cohen_test_save_dir,new_name+\".png\")\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    cv2.imwrite(ann_out_path,ann)\n",
    "    cv2.imwrite(img_out_path,img)\n",
    "    \n",
    "    disp = label2rgb(ann,img, bg_label = 0)\n",
    "  \n",
    "    \n",
    "    if show == True:\n",
    "        plt.figure(figsize = (5,5))\n",
    "        print(ann_file)\n",
    "        print(img_file)\n",
    "        print('saving to..')\n",
    "        print(ann_out_path)\n",
    "        print(img_out_path)\n",
    "        plt.imshow(disp)\n",
    "        plt.axis('off')\n",
    "        plt.show()\n",
    "    \n",
    "    \n",
    "        \n",
    "   \n",
    "   \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "for img_file in os.listdir(NIH_img):\n",
    "    ann_file = img_file.split('.')[0]+'_mask.png'\n",
    "    \n",
    "    img = cv2.imread(os.path.join(NIH_img,img_file))\n",
    "    ann = cv2.imread(os.path.join(NIH_ann,ann_file))\n",
    "    \n",
    "    if show==True:\n",
    "        plt.figure()\n",
    "        plt.subplot(1,2,1)\n",
    "        plt.imshow(img)\n",
    "        plt.subplot(1,2,2)\n",
    "        plt.imshow(ann)\n",
    "        plt.show()\n",
    "    \n",
    "    ann_outpath = os.path.join(NIH_test_save_dir, ann_file)\n",
    "    img_outpath = os.path.join(NIH_test_save_dir, img_file)\n",
    "    \n",
    "    cv2.imwrite(ann_outpath, ann)\n",
    "    cv2.imwrite(img_outpath, img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-20-c1f134326479>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     23\u001b[0m             \u001b[0mimg_path\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mXLSor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtag\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m\".png\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m             \u001b[0;31m#img=cv2.imread(\"XLSor_data/data/Augmentation/\"+img_file,0)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 25\u001b[0;31m             \u001b[0mimg\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcv2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg_path\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     26\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "#now lets do the training data\n",
    "from skimage.segmentation import mark_boundaries\n",
    "os.listdir(XLSor)\n",
    "\n",
    "ann_files = []\n",
    "img_files = []\n",
    "\n",
    "count = 0\n",
    "for file in os.listdir(XLSor):\n",
    "    if \"mask\" in file:\n",
    "        \n",
    "        ann_path = os.path.join(XLSor,file)\n",
    "        ann=cv2.imread(ann_path,0)\n",
    "        #ann=cv2.imread(\"XLSor_data/data/Augmentation/\"+file,0)\n",
    "        tag=file.split(\"mask\")[0]\n",
    "        \n",
    "        for i in range(6):\n",
    "            \n",
    "           \n",
    "                \n",
    "            #img_file = tag+str(i)+\".png\"\n",
    "            \n",
    "            img_path=os.path.join(XLSor, tag+str(i)+\".png\")\n",
    "            #img=cv2.imread(\"XLSor_data/data/Augmentation/\"+img_file,0)\n",
    "            img=cv2.imread(img_path,0)\n",
    "            \n",
    "            \n",
    "            \n",
    "            \n",
    "            new_tag = \"IMAGE_\"+str(count)\n",
    "            count+=1\n",
    "            img_out_path = os.path.join(train_save_dir, new_tag+'.png')\n",
    "            ann_out_path = os.path.join(train_save_dir, new_tag+'_mask.png')\n",
    "            \n",
    "            cv2.imwrite(img_out_path,img)\n",
    "            cv2.imwrite(ann_out_path,ann)\n",
    "            \n",
    "            disp = label2rgb(ann,img, bg_label = 0)\n",
    "            #disp = mark_boundaries(img,ann)\n",
    "            if show==True:\n",
    "                print(img_path)\n",
    "                print(ann_path)\n",
    "                print(img_out_path)\n",
    "                print(ann_out_path)\n",
    "\n",
    "                plt.figure(figsize = (5,5))\n",
    "                plt.imshow(disp)\n",
    "                plt.axis('off')\n",
    "                plt.show()\n",
    "     \n",
    "            \n",
    "            \n",
    "            \n",
    "            \n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "open_CV",
   "language": "python",
   "name": "open_cv"
  },
  "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}