process_abus23.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = \"datasets/Train\"\n",
    "label_file = \"labels.csv\"\n",
    "dataset_name = \"abus23_25\"\n",
    "\n",
    "output_folder = f\"datasets/{dataset_name}_png\"\n",
    "slice_min_lesion_px = 25\n",
    "\n",
    "use_classes = False  #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": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def LoadAbus23(dataset_path, label_file):\n",
    "    dataset = pd.read_csv(os.path.join(dataset_path, label_file))\n",
    "    print(\"Dataset columns:\", dataset.columns)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset columns: Index(['case_id', 'label', 'data_path', 'mask_path'], dtype='object')\n",
      "Dataset classes: ['M', 'B']\n"
     ]
    }
   ],
   "source": [
    "# Load dataset\n",
    "dataset = LoadAbus23(dataset_path, label_file)\n",
    "\n",
    "\n",
    "# Get class index\n",
    "if 'label' in dataset.columns:\n",
    "    classes = dataset['label'].unique().tolist()\n",
    "    print(\"Dataset classes:\", classes)\n",
    "else:\n",
    "    print(\"No classes found!\")\n",
    "    classes = []"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define function for reading NRRD files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ReadDCM(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": 6,
   "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": [
    "### Create 8-bit slices and masks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_folder = os.path.join(output_folder, \"images\")\n",
    "masks_folder = os.path.join(output_folder, \"masks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [05:23<00:00,  3.23s/it]\n"
     ]
    }
   ],
   "source": [
    "os.makedirs(masks_folder, exist_ok=True)\n",
    "os.makedirs(images_folder, exist_ok=True)\n",
    "\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",
    "    data, metadata = ReadDCM(os.path.join(dataset_path, row.data_path.replace('\\\\','/')))\n",
    "    mask, _ = ReadDCM(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",
    "\n",
    "    # Chech the labels\n",
    "    if len(classes) > 0:\n",
    "        label = row.label\n",
    "        assert label in classes\n",
    "    else:\n",
    "        label = 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",
    "        # Binarize mask_Slice\n",
    "        mask_slice = (mask_slice > 0)\n",
    "        if np.sum(mask_slice) < slice_min_lesion_px:\n",
    "            continue\n",
    "        \n",
    "        mask_slice = mask_slice.astype(np.uint8) * 255\n",
    "        \n",
    "        # Normalize to 8-bits\n",
    "        data_slice = normalize_8bits(data_slice)\n",
    "        \n",
    "        \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+1:0>3}\"\n",
    "        \n",
    "        # Save image\n",
    "        image_out = os.path.join(images_folder, f\"{image_name}_{label}.png\")\n",
    "        cv2.imwrite(image_out, data_slice)\n",
    "    \n",
    "        # Save mask\n",
    "        image_out = os.path.join(masks_folder, f\"{image_name}_{label}.png\")\n",
    "        cv2.imwrite(image_out,  mask_slice)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create YOLO data\n",
    "\n",
    "Format YOLOv8 segement:\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\n",
    "\n",
    "Format YOLOv8 detect:\n",
    "- https://docs.ultralytics.com/datasets/detect/\n",
    "- `<object-class> <x> <y> <width> <height>`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "segmentation_data_path = os.path.join(output_folder, f\"yolo_seg_data\")\n",
    "detection_data_path = os.path.join(output_folder, f\"yolo_det_data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Classes are not used\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 11/3170 [00:00<00:30, 105.18it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 094_248_M.png\n",
      "More than 1 coutour found for image 093_193_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 30/3170 [00:00<00:20, 150.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 007_274_B.png\n",
      "More than 1 coutour found for image 084_215_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▎         | 117/3170 [00:00<00:15, 202.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 020_337_M.png\n",
      "More than 1 coutour found for image 036_265_M.png\n",
      "More than 1 coutour found for image 093_196_M.png\n",
      "More than 1 coutour found for image 026_177_B.png\n",
      "More than 1 coutour found for image 084_211_B.png\n",
      "More than 1 coutour found for image 087_226_M.png\n",
      "More than 1 coutour found for image 068_199_M.png\n",
      "More than 1 coutour found for image 087_225_M.png\n",
      "More than 1 coutour found for image 024_145_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  6%|▌         | 181/3170 [00:00<00:14, 206.42it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 019_241_M.png\n",
      "More than 1 coutour found for image 066_234_M.png\n",
      "More than 1 coutour found for image 019_200_M.png\n",
      "More than 1 coutour found for image 012_193_M.png\n",
      "More than 1 coutour found for image 022_275_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 223/3170 [00:01<00:14, 205.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 084_212_B.png\n",
      "More than 1 coutour found for image 022_301_M.png\n",
      "More than 1 coutour found for image 022_265_M.png\n",
      "More than 1 coutour found for image 084_221_B.png\n",
      "More than 1 coutour found for image 022_300_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|▊         | 264/3170 [00:01<00:14, 195.53it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 012_192_M.png\n",
      "More than 1 coutour found for image 077_237_M.png\n",
      "More than 1 coutour found for image 046_158_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|▉         | 305/3170 [00:01<00:14, 199.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 074_158_B.png\n",
      "More than 1 coutour found for image 047_160_M.png\n",
      "More than 1 coutour found for image 059_233_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 11%|█         | 348/3170 [00:01<00:14, 200.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 046_146_B.png\n",
      "More than 1 coutour found for image 026_178_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 13%|█▎        | 412/3170 [00:02<00:13, 203.07it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 068_200_M.png\n",
      "More than 1 coutour found for image 046_155_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 14%|█▍        | 453/3170 [00:02<00:14, 187.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 012_190_M.png\n",
      "More than 1 coutour found for image 007_270_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 16%|█▌        | 492/3170 [00:02<00:15, 178.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 046_175_B.png\n",
      "More than 1 coutour found for image 024_164_M.png\n",
      "More than 1 coutour found for image 068_178_M.png\n",
      "More than 1 coutour found for image 057_199_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 18%|█▊        | 565/3170 [00:02<00:15, 171.68it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 070_155_M.png\n",
      "More than 1 coutour found for image 079_138_B.png\n",
      "More than 1 coutour found for image 029_160_M.png\n",
      "More than 1 coutour found for image 016_190_M.png\n",
      "More than 1 coutour found for image 019_199_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|█▉        | 622/3170 [00:03<00:13, 182.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 066_209_M.png\n",
      "More than 1 coutour found for image 034_250_M.png\n",
      "More than 1 coutour found for image 055_252_B.png\n",
      "More than 1 coutour found for image 059_234_M.png\n",
      "More than 1 coutour found for image 084_209_B.png\n",
      "More than 1 coutour found for image 012_251_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 21%|██        | 664/3170 [00:03<00:13, 192.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 012_189_M.png\n",
      "More than 1 coutour found for image 026_161_B.png\n",
      "More than 1 coutour found for image 039_196_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 22%|██▏       | 684/3170 [00:03<00:13, 189.86it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 012_201_M.png\n",
      "More than 1 coutour found for image 022_310_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 23%|██▎       | 740/3170 [00:03<00:14, 167.31it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 094_223_M.png\n",
      "More than 1 coutour found for image 079_136_B.png\n",
      "More than 1 coutour found for image 084_216_B.png\n",
      "More than 1 coutour found for image 056_220_M.png\n",
      "More than 1 coutour found for image 077_235_M.png\n",
      "More than 1 coutour found for image 093_200_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 25%|██▍       | 779/3170 [00:04<00:13, 177.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 074_166_B.png\n",
      "More than 1 coutour found for image 055_255_B.png\n",
      "More than 1 coutour found for image 070_156_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██▋       | 842/3170 [00:04<00:12, 193.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 084_210_B.png\n",
      "More than 1 coutour found for image 046_152_B.png\n",
      "More than 1 coutour found for image 068_174_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██▊       | 883/3170 [00:04<00:11, 197.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 039_197_B.png\n",
      "More than 1 coutour found for image 066_235_M.png\n",
      "More than 1 coutour found for image 093_199_M.png\n",
      "More than 1 coutour found for image 059_214_M.png\n",
      "More than 1 coutour found for image 006_175_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 29%|██▉       | 924/3170 [00:04<00:11, 198.61it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 074_157_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 31%|███       | 982/3170 [00:05<00:12, 170.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 068_198_M.png\n",
      "More than 1 coutour found for image 021_261_M.png\n",
      "More than 1 coutour found for image 007_271_B.png\n",
      "More than 1 coutour found for image 021_262_M.png\n",
      "More than 1 coutour found for image 012_191_M.png\n",
      "More than 1 coutour found for image 066_236_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 32%|███▏      | 1018/3170 [00:05<00:12, 166.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 094_246_M.png\n",
      "More than 1 coutour found for image 026_176_B.png\n",
      "More than 1 coutour found for image 046_177_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 34%|███▍      | 1071/3170 [00:05<00:12, 165.16it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 055_254_B.png\n",
      "More than 1 coutour found for image 076_221_M.png\n",
      "More than 1 coutour found for image 022_269_M.png\n",
      "More than 1 coutour found for image 037_245_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 35%|███▍      | 1106/3170 [00:06<00:12, 165.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 084_213_B.png\n",
      "More than 1 coutour found for image 079_140_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 37%|███▋      | 1158/3170 [00:06<00:11, 168.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 099_294_M.png\n",
      "More than 1 coutour found for image 036_266_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 38%|███▊      | 1193/3170 [00:06<00:12, 152.47it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 093_198_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 39%|███▉      | 1249/3170 [00:06<00:11, 172.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 000_262_M.png\n",
      "More than 1 coutour found for image 007_273_B.png\n",
      "More than 1 coutour found for image 066_240_M.png\n",
      "More than 1 coutour found for image 087_229_M.png\n",
      "More than 1 coutour found for image 072_237_M.png\n",
      "More than 1 coutour found for image 093_194_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 42%|████▏     | 1326/3170 [00:07<00:10, 179.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 022_274_M.png\n",
      "More than 1 coutour found for image 006_172_M.png\n",
      "More than 1 coutour found for image 034_254_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 1368/3170 [00:07<00:09, 192.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 097_233_B.png\n",
      "More than 1 coutour found for image 066_244_M.png\n",
      "More than 1 coutour found for image 012_188_M.png\n",
      "More than 1 coutour found for image 012_199_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|████▍     | 1409/3170 [00:07<00:09, 193.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 056_222_M.png\n",
      "More than 1 coutour found for image 059_333_M.png\n",
      "More than 1 coutour found for image 026_179_B.png\n",
      "More than 1 coutour found for image 036_268_M.png\n",
      "More than 1 coutour found for image 012_204_M.png\n",
      "More than 1 coutour found for image 046_173_B.png\n",
      "More than 1 coutour found for image 066_210_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 46%|████▋     | 1470/3170 [00:08<00:08, 199.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 084_214_B.png\n",
      "More than 1 coutour found for image 088_179_M.png\n",
      "More than 1 coutour found for image 093_195_M.png\n",
      "More than 1 coutour found for image 012_200_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 48%|████▊     | 1512/3170 [00:08<00:08, 203.76it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 046_157_B.png\n",
      "More than 1 coutour found for image 072_236_M.png\n",
      "More than 1 coutour found for image 059_217_M.png\n",
      "More than 1 coutour found for image 059_337_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 49%|████▉     | 1553/3170 [00:08<00:08, 191.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 093_231_M.png\n",
      "More than 1 coutour found for image 008_202_M.png\n",
      "More than 1 coutour found for image 059_336_M.png\n",
      "More than 1 coutour found for image 084_223_B.png\n",
      "More than 1 coutour found for image 076_222_M.png\n",
      "More than 1 coutour found for image 022_267_M.png\n",
      "More than 1 coutour found for image 034_249_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 51%|█████     | 1614/3170 [00:08<00:07, 195.21it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 024_163_M.png\n",
      "More than 1 coutour found for image 094_249_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 53%|█████▎    | 1675/3170 [00:09<00:07, 190.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 046_156_B.png\n",
      "More than 1 coutour found for image 068_179_M.png\n",
      "More than 1 coutour found for image 061_177_B.png\n",
      "More than 1 coutour found for image 059_232_M.png\n",
      "More than 1 coutour found for image 076_223_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 55%|█████▍    | 1738/3170 [00:09<00:07, 199.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 024_158_M.png\n",
      "More than 1 coutour found for image 072_235_M.png\n",
      "More than 1 coutour found for image 019_201_M.png\n",
      "More than 1 coutour found for image 066_242_M.png\n",
      "More than 1 coutour found for image 034_252_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 56%|█████▌    | 1780/3170 [00:09<00:07, 198.55it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 070_154_M.png\n",
      "More than 1 coutour found for image 068_177_M.png\n",
      "More than 1 coutour found for image 022_268_M.png\n",
      "More than 1 coutour found for image 059_203_M.png\n",
      "More than 1 coutour found for image 093_197_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 57%|█████▋    | 1821/3170 [00:09<00:06, 197.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 012_246_M.png\n",
      "More than 1 coutour found for image 017_213_B.png\n",
      "More than 1 coutour found for image 007_272_B.png\n",
      "More than 1 coutour found for image 022_264_M.png\n",
      "More than 1 coutour found for image 012_203_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 59%|█████▉    | 1881/3170 [00:10<00:06, 192.92it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 020_312_M.png\n",
      "More than 1 coutour found for image 077_238_M.png\n",
      "More than 1 coutour found for image 012_197_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 61%|██████    | 1923/3170 [00:10<00:06, 199.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 088_177_M.png\n",
      "More than 1 coutour found for image 036_263_M.png\n",
      "More than 1 coutour found for image 029_156_M.png\n",
      "More than 1 coutour found for image 066_243_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 63%|██████▎   | 1986/3170 [00:10<00:05, 201.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 059_213_M.png\n",
      "More than 1 coutour found for image 022_270_M.png\n",
      "More than 1 coutour found for image 036_264_M.png\n",
      "More than 1 coutour found for image 045_158_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 64%|██████▍   | 2029/3170 [00:10<00:05, 205.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 029_157_M.png\n",
      "More than 1 coutour found for image 061_159_B.png\n",
      "More than 1 coutour found for image 056_221_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 65%|██████▌   | 2071/3170 [00:11<00:05, 198.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 000_265_M.png\n",
      "More than 1 coutour found for image 046_150_B.png\n",
      "More than 1 coutour found for image 059_210_M.png\n",
      "More than 1 coutour found for image 070_153_M.png\n",
      "More than 1 coutour found for image 012_194_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 67%|██████▋   | 2133/3170 [00:11<00:05, 200.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 045_159_M.png\n",
      "More than 1 coutour found for image 029_158_M.png\n",
      "More than 1 coutour found for image 094_247_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 69%|██████▊   | 2174/3170 [00:11<00:05, 188.76it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 088_178_M.png\n",
      "More than 1 coutour found for image 098_230_M.png\n",
      "More than 1 coutour found for image 022_309_M.png\n",
      "More than 1 coutour found for image 019_197_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|██████▉   | 2212/3170 [00:11<00:05, 177.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 093_201_M.png\n",
      "More than 1 coutour found for image 020_311_M.png\n",
      "More than 1 coutour found for image 066_239_M.png\n",
      "More than 1 coutour found for image 059_211_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 71%|███████   | 2248/3170 [00:12<00:05, 173.06it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 034_251_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 73%|███████▎  | 2302/3170 [00:12<00:05, 163.04it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 046_174_B.png\n",
      "More than 1 coutour found for image 012_196_M.png\n",
      "More than 1 coutour found for image 087_243_M.png\n",
      "More than 1 coutour found for image 035_226_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 75%|███████▍  | 2364/3170 [00:12<00:04, 188.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 077_236_M.png\n",
      "More than 1 coutour found for image 012_195_M.png\n",
      "More than 1 coutour found for image 017_215_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 77%|███████▋  | 2428/3170 [00:13<00:03, 201.99it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 017_214_B.png\n",
      "More than 1 coutour found for image 097_232_B.png\n",
      "More than 1 coutour found for image 003_299_M.png\n",
      "More than 1 coutour found for image 085_229_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 78%|███████▊  | 2469/3170 [00:13<00:03, 197.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 035_225_M.png\n",
      "More than 1 coutour found for image 022_276_M.png\n",
      "More than 1 coutour found for image 098_229_M.png\n",
      "More than 1 coutour found for image 093_233_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|███████▉  | 2529/3170 [00:13<00:03, 196.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 055_253_B.png\n",
      "More than 1 coutour found for image 032_162_M.png\n",
      "More than 1 coutour found for image 026_167_B.png\n",
      "More than 1 coutour found for image 022_271_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 81%|████████  | 2571/3170 [00:13<00:02, 201.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 026_180_B.png\n",
      "More than 1 coutour found for image 000_263_M.png\n",
      "More than 1 coutour found for image 022_266_M.png\n",
      "More than 1 coutour found for image 093_192_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 82%|████████▏ | 2613/3170 [00:13<00:02, 200.57it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 022_314_M.png\n",
      "More than 1 coutour found for image 059_335_M.png\n",
      "More than 1 coutour found for image 022_263_M.png\n",
      "More than 1 coutour found for image 070_157_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 84%|████████▍ | 2673/3170 [00:14<00:02, 185.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 084_222_B.png\n",
      "More than 1 coutour found for image 070_158_M.png\n",
      "More than 1 coutour found for image 029_159_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▋ | 2735/3170 [00:14<00:02, 195.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 066_241_M.png\n",
      "More than 1 coutour found for image 019_196_M.png\n",
      "More than 1 coutour found for image 046_149_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 88%|████████▊ | 2777/3170 [00:14<00:01, 201.15it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 059_338_M.png\n",
      "More than 1 coutour found for image 045_160_M.png\n",
      "More than 1 coutour found for image 099_295_M.png\n",
      "More than 1 coutour found for image 006_171_M.png\n",
      "More than 1 coutour found for image 022_302_M.png\n",
      "More than 1 coutour found for image 076_224_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 89%|████████▉ | 2820/3170 [00:15<00:01, 198.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 046_153_B.png\n",
      "More than 1 coutour found for image 059_212_M.png\n",
      "More than 1 coutour found for image 059_206_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 91%|█████████ | 2881/3170 [00:15<00:01, 197.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 036_267_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 92%|█████████▏| 2922/3170 [00:15<00:01, 200.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 034_253_M.png\n",
      "More than 1 coutour found for image 059_215_M.png\n",
      "More than 1 coutour found for image 079_139_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 95%|█████████▍| 3008/3170 [00:15<00:00, 204.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 072_218_M.png\n",
      "More than 1 coutour found for image 000_264_M.png\n",
      "More than 1 coutour found for image 012_198_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 96%|█████████▌| 3050/3170 [00:16<00:00, 201.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 077_243_M.png\n",
      "More than 1 coutour found for image 000_261_M.png\n",
      "More than 1 coutour found for image 020_313_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 98%|█████████▊| 3091/3170 [00:16<00:00, 189.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 046_176_B.png\n",
      "More than 1 coutour found for image 022_260_M.png\n",
      "More than 1 coutour found for image 079_137_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 99%|█████████▊| 3130/3170 [00:16<00:00, 180.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 057_198_B.png\n",
      "More than 1 coutour found for image 093_232_M.png\n",
      "More than 1 coutour found for image 019_202_M.png\n",
      "More than 1 coutour found for image 094_244_M.png\n",
      "More than 1 coutour found for image 022_311_M.png\n",
      "More than 1 coutour found for image 012_202_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3170/3170 [00:16<00:00, 188.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "More than 1 coutour found for image 017_212_B.png\n",
      "More than 1 coutour found for image 089_182_M.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import shutil\n",
    "\n",
    "# Create folders\n",
    "os.makedirs(segmentation_data_path, exist_ok=True)\n",
    "os.makedirs(detection_data_path, exist_ok=True)\n",
    "        \n",
    "# List of image cases\n",
    "list_cases = os.listdir(images_folder)\n",
    "\n",
    "# Get classes\n",
    "if not use_classes or any(len(item.split('_')) != 3 for item in list_cases):\n",
    "    use_classes = False\n",
    "    print(\"Classes are not used\")\n",
    "else:\n",
    "    classes = sorted(set([ item[:-4].split('_')[-1] for item in list_cases]))\n",
    "    print(classes)\n",
    "    \n",
    "# For each image\n",
    "for image in tqdm(list_cases):\n",
    "    \n",
    "    # Check image name format\n",
    "    assert len(image[:-4].split(\"_\")) == 3\n",
    "    \n",
    "    # Get class label\n",
    "    if use_classes:\n",
    "        label = image[:-4].split(\"_\")[-1]\n",
    "        label_index = classes.index(label)\n",
    "    else:\n",
    "        label_index = 0\n",
    "    \n",
    "    # Read mask\n",
    "    mask = cv2.imread(os.path.join(masks_folder, image), cv2.IMREAD_GRAYSCALE)\n",
    "    mask_size = mask.shape[1], mask.shape[0]\n",
    "    \n",
    "    # Get contours\n",
    "    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    contours_list = [contour.squeeze().tolist() for contour in contours]\n",
    "    \n",
    "    if len(contours_list) > 1:\n",
    "        print(f\"More than 1 coutour found for image {image}\")\n",
    "\n",
    "    # Copy image files\n",
    "    shutil.copyfile(os.path.join(images_folder, image), os.path.join(detection_data_path, image))\n",
    "    shutil.copyfile(os.path.join(images_folder, image), os.path.join(segmentation_data_path, image))\n",
    "\n",
    "    # Oepn label files\n",
    "    det_fp = open(os.path.join(detection_data_path, f\"{image[:-4]}.txt\"), \"w\")\n",
    "    seg_fp = open(os.path.join(segmentation_data_path, f\"{image[:-4]}.txt\"), \"w\")\n",
    "        \n",
    "    # For each countour\n",
    "    for contour in contours_list:\n",
    "        \n",
    "        # Get lesion countour data\n",
    "        x, y = [point[0] for point in contour], [point[1] for point in contour] #TODO: use numpy\n",
    "        x1, x2, y1, y2 = min(x), max(x), min(y), max(y)\n",
    "        center_x, center_y = (x1+x2)/(2*mask_size[0]), (y1+y2)/(2*mask_size[1])\n",
    "        width, height = (x2-x1)/mask_size[0], (y2-y1)/mask_size[1]\n",
    "        \n",
    "        # Check lesion size\n",
    "        if width < 0.001 or height < 0.001:\n",
    "            print(f\"Small lesion found in image ({image})\")\n",
    "            continue\n",
    "        \n",
    "        # Write labels\n",
    "        det_str = f\"{center_x:0.6f} {center_y:0.6f} {width:0.6f} {height:0.6f}\"\n",
    "        det_fp.write(f\"{label_index} {det_str}\\n\")   \n",
    "        seg_str = \" \".join([f\"{point[0]/mask_size[0]:0.6f} {point[1]/mask_size[1]:0.6f}\" for point in contour])\n",
    "        seg_fp.write(f\"{label_index} {seg_str}\\n\")\n",
    "        \n",
    "    # Close files\n",
    "    det_fp.close()\n",
    "    seg_fp.close()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create YOLO Train/Val split data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0: 100}\n",
      "['064', '040', '081', '077', '055', '093', '054', '033', '083', '018', '051', '068', '036', '084', '065', '030', '058', '000', '004', '037', '002', '060', '046', '078', '088', '057', '052', '034', '008', '003', '095', '011', '020', '023', '049', '059', '076', '098', '028', '086', '048', '006', '045', '026', '001', '073', '017', '056', '024', '032', '085', '019', '094', '005', '069', '071', '063', '035', '042', '043', '029', '047', '013', '079', '041', '027', '015', '012', '021', '070', '050', '097', '096', '092', '072', '016', '091', '082', '044', '038'] 80\n",
      "['022', '061', '066', '080', '039', '014', '009', '075', '099', '074', '007', '031', '089', '053', '062', '025', '067', '090', '010', '087'] 20\n"
     ]
    }
   ],
   "source": [
    "# Get patients list\n",
    "list_items = [item[:-4] for item in os.listdir(detection_data_path) if item.endswith(\".png\")]\n",
    "patients = list(set([ item.split(\"_\")[0] for item in list_items]))\n",
    "\n",
    "# Get lesion type per patient\n",
    "if use_classes:\n",
    "    patients_type = {}     \n",
    "    for item in list_items:\n",
    "        patient = item.split(\"_\")[0] \n",
    "        label = item.split(\"_\")[-1]\n",
    "        if label not in patients_type:\n",
    "            patients_type[label] = [patient]\n",
    "        elif patient not in patients_type[label]:\n",
    "            patients_type[label].append(patient)\n",
    "else:\n",
    "    patients_type = {0: patients}\n",
    "    \n",
    "print({k: len(v) for k, v in patients_type.items()})\n",
    "    \n",
    "# Distribute patients in Train/Val using val_frac with balanced lesion types\n",
    "train_p = []\n",
    "val_p = []\n",
    "for k in patients_type.keys():\n",
    "    num_train = int(len(patients_type[k])*(1-val_frac))\n",
    "    train_type_p, val_type_p = patients_type[k][:num_train], patients_type[k][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 = [ f\"{item}.png\" for item in list_items if item.split(\"_\")[0] in train_p]\n",
    "val = [ f\"{item}.png\" for item in list_items if item.split(\"_\")[0] in val_p]\n",
    "\n",
    "import random\n",
    "random.seed(0)\n",
    "random.shuffle(train)\n",
    "random.shuffle(val)\n",
    "\n",
    "# Generate train.txt file\n",
    "train_det = os.path.join(output_folder, f\"train_det.txt\" ) \n",
    "with open(train_det, \"w\") as fp:\n",
    "    fp.writelines([os.path.join(detection_data_path, t) + '\\n' for t in train])\n",
    "    \n",
    "train_seg = os.path.join(output_folder, f\"train_seg.txt\" ) \n",
    "with open(train_seg, \"w\") as fp:\n",
    "    fp.writelines([os.path.join(segmentation_data_path, t) + '\\n' for t in train])\n",
    "    \n",
    "# Generate val.txt file\n",
    "val_det = os.path.join(output_folder, f\"val_det.txt\")\n",
    "with open(val_det, \"w\") as fp:\n",
    "    fp.writelines([os.path.join(detection_data_path, v) + '\\n' for v in val])\n",
    "    \n",
    "val_seg = os.path.join(output_folder, f\"val_seg.txt\")\n",
    "with open(val_seg, \"w\") as fp:\n",
    "    fp.writelines([os.path.join(segmentation_data_path, 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
}