{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b5d65232",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import ast\n",
    "import SimpleITK as sitk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5f398088",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = \"Data/MAMA-MIA\"\n",
    "data_dir= os.path.join(root_dir, \"images\")\n",
    "output_dir = os.path.join(root_dir, \"4D_images\")\n",
    "segmentation_dir = os.path.join(root_dir, \"segmentations\", \"expert\")\n",
    "clinical_data = pd.read_excel(os.path.join(root_dir, \"clinical_and_imaging_info.xlsx\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bdf2e5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dict = {}\n",
    "for sub in os.listdir(data_dir):\n",
    "    if os.path.isdir(os.path.join(data_dir, sub)):\n",
    "        data_dict[sub] = {\"images\": [], \"timepoints\": []}\n",
    "        for img in os.listdir(os.path.join(data_dir, sub)):\n",
    "            if img.endswith(\".nii.gz\"):\n",
    "                data_dict[sub][\"images\"].append(os.path.join(data_dir, sub, img))\n",
    "                \n",
    "        data_dict[sub][\"images\"].sort()\n",
    "        \n",
    "        acq_times = clinical_data[clinical_data[\"patient_id\"] == sub][\"acquisition_times\"].values[0]\n",
    "        if isinstance(acq_times, str):\n",
    "            data_dict[sub][\"timepoints\"] = ast.literal_eval(acq_times)\n",
    "        else:\n",
    "            data_dict[sub][\"timepoints\"] = acq_times\n",
    "        data_dict[sub][\"segmentation\"] = os.path.join(segmentation_dir, f\"{sub}.nii.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a64d3d85",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sub in data_dict.keys():\n",
    "    if not os.path.exists(os.path.join(output_dir, sub, f\"{sub}.nii.gz\")):\n",
    "        image_list = []\n",
    "        for img in data_dict[sub][\"images\"]:\n",
    "            sitk_img = sitk.ReadImage(img)\n",
    "            image_list.append(sitk_img)\n",
    "        sitk_4D_image = sitk.JoinSeries(image_list)\n",
    "        os.makedirs(os.path.join(output_dir, sub), exist_ok=True)\n",
    "        sitk.WriteImage(sitk_4D_image, os.path.join(output_dir, sub, f\"{sub}.nii.gz\"))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}