diff --git a/README.md b/README.md
index db7b691a7e440e4c6c921fcff7c043242c7cdf3f..f1d7a38cfe45949db7c2575eb92556ca5db6111b 100644
--- a/README.md
+++ b/README.md
@@ -3,95 +3,269 @@
## *μPIX : Leveraging Generative AI for Enhanced, Personalized and Sustainable Microscopy*
### Gabriel Bon, Daniel Sapede, Cédric Matthews, and Fabrice Daian
+----
+----
+
| | Paper |
|--------------------------|--------------------------------------------------------------------------------------------------------------------------------|
| Preprint | [](https://doi.org/10.1101/2024.10.25.620201) |
| Paper | To be published |
-| | Dataset |
-|--------------------------|--------------------------------------------------------------------------------------------------------------------------------|
-| CSBDeep | [CSBDeep Directory](./data/csbdeep) |
-| Gastruloids | [Gastruloids Directory](./data/gastruloids) |
-| Metrology | [Metrology Directory](./data/metrology) |
-
-
-| | Code Requirements |
-|--------------------------|--------------------------------------------------------------------------------------------------------------------------------|
-| Python | 3.11 |
-| Tensorflow | 2.15 |
-| CUDA | 11.8 |
-| cuDNN | 8 |
+| | Dataset | pre-trained µPIX Models |
+|--------------------------|--------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------|
+| CSBDeep | [download](https://sync.lis-lab.fr/index.php/s/sy3SrGgqNafbP5X/download) | [download](https://sync.lis-lab.fr/index.php/s/9SxkR2QH3Z79Bbc/download) |
+| Gastruloids | [download](https://sync.lis-lab.fr/index.php/s/kDyLKjCt48BfNMK/download) | [download](https://sync.lis-lab.fr/index.php/s/n2oJAw4LobTCeNZ/download) |
+| Metrology | [download](https://sync.lis-lab.fr/index.php/s/mYDRTeAQxMxNPPJ/download) | [download](https://sync.lis-lab.fr/index.php/s/degZsCxN7ZXxeB6/download) |
-| | Models |
+| Demo Notebooks | |
|--------------------------|--------------------------------------------------------------------------------------------------------------------------------|
-µPIX model on Metrology Dataset | [mupixmetro]() |
-µPIX model on CSBDeep Dataset | [mupixcsb]() |
-µPIX model on Gastruloids Dataset | [mupixgastruloid]() |
-
+| Inference on Metrology dataset using pre-trained µPIX model | [](https://colab.research.google.com/drive/1EbV04iv141q7ebVUOErJupYboaCV40OD?usp=drive_link) |
-| | Demos |
+| | Code & System Requirements |
|--------------------------|--------------------------------------------------------------------------------------------------------------------------------|
-| Train a µPix model from scratch on a custom dataset | [](-) |
-| Fine-tune a µPix model | [](-) |
-| Using a µPIX model on inference on new images | [](-) |
-
-
-
-
-
-## 1 - Setup µPIX environment
-
-To setup µPIX environment, you have the choice on relying either on a Conda environment or a Docker environment.
+| Python | 3.11 |
+| Tensorflow | 2.14 |
+| CUDA | 12.2 |
+| OS | Linux |
+----
+----
-### a) Setup a Conda Environement
+## Install required µPIX Environment and Sources
-If you don't have alread miniconda installed, please refere to the [officical miniconda documentation](https://docs.anaconda.com/miniconda/install/)
+If you don't have alread miniconda installed on your system, please refere to the [officical miniconda documentation](https://docs.anaconda.com/miniconda/install/) and then :
-1) Create a conda environment:
+Create a miniconda environment and source it:
```bash
conda create -n mupix_env python=3.11
```
-and
```bash
source ~/.bashrc
```
-and finally activate the environment:
+Finally activate the environment:
```bash
conda activate mupix_env
```
-2) Clone this repository to download µPIX scripts and datasets :
+Clone this repository to download µPIX sources :
```bash
git clone https://gitlab.lis-lab.fr/sicomp/mupix
```
-then:
+and then proceed to the installation of required Python packages:
+
+```bash
+pip install -r mupix/requirements.txt
+```
+
+## 2 - Train a µPIX model from scratch using a custom dataset
+
+Once the environment is installed, you can start using µPIX scripts.
+
+Move to the `mupix` directory:
```bash
cd mupix
```
+As an example, we will train a µPIX model from scratch using the Metrology dataset.
+
+
+### 2.1 - Download/Prepare the dataset
+First download the dataset manually ([here](https://sync.lis-lab.fr/index.php/s/mYDRTeAQxMxNPPJ/download)), and unzip it.
+
+Alternatively, if you have ```curl``` installed on your system, you can download and extract the dataset using:
+
+```bash
+curl -o metrology.zip https://sync.lis-lab.fr/index.php/s/mYDRTeAQxMxNPPJ/download -q
+```
+and
```bash
-pip install -r requirements.txt
+unzip -qq metrology.zip
```
+This will create a ```metrology``` directory containing two subdirectories: ```train``` and ```test```.
-### b) Using the µPIX Docker container
+- ```train``` (Training data) directory contains two folders:
+- - ```GT``` folder contains clean images (Ground Truth)
+- - ```low``` folder contains the corresponding noisy images
-If you don't have Docker already installed on your computer, please use our help to guide through the [installation of Docker on your machine](./docker_install.md).
+**Important: Each image in the GT/ folder must have a corresponding noisy image in low/ with the same filename, both in TIF format.**
+Example: if it exist an image called ```image0.tif``` in the ```GT``` folder, it must have the corresponding ```image0.tif``` noisy image in the ```low``` folder.
-We have developped a "ready-to-use" Docker image ([repo on dockerhub](https://hub.docker.com/r/gabrielbon/mupix)) containing a ready-to-use µPIX environment.
+- ```test``` (Testing data) directory contains all the images that will be used for testing/inference the model.
-To get the image, open a Terminal and type:
+### 2.2 - Create a µPIX experiment
-```bash
-docker pull gabrielbon/mupix:latest
+The ```new_experiment.py``` script allows you to set up a new experiment by specifying the required paths and parameters to the dataset.
+
+#### Usage
+```
+python new_experiment.py --experiment_name EXPERIMENT_NAME
+ --experiment_path EXPERIMENT_PATH
+ --clean_data_path CLEAN_DATA_PATH
+ --noisy_data_path NOISY_DATA_PATH
+ [--test_data_path TEST_DATA_PATH]
```
-Finally, you can clone this repository to install µPIX scripts and datasets:
+| Argument | Description |
+|--------------------------|-------------|
+| `--experiment_name` | Name of the µPIX experiment. |
+| `--experiment_path` | Path where the µPIX experiment will be saved. |
+| `--clean_data_path` | Path to the clean dataset (Train). |
+| `--noisy_data_path` | Path to the noisy dataset (Train). |
+| `--test_data_path` *(optional)* | Path to the test dataset (if available). |
+
+For our example, here is the command you need to type to create a µPIX experiment for the ```metrology``` dataset:
+
```bash
-git clone https://gitlab.lis-lab.fr/sicomp/mupix
+python new_experiment.py --experiment_name metrology_experiment
+ --experiment_path ./experiments
+ --clean_data_path ./metrology/train/GT
+ --noisy_data_path ./metrology/train/low
+ --test_data_path ./metrology/test
+```
+If everything works correctly you should see:
```
+Experiment 'metrology_experiment' created successfully at ./experiments/metrology_experiment/
+```
+
+This command has created a new `experiments/metrology_experiment` directory containing:
+- `hyperparameters.json`: default parameters used by µPIX
+- `results` directory: will contains training logs you can monitor during the training (```logs.txt```), model checkpoints (```networks```), image generated during the training (```images```) of a the µPIX model.
+- `predictions` directory: will contains the generated image on the test dataset.
+
+
+For information, here is the ```hyperparameters.json``` file created by default by µPIX for our example experiment that could be freely modifiable
+
+```json
+{
+ "learning_rate_generator": 0.0001,
+ "learning_rate_discriminator": 0.0001,
+ "batch_size": 16,
+ "num_epochs": 100,
+ "loss_weight": 10,
+ "tile_size": 256,
+ "patience": 20,
+ "valid_size": 0.1,
+ "seed": 42,
+ "data_paths": {
+ "clean": "./data/metrology/GT/",
+ "noisy": "./data/metrology/low",
+ "test": null
+ }
+}
+```
+
+| Hyperparameter | Type | Description |
+|--------------------------------|---------|-------------|
+| `learning_rate_generator` | `float` | Learning rate for the generator. |
+| `learning_rate_discriminator` | `float` | Learning rate for the discriminator. |
+| `batch_size` | `int` | Number of samples per batch. |
+| `num_epochs` | `int` | Total number of training epochs. |
+| `loss_weight` | `int` | Weight factor for µPIX loss calculation. |
+| `tile_size` | `int` | Size of image tiles used for training. |
+| `patience` | `int` | Number of epochs to wait before triggering µPIX early stopping if no improvement. |
+| `valid_size` | `float` | Proportion of the dataset used for validation. |
+| `seed` | `int` | Random seed for reproducibility. |
+| `data_paths.clean` | `str` | Path to the clean dataset. |
+| `data_paths.noisy` | `str` | Path to the noisy dataset. |
+| `data_paths.test` | `str or null` | Path to the test dataset (if available). |
+
+
+### 2.3 - Training
+
+The `mupixtraining.py` script allows you to train a µPIX model by specifying the required path to an existing experiment.
+
+#### Usage
+```
+python mupixtraining.py --experiment_path EXPERIMENT_PATH [--retrain]
+```
+
+#### Details
+
+| Argument | Description |
+|--------------------------|-------------|
+| `--experiment_path` | Path to the previously created experiment. |
+| `--retrain` *(optional)* | Use this flag to continue training an existing µPIX model located inside the experiment path. |
+
+
+For our example, here is the command you need to type to start the training of a µPIX model for the ```metrology``` dataset:
+
+```
+python mupixtraining.py --experiment_path ./experiments/metrology_experiment
+```
+
+You should see the training starting:
+
+```
+Segmentation Models: using `tf.keras` framework.
+Input Path: ('./dataset_metrology/train/low', './dataset_metrology/train/GT')
+Output Path: ./experiments/metrology/results
+Retraining: False
+Data loading...
+85 images loaded
+2212 tiles extracted for training
+303 tiles extracted for validation
+Data loaded !
+Training from scratch...
+2025-02-27 12:27:50.608620: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:894] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
+2025-02-27 12:27:50.608906: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2211] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
+Skipping registering GPU devices...
+/home/fabrice/miniconda3/envs/mupix_env/lib/python3.11/site-packages/keras/src/initializers/initializers.py:120: UserWarning: The initializer RandomNormal is unseeded and being called multiple times, which will return identical values each time (even if the initializer is unseeded). Please update your code to provide a seed to the initializer, or avoid using the same initializer instance more than once.
+ warnings.warn(
+3645.2395 4581.112
+10/10 [==============================] - 6s 509ms/step
+Epoch 1
+Val> mse[3.827e+08], ssim[-1.143e-04]
+1/1 [==============================] - 0s 399ms/step
+Step>1, Generator loss : 1.265e+05
+1/1 [==============================] - 0s 330ms/step
+Step>2, Generator loss : 1.883e+05
+1/1 [==============================] - 0s 354ms/step
+Step>3, Generator loss : 1.941e+05
+1/1 [==============================] - 0s 363ms/step
+...
+...
+...
+```
+By default, the model is trained for 100 epochs (see ```hyperparameters.json```) but contains an ```EarlyStopping``` mechanisms (See the paper methods section pour details) governed by the ```patience``` parameters (see ```hyperparameters.json```). The number of steps per epochs is calculated by dividing the number of training tiles divided by the ```batch_size```.
+
+You can stop the training anytime. The best checkpoints of your model is available inside ```experiments/metrology_experiment/results/networks/```
+
+
+
+### 2.4 - Inference on the test dataset
+
+The `mupixinfer.py` script allows you to use a pre-trained µPIX model to denoise a dataset located in the `test` directory inside the experiment path.
+
+#### Usage
+
+```
+python mupixinfer.py --experiment_path EXPERIMENT_PATH
+```
+
+For our example, here is the command you need to type to start the inference using the newly trained µPIX model best (checkpoint) for the ```metrology``` test dataset:
+
+```
+python mupixinfer.py --experiment_path ./experiments/metrology_experiment/
+```
+
+
+
+
+
+
+
+
+
+1 - Train a µPIX model from scratch
+
+2 - Use a pre-trained µPIX model to denoise images
+
+3 - Fine-Tune a pre-trained µPIX model on a new dataset
+
+
+
## 2 - Create a µPIX experiments
@@ -240,4 +414,25 @@ docker run -it --entrypoint bin/bash -v /path/to/your/file/mupix:/workspace -v /
Given you are working in the container, you will have to specify the path do data differently
## 3 - License
-This code repository is release under the [CC BY-NS-SA 4.0](https://gitlab.lis-lab.fr/sicomp/mupix/-/blob/main/LICENSE?ref_type=heads)
\ No newline at end of file
+This code repository is release under the [CC BY-NS-SA 4.0](https://gitlab.lis-lab.fr/sicomp/mupix/-/blob/main/LICENSE?ref_type=heads)
+
+
+
+
+### b) Using the µPIX Docker container
+
+If you don't have Docker already installed on your computer, please use our help to guide through the [installation of Docker on your machine](./docker_install.md).
+
+
+We have developped a "ready-to-use" Docker image ([repo on dockerhub](https://hub.docker.com/r/gabrielbon/mupix)) containing a ready-to-use µPIX environment.
+
+To get the image, open a Terminal and type:
+
+```bash
+docker pull gabrielbon/mupix:latest
+```
+
+Finally, you can clone this repository to install µPIX scripts and datasets:
+```bash
+git clone https://gitlab.lis-lab.fr/sicomp/mupix
+```
diff --git a/requirements.txt b/requirements.txt
index 4299c69272303776cdbd156156b3d5d200930cb5..bd1cd4053bdd5ded2b99663c88d143af76062041 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,5 @@
absl-py==2.1.0
-albucore==0.0.23
-albumentations==2.0.4
+albumentations==1.3.0
annotated-types==0.7.0
astunparse==1.6.3
cachetools==5.5.2