data_preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing necessary libraries\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import pandas as pd\n",
    "import random\n",
    "import copy\n",
    "import shutil\n",
    "import numpy as np\n",
    "import ipywidgets as widgets\n",
    "import wandb\n",
    "from ultralytics import YOLO\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Login into Weights and Bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WANDB_API_KEY = os.getenv(\"WANDB_API_KEY\")\n",
    "wandb.login(key=WANDB_API_KEY)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the images data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining input images path, output path, masks path and labels path\n",
    "path='dataset'\n",
    "op_path='data'\n",
    "rawimages_path=os.path.join(op_path, 'raw_images')\n",
    "maskimages_path=os.path.join(op_path, 'mask_images')\n",
    "labels_path=os.path.join(op_path, 'labels')\n",
    "# os.makedirs(rawimages_path)\n",
    "# os.makedirs(maskimages_path)\n",
    "# os.makedirs(labels_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining output train,val and test paths\n",
    "imgtrainpath = os.path.join(op_path,'images','train')\n",
    "imgvalpath=os.path.join(op_path,'images','validation')\n",
    "imgtestpath=os.path.join(op_path,'images','test')\n",
    "\n",
    "labeltrainpath=os.path.join(op_path,'labels','train')\n",
    "labelvalpath=os.path.join(op_path,'labels','validation')\n",
    "labeltestpath=os.path.join(op_path,'labels','test')\n",
    "\n",
    "os.makedirs(imgtrainpath)\n",
    "os.makedirs(imgvalpath)\n",
    "os.makedirs(imgtestpath)\n",
    "\n",
    "# os.makedirs(labeltrainpath)\n",
    "# os.makedirs(labelvalpath)\n",
    "# os.makedirs(labeltestpath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we transfer raw images and their corresponding color masks to their relevant output directories. \n",
    "\n",
    "Also we calculate the total number of sub-directories, total images, raw images and mask images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "from collections import Counter\n",
    "\n",
    "def process_images(root_path, rawimages_path, maskimages_path):\n",
    "    stats = Counter()\n",
    "\n",
    "    for directory in os.scandir(root_path):\n",
    "        if not directory.is_dir():\n",
    "            continue\n",
    "        \n",
    "        stats['directories'] += 1\n",
    "        \n",
    "        for sub_dir in os.scandir(directory.path):\n",
    "            if not sub_dir.is_dir():\n",
    "                continue\n",
    "            \n",
    "            stats['sub_directories'] += 1\n",
    "            \n",
    "            for image in os.scandir(sub_dir.path):\n",
    "                if not image.is_file():\n",
    "                    continue\n",
    "                \n",
    "                stats['total_images'] += 1\n",
    "                newname = f\"{sub_dir.name}_{image.name}\"\n",
    "                \n",
    "                if 'mask' not in image.name:\n",
    "                    process_image(image.path, rawimages_path, newname)\n",
    "                    stats['raw_images'] += 1\n",
    "                elif 'color_mask' in image.name:\n",
    "                    process_image(image.path, maskimages_path, newname)\n",
    "                    stats['masks'] += 1\n",
    "\n",
    "    return stats\n",
    "\n",
    "def process_image(src_path, dest_dir, new_name):\n",
    "    try:\n",
    "        dest_path = os.path.join(dest_dir, new_name)\n",
    "        shutil.copy2(src_path, dest_path)\n",
    "    except (IOError, OSError) as e:\n",
    "        print(f\"Error processing {src_path}: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_path = path\n",
    "rawimages_path = rawimages_path\n",
    "maskimages_path = maskimages_path\n",
    "\n",
    "stats = process_images(root_path, rawimages_path, maskimages_path)\n",
    "\n",
    "print(f\"Total directories: {stats['directories']}\")\n",
    "print(f\"Total sub-directories: {stats['sub_directories']}\")\n",
    "print(f\"Total images: {stats['total_images']}\")\n",
    "print(f\"Total raw images: {stats['raw_images']}\")\n",
    "print(f\"Total masks: {stats['masks']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see here that we have equal number of raw images and masks. Let's verify if all the raw images and masks have been copied"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# checking if all raw images and masks have been copied successfully\n",
    "len(os.listdir(rawimages_path)), len(os.listdir(maskimages_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Checking first five raw images\n",
    "os.listdir(rawimages_path)[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Checking first five masks\n",
    "os.listdir(maskimages_path)[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the raw images and masks are not sorted in order. So we need to arrange them in order to visualize them as image-mask pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sorting raw images and masks\n",
    "rawimages_list=sorted(os.listdir(rawimages_path))\n",
    "maskimages_list=sorted(os.listdir(maskimages_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Checking first five sorted raw images\n",
    "rawimages_list[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Checking first five sorted masks\n",
    "maskimages_list[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "raw images and masks are now sorted."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualising images and masks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "define the color-class mapping corresponding to the class labels and the raw images given in the problem description."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, Tuple, NamedTuple\n",
    "\n",
    "class ColorClass(NamedTuple):\n",
    "    color: Tuple[int, int, int]\n",
    "    name: str\n",
    "\n",
    "class SegmentationClasses:\n",
    "    def __init__(self):\n",
    "        # the colour scheme is given based on the documentation\n",
    "        self.class_info: Dict[int, ColorClass] = {\n",
    "            0: ColorClass((127, 127, 127), 'Black Background'),\n",
    "            1: ColorClass((210, 140, 140), 'Abdominal Wall'),\n",
    "            2: ColorClass((255, 114, 114), 'Liver'),\n",
    "            3: ColorClass((231, 70, 156), 'Gastrointestinal Tract'),\n",
    "            4: ColorClass((186, 183, 75), 'Fat'),\n",
    "            5: ColorClass((170, 255, 0), 'Grasper'),\n",
    "            6: ColorClass((255, 85, 0), 'Connective Tissue'),\n",
    "            7: ColorClass((255, 0, 0), 'Blood'),\n",
    "            8: ColorClass((255, 255, 0), 'Cystic Duct'),\n",
    "            9: ColorClass((169, 255, 184), 'L-hook Electrocautery'),\n",
    "            10: ColorClass((255, 160, 165), 'Gallbladder'),\n",
    "            11: ColorClass((0, 50, 128), 'Hepatic Vein'),\n",
    "            12: ColorClass((111, 74, 0), 'Liver Ligament')\n",
    "        }\n",
    "        \n",
    "        self.color_to_class: Dict[Tuple[int, int, int], int] = {\n",
    "            info.color: class_index for class_index, info in self.class_info.items()\n",
    "        }\n",
    "        \n",
    "        self.name_to_class: Dict[str, int] = {\n",
    "            info.name: class_index for class_index, info in self.class_info.items()\n",
    "        }\n",
    "    \n",
    "    def get_class_from_color(self, color: Tuple[int, int, int]) -> int:\n",
    "        return self.color_to_class.get(color, -1)  # Returns -1 if color not found\n",
    "    \n",
    "    def get_color_from_class(self, class_index: int) -> Tuple[int, int, int]:\n",
    "        return self.class_info[class_index].color if class_index in self.class_info else (0, 0, 0)\n",
    "    \n",
    "    def get_name_from_class(self, class_index: int) -> str:\n",
    "        return self.class_info[class_index].name if class_index in self.class_info else \"Unknown\"\n",
    "    \n",
    "    def get_class_from_name(self, name: str) -> int:\n",
    "        return self.name_to_class.get(name, -1)  # Returns -1 if name not found\n",
    "    \n",
    "    def get_color_name(self, class_index: int) -> str:\n",
    "        if class_index in self.class_info:\n",
    "            color = self.class_info[class_index].color\n",
    "            return f\"RGB{color}\"\n",
    "        return \"Unknown\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Usage\n",
    "segmentation_classes = SegmentationClasses()\n",
    "\n",
    "# Examples\n",
    "print(segmentation_classes.get_class_from_color((255, 0, 0)))  # Should print 7 (Blood)\n",
    "print(segmentation_classes.get_color_from_class(2))  # Should print (255, 114, 114) (Liver)\n",
    "print(segmentation_classes.get_name_from_class(5))  # Should print \"Grasper\"\n",
    "print(segmentation_classes.get_class_from_name(\"Liver\"))  # Should print 2\n",
    "print(segmentation_classes.get_color_name(3))  # Should print \"RGB(231, 70, 156)\" (Gastrointestinal Tract)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import numpy as np\n",
    "from typing import Tuple, List\n",
    "\n",
    "def load_image(path: str) -> np.ndarray:\n",
    "    \"\"\"Load and convert an image to RGB.\"\"\"\n",
    "    return cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)\n",
    "\n",
    "def get_corresponding_mask_name(image_name: str) -> str:\n",
    "    \"\"\"Generate the corresponding mask name for a given image name.\"\"\"\n",
    "    return image_name.replace('_endo.png', '_endo_color_mask.png')\n",
    "\n",
    "def plot_image_and_mask(rawimages_path: str, maskimages_path: str, figsize: Tuple[int, int] = (20, 10)) -> None:\n",
    "    \"\"\"\n",
    "    Plot a random pair of raw image and its corresponding mask.\n",
    "    \n",
    "    Args:\n",
    "    rawimages_path (str): Path to the directory containing raw images.\n",
    "    maskimages_path (str): Path to the directory containing mask images.\n",
    "    figsize (Tuple[int, int]): Figure size for the plot. Default is (20, 10).\n",
    "    \"\"\"\n",
    "    rawimages_list = [f for f in os.listdir(rawimages_path) if f.endswith('_endo.png')]\n",
    "    \n",
    "    if not rawimages_list:\n",
    "        print(\"No images found in the raw images directory.\")\n",
    "        return\n",
    "    \n",
    "    img_name = random.choice(rawimages_list)\n",
    "    mask_name = get_corresponding_mask_name(img_name)\n",
    "    \n",
    "    img_path = os.path.join(rawimages_path, img_name)\n",
    "    mask_path = os.path.join(maskimages_path, mask_name)\n",
    "    \n",
    "    if not os.path.exists(mask_path):\n",
    "        print(f\"Corresponding mask not found for {img_name}\")\n",
    "        return\n",
    "    \n",
    "    fig, axes = plt.subplots(1, 2, figsize=figsize)\n",
    "    fig.suptitle(\"Raw Image vs Mask\", fontsize=16)\n",
    "    \n",
    "    for ax, path, title in zip(axes, [img_path, mask_path], [img_name, mask_name]):\n",
    "        image = load_image(path)\n",
    "        ax.imshow(image)\n",
    "        ax.set_title(title, fontsize=14)\n",
    "        ax.axis('off')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "def plot_multiple_pairs(rawimages_path: str, maskimages_path: str, num_pairs: int = 3, figsize: Tuple[int, int] = (20, 20)) -> None:\n",
    "    \"\"\"\n",
    "    Plot multiple random pairs of raw images and their corresponding masks.\n",
    "    \n",
    "    Args:\n",
    "    rawimages_path (str): Path to the directory containing raw images.\n",
    "    maskimages_path (str): Path to the directory containing mask images.\n",
    "    num_pairs (int): Number of image-mask pairs to plot. Default is 3.\n",
    "    figsize (Tuple[int, int]): Figure size for the plot. Default is (20, 20).\n",
    "    \"\"\"\n",
    "    rawimages_list = [f for f in os.listdir(rawimages_path) if f.endswith('_endo.png')]\n",
    "    \n",
    "    if not rawimages_list:\n",
    "        print(\"No images found in the raw images directory.\")\n",
    "        return\n",
    "    \n",
    "    fig, axes = plt.subplots(num_pairs, 2, figsize=figsize)\n",
    "    fig.suptitle(f\"{num_pairs} Random Image-Mask Pairs\", fontsize=16)\n",
    "    \n",
    "    for row in axes:\n",
    "        img_name = random.choice(rawimages_list)\n",
    "        mask_name = get_corresponding_mask_name(img_name)\n",
    "        \n",
    "        img_path = os.path.join(rawimages_path, img_name)\n",
    "        mask_path = os.path.join(maskimages_path, mask_name)\n",
    "        \n",
    "        if not os.path.exists(mask_path):\n",
    "            print(f\"Corresponding mask not found for {img_name}\")\n",
    "            continue\n",
    "        \n",
    "        for ax, path in zip(row, [img_path, mask_path]):\n",
    "            image = load_image(path)\n",
    "            ax.imshow(image)\n",
    "            ax.set_title(os.path.basename(path), fontsize=10)\n",
    "            ax.axis('off')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage:\n",
    "plot_image_and_mask(rawimages_path, maskimages_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_multiple_pairs(rawimages_path, maskimages_path, num_pairs=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Masks visualisation with contours"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will find unique colors in mask and draw segmentation contours for any color present in it. Since we don't have pure black color (RGB-0,0,0) for any of the classes in masks, we will use it to draw contours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_contour_for_one_color_on_mask(\n",
    "    rawimages_path: str,\n",
    "    rawimages_list: List[str],\n",
    "    maskimages_path: str,\n",
    "    maskimages_list: List[str]\n",
    ") -> None:\n",
    "    segmentation_classes = SegmentationClasses()\n",
    "    \n",
    "    # Choose a random image\n",
    "    k = random.randint(0, len(rawimages_list) - 1)\n",
    "    \n",
    "    # Load images\n",
    "    img_path = os.path.join(rawimages_path, rawimages_list[k])\n",
    "    mask_path = os.path.join(maskimages_path, maskimages_list[k])\n",
    "    \n",
    "    img = cv2.imread(img_path)\n",
    "    mask = cv2.imread(mask_path)\n",
    "    \n",
    "    if img is None or mask is None:\n",
    "        raise ValueError(\"Failed to load image or mask\")\n",
    "    \n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    # Get unique colors in mask\n",
    "    unique_colors = np.unique(mask.reshape(-1, 3), axis=0)\n",
    "    defined_colors = [color for color in unique_colors if tuple(color) in segmentation_classes.color_to_class]\n",
    "    \n",
    "    if not defined_colors:\n",
    "        raise ValueError(\"No defined colors found in the mask\")\n",
    "    \n",
    "    # Select a random color\n",
    "    color = random.choice(defined_colors)\n",
    "    \n",
    "    # Get class information\n",
    "    class_index = segmentation_classes.get_class_from_color(tuple(color))\n",
    "    class_name = segmentation_classes.get_name_from_class(class_index)\n",
    "    \n",
    "    # Create mask for the selected color and find contours\n",
    "    color_mask = cv2.inRange(mask, color, color)\n",
    "    contours, _ = cv2.findContours(color_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    \n",
    "    # Draw contours on a copy of the mask\n",
    "    mask_with_contour = mask.copy()\n",
    "    cv2.drawContours(mask_with_contour, contours, -1, (0, 0, 0), 4)\n",
    "    \n",
    "    # Plotting\n",
    "    fig, axes = plt.subplots(1, 3, figsize=(30, 10))\n",
    "    images = [img, mask, mask_with_contour]\n",
    "    titles = [\n",
    "        os.path.basename(img_path),\n",
    "        'Mask',\n",
    "        f'Mask with contour on {class_name}'\n",
    "    ]\n",
    "    \n",
    "    for ax, image, title in zip(axes, images, titles):\n",
    "        ax.imshow(image)\n",
    "        ax.set_title(title, fontsize=20)\n",
    "        ax.axis('off')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_contour_for_one_color_on_mask(rawimages_path, rawimages_list, maskimages_path, maskimages_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing text file for masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "def process_mask_images(maskimages_list: List[str], maskimages_path: str, labels_path: str, segmentation_classes: SegmentationClasses) -> int:\n",
    "    total_labels = 0\n",
    "\n",
    "    for img in tqdm(maskimages_list, desc=\"Processing mask images\"):\n",
    "        # Extract shortened mask name\n",
    "        parts = img.split('_')\n",
    "        endo_index = parts.index('endo')\n",
    "        newname = '_'.join(parts[:endo_index+1])\n",
    "        \n",
    "        # Read the image\n",
    "        image_path = os.path.join(maskimages_path, img)\n",
    "        image = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)\n",
    "        \n",
    "        # Get unique colors present in mask\n",
    "        unique_colors = np.unique(image.reshape(-1, 3), axis=0)\n",
    "        \n",
    "        # Filter colors defined in segmentation_classes\n",
    "        unique_colors_defined = [color for color in unique_colors if tuple(color) in segmentation_classes.color_to_class]\n",
    "        \n",
    "        H, W, _ = image.shape\n",
    "        class_contour_mapping = {}\n",
    "\n",
    "        for color in unique_colors_defined:\n",
    "            class_code = segmentation_classes.get_class_from_color(tuple(color))\n",
    "            mask = cv2.inRange(image, color, color)\n",
    "            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)\n",
    "            class_contour_mapping[class_code] = contours\n",
    "            \n",
    "        # Write label text file    \n",
    "        write_polygon_file(class_contour_mapping, H, W, labels_path, newname)\n",
    "        total_labels += 1\n",
    "    \n",
    "    return total_labels\n",
    "\n",
    "def write_polygon_file(class_contour_mapping: Dict[int, List], H: int, W: int, output_path: str, img_name: str):\n",
    "    coordinates = {}\n",
    "    min_contour_area = 20\n",
    "\n",
    "    for class_code, contours in class_contour_mapping.items():\n",
    "        polygons = []\n",
    "        for cnt in contours:\n",
    "            if cv2.contourArea(cnt) > min_contour_area:\n",
    "                polygon = np.round(cnt.squeeze() / [W, H], 4).flatten().tolist()\n",
    "                polygons.append(polygon)\n",
    "        if polygons:\n",
    "            coordinates[class_code] = polygons\n",
    "\n",
    "    output_file = os.path.join(output_path, f\"{img_name}.txt\")\n",
    "    with open(output_file, 'w') as f:\n",
    "        for class_code, polygons in coordinates.items():\n",
    "            for polygon in polygons:\n",
    "                polygon_str = ' '.join(map(str, polygon))\n",
    "                f.write(f\"{class_code} {polygon_str}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_labels = process_mask_images(maskimages_list, maskimages_path, labels_path, segmentation_classes)\n",
    "print(f\"Total number of labels created: {total_labels}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split the data into Training, Validation and Test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Tuple\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "def create_dataset(images_list: List[str], train_size: float = 0.8, val_size: float = 0.1, \n",
    "                   test_size: float = 0.1, random_state: int = 42) -> Tuple[List[str], List[str], List[str]]:\n",
    "    \"\"\"\n",
    "    Create training, validation, and test datasets from a list of images.\n",
    "    \n",
    "    Args:\n",
    "    images_list (List[str]): List of image file names.\n",
    "    train_size (float): Proportion of the dataset to include in the train split.\n",
    "    val_size (float): Proportion of the dataset to include in the validation split.\n",
    "    test_size (float): Proportion of the dataset to include in the test split.\n",
    "    random_state (int): Controls the shuffling applied to the data before applying the split.\n",
    "    \n",
    "    Returns:\n",
    "    Tuple[List[str], List[str], List[str]]: Train, validation, and test image lists.\n",
    "    \"\"\"\n",
    "    if not np.isclose(train_size + val_size + test_size, 1.0):\n",
    "        raise ValueError(\"The sum of train_size, val_size, and test_size should be 1.0\")\n",
    "    \n",
    "    # First, split off the test set\n",
    "    train_val, test_images = train_test_split(images_list, test_size=test_size, random_state=random_state)\n",
    "    \n",
    "    # Then split the remaining data into train and validation sets\n",
    "    relative_val_size = val_size / (train_size + val_size)\n",
    "    train_images, val_images = train_test_split(train_val, test_size=relative_val_size, random_state=random_state)\n",
    "    \n",
    "    return train_images, val_images, test_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_images, val_images, test_images = create_dataset(rawimages_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print dataset sizes\n",
    "print(f\"Training set size: {len(train_images)}\")\n",
    "print(f\"Validation set size: {len(val_images)}\")\n",
    "print(f\"Test set size: {len(test_images)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Optional: Verify the proportions\n",
    "total = len(rawimages_list)\n",
    "print(f\"Training set proportion: {len(train_images)/total:.2f}\")\n",
    "print(f\"Validation set proportion: {len(val_images)/total:.2f}\")\n",
    "print(f\"Test set proportion: {len(test_images)/total:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the dataset is bifurcated, we will write a function to create names of label files corresponding to the names of image files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_label_filenames(image_files: List[str], label_extension: str = '.txt') -> List[str]:\n",
    "    \"\"\"\n",
    "    Create label filenames from image filenames by changing the extension.\n",
    "    \n",
    "    Args:\n",
    "    image_files (List[str]): List of image file names.\n",
    "    label_extension (str): The extension to use for label files. Default is '.txt'.\n",
    "    \n",
    "    Returns:\n",
    "    List[str]: List of label file names.\n",
    "    \"\"\"\n",
    "    return [os.path.splitext(file)[0] + label_extension for file in image_files]\n",
    "\n",
    "def create_dataset_labels(train_images: List[str], val_images: List[str], test_images: List[str], \n",
    "                          label_extension: str = '.txt') -> Tuple[List[str], List[str], List[str]]:\n",
    "    \"\"\"\n",
    "    Create label filenames for train, validation, and test datasets.\n",
    "    \n",
    "    Args:\n",
    "    train_images (List[str]): List of training image file names.\n",
    "    val_images (List[str]): List of validation image file names.\n",
    "    test_images (List[str]): List of test image file names.\n",
    "    label_extension (str): The extension to use for label files. Default is '.txt'.\n",
    "    \n",
    "    Returns:\n",
    "    Tuple[List[str], List[str], List[str]]: Train, validation, and test label file lists.\n",
    "    \"\"\"\n",
    "    train_labels = create_label_filenames(train_images, label_extension)\n",
    "    val_labels = create_label_filenames(val_images, label_extension)\n",
    "    test_labels = create_label_filenames(test_images, label_extension)\n",
    "    \n",
    "    return train_labels, val_labels, test_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_extension = '.txt'\n",
    "train_labels, val_labels, test_labels = create_dataset_labels(train_images, val_images, test_images, label_extension)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print dataset sizes\n",
    "print(f\"Training set size: {len(train_labels)}\")\n",
    "print(f\"Validation set size: {len(val_labels)}\")\n",
    "print(f\"Test set size: {len(test_labels)}\")\n",
    "\n",
    "# Optional: Verify that the number of labels matches the number of images\n",
    "assert len(train_labels) == len(train_images), \"Mismatch in training set size\"\n",
    "assert len(val_labels) == len(val_images), \"Mismatch in validation set size\"\n",
    "assert len(test_labels) == len(test_images), \"Mismatch in test set size\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "move images and labels to their respective directories\n",
    "\n",
    "resize the images to a fixed size so that all images have exactly same size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import cv2\n",
    "from typing import List, Tuple\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "from tqdm import tqdm\n",
    "\n",
    "def resize_and_move_image(args: Tuple[str, str, str, int]) -> None:\n",
    "    \"\"\"\n",
    "    Resize an image and move it to the destination path.\n",
    "    \n",
    "    Args:\n",
    "    args (Tuple[str, str, str, int]): (file, source_path, destination_path, image_size)\n",
    "    \"\"\"\n",
    "    file, source_path, destination_path, image_size = args\n",
    "    filepath = os.path.join(source_path, file)\n",
    "    finalimage_path = os.path.join(destination_path, file)\n",
    "    img_resized = cv2.resize(cv2.imread(filepath), (image_size, image_size))\n",
    "    cv2.imwrite(finalimage_path, img_resized)\n",
    "\n",
    "def move_file(args: Tuple[str, str, str]) -> None:\n",
    "    \"\"\"\n",
    "    Move a file from source to destination.\n",
    "    \n",
    "    Args:\n",
    "    args (Tuple[str, str, str]): (file, source_path, destination_path)\n",
    "    \"\"\"\n",
    "    file, source_path, destination_path = args\n",
    "    filepath = os.path.join(source_path, file)\n",
    "    shutil.move(filepath, destination_path)\n",
    "\n",
    "def process_dataset(image_list: List[str], label_list: List[str], \n",
    "                    image_source_path: str, label_source_path: str,\n",
    "                    image_dest_path: str, label_dest_path: str,\n",
    "                    image_size: int = 640, max_workers: int = 4) -> None:\n",
    "    \"\"\"\n",
    "    Process a dataset by resizing and moving images, and moving labels.\n",
    "    \n",
    "    Args:\n",
    "    image_list (List[str]): List of image file names.\n",
    "    label_list (List[str]): List of label file names.\n",
    "    image_source_path (str): Source path for images.\n",
    "    label_source_path (str): Source path for labels.\n",
    "    image_dest_path (str): Destination path for images.\n",
    "    label_dest_path (str): Destination path for labels.\n",
    "    image_size (int): Size to resize images to. Default is 640.\n",
    "    max_workers (int): Maximum number of worker threads. Default is 4.\n",
    "    \"\"\"\n",
    "    os.makedirs(image_dest_path, exist_ok=True)\n",
    "    os.makedirs(label_dest_path, exist_ok=True)\n",
    "\n",
    "    with ThreadPoolExecutor(max_workers=max_workers) as executor:\n",
    "        # Process images\n",
    "        image_args = [(file, image_source_path, image_dest_path, image_size) for file in image_list]\n",
    "        list(tqdm(executor.map(resize_and_move_image, image_args), total=len(image_list), desc=\"Processing images\"))\n",
    "\n",
    "        # Process labels\n",
    "        label_args = [(file, label_source_path, label_dest_path) for file in label_list]\n",
    "        list(tqdm(executor.map(move_file, label_args), total=len(label_list), desc=\"Processing labels\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_size = 640\n",
    "max_workers = 4  # Adjust based on your system's capabilities\n",
    "\n",
    "# Process training set\n",
    "process_dataset(train_images, train_labels, rawimages_path, labels_path, \n",
    "                imgtrainpath, labeltrainpath, image_size, max_workers)\n",
    "\n",
    "# Process validation set\n",
    "process_dataset(val_images, val_labels, rawimages_path, labels_path, \n",
    "                imgvalpath, labelvalpath, image_size, max_workers)\n",
    "\n",
    "# Process test set\n",
    "process_dataset(test_images, test_labels, rawimages_path, labels_path, \n",
    "                imgtestpath, labeltestpath, image_size, max_workers)\n",
    "\n",
    "print(\"All datasets processed successfully.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have the list of test masks handy as it will be required later for visualisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import os\n",
    "\n",
    "def create_mask_list(image_list: List[str], mask_suffix: str = '_color_mask.png') -> List[str]:\n",
    "    \"\"\"\n",
    "    Create a list of mask filenames from a list of image filenames.\n",
    "    \n",
    "    Args:\n",
    "    image_list (List[str]): List of image filenames.\n",
    "    mask_suffix (str): Suffix to append to the base filename to create the mask filename.\n",
    "                       Default is '_color_mask.png'.\n",
    "    \n",
    "    Returns:\n",
    "    List[str]: List of corresponding mask filenames.\n",
    "    \"\"\"\n",
    "    return [os.path.splitext(image)[0] + mask_suffix for image in image_list]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Usage\n",
    "mask_suffix = '_color_mask.png'\n",
    "test_masks = create_mask_list(test_images, mask_suffix)\n",
    "\n",
    "# Print some information\n",
    "print(f\"Number of test images: {len(test_images)}\")\n",
    "print(f\"Number of test masks: {len(test_masks)}\")\n",
    "\n",
    "# Optional: Print a few examples to verify\n",
    "print(\"\\nExample mappings (image -> mask):\")\n",
    "for i in range(min(5, len(test_images))):\n",
    "    print(f\"{test_images[i]} -> {test_masks[i]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Upload the data into Google Cloud Storage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "setup the GCP credential using service account"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Authenticate with GCS\n",
    "from google.cloud import storage\n",
    "\n",
    "from google.oauth2 import service_account\n",
    "\n",
    "credentials = service_account.Credentials.from_service_account_file('gcp-credential.json')\n",
    "\n",
    "# Use the credentials to create a GCS client\n",
    "client = storage.Client(credentials=credentials)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify the client\n",
    "buckets = client.list_buckets()\n",
    "\n",
    "print(\"Buckets:\")\n",
    "for bucket in buckets:\n",
    "    print(bucket.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def upload_folder(bucket_name, local_folder, remote_folder):\n",
    "    \"\"\"Uploads a local folder and its subfolders to a GCS bucket.\n",
    "\n",
    "    Args:\n",
    "        bucket_name (str): The name of the GCS bucket.\n",
    "        local_folder (str): The path to the local folder to upload.\n",
    "        remote_folder (str): The desired prefix for the uploaded files in GCS.\n",
    "    \"\"\"\n",
    "\n",
    "    bucket = client.get_bucket(bucket_name)\n",
    "\n",
    "    for root, dirs, files in os.walk(local_folder):\n",
    "        for file in files:\n",
    "            local_path = os.path.join(root, file)\n",
    "            remote_path = os.path.join(remote_folder, root[len(local_folder) + 1:], file)\n",
    "            blob = bucket.blob(remote_path)\n",
    "            blob.upload_from_filename(local_path)\n",
    "\n",
    "            print(f\"Uploaded {local_path} to {blob.public_url}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket_name = 'endoinsight-ai'\n",
    "local_folder = 'data/labels'\n",
    "remote_folder = 'data/labels'  # Prefix for uploaded files in GCS\n",
    "\n",
    "upload_folder(bucket_name, local_folder, remote_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "folders = ['data/labels', 'data/images']\n",
    "\n",
    "for folder in folders:\n",
    "    file_count = sum(1 for _ in bucket.list_blobs(prefix=folder))\n",
    "    print(f\"Number of files in the {folder} folder: {file_count}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Config File for YOLOv8-Seg Model Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The config file is required to use YOLOv8 model. The names of classes present in dataset and the directories for the training, validation and test datasets are indicated in the config file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining newline variable for config file\n",
    "newline='\\n'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the config file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining lines of config text file\n",
    "\n",
    "ln_1='# Train/val/test sets'+newline # starting with a comment line\n",
    "\n",
    "# train, val and test path declaration\n",
    "ln_2='train: ' +\"'\"+imgtrainpath+\"'\"+newline\n",
    "ln_3='val: ' +\"'\" + imgvalpath+\"'\"+newline\n",
    "ln_4='test: ' +\"'\" + imgtestpath+\"'\"+newline\n",
    "ln_5=newline\n",
    "\n",
    "# names of the classes declaration\n",
    "ln_6='# Classes'+newline\n",
    "ln_7='names:'+newline\n",
    "ln_8='  0: Black Background'+newline\n",
    "ln_9='  1: Abdominal Wall'+newline\n",
    "ln_10='  2: Liver'+newline\n",
    "ln_11='  3: Gastrointestinal Tract'+newline\n",
    "ln_12='  4: Fat'+newline\n",
    "ln_13='  5: Grasper'+newline\n",
    "ln_14='  6: Connective Tissue'+newline\n",
    "ln_15='  7: Blood'+newline\n",
    "ln_16='  8: Cystic Duct'+newline\n",
    "ln_17='  9: L-hook Electrocautery'+newline\n",
    "ln_18='  10: Gallbladder'+newline\n",
    "ln_19='  11: Hepatic Vein'+newline\n",
    "ln_20='  12: Liver Ligament'\n",
    "\n",
    "#listing all config lines\n",
    "config_lines=[ln_1, ln_2, ln_3, ln_4, ln_5, ln_6, ln_7, ln_8, ln_9, ln_10, ln_11, ln_12,\n",
    "             ln_13, ln_14, ln_15, ln_16, ln_17, ln_18, ln_19, ln_20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating path for config file\n",
    "config_path=os.path.join(op_path, 'config.yaml')\n",
    "config_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Writing config file\n",
    "with open(config_path, 'w') as f:\n",
    "    f.writelines(config_lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the config file\n",
    "with open(config_path, 'r') as f:\n",
    "    config_file = f.read()\n",
    "\n",
    "print(config_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "now we are ready to do model training"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}