From fa20a4ce6e739f42ee9bb9bf8357df66df9b2c0a Mon Sep 17 00:00:00 2001 From: Sajjad-scbas <106543263+Sajjad-scbas@users.noreply.github.com> Date: Mon, 27 Oct 2025 19:08:13 +0100 Subject: [PATCH 1/2] fix google colab --- predict_zaugnet.py | 2 +- predict_zaugnet_plus.py | 2 +- zaugnet_colab.ipynb | 118 ++++++++++------------------------------ 3 files changed, 31 insertions(+), 91 deletions(-) diff --git a/predict_zaugnet.py b/predict_zaugnet.py index 8160b4d..b614a08 100644 --- a/predict_zaugnet.py +++ b/predict_zaugnet.py @@ -21,7 +21,7 @@ def load_model(cfg, dataset, model_name): zaug = ZAugGenerator(cfg) zaug.set_multiple_gpus() - zaug.load_state_dict(torch.load(path_model)) + zaug.load_state_dict(torch.load(path_model, map_location=f"cuda:{cfg.device_ids[0]}")) zaug.eval() return zaug diff --git a/predict_zaugnet_plus.py b/predict_zaugnet_plus.py index d2a9c6c..a312a87 100644 --- a/predict_zaugnet_plus.py +++ b/predict_zaugnet_plus.py @@ -19,7 +19,7 @@ def load_model(cfg, dataset, model_name): print(f"The model used for this prediction : {path_model}") zaug = ZAugGenerator(cfg) zaug.set_multiple_gpus() - zaug.load_state_dict(torch.load(path_model)) + zaug.load_state_dict(torch.load(path_model, map_location=f"cuda:{cfg.device_ids[0]}")) zaug.eval() return zaug diff --git a/zaugnet_colab.ipynb b/zaugnet_colab.ipynb index 0d1d08d..b8293b8 100644 --- a/zaugnet_colab.ipynb +++ b/zaugnet_colab.ipynb @@ -8,7 +8,7 @@ "source": [ "# **ZAugNet Colab Notebook**\n", "\n", - "This Colab Notebook is related to the paper [Self-Supervised Z-Slice Augmentation for 3D Bio-Imaging via Knowledge Distillation](https://arxiv.org/abs/2503.04843).\n", + "This Colab Notebook is related to the paper [ZAugNet for Z-Slice Augmentation in Bio-Imaging](https://arxiv.org/).\n", "\n", "It provides an intuitive interface for training and/or testing ZAugNet and ZAugNet+, eliminating the need to interact with the code directly.\n", "\n" @@ -62,103 +62,45 @@ ] }, { - "cell_type": "markdown", + "cell_type": "code", + "execution_count": null, "metadata": { - "id": "6X_xvrT2sF8t" + "id": "u-6NjOJ0fruD" }, + "outputs": [], "source": [ - "### **1 - Download ZAugNet and install the requirements**\n" + "# shutil.copytree('./gdrive/MyDrive/zaugnet/', './zaugnet')" ] }, { "cell_type": "markdown", "metadata": { - "id": "F5fBoJVbitFC" - }, - "source": [ - "* Download ZAugNet from GitHub:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "5wsVUeAfimGr" + "id": "6X_xvrT2sF8t" }, - "outputs": [], "source": [ - "!git clone https://github.com/VirtualEmbryo/ZAugNet.git" -<<<<<<< HEAD + "### **1 - Download ZAugNet and install the requirements**\n" ] }, { "cell_type": "markdown", "metadata": { - "id": "WdKYdfzX2nge" + "id": "F5fBoJVbitFC" }, "source": [ - "* **(Optional)** Download train/test data images and pre-trained models:" + "* Download ZAugNet from GitHub:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { - "id": "osmjfhvc23lq" + "id": "5wsVUeAfimGr" }, "outputs": [], "source": [ - "# Download data and pre-trained models from Zenodo\n", - "!curl \"https://zenodo.org/records/14961732/files/zenodo.zip?download=1\" --output /content/ZAugNet/zaugnet_data_and_models.zip\n", - "\n", - "# Unzip the folder\n", - "import shutil\n", - "shutil.unpack_archive(\"/content/ZAugNet/zaugnet_data_and_models.zip\", \"/content/ZAugNet/zaugnet_data_and_models\")\n", - "\n", - "# Create data folder and move files\n", - "!mkdir /content/ZAugNet/data\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/ascidians /content/ZAugNet/data/ascidians\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/filaments /content/ZAugNet/data/filaments\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/humans /content/ZAugNet/data/humans\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/nuclei /content/ZAugNet/data/nuclei\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/results /content/ZAugNet/results" -======= ->>>>>>> 24dd0f40c0f345e5b34afd3b910460dc071e1e15 + "!git clone https://github.com/apasqui/zaugnet.git" ] }, - { - "cell_type": "markdown", - "source": [ - "* **(Optional)** Download train/test data images and pre-trained models:" - ], - "metadata": { - "id": "WdKYdfzX2nge" - } - }, - { - "cell_type": "code", - "source": [ - "# Download data and pre-trained models from Zenodo\n", - "!curl \"https://zenodo.org/records/14961732/files/zenodo.zip?download=1\" --output /content/ZAugNet/zaugnet_data_and_models.zip\n", - "\n", - "# Unzip the folder\n", - "import shutil\n", - "shutil.unpack_archive(\"/content/ZAugNet/zaugnet_data_and_models.zip\", \"/content/ZAugNet/zaugnet_data_and_models\")\n", - "\n", - "# Create data folder and move files\n", - "!mkdir /content/ZAugNet/data\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/ascidians /content/ZAugNet/data/ascidians\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/filaments /content/ZAugNet/data/filaments\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/humans /content/ZAugNet/data/humans\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/nuclei /content/ZAugNet/data/nuclei\n", - "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/results /content/ZAugNet/results" - ], - "metadata": { - "id": "osmjfhvc23lq" - }, - "execution_count": null, - "outputs": [] - }, { "cell_type": "markdown", "metadata": { @@ -176,7 +118,7 @@ }, "outputs": [], "source": [ - "os.chdir('./ZAugNet')\n", + "os.chdir('./zaugnet')\n", "%pip install -r requirements.txt" ] }, @@ -186,16 +128,14 @@ "id": "DcTSFrqDfN84" }, "source": [ - "* Insert paths to train and test folders.\n", - "\n", - "You can upload your .tif images on Google Drive and use the path ./gdrive/MyDrive/path_folder.\n", + "* Upload your train and test folders containing the .tif images in your Google Drive account.\n", "\n", - "Or you can use our data and use the path /content/ZAugNet/data/ascidiands for example." + "* Insert paths to train and test folders:" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 9, "metadata": { "id": "Oo0i4L3-Jinx" }, @@ -203,16 +143,18 @@ "source": [ "os.chdir('/content')\n", "# Put the adresse of the tif files for training and test data in google drive\n", - "train_path = \"\" #@param {type:\"string\"}\n", - "test_path = \"\" #@param {type:\"string\"}\n", + "train_path = \"./gdrive/MyDrive/\" #@param {type:\"string\"}\n", + "test_path = \"./gdrive/MyDrive/\" #@param {type:\"string\"}\n", "\n", "# Change the directory and create necessary folders\n", - "if not os.path.exists('./ZAugNet/data/train') :\n", - " shutil.copytree(train_path, './ZAugNet/data/train/')\n", - "if not os.path.exists('./ZAugNet/data/test') :\n", - " shutil.copytree(test_path, './ZAugNet/data/test/')\n", + "if not os.path.exists('./zaugnet/data/train') :\n", + " #os.makedirs('./zaugnet/data/train/')\n", + " shutil.copytree(train_path, './zaugnet/data/train/')\n", + "if not os.path.exists('./zaugnet/data/test') :\n", + " #os.makedirs('./zaugnet/data/test/')\n", + " shutil.copytree(test_path, './zaugnet/data/test/')\n", "\n", - "os.chdir('./ZAugNet')" + "os.chdir('./zaugnet')" ] }, { @@ -309,7 +251,7 @@ "source": [ "* To train a new ZAugNet or ZAugNet+ model, execute the following block.\n", "\n", - " The trained model will be saved in /content/ZAugNet/results/\\/ as `.pt` file.\n", + " The trained model will be saved in /content/zaugnet/results/\\/ as `.pt` file.\n", "\n" ] }, @@ -363,7 +305,7 @@ "source": [ "* To predict using a ZAugNet or ZAugNet+ model in the middle execute the following block.\n", "\n", - " The predicted images will be saved in /content/ZAugNet/predictions/ as a `.tif` files." + " The predicted images will be saved in /content/zaugnet/predictions/ as a `.tif` files." ] }, { @@ -413,9 +355,7 @@ }, { "cell_type": "markdown", - "metadata": { - "id": "mWpUvSG_2Xqr" - }, + "metadata": {}, "source": [ "You can use ZAugNet+ to predict at a specefic distance between two consecutive original z-slices $I_0$ and $I_1$ using DPM" ] @@ -456,4 +396,4 @@ }, "nbformat": 4, "nbformat_minor": 0 -} \ No newline at end of file +} From deb9e61549502c77314d3373c80a6100bfa8ff75 Mon Sep 17 00:00:00 2001 From: Sajjad-scbas <106543263+Sajjad-scbas@users.noreply.github.com> Date: Mon, 27 Oct 2025 19:09:12 +0100 Subject: [PATCH 2/2] fix Google Colab --- zaugnet_colab.ipynb | 80 +++++++++++++++++++++++++++++---------------- 1 file changed, 52 insertions(+), 28 deletions(-) diff --git a/zaugnet_colab.ipynb b/zaugnet_colab.ipynb index b8293b8..6690f23 100644 --- a/zaugnet_colab.ipynb +++ b/zaugnet_colab.ipynb @@ -8,7 +8,7 @@ "source": [ "# **ZAugNet Colab Notebook**\n", "\n", - "This Colab Notebook is related to the paper [ZAugNet for Z-Slice Augmentation in Bio-Imaging](https://arxiv.org/).\n", + "This Colab Notebook is related to the paper [Self-Supervised Z-Slice Augmentation for 3D Bio-Imaging via Knowledge Distillation](https://arxiv.org/abs/2503.04843).\n", "\n", "It provides an intuitive interface for training and/or testing ZAugNet and ZAugNet+, eliminating the need to interact with the code directly.\n", "\n" @@ -62,43 +62,65 @@ ] }, { - "cell_type": "code", - "execution_count": null, + "cell_type": "markdown", "metadata": { - "id": "u-6NjOJ0fruD" + "id": "6X_xvrT2sF8t" }, - "outputs": [], "source": [ - "# shutil.copytree('./gdrive/MyDrive/zaugnet/', './zaugnet')" + "### **1 - Download ZAugNet and install the requirements**\n" ] }, { "cell_type": "markdown", "metadata": { - "id": "6X_xvrT2sF8t" + "id": "F5fBoJVbitFC" }, "source": [ - "### **1 - Download ZAugNet and install the requirements**\n" + "* Download ZAugNet from GitHub:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5wsVUeAfimGr" + }, + "outputs": [], + "source": [ + "!git clone https://github.com/VirtualEmbryo/ZAugNet.git" ] }, { "cell_type": "markdown", "metadata": { - "id": "F5fBoJVbitFC" + "id": "WdKYdfzX2nge" }, "source": [ - "* Download ZAugNet from GitHub:" + "* **(Optional)** Download train/test data images and pre-trained models:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { - "id": "5wsVUeAfimGr" + "id": "osmjfhvc23lq" }, "outputs": [], "source": [ - "!git clone https://github.com/apasqui/zaugnet.git" + "# Download data and pre-trained models from Zenodo\n", + "!curl \"https://zenodo.org/records/14961732/files/zenodo.zip?download=1\" --output /content/ZAugNet/zaugnet_data_and_models.zip\n", + "\n", + "# Unzip the folder\n", + "import shutil\n", + "shutil.unpack_archive(\"/content/ZAugNet/zaugnet_data_and_models.zip\", \"/content/ZAugNet/zaugnet_data_and_models\")\n", + "\n", + "# Create data folder and move files\n", + "!mkdir /content/ZAugNet/data\n", + "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/ascidians /content/ZAugNet/data/ascidians\n", + "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/filaments /content/ZAugNet/data/filaments\n", + "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/humans /content/ZAugNet/data/humans\n", + "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/nuclei /content/ZAugNet/data/nuclei\n", + "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/results /content/ZAugNet/results" ] }, { @@ -118,7 +140,7 @@ }, "outputs": [], "source": [ - "os.chdir('./zaugnet')\n", + "os.chdir('./ZAugNet')\n", "%pip install -r requirements.txt" ] }, @@ -128,14 +150,16 @@ "id": "DcTSFrqDfN84" }, "source": [ - "* Upload your train and test folders containing the .tif images in your Google Drive account.\n", + "* Insert paths to train and test folders.\n", + "\n", + "You can upload your .tif images on Google Drive and use the path ./gdrive/MyDrive/path_folder.\n", "\n", - "* Insert paths to train and test folders:" + "Or you can use our data and use the path /content/ZAugNet/data/ascidiands for example." ] }, { "cell_type": "code", - "execution_count": 9, + "execution_count": null, "metadata": { "id": "Oo0i4L3-Jinx" }, @@ -143,18 +167,16 @@ "source": [ "os.chdir('/content')\n", "# Put the adresse of the tif files for training and test data in google drive\n", - "train_path = \"./gdrive/MyDrive/\" #@param {type:\"string\"}\n", - "test_path = \"./gdrive/MyDrive/\" #@param {type:\"string\"}\n", + "train_path = \"\" #@param {type:\"string\"}\n", + "test_path = \"\" #@param {type:\"string\"}\n", "\n", "# Change the directory and create necessary folders\n", - "if not os.path.exists('./zaugnet/data/train') :\n", - " #os.makedirs('./zaugnet/data/train/')\n", - " shutil.copytree(train_path, './zaugnet/data/train/')\n", - "if not os.path.exists('./zaugnet/data/test') :\n", - " #os.makedirs('./zaugnet/data/test/')\n", - " shutil.copytree(test_path, './zaugnet/data/test/')\n", + "if not os.path.exists('./ZAugNet/data/train') :\n", + " shutil.copytree(train_path, './ZAugNet/data/train/')\n", + "if not os.path.exists('./ZAugNet/data/test') :\n", + " shutil.copytree(test_path, './ZAugNet/data/test/')\n", "\n", - "os.chdir('./zaugnet')" + "os.chdir('./ZAugNet')" ] }, { @@ -251,7 +273,7 @@ "source": [ "* To train a new ZAugNet or ZAugNet+ model, execute the following block.\n", "\n", - " The trained model will be saved in /content/zaugnet/results/\\/ as `.pt` file.\n", + " The trained model will be saved in /content/ZAugNet/results/\\/ as `.pt` file.\n", "\n" ] }, @@ -305,7 +327,7 @@ "source": [ "* To predict using a ZAugNet or ZAugNet+ model in the middle execute the following block.\n", "\n", - " The predicted images will be saved in /content/zaugnet/predictions/ as a `.tif` files." + " The predicted images will be saved in /content/ZAugNet/predictions/ as a `.tif` files." ] }, { @@ -355,7 +377,9 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "id": "mWpUvSG_2Xqr" + }, "source": [ "You can use ZAugNet+ to predict at a specefic distance between two consecutive original z-slices $I_0$ and $I_1$ using DPM" ]