Tooth_Segmentation_Corrected.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3a4ad7fd",
   "metadata": {},
   "source": [
    "# Tooth Segmentation with 3D U-Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9bfa4f05",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Import necessary libraries\n",
    "import os\n",
    "import numpy as np\n",
    "import pydicom\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import layers, models\n",
    "from sklearn.model_selection import train_test_split\n",
    "import time\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0793fbd",
   "metadata": {},
   "source": [
    "### Step 1: Load and Normalize DICOM Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a9236917",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Function to load DICOM images and stack them into a 3D volume\n",
    "def load_dicom_images(dicom_folder_path):\n",
    "    dicom_files = [os.path.join(dicom_folder_path, f) for f in os.listdir(dicom_folder_path) if f.endswith('.dcm')]\n",
    "    dicom_files = sorted(dicom_files)  # Ensure consistent order\n",
    "    images = [pydicom.dcmread(file).pixel_array for file in dicom_files]\n",
    "    volume = np.stack(images, axis=-1)\n",
    "    volume = np.array(volume, dtype=np.float32)\n",
    "    return volume\n",
    "\n",
    "# Function to normalize the DICOM volumes\n",
    "def normalize(volume, min_bound=-1000, max_bound=400):\n",
    "    volume = (volume - min_bound) / (max_bound - min_bound)\n",
    "    volume = np.clip(volume, 0, 1)\n",
    "    return volume\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0646757e",
   "metadata": {},
   "source": [
    "### Step 2: Load and Prepare Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "54a46004",
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[WinError 3] The system cannot find the path specified: 'D:/DIcom gans/Data/labels'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[6], line 37\u001b[0m\n\u001b[0;32m     34\u001b[0m load_and_append_dicom(input_folder_path, input_volumes, is_normalized\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m     36\u001b[0m \u001b[38;5;66;03m# Load label volumes (assumed not to be normalized)\u001b[39;00m\n\u001b[1;32m---> 37\u001b[0m \u001b[43mload_and_append_dicom\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlabel_folder_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget_volumes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mis_normalized\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[0;32m     39\u001b[0m \u001b[38;5;66;03m# Convert lists to numpy arrays for model compatibility\u001b[39;00m\n\u001b[0;32m     40\u001b[0m input_volumes \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(input_volumes)\n",
      "Cell \u001b[1;32mIn[6], line 17\u001b[0m, in \u001b[0;36mload_and_append_dicom\u001b[1;34m(folder_path, volume_list, is_normalized)\u001b[0m\n\u001b[0;32m     13\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mload_and_append_dicom\u001b[39m(folder_path, volume_list, is_normalized\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\n\u001b[0;32m     14\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m     15\u001b[0m \u001b[38;5;124;03m    Loads DICOM images from a folder or single directory and appends normalized volume to list.\u001b[39;00m\n\u001b[0;32m     16\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m---> 17\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(f\u001b[38;5;241m.\u001b[39mendswith(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m.dcm\u001b[39m\u001b[38;5;124m'\u001b[39m) \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m \u001b[43mos\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlistdir\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfolder_path\u001b[49m\u001b[43m)\u001b[49m):\n\u001b[0;32m     18\u001b[0m         \u001b[38;5;66;03m# If .dcm files are directly inside the folder\u001b[39;00m\n\u001b[0;32m     19\u001b[0m         dicom_volume \u001b[38;5;241m=\u001b[39m load_dicom_images(folder_path)\n\u001b[0;32m     20\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m is_normalized:\n",
      "\u001b[1;31mFileNotFoundError\u001b[0m: [WinError 3] The system cannot find the path specified: 'D:/DIcom gans/Data/labels'"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "# Paths to input and label data\n",
    "input_folder_path = \"D:/DIcom gans/Data/raw\"  # Path to folder containing raw DICOM input volumes\n",
    "label_folder_path = \"D:/DIcom gans/Data/labels\"  # Path to folder containing ground truth labels\n",
    "\n",
    "# Initialize lists to store input and label volumes\n",
    "input_volumes = []\n",
    "target_volumes = []\n",
    "\n",
    "# Load and normalize input volumes\n",
    "def load_and_append_dicom(folder_path, volume_list, is_normalized=True):\n",
    "    \"\"\"\n",
    "    Loads DICOM images from a folder or single directory and appends normalized volume to list.\n",
    "    \"\"\"\n",
    "    if any(f.endswith('.dcm') for f in os.listdir(folder_path)):\n",
    "        # If .dcm files are directly inside the folder\n",
    "        dicom_volume = load_dicom_images(folder_path)\n",
    "        if is_normalized:\n",
    "            dicom_volume = normalize(dicom_volume)\n",
    "        volume_list.append(dicom_volume)\n",
    "    else:\n",
    "        # If subfolders are present, process each subfolder separately\n",
    "        for subfolder in os.listdir(folder_path):\n",
    "            subfolder_path = os.path.join(folder_path, subfolder)\n",
    "            if os.path.isdir(subfolder_path):\n",
    "                dicom_volume = load_dicom_images(subfolder_path)\n",
    "                if is_normalized:\n",
    "                    dicom_volume = normalize(dicom_volume)\n",
    "                volume_list.append(dicom_volume)\n",
    "\n",
    "# Load input volumes\n",
    "load_and_append_dicom(input_folder_path, input_volumes, is_normalized=True)\n",
    "\n",
    "# Load label volumes (assumed not to be normalized)\n",
    "load_and_append_dicom(label_folder_path, target_volumes, is_normalized=False)\n",
    "\n",
    "# Convert lists to numpy arrays for model compatibility\n",
    "input_volumes = np.array(input_volumes)\n",
    "target_volumes = np.array(target_volumes)\n",
    "\n",
    "# Expand dimensions to add a channel axis (required for U-Net input)\n",
    "input_volumes = np.expand_dims(input_volumes, axis=-1)\n",
    "target_volumes = np.expand_dims(target_volumes, axis=-1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "375fcfbf",
   "metadata": {},
   "source": [
    "### Step 3: Split Data into Training and Validation Sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "948dc806",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Split data into training and validation sets\n",
    "train_X, val_X, train_y, val_y = train_test_split(input_volumes, target_volumes, test_size=0.2, random_state=42)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d9d84a9",
   "metadata": {},
   "source": [
    "### Step 4: Define U-Net Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a515b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Define U-Net model for 3D segmentation\n",
    "def unet_model(input_shape=(128, 128, 64, 1)):\n",
    "    inputs = layers.Input(input_shape)\n",
    "    \n",
    "    # Encoder\n",
    "    c1 = layers.Conv3D(32, 3, activation=\"relu\", padding=\"same\")(inputs)\n",
    "    c1 = layers.BatchNormalization()(c1)\n",
    "    c1 = layers.Conv3D(32, 3, activation=\"relu\", padding=\"same\")(c1)\n",
    "    c1 = layers.BatchNormalization()(c1)\n",
    "    p1 = layers.MaxPooling3D((2, 2, 2))(c1)\n",
    "\n",
    "    c2 = layers.Conv3D(64, 3, activation=\"relu\", padding=\"same\")(p1)\n",
    "    c2 = layers.BatchNormalization()(c2)\n",
    "    c2 = layers.Conv3D(64, 3, activation=\"relu\", padding=\"same\")(c2)\n",
    "    c2 = layers.BatchNormalization()(c2)\n",
    "    p2 = layers.MaxPooling3D((2, 2, 2))(c2)\n",
    "    p2 = layers.Dropout(0.3)(p2)\n",
    "\n",
    "    c3 = layers.Conv3D(128, 3, activation=\"relu\", padding=\"same\")(p2)\n",
    "    c3 = layers.BatchNormalization()(c3)\n",
    "    c3 = layers.Conv3D(128, 3, activation=\"relu\", padding=\"same\")(c3)\n",
    "    c3 = layers.BatchNormalization()(c3)\n",
    "\n",
    "    # Decoder\n",
    "    u1 = layers.Conv3DTranspose(64, 3, strides=(2, 2, 2), padding=\"same\")(c3)\n",
    "    u1 = layers.concatenate([u1, c2])\n",
    "    c4 = layers.Conv3D(64, 3, activation=\"relu\", padding=\"same\")(u1)\n",
    "    c4 = layers.BatchNormalization()(c4)\n",
    "    c4 = layers.Conv3D(64, 3, activation=\"relu\", padding=\"same\")(c4)\n",
    "    c4 = layers.BatchNormalization()(c4)\n",
    "\n",
    "    u2 = layers.Conv3DTranspose(32, 3, strides=(2, 2, 2), padding=\"same\")(c4)\n",
    "    u2 = layers.concatenate([u2, c1])\n",
    "    c5 = layers.Conv3D(32, 3, activation=\"relu\", padding=\"same\")(u2)\n",
    "    c5 = layers.BatchNormalization()(c5)\n",
    "    c5 = layers.Conv3D(32, 3, activation=\"relu\", padding=\"same\")(c5)\n",
    "    c5 = layers.BatchNormalization()(c5)\n",
    "\n",
    "    # Output Layer\n",
    "    outputs = layers.Conv3D(1, 1, activation=\"sigmoid\")(c5)\n",
    "\n",
    "    model = models.Model(inputs, outputs)\n",
    "    return model\n",
    "\n",
    "# Instantiate the model\n",
    "model = unet_model(input_shape=train_X.shape[1:])\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22a8dfba",
   "metadata": {},
   "source": [
    "### Step 5: Compile the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17b02e69",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Compile the model\n",
    "model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),\n",
    "              loss='binary_crossentropy',\n",
    "              metrics=['accuracy'])\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca1f0fa",
   "metadata": {},
   "source": [
    "### Step 6: Train the Model with Time Estimation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdbe005a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Callbacks\n",
    "checkpoint_cb = tf.keras.callbacks.ModelCheckpoint(\"best_model.h5\", save_best_only=True)\n",
    "early_stopping_cb = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)\n",
    "\n",
    "# Measure time for the first epoch to estimate total time\n",
    "start_time = time.time()\n",
    "\n",
    "# Train the model\n",
    "history = model.fit(train_X, train_y,\n",
    "                    validation_data=(val_X, val_y),\n",
    "                    epochs=50,\n",
    "                    batch_size=1,\n",
    "                    callbacks=[checkpoint_cb, early_stopping_cb],\n",
    "                    verbose=1)\n",
    "\n",
    "# Calculate time estimation for total training\n",
    "elapsed_time = time.time() - start_time\n",
    "print(f\"Time for one epoch: {elapsed_time:.2f} seconds\")\n",
    "estimated_total_time = elapsed_time * 50  # Adjust based on number of epochs\n",
    "print(f\"Estimated total training time: {estimated_total_time / 60:.2f} minutes\")\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}