1_stratification_data_split.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "formed-stockholm",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import shutil\n",
    "import math\n",
    "import pydicom\n",
    "import cv2\n",
    "import random\n",
    "random.seed(10)\n",
    "\n",
    "cmmd_manifest_directory = \"/media/craig/Larry/python/manifest-1616439774456/\"\n",
    "parent_dir = (cmmd_manifest_directory[:-23]) #e.g. \"/media/craig/Larry/python\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cellular-criterion",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>subject_id</th>\n",
       "      <th>leftright</th>\n",
       "      <th>age</th>\n",
       "      <th>abnormality</th>\n",
       "      <th>classification</th>\n",
       "      <th>subtype</th>\n",
       "      <th>file_location</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>D1-0001</td>\n",
       "      <td>R</td>\n",
       "      <td>44</td>\n",
       "      <td>calcification</td>\n",
       "      <td>Benign</td>\n",
       "      <td>NaN</td>\n",
       "      <td>./CMMD/D1-0001/07-18-2010-NA-NA-79377/1.000000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>D1-0001</td>\n",
       "      <td>R</td>\n",
       "      <td>44</td>\n",
       "      <td>calcification</td>\n",
       "      <td>Benign</td>\n",
       "      <td>NaN</td>\n",
       "      <td>./CMMD/D1-0001/07-18-2010-NA-NA-79377/1.000000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>D1-0002</td>\n",
       "      <td>L</td>\n",
       "      <td>40</td>\n",
       "      <td>calcification</td>\n",
       "      <td>Benign</td>\n",
       "      <td>NaN</td>\n",
       "      <td>./CMMD/D1-0002/07-18-2010-NA-NA-49231/1.000000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>D1-0002</td>\n",
       "      <td>L</td>\n",
       "      <td>40</td>\n",
       "      <td>calcification</td>\n",
       "      <td>Benign</td>\n",
       "      <td>NaN</td>\n",
       "      <td>./CMMD/D1-0002/07-18-2010-NA-NA-49231/1.000000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>D1-0003</td>\n",
       "      <td>L</td>\n",
       "      <td>39</td>\n",
       "      <td>calcification</td>\n",
       "      <td>Benign</td>\n",
       "      <td>NaN</td>\n",
       "      <td>./CMMD/D1-0003/07-18-2011-NA-NA-25491/1.000000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3363</th>\n",
       "      <td>D2-0564</td>\n",
       "      <td>R</td>\n",
       "      <td>30</td>\n",
       "      <td>mass</td>\n",
       "      <td>Malignant</td>\n",
       "      <td>Luminal B</td>\n",
       "      <td>./CMMD/D2-0564/07-18-2011-NA-NA-19553/1.000000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3364</th>\n",
       "      <td>D2-0565</td>\n",
       "      <td>L</td>\n",
       "      <td>75</td>\n",
       "      <td>mass</td>\n",
       "      <td>Malignant</td>\n",
       "      <td>NaN</td>\n",
       "      <td>./CMMD/D2-0565/07-18-2011-NA-NA-72392/1.000000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3365</th>\n",
       "      <td>D2-0565</td>\n",
       "      <td>L</td>\n",
       "      <td>75</td>\n",
       "      <td>mass</td>\n",
       "      <td>Malignant</td>\n",
       "      <td>NaN</td>\n",
       "      <td>./CMMD/D2-0565/07-18-2011-NA-NA-72392/1.000000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3366</th>\n",
       "      <td>D2-0565</td>\n",
       "      <td>R</td>\n",
       "      <td>75</td>\n",
       "      <td>mass</td>\n",
       "      <td>Malignant</td>\n",
       "      <td>Luminal A</td>\n",
       "      <td>./CMMD/D2-0565/07-18-2011-NA-NA-72392/1.000000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3367</th>\n",
       "      <td>D2-0565</td>\n",
       "      <td>R</td>\n",
       "      <td>75</td>\n",
       "      <td>mass</td>\n",
       "      <td>Malignant</td>\n",
       "      <td>Luminal A</td>\n",
       "      <td>./CMMD/D2-0565/07-18-2011-NA-NA-72392/1.000000...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3368 rows × 7 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     subject_id leftright  age    abnormality classification    subtype  \\\n",
       "0       D1-0001         R   44  calcification         Benign        NaN   \n",
       "1       D1-0001         R   44  calcification         Benign        NaN   \n",
       "2       D1-0002         L   40  calcification         Benign        NaN   \n",
       "3       D1-0002         L   40  calcification         Benign        NaN   \n",
       "4       D1-0003         L   39  calcification         Benign        NaN   \n",
       "...         ...       ...  ...            ...            ...        ...   \n",
       "3363    D2-0564         R   30           mass      Malignant  Luminal B   \n",
       "3364    D2-0565         L   75           mass      Malignant        NaN   \n",
       "3365    D2-0565         L   75           mass      Malignant        NaN   \n",
       "3366    D2-0565         R   75           mass      Malignant  Luminal A   \n",
       "3367    D2-0565         R   75           mass      Malignant  Luminal A   \n",
       "\n",
       "                                          file_location  \n",
       "0     ./CMMD/D1-0001/07-18-2010-NA-NA-79377/1.000000...  \n",
       "1     ./CMMD/D1-0001/07-18-2010-NA-NA-79377/1.000000...  \n",
       "2     ./CMMD/D1-0002/07-18-2010-NA-NA-49231/1.000000...  \n",
       "3     ./CMMD/D1-0002/07-18-2010-NA-NA-49231/1.000000...  \n",
       "4     ./CMMD/D1-0003/07-18-2011-NA-NA-25491/1.000000...  \n",
       "...                                                 ...  \n",
       "3363  ./CMMD/D2-0564/07-18-2011-NA-NA-19553/1.000000...  \n",
       "3364  ./CMMD/D2-0565/07-18-2011-NA-NA-72392/1.000000...  \n",
       "3365  ./CMMD/D2-0565/07-18-2011-NA-NA-72392/1.000000...  \n",
       "3366  ./CMMD/D2-0565/07-18-2011-NA-NA-72392/1.000000...  \n",
       "3367  ./CMMD/D2-0565/07-18-2011-NA-NA-72392/1.000000...  \n",
       "\n",
       "[3368 rows x 7 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(\"./CMMD_metadata_subset.csv\")\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "severe-render",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building benign\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1106/1106 [00:21<00:00, 52.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building malignant\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2262/2262 [00:43<00:00, 52.20it/s]\n"
     ]
    }
   ],
   "source": [
    "#Create dataframe which excludes all non benign classifications\n",
    "benign_df = df.loc[df['classification'] == 'Benign']\n",
    "\n",
    "#Create dataframe which excludes all non malignant classifications\n",
    "malignant_df = df.loc[df['classification'] == 'Malignant']\n",
    "\n",
    "cmmd_dir = parent_dir\n",
    "benign_loc = cmmd_dir+\"cmmd_data/benign/\"\n",
    "malignant_loc = cmmd_dir+\"cmmd_data/malignant/\"\n",
    "\n",
    "#create directory if doesnt exist\n",
    "Path(benign_loc).mkdir(parents=True, exist_ok=True)\n",
    "#create directory if doesnt exist\n",
    "Path(malignant_loc).mkdir(parents=True, exist_ok=True)      \n",
    "\n",
    "matches = [\"1-3.dcm\", \"1-4.dcm\"]\n",
    "\n",
    "def create_benign_malignant(df, dest_folder):\n",
    "    \n",
    "    for index, row in tqdm(df.iterrows(), total=df.shape[0]):  \n",
    "        src = cmmd_dir+\"manifest-1616439774456/\"+row['file_location']\n",
    "        basename = os.path.basename(src) #<- basename = file name + extension\n",
    "        if any(x in basename for x in matches): # Check for a 3rd or 4th in path\n",
    "            #append \"_b\" to subject ID to show this is a second case for the same patient\n",
    "            dest = dest_folder+row['subject_id']+\"_b/\"+basename \n",
    "            Path(dest_folder+row['subject_id']+\"_b/\").mkdir(parents=True, exist_ok=True)\n",
    "        else:\n",
    "            dest = dest_folder+row['subject_id']+\"/\"+basename\n",
    "            Path(dest_folder+row['subject_id']+\"/\").mkdir(parents=True, exist_ok=True)\n",
    "        shutil.copyfile(src, dest)\n",
    "\n",
    "print(\"Building benign\")\n",
    "create_benign_malignant(benign_df, benign_loc)\n",
    "print(\"Building malignant\")\n",
    "create_benign_malignant(malignant_df, malignant_loc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "brief-yield",
   "metadata": {},
   "outputs": [],
   "source": [
    "#method to move 20% of a directory into another location\n",
    "#splits a dataset into train/test and/or train/validate\n",
    "def create_test_dataset(data_location, destination):\n",
    "    count = (len(os.listdir(data_location))/5)\n",
    "    count = (math.ceil(count))\n",
    "    test_set = random.sample(os.listdir(data_location), count)\n",
    "\n",
    "    for i in tqdm(range(len(test_set))):\n",
    "\n",
    "        shutil.move(data_location+test_set[i], destination)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fallen-logging",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating 20% test split for benign set...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 111/111 [00:00<00:00, 4665.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating 20% test split for malignant set...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 227/227 [00:00<00:00, 4323.53it/s]\n"
     ]
    }
   ],
   "source": [
    "benign_testset_location = cmmd_dir+\"TEST/benign/\"\n",
    "malignant_testset_location = cmmd_dir+\"TEST/malignant/\"\n",
    "#create directory if doesnt exist\n",
    "Path(benign_testset_location).mkdir(parents=True, exist_ok=True)\n",
    "#create directory if doesnt exist\n",
    "Path(malignant_testset_location).mkdir(parents=True, exist_ok=True)   \n",
    "    \n",
    "print(\"Creating 20% test split for benign set...\")        \n",
    "create_test_dataset(benign_loc, benign_testset_location)\n",
    "\n",
    "print(\"Creating 20% test split for malignant set...\")  \n",
    "create_test_dataset(malignant_loc, malignant_testset_location)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "sensitive-brief",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 89/89 [00:00<00:00, 4007.61it/s]\n",
      "100%|██████████| 181/181 [00:00<00:00, 3552.30it/s]\n"
     ]
    }
   ],
   "source": [
    "benign_valset_location = cmmd_dir+\"VAL/benign/\"\n",
    "malignant_valset_location = cmmd_dir+\"VAL/malignant/\"\n",
    "#create directory if doesnt exist\n",
    "Path(benign_valset_location).mkdir(parents=True, exist_ok=True)\n",
    "#create directory if doesnt exist\n",
    "Path(malignant_valset_location).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "create_test_dataset(benign_loc, benign_valset_location)\n",
    "create_test_dataset(malignant_loc, malignant_valset_location)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "annual-marina",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/media/craig/Larry/python/TRAIN'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "shutil.move(cmmd_dir+\"cmmd_data\", cmmd_dir+\"TRAIN\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "academic-component",
   "metadata": {},
   "outputs": [],
   "source": [
    "def move_dcm_from_subdir(source, destination):\n",
    "    Path(destination).mkdir(parents=True, exist_ok=True)\n",
    "    files_list = os.listdir(source)\n",
    "    j=1\n",
    "    for files in files_list:\n",
    "        files_list2 = os.listdir(source+files)\n",
    "        for x in files_list2:\n",
    "            shutil.move(source+files+\"/\"+x, destination+str(j)+\".dcm\")\n",
    "            j+=1\n",
    "        \n",
    "move_dcm_from_subdir(source=cmmd_dir+\"TRAIN/benign/\",\n",
    "                     destination = cmmd_dir+\"cmmd_data/TRAIN/benign/\")\n",
    "move_dcm_from_subdir(source=cmmd_dir+\"TRAIN/malignant/\",\n",
    "                     destination = cmmd_dir+\"cmmd_data/TRAIN/malignant/\")\n",
    "\n",
    "move_dcm_from_subdir(source=cmmd_dir+\"VAL/benign/\",\n",
    "                     destination = cmmd_dir+\"cmmd_data/VAL/benign/\")\n",
    "move_dcm_from_subdir(source=cmmd_dir+\"VAL/malignant/\",\n",
    "                     destination = cmmd_dir+\"cmmd_data/VAL/malignant/\")\n",
    "\n",
    "move_dcm_from_subdir(source=cmmd_dir+\"TEST/benign/\",\n",
    "                     destination = cmmd_dir+\"cmmd_data/TEST/benign/\")\n",
    "move_dcm_from_subdir(source=cmmd_dir+\"TEST/malignant/\",\n",
    "                     destination = cmmd_dir+\"cmmd_data/TEST/malignant/\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "tested-dylan",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rm_dir(directiory):\n",
    "    ## Try to remove tree\n",
    "    try:\n",
    "        shutil.rmtree(directiory)\n",
    "    except OSError as e:\n",
    "        print (\"Error: %s - %s.\" % (e.filename, e.strerror))\n",
    "\n",
    "rm_dir(cmmd_dir+\"TRAIN\")\n",
    "rm_dir(cmmd_dir+\"TEST\")\n",
    "rm_dir(cmmd_dir+\"VAL\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "several-equality",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working on benign for \n",
      "/media/craig/Larry/python/cmmd_data/PNG/TRAIN/\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 706/706 [00:20<00:00, 35.20it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working on malignant for \n",
      "/media/craig/Larry/python/cmmd_data/PNG/TRAIN/\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1446/1446 [00:42<00:00, 34.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working on benign for \n",
      "/media/craig/Larry/python/cmmd_data/PNG/TEST/\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 222/222 [00:06<00:00, 34.39it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working on malignant for \n",
      "/media/craig/Larry/python/cmmd_data/PNG/TEST/\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 454/454 [00:13<00:00, 33.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working on benign for \n",
      "/media/craig/Larry/python/cmmd_data/PNG/VAL/\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 178/178 [00:05<00:00, 32.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working on malignant for \n",
      "/media/craig/Larry/python/cmmd_data/PNG/VAL/\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 362/362 [00:11<00:00, 32.18it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data converted to PNG and filesystem is tidied.\n"
     ]
    }
   ],
   "source": [
    "def convert_dicom_to_png(input_dir, output_dir):\n",
    "    \n",
    "    if not os.path.exists(output_dir): #if file doesnt exist, create it\n",
    "        Path(output_dir).mkdir(parents=True, exist_ok=True)\n",
    "        \n",
    "    classification_list = [ classification for classification in  os.listdir(input_dir)]\n",
    "    for classification in classification_list:\n",
    "        print(\"Working on \"+classification+\" for \\n\"+output_dir)\n",
    "        dicom_list = [ dcm_image for dcm_image in  os.listdir(input_dir+classification)]\n",
    "        for dcm_image in tqdm(dicom_list):\n",
    "            ds = pydicom.read_file(input_dir+classification+\"/\"+dcm_image) # read dicom image\n",
    "            img = ds.pixel_array # get image array\n",
    "            if not os.path.exists(output_dir + classification): #if file doesnt exist, create it\n",
    "                Path(output_dir + classification).mkdir(parents=True, exist_ok=True)\n",
    "            cv2.imwrite(output_dir + classification + \"/img_\" +dcm_image.replace('.dcm','.png'),img) # write png image\n",
    "\n",
    "\n",
    "convert_dicom_to_png(input_dir = cmmd_dir + 'cmmd_data/TRAIN/',\n",
    "                     output_dir = cmmd_dir + 'cmmd_data/PNG/TRAIN/')    \n",
    "    \n",
    "convert_dicom_to_png(input_dir = cmmd_dir + 'cmmd_data/TEST/',\n",
    "                     output_dir = cmmd_dir + 'cmmd_data/PNG/TEST/')    \n",
    "\n",
    "convert_dicom_to_png(input_dir = cmmd_dir + 'cmmd_data/VAL/',\n",
    "                     output_dir = cmmd_dir + 'cmmd_data/PNG/VAL/') \n",
    "\n",
    "#Tidy up directory\n",
    "rm_dir(cmmd_dir+\"cmmd_data/TRAIN/\")\n",
    "rm_dir(cmmd_dir+\"cmmd_data/TEST/\")\n",
    "rm_dir(cmmd_dir+\"cmmd_data/VAL/\")\n",
    "rm_dir(cmmd_manifest_directory)\n",
    "print(\"Data converted to PNG and filesystem is tidied.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}