prepare_data.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "YOLOv8: https://github.com/ultralytics/ultralytics"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "import pandas as pd\n",
    "from typing import Tuple, Dict, Any\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import cv2\n",
    "import os"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Basic parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = \"datasets/Train\"\n",
    "labels_file = \"labels.csv\"\n",
    "\n",
    "output_folder = \"datasets/train_png2\"\n",
    "use_classes = True  #If false, only one class is used (0)\n",
    "val_frac = 0.2 # Validation fraction of the data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load dataset info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset columns: Index(['case_id', 'label', 'data_path', 'mask_path'], dtype='object')\n"
     ]
    }
   ],
   "source": [
    "# Load dataset\n",
    "dataset = pd.read_csv(os.path.join(dataset_path, labels_file))\n",
    "print(\"Dataset columns:\", dataset.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset classes: ['M', 'B']\n"
     ]
    }
   ],
   "source": [
    "# Get class index\n",
    "classes = dataset['label'].unique().tolist()\n",
    "print(\"Dataset classes:\", classes)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define function for reading NRRD files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ReadNRRD(filename: str) -> Tuple[sitk.Image, Dict[str, Any]]:\n",
    "    reader = sitk.ImageFileReader()\n",
    "    reader.SetFileName(filename)\n",
    "    reader.LoadPrivateTagsOn()\n",
    "    reader.ReadImageInformation()\n",
    "\n",
    "    image = reader.Execute()\n",
    "    metadata = {}\n",
    "    for key in reader.GetMetaDataKeys():\n",
    "        if reader.HasMetaDataKey(key):\n",
    "            metadata[key] = reader.GetMetaData(key)     \n",
    "            \n",
    "    return image, metadata"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define function for 8-bits normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_8bits(image: np.ndarray):\n",
    "    return (255.0 *(image - image.min()) / (image.max() - image.min())).astype(np.uint8)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save 8-bit slices and labels\n",
    "\n",
    "Format YOLOv8:\n",
    "- https://docs.ultralytics.com/datasets/segment/\n",
    "- `<class-index> <x1> <y1> <x2> <y2> ... <xn> <yn>`\n",
    "- Others: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/#13-prepare-dataset-for-yolov5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 2/100 [00:04<03:31,  2.16s/it]"
     ]
    }
   ],
   "source": [
    "os.makedirs(output_folder, exist_ok=True)\n",
    "\n",
    "# Create slice from data\n",
    "for _, row in tqdm(dataset.iterrows(), total=len(dataset)):\n",
    "    \n",
    "    # Load data and GT\n",
    "    id = row.case_id\n",
    "    label = row.label\n",
    "    data, metadata = ReadNRRD(os.path.join(dataset_path, row.data_path.replace('\\\\','/')))\n",
    "    mask, _ = ReadNRRD(os.path.join(dataset_path, row.mask_path.replace('\\\\','/')))\n",
    "    \n",
    "    # Check image and mask size and get numpy arrays\n",
    "    assert data.GetSize() == mask.GetSize()\n",
    "    data_array = sitk.GetArrayFromImage(data)\n",
    "    mask_array = sitk.GetArrayFromImage(mask)\n",
    "    image_size = data_array.shape[2], data_array.shape[1]\n",
    "    \n",
    "    # Chech the labels\n",
    "    assert label in classes\n",
    "    \n",
    "    if use_classes: # Multiples data classes:\n",
    "        data_class = classes.index(label)\n",
    "    else:           # Single data class:\n",
    "        data_class = 0\n",
    "    \n",
    "    # For each slice\n",
    "    for idx in range(len(data_array)): #first dimension is z in numpy (z,y,x)\n",
    "        data_slice = data_array[idx, ...]\n",
    "        mask_slice = mask_array[idx, ...]\n",
    "        \n",
    "        # Use only slices with mask data\n",
    "        if np.sum(mask_slice) == 0:\n",
    "            continue\n",
    "        \n",
    "        # Normalize to 8-bits\n",
    "        data_slice = normalize_8bits(data_slice)\n",
    "        \n",
    "        # Get contours from the mask\n",
    "        contours, _ = cv2.findContours(mask_slice, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "        contours_list = [contour.squeeze().tolist() for contour in contours]\n",
    "        \n",
    "        image_name = f\"{id:0>3}_{idx:0>3}\"\n",
    "        \n",
    "        # Save image\n",
    "        image_out = os.path.join(output_folder, f\"{image_name}.png\")\n",
    "        cv2.imwrite(image_out, data_slice)\n",
    "    \n",
    "        # Save  label\n",
    "        label_out = os.path.join(output_folder, f\"{image_name}.txt\")\n",
    "        with open(label_out, \"w\") as fp:\n",
    "            for contour in contours_list:\n",
    "                contour_str = \" \".join([f\"{point[0]/image_size[0]:0.6f} {point[1]/image_size[1]:0.6f}\" for point in contour])\n",
    "                fp.write(f\"{data_class} {contour_str}\\n\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Train and Val files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0: ['066', '006', '037', '019', '086', '045', '032', '068', '012', '022', '093', '088', '033', '024', '059', '089', '058', '023', '036', '009', '015', '035', '077', '094', '087', '075', '021', '008', '056', '002', '060', '076', '092', '029', '003', '067', '020', '063', '098', '044', '038', '090', '054', '000', '070', '016', '069', '080', '034', '030', '081', '072', '047', '099', '040', '028', '085', '010'], 1: ['082', '026', '043', '079', '052', '074', '031', '097', '084', '039', '046', '096', '071', '048', '014', '061', '041', '065', '018', '017', '025', '007', '064', '095', '073', '049', '004', '005', '013', '050', '057', '062', '055', '051', '011', '053', '091', '083', '042', '078', '027', '001']}\n",
      "58 42\n",
      "['066', '006', '037', '019', '086', '045', '032', '068', '012', '022', '093', '088', '033', '024', '059', '089', '058', '023', '036', '009', '015', '035', '077', '094', '087', '075', '021', '008', '056', '002', '060', '076', '092', '029', '003', '067', '020', '063', '098', '044', '038', '090', '054', '000', '070', '016', '082', '026', '043', '079', '052', '074', '031', '097', '084', '039', '046', '096', '071', '048', '014', '061', '041', '065', '018', '017', '025', '007', '064', '095', '073', '049', '004', '005', '013', '050', '057', '062', '055'] 79\n",
      "['069', '080', '034', '030', '081', '072', '047', '099', '040', '028', '085', '010', '051', '011', '053', '091', '083', '042', '078', '027', '001'] 21\n"
     ]
    }
   ],
   "source": [
    "# Get patients list\n",
    "list_items = [item[:-4] for item in os.listdir(output_folder) if item.endswith(\".png\")]\n",
    "patients = list(set([ item.split(\"_\")[0] for item in list_items]))\n",
    "\n",
    "# Get lesion type per patient\n",
    "patients_type = {i: [] for i in range(len(classes))}\n",
    "for item in list_items:\n",
    "    p = item.split(\"_\")[0]\n",
    "    item_path = os.path.join(output_folder, f\"{item}.txt\")\n",
    "    with open(item_path) as fp:\n",
    "        class_type = int(fp.readline().split(\" \")[0])\n",
    "    if p not in patients_type[class_type]:\n",
    "        patients_type[class_type].append(p)\n",
    "    \n",
    "print(patients_type)\n",
    "print(len(patients_type[0]),len(patients_type[1]) )\n",
    "    \n",
    "# Distribute patients in Train/Val using val_frac with balanced lesion types\n",
    "train_p = []\n",
    "val_p = []\n",
    "for i in range(len(classes)):\n",
    "    num_train = int(len(patients_type[i])*(1-val_frac))\n",
    "    train_type_p, val_type_p = patients_type[i][:num_train], patients_type[i][num_train:]\n",
    "    train_p += train_type_p\n",
    "    val_p += val_type_p\n",
    "    \n",
    "print(train_p, len(train_p))\n",
    "print(val_p, len(val_p))\n",
    "    \n",
    "# Get distributed patient images for train and val\n",
    "train = [ os.path.join(output_folder, f\"{item}.png\") for item in list_items if item.split(\"_\")[0] in train_p]\n",
    "val = [ os.path.join(output_folder, f\"{item}.png\") for item in list_items if item.split(\"_\")[0] in val_p]\n",
    "\n",
    "# Generate train.txt file\n",
    "train_file = os.path.join(output_folder, \"../train.txt\")\n",
    "with open(train_file, \"w\") as fp:\n",
    "    fp.writelines([t + '\\n' for t in train])\n",
    "    \n",
    "# Generate val.txt file\n",
    "val_file = os.path.join(output_folder, \"../val.txt\")\n",
    "with open(val_file, \"w\") as fp:\n",
    "    fp.writelines([v + '\\n' for v in val])\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310_1",
   "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.10.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "12f3aea31d254c34d7f2ef2d829cf85510b91fc37a99fa9bf72c9e4d79053620"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}