[UNet3+/TF2] Initial contribution (#1267)

TensorFlow2/Segmentation/Contrib/UNet3P/.gitignore

@@ -0,0 +1,18 @@
+.idea
+__pycache__
+
+checkpoint/tb_logs/*
+checkpoint/*.hdf5
+checkpoint/*.csv
+!checkpoint/tb_logs/.gitkeep
+
+#data/*
+/data/**/*.png
+/data/**/*.jpg
+/data/**/*.nii
+!data/**/.gitkeep
+
+data_preparation/verify_preprocess_data.ipynb
+old_data_preperation/
+others/
+**/outputs

TensorFlow2/Segmentation/Contrib/UNet3P/Dockerfile

@@ -0,0 +1,15 @@
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/tensorflow:22.12-tf2-py3
+FROM ${FROM_IMAGE_NAME}
+
+ADD . /workspace/unet3p
+WORKDIR /workspace/unet3p
+
+RUN pip install -r requirements.txt
+
+#For opencv, inside docker run these commands
+RUN apt-get update
+RUN apt-get install ffmpeg libsm6 libxext6  -y
+
+# reinstall jupyterlab
+RUN pip uninstall jupyterlab -y
+RUN pip install jupyterlab

TensorFlow2/Segmentation/Contrib/UNet3P/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2023 Hamid Ali
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

TensorFlow2/Segmentation/Contrib/UNet3P/README.md

@@ -0,0 +1,319 @@
+# UNet 3+: A Full-Scale Connected UNet for Medical Image Segmentation
+
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/unet-3-a-full-scale-connected-unet-for/medical-image-segmentation-on-lits2017)](https://paperswithcode.com/sota/medical-image-segmentation-on-lits2017?p=unet-3-a-full-scale-connected-unet-for)
+
+This repository provides a script and recipe to train UNet3+ to achieve state of the art accuracy.
+
+**The code and associated performance metrics were contributed by the community and are not maintained by NVIDIA.**
+
+## Table of Contents
+
+- [UNet 3+](https://arxiv.org/abs/2004.08790) for Image Segmentation in Tensorflow 2.0.
+    - [Table of Contents](#table-of-contents)
+    - [Feature Support Matrix](#feature-support-matrix)
+    - [Installation](#installation)
+    - [Code Structure](#code-structure)
+    - [Config](#config)
+    - [Data Preparation](#data-preparation)
+    - [Models](#models)
+    - [Performance](#performance)
+    - [Inference Demo](#inference-demo)
+    - [Known issues](#known-issues)
+    - [Release notes](#release-notes)
+
+## Feature Support Matrix
+
+The following features are supported by our code base:
+
+|                  Feature                   | UNet3+ Supports |
+|:------------------------------------------:|:---------------:|
+|                    DALI                    |     ✓     |
+|       TensorFlow Multi-GPU Training        |     ✓     |
+| TensorFlow Automatic Mixed Precision(AMP)  |     ✓     |
+| TensorFlow Accelerated Linear Algebra(XLA) |     ✓     |
+
+#### [NVIDIA DALI](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/index.html)
+
+The NVIDIA Data Loading Library (DALI) is a library for data loading and
+pre-processing to accelerate deep learning applications. It provides a
+collection of highly optimized building blocks for loading and processing
+image, video and audio data. It can be used as a portable drop-in
+replacement for built in data loaders and data iterators in popular deep
+learning frameworks.
+
+#### [TensorFlow Multi-GPU Training](https://www.tensorflow.org/guide/distributed_training)
+
+Distribute training across multiple GPUs, multiple machines, or TPUs.
+
+#### [TensorFlow Automatic Mixed Precision(AMP)](https://www.tensorflow.org/guide/mixed_precision)
+
+Mixed precision is the use of both 16-bit and 32-bit floating-point types in a model during training to make it run
+faster and use less memory. By keeping certain parts of the model in the 32-bit types for numeric stability, the model
+will have a lower step time and train equally as well in terms of the evaluation metrics such as accuracy.
+
+#### [TensorFlow Accelerated Linear Algebra(XLA)](https://www.tensorflow.org/xla)
+
+In a TensorFlow program, all of the operations are executed individually by the TensorFlow executor. Each TensorFlow
+operation has a precompiled GPU kernel implementation that the executor dispatches to.
+XLA provides an alternative mode of running models: it compiles the TensorFlow graph into a sequence of computation
+kernels generated specifically for the given model. Because these kernels are unique to the model, they can exploit
+model-specific information for optimization.
+
+For details on how to enable these features while training and evaluation see [Benchmarking](#benchmarking) section.
+
+## Installation
+
+* Clone code
+
+```
+git clone https://github.com/hamidriasat/NVIDIA-DeepLearningExamples.git
+cd NVIDIA-DeepLearningExamples/TensorFlow2/Segmentation/UNet3P/
+```
+
+* Build the UNet3P TensorFlow NGC container
+  From `Dockerfile` this will create a docker image with name `unet3p`. This image will contain all the components
+  required to successfully run the UNet3+ code.
+
+```
+docker build -t unet3p .
+```
+
+The NGC container contains all the components optimized for usage on NVIDIA hardware.
+
+* Start an interactive session in the NGC container
+
+To run preprocessing/training/inference, following command will launch the container and mount the current directory
+to `/workspace/unet3p` as a volume in the container
+
+```
+docker run --rm -it --shm-size=1g --ulimit memlock=-1 --pids-limit=8192 --gpus all -p 5012:8888 -v $PWD/:/workspace/unet3p --name unet3p unet3p:latest /bin/bash
+```
+
+Here we are mapping external port `5012` to `8888` inside docker. This will be used for visualization purpose.
+
+## Code Structure
+
+- **callbacks**: Custom callbacks to monitor training time, latency and throughput
+- **checkpoint**: Model checkpoint and logs directory
+- **configs**: Configuration file (see [Config](#config) for more details)
+- **data**: Dataset files (see [Data Preparation](#data-preparation) for more details)
+- **data_generators**: Data loaders for UNet3+ (see [Data Generators](#data-generators) for more details)
+- **data_preparation**: For LiTS data preparation and data verification
+- **figures**: Model architecture image
+- **losses**: Implementations of UNet3+ hybrid loss function and dice coefficient
+- **models**: Unet3+ model files (see [Models](#models) for more details)
+- **utils**: Generic utility functions
+- **benchmark_inference.py**: Benchmark script to output model throughput and latency while inference
+- **evaluate.py**: Evaluation script to validate accuracy on trained model
+- **predict.ipynb**: Prediction file used to visualize model output inside notebook(helpful for remote server
+  visualization)
+- **predict.py**: Prediction script used to visualize model output
+- **train.py**: Training script
+
+## Data Preparation
+
+- This code can be used to reproduce UNet3+ paper results
+  on [LiTS - Liver Tumor Segmentation Challenge](https://competitions.codalab.org/competitions/15595).
+- You can also use it to train UNet3+ on custom dataset.
+
+For dataset preparation read [here](data_preparation/README.md).
+
+## Config
+
+Configurations are passed through `yaml` file. For more details on config file read [here](configs/).
+
+## Data Generators
+
+We support two types of data loaders. `NVIDIA DALI` and `TensorFlow Sequence`
+generators. For more details on supported generator types read [here](data_generators/).
+
+## Models
+
+UNet 3+ is latest from Unet family, proposed for semantic image segmentation. it takes advantage of full-scale skip
+connections and deep supervisions.The full-scale skip connections incorporate low-level details with high-level
+semantics from feature maps in different scales; while the deep supervision learns hierarchical representations from the
+full-scale aggregated feature maps.
+
+![alt text](figures/unet3p_architecture.png)
+
+Figure 1. UNet3+ architecture diagram from [original paper](https://arxiv.org/abs/2004.08790).
+
+This repo contains all three versions of UNet3+.
+
+| #   |                          Description                          |                            Model Name                             | Training Supported |
+|:----|:-------------------------------------------------------------:|:-----------------------------------------------------------------:|:------------------:|
+| 1   |                       UNet3+ Base model                       |                 [unet3plus](models/unet3plus.py)                  |      ✓       |
+| 2   |                 UNet3+ with Deep Supervision                  |     [unet3plus_deepsup](models/unet3plus_deep_supervision.py)     |      ✓       |
+| 3   | UNet3+ with Deep Supervision and Classification Guided Module | [unet3plus_deepsup_cgm](models/unet3plus_deep_supervision_cgm.py) |      ✗       |
+
+Available backbones are `unet3plus`, `vgg16` and  `vgg19`. All backbones are untrained networks.
+
+In our case all results are reported using `vgg19` backbone and  `unet3plus` variant.
+
+[Here](losses/unet_loss.py) you can find UNet3+ hybrid loss.
+
+## Performance
+
+### Benchmarking
+
+The following section shows how to run benchmarks to measure the model performance in training and inference modes.
+
+#### Training performance benchmark
+
+Run the `python train.py` script with the required model configurations to print training benchmark results for each
+model
+configuration. At the end of the training, a line reporting the training throughput and latency will be printed.
+
+To calculate dice score on trained model call `python evaluate.py` with required parameters.
+
+##### Example 1
+
+To train base model `unet3plus` with `vgg19` backbone on `single GPU`
+using `TensorFlow Sequence Generator` without `Automatic Mixed Precision(AMP)` and `Accelerated Linear Algebra(XLA)` run
+
+```
+python train.py MODEL.TYPE=unet3plus MODEL.BACKBONE.TYPE=vgg19 \
+USE_MULTI_GPUS.VALUE=False \
+DATA_GENERATOR_TYPE=TF_GENERATOR \
+OPTIMIZATION.AMP=False OPTIMIZATION.XLA=False
+```
+
+##### Example 2
+
+To train base model `unet3plus` with `vgg19` backbone on `multiple GPUs`
+using `TensorFlow Sequence Generator` without `Automatic Mixed Precision(AMP)` and `Accelerated Linear Algebra(XLA)` run
+
+```
+python train.py MODEL.TYPE=unet3plus MODEL.BACKBONE.TYPE=vgg19 \
+USE_MULTI_GPUS.VALUE=True USE_MULTI_GPUS.GPU_IDS=-1 \
+DATA_GENERATOR_TYPE=TF_GENERATOR \
+OPTIMIZATION.AMP=False OPTIMIZATION.XLA=False
+```
+
+##### Example 3
+
+To train base model `unet3plus` with `vgg19` backbone on `multiple GPUs`
+using `NVIDIA DALI Generator` with `Automatic Mixed Precision(AMP)` and `Accelerated Linear Algebra(XLA)` run
+
+```
+python train.py MODEL.TYPE=unet3plus MODEL.BACKBONE.TYPE=vgg19 \
+USE_MULTI_GPUS.VALUE=True USE_MULTI_GPUS.GPU_IDS=-1 \
+DATA_GENERATOR_TYPE=DALI_GENERATOR \
+OPTIMIZATION.AMP=True OPTIMIZATION.XLA=True
+```
+
+To evaluate/calculate dice accuracy of model pass same parameters to `evaluate.py` file.
+
+Please check [Config](configs/config.yaml) file for more details about default training parameters.
+
+#### Inference performance benchmark
+
+To benchmark inference time, run the `python benchmark_inference.py` script with the required model configurations to
+print
+inference benchmark results for each model configuration. At the end, a line reporting the inference throughput and
+latency will be printed.
+
+For inference run without `data generator` and `GPUs` details but with `batch size`, `warmup_steps`
+and `bench_steps`
+parameters.
+
+```
+python benchmark_inference.py MODEL.TYPE=unet3plus MODEL.BACKBONE.TYPE=vgg19 \
+HYPER_PARAMETERS.BATCH_SIZE=16 \
+OPTIMIZATION.AMP=False OPTIMIZATION.XLA=False \
++warmup_steps=50 +bench_steps=100
+```
+
+Each of these scripts will by default run a warm-up for 50 iterations and then start benchmarking for another 100
+steps.
+You can adjust these settings with `+warmup_steps` and `+bench_steps` parameters.
+
+### Results
+
+The following section provide details of results that are achieved in different settings of model training and
+inference.
+
+**These results were contributed by the community and are not maintained by NVIDIA.**
+
+#### Training accuracy results
+
+###### Training accuracy: NVIDIA DGX A100 (8xA100 80G)
+
+| #GPU | Generator |   XLA   |   AMP   | Training Time<br/>HH:MM:SS &darr; | Latency Avg [ms] &darr; | Throughput Avg [img/s] &uarr; | Speed Up  | Dice Score |
+|:----:|:---------:|:-------:|:-------:|:---------------------------------:|:-----------------------:|:-----------------------------:|:---------:|:----------:|
+|  1   |    TF     | &cross; | &cross; |            51:38:24.24            |         616.14          |             25.97             |    ---    |  0.96032   |
+|  8   |    TF     | &cross; | &cross; |             11:30:45              |         999.39          |            128.08             | 1x (base) |  0.95224   |
+|  8   |   DALI    | &cross; | &cross; |              6:23:43              |         614.26          |            208.38             |   1.8x    |  0.94566   |
+|  8   |   DALI    | &check; | &cross; |              7:33:15              |         692.71          |            184.78             |   1.5x    |  0.94806   |
+|  8   |   DALI    | &cross; | &check; |              3:49:55              |         357.34          |             358.2             |    3x     |  0.94786   |
+|  8   |   DALI    | &check; | &check; |              3:14:24              |         302.83          |            422.68             |   3.5x    |   0.9474   |
+
+Latency is reported in milliseconds per batch whereas throughput is reported in images per second.
+Speed Up comparison is efficiency achieved in terms of training time between different runs.
+
+Note: Training time includes time to load cuDNN in first iteration and the first epoch which take little longer as
+compared to later epochs because in first epoch tensorflow optimizes the training graph. In terms of latency and
+throughput it does not matter much because we have trained networks for 100 epochs which normalizes this during
+averaging.
+
+#### Inference performance results
+
+###### Inference performance: NVIDIA DGX A100 (1xA100 80G)
+
+| Batch Size |   XLA   |   AMP   | Latency Avg [ms] &darr; | Throughput Avg [img/s] &uarr; |
+|:----------:|:-------:|:-------:|:-----------------------:|:-----------------------------:|
+|     1      | &cross; | &cross; |          59.54          |             16.79             |
+|     1      | &check; | &cross; |          70.59          |             14.17             |
+|     1      | &cross; | &check; |          56.17          |             17.80             |
+|     1      | &check; | &check; |          55.54          |             18.16             |
+|     16     | &cross; | &cross; |         225.59          |             70.93             |
+|     16     | &check; | &cross; |         379.93          |             43.15             |
+|     16     | &cross; | &check; |         184.98          |             87.02             |
+|     16     | &check; | &check; |         153.65          |             103.6             |
+
+Inference results are tested on single gpu. Here data generator type does not matter because only prediction time
+is calculated and averaged between 5 runs.
+
+## Inference Demo
+
+Model output can be visualized from Jupyter Notebook. Use below command to start Jupyter Lab on port `8888`.
+
+```
+jupyter lab --no-browser --allow-root --ip=0.0.0.0 --port=8888
+```
+
+While starting container we mapped system port `5012` to `8888` inside docker.
+> Note: Make sure you have server network ip and port access in case you are working with remote sever.
+
+Now in browser go to link `http://<server ip here>:5012/` to access Jupyter Lab.
+Open [predict.ipynb](predict.ipynb) notebook and rerun the whole notebook to visualize model output.
+
+There are two options for visualization, you can
+
+1. Visualize from directory
+
+It's going to make prediction and show all images from given directory. Useful for detailed evaluation.
+
+2. Visualize from list
+
+It's going to make prediction on elements of given list. Use for testing on specific cases.
+
+For custom data visualization set `SHOW_CENTER_CHANNEL_IMAGE=False`. This should set True for only UNet3+ LiTS data.
+
+For further details on visualization options see [predict.ipynb](predict.ipynb) notebook.
+
+## Known issues
+
+There are no known issues in this release.
+
+## Release notes
+
+### Changelog
+
+Feb 2023
+
+- Initial release
+
+We appreciate any feedback so reporting problems, and asking questions are welcomed here.
+
+Licensed under [Apache-2.0 License](LICENSE)

TensorFlow2/Segmentation/Contrib/UNet3P/benchmark_inference.py

@@ -0,0 +1,101 @@
+"""
+Script to benchmark model throughput and latency
+"""
+import os
+import numpy as np
+from tqdm import tqdm
+from timeit import default_timer as timer
+import hydra
+from omegaconf import DictConfig
+import tensorflow as tf
+from tensorflow.keras import mixed_precision
+
+from data_generators import tf_data_generator
+from utils.general_utils import join_paths, suppress_warnings
+from utils.images_utils import postprocess_mask
+from models.model import prepare_model
+
+
+def benchmark_time(cfg: DictConfig):
+    """
+    Output throughput and latency
+    """
+
+    # suppress TensorFlow and DALI warnings
+    suppress_warnings()
+
+    if cfg.OPTIMIZATION.AMP:
+        print("Enabling Automatic Mixed Precision(AMP)")
+        policy = mixed_precision.Policy('mixed_float16')
+        mixed_precision.set_global_policy(policy)
+
+    if cfg.OPTIMIZATION.XLA:
+        print("Enabling Accelerated Linear Algebra(XLA)")
+        tf.config.optimizer.set_jit(True)
+
+    # data generator
+    val_generator = tf_data_generator.DataGenerator(cfg, mode="VAL")
+    validation_steps = val_generator.__len__()
+
+    warmup_steps, bench_steps = 50, 100
+    if "warmup_steps" in cfg.keys():
+        warmup_steps = cfg.warmup_steps
+    if "bench_steps" in cfg.keys():
+        bench_steps = cfg.bench_steps
+    validation_steps = min(validation_steps, (warmup_steps + bench_steps))
+
+    progress_bar = tqdm(total=validation_steps)
+
+    # create model
+    model = prepare_model(cfg)
+
+    # weights model path
+    checkpoint_path = join_paths(
+        cfg.WORK_DIR,
+        cfg.CALLBACKS.MODEL_CHECKPOINT.PATH,
+        f"{cfg.MODEL.WEIGHTS_FILE_NAME}.hdf5"
+    )
+
+    assert os.path.exists(checkpoint_path), \
+        f"Model weight's file does not exist at \n{checkpoint_path}"
+
+    # load model weights
+    model.load_weights(checkpoint_path, by_name=True, skip_mismatch=True)
+    # model.summary()
+
+    time_taken = []
+    # for each batch
+    for i, (batch_images, batch_mask) in enumerate(val_generator):
+
+        start_time = timer()
+        # make prediction on batch
+        batch_predictions = model.predict_on_batch(batch_images)
+        if len(model.outputs) > 1:
+            batch_predictions = batch_predictions[0]
+
+        # do postprocessing on predicted mask
+        batch_predictions = postprocess_mask(batch_predictions, cfg.OUTPUT.CLASSES)
+
+        time_taken.append(timer() - start_time)
+
+        progress_bar.update(1)
+        if i >= validation_steps:
+            break
+    progress_bar.close()
+
+    mean_time = np.mean(time_taken[warmup_steps:])  # skipping warmup_steps
+    throughput = (cfg.HYPER_PARAMETERS.BATCH_SIZE / mean_time)
+    print(f"Latency: {round(mean_time * 1e3, 2)} msec")
+    print(f"Throughput/FPS: {round(throughput, 2)} samples/sec")
+
+
[email protected](version_base=None, config_path="configs", config_name="config")
+def main(cfg: DictConfig):
+    """
+    Read config file and pass to benchmark_time method
+    """
+    benchmark_time(cfg)
+
+
+if __name__ == "__main__":
+    main()

TensorFlow2/Segmentation/Contrib/UNet3P/callbacks/timing_callback.py

@@ -0,0 +1,30 @@
+import sys
+from timeit import default_timer as timer
+import tensorflow as tf
+
+
+class TimingCallback(tf.keras.callbacks.Callback):
+    """
+    Custom callback to note training time, latency and throughput
+    """
+
+    def __init__(self, ):
+        super(TimingCallback, self).__init__()
+        self.train_start_time = None
+        self.train_end_time = None
+        self.batch_time = []
+        self.batch_start_time = None
+
+    def on_train_begin(self, logs: dict):
+        tf.print("Training starting time noted.", output_stream=sys.stdout)
+        self.train_start_time = timer()
+
+    def on_train_end(self, logs: dict):
+        tf.print("Training ending time noted.", output_stream=sys.stdout)
+        self.train_end_time = timer()
+
+    def on_train_batch_begin(self, batch: int, logs: dict):
+        self.batch_start_time = timer()
+
+    def on_train_batch_end(self, batch: int, logs: dict):
+        self.batch_time.append(timer() - self.batch_start_time)

TensorFlow2/Segmentation/Contrib/UNet3P/configs/README.md

@@ -0,0 +1,86 @@
+Here we provide **overview** of our config file and how you can use your own custom settings's for training and
+evaluation.
+
+We are using [Hydra](https://hydra.cc/) for passing configurations. Hydra is a framework for elegantly configuring
+complex applications. In Hydra you can easily [extend](https://hydra.cc/docs/patterns/extending_configs/)
+and [interpolate](https://hydra.cc/docs/advanced/override_grammar/basic/#primitives) `yaml` config files.
+
+#### Override Hydra config from command line
+
+[Here](https://hydra.cc/docs/1.0/advanced/override_grammar/basic/)  you can read how to pass or override configurations
+through command line. Overall to
+
+###### Override higher level attribute
+
+Directly access the key and override its value
+
+- For instance to override Data generator pass `DATA_GENERATOR_TYPE=DALI_GENERATOR`
+
+###### Override nested attribute
+
+Use `.` to access nested keys
+
+- For instance to override model type `MODEL.TYPE=unet3plus`
+- To override model backbone `MODEL.BACKBONE.TYPE=vgg19`
+
+To add new element from command line add `+` before attribute name. E.g. `+warmup_steps=50` because warm steps is not
+added in config file.
+
+> Note: Don't add space between list elements, it will create problem with Hydra.
+
+Most of the configurations attributes in our [config](./../configs/config.yaml) are self-explanatory. However, for some
+attributes additional comments are added.
+
+You can override configurations from command line too, but it's **advisable to override them from config file** because
+it's
+easy.
+
+By default, hydra stores a log file of each run in a separate directory. We have disabled it in our case,
+if you want to enable them to keep record of each run configuration's then comment out the settings at the end of config
+file.
+
+```yaml
+# project root working directory, automatically read by hydra (.../UNet3P)
+WORK_DIR: ${hydra:runtime.cwd}
+DATA_PREPARATION:
+  # unprocessed LiTS scan data paths, for custom data training skip this section details 
+  SCANS_TRAIN_DATA_PATH: "/data/Training Batch 2/"
+  ...
+DATASET:
+  # training data paths, should be relative from project root path
+  TRAIN:
+    IMAGES_PATH: "/data/train/images"
+  ...
+MODEL:
+  # available variants are unet3plus, unet3plus_deepsup, unet3plus_deepsup_cgm
+  TYPE: "unet3plus"
+  BACKBONE:
+  ...
+...
+DATA_GENERATOR_TYPE: "DALI_GENERATOR"  # options are TF_GENERATOR or DALI_GENERATOR
+SHOW_CENTER_CHANNEL_IMAGE: True  # only true for UNet3+. for custom dataset it should be False
+# Model input shape
+INPUT:
+  HEIGHT: 320
+  ...
+# Model output classes
+OUTPUT:
+  CLASSES: 2
+HYPER_PARAMETERS:
+  EPOCHS: 5
+  BATCH_SIZE: 2  # specify per gpu batch size
+  ...
+CALLBACKS:
+  TENSORBOARD:
+  ...
+PREPROCESS_DATA:
+  RESIZE:
+    VALUE: False  # if True, resize to input height and width
+    ...
+USE_MULTI_GPUS:
+  ...
+# to stop hydra from storing logs files
+defaults:
+  ...
+
+```

TensorFlow2/Segmentation/Contrib/UNet3P/configs/config.yaml

@@ -0,0 +1,118 @@
+# project root working directory, automatically read by hydra (.../UNet3P)
+WORK_DIR: ${hydra:runtime.cwd}
+
+DATA_PREPARATION:
+  # unprocessed LiTS scan data paths, for custom data training skip this section details
+  SCANS_TRAIN_DATA_PATH: "/data/Training Batch 2/"
+  SCANS_VAL_DATA_PATH: "/data/Training Batch 1/"
+
+  # Resize scans to model input size
+  RESIZED_HEIGHT: ${INPUT.HEIGHT}
+  RESIZED_WIDTH: ${INPUT.WIDTH}
+
+  # Clip scans value in given range
+  SCAN_MIN_VALUE: -200
+  SCAN_MAX_VALUE: 250
+
+DATASET:
+  # paths should be relative from project root path
+  TRAIN:
+    IMAGES_PATH: "/data/train/images"
+    MASK_PATH: "/data/train/mask"
+  VAL:
+    IMAGES_PATH: "/data/val/images"
+    MASK_PATH: "/data/val/mask"
+
+
+MODEL:
+  # available variants are unet3plus, unet3plus_deepsup, unet3plus_deepsup_cgm
+  TYPE: "unet3plus"
+  WEIGHTS_FILE_NAME: model_${MODEL.TYPE}
+  BACKBONE:
+    # available variants are unet3plus, vgg16, vgg19
+    TYPE: "vgg19"
+
+DATA_GENERATOR_TYPE: "DALI_GENERATOR"  # options are TF_GENERATOR or DALI_GENERATOR
+SEED: 5  # for result's reproducibility
+VERBOSE: 1  # For logs printing details, available options are 0, 1, 2
+DATALOADER_WORKERS: 3  # number of workers used for data loading
+SHOW_CENTER_CHANNEL_IMAGE: True  # only true for UNet3+ for custom dataset it should be False
+
+# Model input shape
+INPUT:
+  HEIGHT: 320
+  WIDTH: 320
+  CHANNELS: 3
+
+# Model output classes
+OUTPUT:
+  CLASSES: 2
+
+
+HYPER_PARAMETERS:
+  EPOCHS: 100
+  BATCH_SIZE: 16  # specify per gpu batch size
+  LEARNING_RATE: 5e-5  # 0.1, 1e-3, 3e-4, 5e-5
+
+
+CALLBACKS:
+  # paths should be relative from project root path
+  TENSORBOARD:
+    PATH: "/checkpoint/tb_logs"
+
+  EARLY_STOPPING:
+    PATIENCE: 100
+
+  MODEL_CHECKPOINT:
+    PATH: "/checkpoint"
+    SAVE_WEIGHTS_ONLY: True
+    SAVE_BEST_ONLY: True
+
+  CSV_LOGGER:
+    PATH: "/checkpoint"
+    APPEND_LOGS: False
+
+
+PREPROCESS_DATA:
+  RESIZE:
+    VALUE: False  # if True, resize to input height and width
+    HEIGHT: ${INPUT.HEIGHT}
+    WIDTH: ${INPUT.WIDTH}
+
+  IMAGE_PREPROCESSING_TYPE: "normalize"
+
+  NORMALIZE_MASK:
+    VALUE: False  # if True, divide mask by given value
+    NORMALIZE_VALUE: 255
+
+  SHUFFLE:
+    TRAIN:
+      VALUE: True
+    VAL:
+      VALUE: False
+
+
+USE_MULTI_GPUS:
+  VALUE: True  # If True use multiple gpus for training
+  # GPU_IDS: Could be integer or list of integers.
+  # In case Integer: if integer value is -1 then it uses all available gpus.
+  # otherwise if positive number, then use given number of gpus.
+  # In case list of Integers: each integer will be considered as gpu id
+  # e.g. [4, 5, 7] means use gpu 5,6 and 8 for training/evaluation
+  GPU_IDS: -1
+
+
+OPTIMIZATION:
+  AMP: True  # Automatic Mixed Precision(AMP)
+  XLA: True  # Accelerated Linear Algebra(XLA)
+
+
+# to stop hydra from storing logs files
+# logs will be stored in outputs directory
+defaults:
+  - _self_
+  - override hydra/hydra_logging: disabled
+  - override hydra/job_logging: disabled
+
+hydra:
+  output_subdir: null

TensorFlow2/Segmentation/Contrib/UNet3P/data_generators/README.md

@@ -0,0 +1,44 @@
+Our code base support two types of data loaders.
+
+- [Tensorflow Sequence Generator](#tensorflow-sequence-generator)
+- [NVIDIA DALI Generator](#nvidia-dali-generator)
+
+## [Tensorflow Sequence Generator](https://www.tensorflow.org/api_docs/python/tf/keras/utils/Sequence)
+
+Sequence data generator is best suited for situations where we need
+advanced control over sample generation or when simple data does not
+fit into memory and must be loaded dynamically.
+
+Our [sequence generator](./../data_generators/tf_data_generator.py) generates
+dataset on multiple cores in real time and feed it right away to deep
+learning model.
+
+## [NVIDIA DALI Generator](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/index.html)
+
+The NVIDIA Data Loading Library (DALI) is a library for data loading and
+pre-processing to accelerate deep learning applications. It provides a
+collection of highly optimized building blocks for loading and processing
+image, video and audio data. It can be used as a portable drop-in
+replacement for built in data loaders and data iterators in popular deep
+learning frameworks.
+
+We've used [DALI Pipeline](./../data_generators/dali_data_generator.py) to directly load
+data on `GPU`, which resulted in reduced latency and training time,
+mitigating bottlenecks, by overlapping training and pre-processing. Our code
+base also support's multi GPU data loading for DALI.
+
+## Use Cases
+
+For training and evaluation you can use both  `TF Sequence` and `DALI` generator with multiple gpus, but for prediction
+and inference benchmark we only support `TF Sequence` generator with single gpu support.
+
+> Reminder: DALI is only supported on Linux platforms. For Windows, you can
+> train using Sequence Generator. The code base will work without DALI
+> installation too.
+
+It's advised to use DALI only when you have large gpu memory to load both model
+and training data at the same time.
+
+Override `DATA_GENERATOR_TYPE` in config to change default generator type. Possible
+options are `TF_GENERATOR` for Sequence generator and `DALI_GENERATOR` for DALI generator.
+ 

TensorFlow2/Segmentation/Contrib/UNet3P/data_generators/dali_data_generator.py

@@ -0,0 +1,264 @@
+"""
+NVIDIA DALI data generator object.
+"""
+import nvidia.dali.fn as fn
+from nvidia.dali import pipeline_def
+import nvidia.dali.types as types
+import nvidia.dali.plugin.tf as dali_tf
+import tensorflow as tf
+from omegaconf import DictConfig
+
+from utils.general_utils import get_data_paths, get_gpus_count
+
+
+def data_generator_pipeline(cfg: DictConfig, mode: str, mask_available: bool):
+    """
+    Returns DALI data pipeline object.
+    """
+    data_paths = get_data_paths(cfg, mode, mask_available)  # get data paths
+    images_paths = data_paths[0]
+    if mask_available:
+        mask_paths = data_paths[1]
+
+    @pipeline_def(batch_size=cfg.HYPER_PARAMETERS.BATCH_SIZE)
+    def single_gpu_pipeline(device):
+        """
+        Returns DALI data pipeline object for single GPU training.
+        """
+        device = 'mixed' if 'gpu' in device.lower() else 'cpu'
+
+        pngs, _ = fn.readers.file(
+            files=images_paths,
+            random_shuffle=cfg.PREPROCESS_DATA.SHUFFLE[mode].VALUE,
+            seed=cfg.SEED
+        )
+        images = fn.decoders.image(pngs, device=device, output_type=types.RGB)
+        if cfg.PREPROCESS_DATA.RESIZE.VALUE:
+            # TODO verify image resizing method
+            images = fn.resize(
+                images,
+                size=[
+                    cfg.PREPROCESS_DATA.RESIZE.HEIGHT,
+                    cfg.PREPROCESS_DATA.RESIZE.WIDTH
+                ]
+            )
+        if cfg.PREPROCESS_DATA.IMAGE_PREPROCESSING_TYPE == "normalize":
+            images = fn.normalize(images, mean=0, stddev=255, )  # axes=(2,)
+
+        if mask_available:
+            labels, _ = fn.readers.file(
+                files=mask_paths,
+                random_shuffle=cfg.PREPROCESS_DATA.SHUFFLE[mode].VALUE,
+                seed=cfg.SEED
+            )
+            labels = fn.decoders.image(
+                labels,
+                device=device,
+                output_type=types.GRAY
+            )
+            if cfg.PREPROCESS_DATA.RESIZE.VALUE:
+                # TODO verify image resizing method
+                labels = fn.resize(
+                    labels,
+                    size=[
+                        cfg.PREPROCESS_DATA.RESIZE.HEIGHT,
+                        cfg.PREPROCESS_DATA.RESIZE.WIDTH
+                    ]
+                )
+            if cfg.PREPROCESS_DATA.NORMALIZE_MASK.VALUE:
+                labels = fn.normalize(
+                    labels,
+                    mean=0,
+                    stddev=cfg.PREPROCESS_DATA.NORMALIZE_MASK.NORMALIZE_VALUE,
+                )
+            if cfg.OUTPUT.CLASSES == 1:
+                labels = fn.cast(labels, dtype=types.FLOAT)
+            else:
+                labels = fn.squeeze(labels, axes=[2])
+                labels = fn.one_hot(labels, num_classes=cfg.OUTPUT.CLASSES)
+
+        if mask_available:
+            return images, labels
+        else:
+            return images,
+
+    @pipeline_def(batch_size=cfg.HYPER_PARAMETERS.BATCH_SIZE)
+    def multi_gpu_pipeline(device, shard_id):
+        """
+        Returns DALI data pipeline object for multi GPU'S training.
+        """
+        device = 'mixed' if 'gpu' in device.lower() else 'cpu'
+        shard_id = 1 if 'cpu' in device else shard_id
+        num_shards = get_gpus_count()
+        # num_shards should be <= #images
+        num_shards = len(images_paths) if num_shards > len(images_paths) else num_shards
+
+        pngs, _ = fn.readers.file(
+            files=images_paths,
+            random_shuffle=cfg.PREPROCESS_DATA.SHUFFLE[mode].VALUE,
+            shard_id=shard_id,
+            num_shards=num_shards,
+            seed=cfg.SEED
+        )
+        images = fn.decoders.image(pngs, device=device, output_type=types.RGB)
+        if cfg.PREPROCESS_DATA.RESIZE.VALUE:
+            # TODO verify image resizing method
+            images = fn.resize(
+                images,
+                size=[
+                    cfg.PREPROCESS_DATA.RESIZE.HEIGHT,
+                    cfg.PREPROCESS_DATA.RESIZE.WIDTH
+                ]
+            )
+        if cfg.PREPROCESS_DATA.IMAGE_PREPROCESSING_TYPE == "normalize":
+            images = fn.normalize(images, mean=0, stddev=255, )  # axes=(2,)
+
+        if mask_available:
+            labels, _ = fn.readers.file(
+                files=mask_paths,
+                random_shuffle=cfg.PREPROCESS_DATA.SHUFFLE[mode].VALUE,
+                shard_id=shard_id,
+                num_shards=num_shards,
+                seed=cfg.SEED
+            )
+            labels = fn.decoders.image(
+                labels,
+                device=device,
+                output_type=types.GRAY
+            )
+            if cfg.PREPROCESS_DATA.RESIZE.VALUE:
+                # TODO verify image resizing method
+                labels = fn.resize(
+                    labels,
+                    size=[
+                        cfg.PREPROCESS_DATA.RESIZE.HEIGHT,
+                        cfg.PREPROCESS_DATA.RESIZE.WIDTH
+                    ]
+                )
+            if cfg.PREPROCESS_DATA.NORMALIZE_MASK.VALUE:
+                labels = fn.normalize(
+                    labels,
+                    mean=0,
+                    stddev=cfg.PREPROCESS_DATA.NORMALIZE_MASK.NORMALIZE_VALUE,
+                )
+            if cfg.OUTPUT.CLASSES == 1:
+                labels = fn.cast(labels, dtype=types.FLOAT)
+            else:
+                labels = fn.squeeze(labels, axes=[2])
+                labels = fn.one_hot(labels, num_classes=cfg.OUTPUT.CLASSES)
+
+        if mask_available:
+            return images, labels
+        else:
+            return images,
+
+    if cfg.USE_MULTI_GPUS.VALUE:
+        return multi_gpu_pipeline
+    else:
+        return single_gpu_pipeline
+
+
+def get_data_shapes(cfg: DictConfig, mask_available: bool):
+    """
+    Returns shapes and dtypes of the outputs.
+    """
+    if mask_available:
+        shapes = (
+            (cfg.HYPER_PARAMETERS.BATCH_SIZE,
+             cfg.INPUT.HEIGHT,
+             cfg.INPUT.WIDTH,
+             cfg.INPUT.CHANNELS),
+            (cfg.HYPER_PARAMETERS.BATCH_SIZE,
+             cfg.INPUT.HEIGHT,
+             cfg.INPUT.WIDTH,
+             cfg.OUTPUT.CLASSES)
+        )
+        dtypes = (
+            tf.float32,
+            tf.float32)
+    else:
+        shapes = (
+            (cfg.HYPER_PARAMETERS.BATCH_SIZE,
+             cfg.INPUT.HEIGHT,
+             cfg.INPUT.WIDTH,
+             cfg.INPUT.CHANNELS),
+        )
+        dtypes = (
+            tf.float32,
+        )
+    return shapes, dtypes
+
+
+def data_generator(cfg: DictConfig,
+                   mode: str,
+                   strategy: tf.distribute.Strategy = None):
+    """
+    Generate batches of data for model by reading images and their
+    corresponding masks using NVIDIA DALI.
+    Works for both single and mult GPU's. In case of multi gpu pass
+    the strategy object too.
+    There are two options you can either pass directory path or list.
+    In case of directory, it should contain relative path of images/mask
+    folder from project root path.
+    In case of list of images, every element should contain absolute path
+    for each image and mask.
+    """
+
+    # check mask are available or not
+    mask_available = False if cfg.DATASET[mode].MASK_PATH is None or str(
+        cfg.DATASET[mode].MASK_PATH).lower() == "none" else True
+
+    # create dali data pipeline
+    data_pipeline = data_generator_pipeline(cfg, mode, mask_available)
+
+    shapes, dtypes = get_data_shapes(cfg, mask_available)
+
+    if cfg.USE_MULTI_GPUS.VALUE:
+        def bound_dataset(input_context):
+            """
+            In case of multi gpu training bound dataset to a device for distributed training.
+            """
+            with tf.device("/gpu:{}".format(input_context.input_pipeline_id)):
+                device_id = input_context.input_pipeline_id
+                return dali_tf.DALIDataset(
+                    pipeline=data_pipeline(
+                        device="gpu",
+                        device_id=device_id,
+                        shard_id=device_id,
+                        num_threads=cfg.DATALOADER_WORKERS
+                    ),
+                    batch_size=cfg.HYPER_PARAMETERS.BATCH_SIZE,
+                    output_shapes=shapes,
+                    output_dtypes=dtypes,
+                    device_id=device_id,
+                )
+
+        # distribute dataset
+        input_options = tf.distribute.InputOptions(
+            experimental_place_dataset_on_device=True,
+            # for older dali versions use experimental_prefetch_to_device
+            # for new dali versions use  experimental_fetch_to_device
+            experimental_fetch_to_device=False,  # experimental_fetch_to_device
+            experimental_replication_mode=tf.distribute.InputReplicationMode.PER_REPLICA)
+
+        # map dataset to given strategy and return it
+        return strategy.distribute_datasets_from_function(bound_dataset, input_options)
+    else:
+        # single gpu pipeline
+        pipeline = data_pipeline(
+            batch_size=cfg.HYPER_PARAMETERS.BATCH_SIZE,
+            num_threads=cfg.DATALOADER_WORKERS,
+            device="gpu",
+            device_id=0
+        )
+
+        # create dataset
+        with tf.device('/gpu:0'):
+            data_generator = dali_tf.DALIDataset(
+                pipeline=pipeline,
+                batch_size=cfg.HYPER_PARAMETERS.BATCH_SIZE,
+                output_shapes=shapes,
+                output_dtypes=dtypes,
+                device_id=0)
+
+        return data_generator

TensorFlow2/Segmentation/Contrib/UNet3P/data_generators/data_generator.py

@@ -0,0 +1,88 @@
+"""
+Data generator
+"""
+import os
+import tensorflow as tf
+from omegaconf import DictConfig
+
+from utils.general_utils import join_paths, get_gpus_count
+from .tf_data_generator import DataGenerator as tf_data_generator
+
+try:
+    from .dali_data_generator import data_generator as dali_data_generator
+except ModuleNotFoundError:
+    print("NVIDIA DALI not installed, please install it."
+          "\nNote: DALI is only available on Linux platform. For Window "
+          "you can use TensorFlow generator for training.")
+
+
+def get_data_generator(cfg: DictConfig,
+                       mode: str,
+                       strategy: tf.distribute.Strategy = None):
+    """
+    Creates and return data generator object based on given type.
+    """
+    if cfg.DATA_GENERATOR_TYPE == "TF_GENERATOR":
+        print(f"Using TensorFlow generator for {mode} data")
+        generator = tf_data_generator(cfg, mode)
+    elif cfg.DATA_GENERATOR_TYPE == "DALI_GENERATOR":
+        print(f"Using NVIDIA DALI generator for {mode} data")
+        if cfg.USE_MULTI_GPUS.VALUE:
+            generator = dali_data_generator(cfg, mode, strategy)
+        else:
+            generator = dali_data_generator(cfg, mode)
+    else:
+        raise ValueError(
+            "Wrong generator type passed."
+            "\nPossible options are TF_GENERATOR and DALI_GENERATOR"
+        )
+    return generator
+
+
+def update_batch_size(cfg: DictConfig):
+    """
+    Scale up batch size to multi gpus in case of TensorFlow generator.
+    """
+    if cfg.DATA_GENERATOR_TYPE == "TF_GENERATOR" and cfg.USE_MULTI_GPUS.VALUE:
+        # change batch size according to available gpus
+        cfg.HYPER_PARAMETERS.BATCH_SIZE = \
+            cfg.HYPER_PARAMETERS.BATCH_SIZE * get_gpus_count()
+
+
+def get_batch_size(cfg: DictConfig):
+    """
+    Return batch size.
+    In case of DALI generator scale up batch size to multi gpus.
+    """
+    if cfg.DATA_GENERATOR_TYPE == "DALI_GENERATOR" and cfg.USE_MULTI_GPUS.VALUE:
+        # change batch size according to available gpus
+        return cfg.HYPER_PARAMETERS.BATCH_SIZE * get_gpus_count()
+    else:
+        return cfg.HYPER_PARAMETERS.BATCH_SIZE
+
+
+def get_iterations(cfg: DictConfig, mode: str):
+    """
+    Return steps per epoch
+    """
+    images_length = len(
+        os.listdir(
+            join_paths(
+                cfg.WORK_DIR,
+                cfg.DATASET[mode].IMAGES_PATH
+            )
+        )
+    )
+
+    if cfg.DATA_GENERATOR_TYPE == "TF_GENERATOR":
+        training_steps = images_length // cfg.HYPER_PARAMETERS.BATCH_SIZE
+    elif cfg.DATA_GENERATOR_TYPE == "DALI_GENERATOR":
+        if cfg.USE_MULTI_GPUS.VALUE:
+            training_steps = images_length // (
+                    cfg.HYPER_PARAMETERS.BATCH_SIZE * get_gpus_count())
+        else:
+            training_steps = images_length // cfg.HYPER_PARAMETERS.BATCH_SIZE
+    else:
+        raise ValueError("Wrong generator type passed.")
+
+    return training_steps

TensorFlow2/Segmentation/Contrib/UNet3P/data_generators/tf_data_generator.py

@@ -0,0 +1,179 @@
+"""
+Tensorflow data generator class.
+"""
+import tensorflow as tf
+import numpy as np
+from omegaconf import DictConfig
+
+from utils.general_utils import get_data_paths
+from utils.images_utils import prepare_image, prepare_mask
+
+
+class DataGenerator(tf.keras.utils.Sequence):
+    """
+    Generate batches of data for model by reading images and their
+    corresponding masks using TensorFlow Sequence Generator.
+    There are two options you can either pass directory path or list.
+    In case of directory, it should contain relative path of images/mask
+    folder from project root path.
+    In case of list of images, every element should contain absolute path
+    for each image and mask.
+    Because this generator is also used for prediction, so during testing you can
+    set mask path to None if mask are not available for visualization.
+    """
+
+    def __init__(self, cfg: DictConfig, mode: str):
+        """
+        Initialization
+        """
+        self.cfg = cfg
+        self.mode = mode
+        self.batch_size = self.cfg.HYPER_PARAMETERS.BATCH_SIZE
+        # set seed for reproducibility
+        np.random.seed(cfg.SEED)
+
+        # check mask are available or not
+        self.mask_available = False if cfg.DATASET[mode].MASK_PATH is None or str(
+            cfg.DATASET[mode].MASK_PATH).lower() == "none" else True
+
+        data_paths = get_data_paths(cfg, mode, self.mask_available)
+
+        self.images_paths = data_paths[0]
+        if self.mask_available:
+            self.mask_paths = data_paths[1]
+
+        # self.images_paths.sort()  # no need for sorting
+
+        self.on_epoch_end()
+
+    def __len__(self):
+        """
+        Denotes the number of batches per epoch
+        """
+        # Tensorflow problem: on_epoch_end is not being called at the end
+        # of each epoch, so forcing on_epoch_end call
+        self.on_epoch_end()
+        return int(
+            np.floor(
+                len(self.images_paths) / self.batch_size
+            )
+        )
+
+    def on_epoch_end(self):
+        """
+        Updates indexes after each epoch
+        """
+        self.indexes = np.arange(len(self.images_paths))
+        if self.cfg.PREPROCESS_DATA.SHUFFLE[self.mode].VALUE:
+            np.random.shuffle(self.indexes)
+
+    def __getitem__(self, index):
+        """
+        Generate one batch of data
+        """
+        # Generate indexes of the batch
+        indexes = self.indexes[
+                  index * self.batch_size:(index + 1) * self.batch_size
+                  ]
+
+        # Generate data
+        return self.__data_generation(indexes)
+
+    def __data_generation(self, indexes):
+        """
+        Generates batch data
+        """
+
+        # create empty array to store batch data
+        batch_images = np.zeros(
+            (
+                self.cfg.HYPER_PARAMETERS.BATCH_SIZE,
+                self.cfg.INPUT.HEIGHT,
+                self.cfg.INPUT.WIDTH,
+                self.cfg.INPUT.CHANNELS
+            )
+        ).astype(np.float32)
+
+        if self.mask_available:
+            batch_masks = np.zeros(
+                (
+                    self.cfg.HYPER_PARAMETERS.BATCH_SIZE,
+                    self.cfg.INPUT.HEIGHT,
+                    self.cfg.INPUT.WIDTH,
+                    self.cfg.OUTPUT.CLASSES
+                )
+            ).astype(np.float32)
+
+        for i, index in enumerate(indexes):
+            # extract path from list
+            img_path = self.images_paths[int(index)]
+            if self.mask_available:
+                mask_path = self.mask_paths[int(index)]
+
+            # prepare image for model by resizing and preprocessing it
+            image = prepare_image(
+                img_path,
+                self.cfg.PREPROCESS_DATA.RESIZE,
+                self.cfg.PREPROCESS_DATA.IMAGE_PREPROCESSING_TYPE,
+            )
+
+            if self.mask_available:
+                # prepare image for model by resizing and preprocessing it
+                mask = prepare_mask(
+                    mask_path,
+                    self.cfg.PREPROCESS_DATA.RESIZE,
+                    self.cfg.PREPROCESS_DATA.NORMALIZE_MASK,
+                )
+
+            # numpy to tensorflow conversion
+            if self.mask_available:
+                image, mask = tf.numpy_function(
+                    self.tf_func,
+                    [image, mask],
+                    [tf.float32, tf.int32]
+                )
+            else:
+                image = tf.numpy_function(
+                    self.tf_func,
+                    [image, ],
+                    [tf.float32, ]
+                )
+
+            # set shape attributes which was lost during Tf conversion
+            image.set_shape(
+                [
+                    self.cfg.INPUT.HEIGHT,
+                    self.cfg.INPUT.WIDTH,
+                    self.cfg.INPUT.CHANNELS
+                ]
+            )
+            batch_images[i] = image
+
+            if self.mask_available:
+                # height x width --> height x width x output classes
+                if self.cfg.OUTPUT.CLASSES == 1:
+                    mask = tf.expand_dims(mask, axis=-1)
+                else:
+                    # convert mask into one hot vectors
+                    mask = tf.one_hot(
+                        mask,
+                        self.cfg.OUTPUT.CLASSES,
+                        dtype=tf.int32
+                    )
+                mask.set_shape(
+                    [
+                        self.cfg.INPUT.HEIGHT,
+                        self.cfg.INPUT.WIDTH,
+                        self.cfg.OUTPUT.CLASSES
+                    ]
+                )
+                batch_masks[i] = mask
+
+        if self.mask_available:
+            return batch_images, batch_masks
+        else:
+            return batch_images,
+
+    @staticmethod
+    def tf_func(*args):
+        return args

TensorFlow2/Segmentation/Contrib/UNet3P/data_preparation/README.md

@@ -0,0 +1,102 @@
+For data two options are available
+
+- [Train on LiTS Data](#lits-liver-tumor-segmentation-challenge)
+- [Train on custom data](#train-on-custom-data)
+
+## LiTS Liver Tumor Segmentation challenge
+
+This dataset consist of 131 Liver CT Scans.
+
+Register [here](https://competitions.codalab.org/competitions/17094) to get dataset access.
+Go to participate → Training Data to get dataset link.
+Download Training Batch 1 and Training Batch 2 zip files and past them under data folder.
+
+`Training Batch 1` size is 3.97GB and `Training Batch 2` zip file size is 11.5GB.
+
+Inside main directory `/workspace/unet3p` run below command to extract zip files
+
+```shell
+bash data_preparation/extract_data.sh
+```
+
+After extraction `Training Batch 1` folder size will be 11.4GB and `Training Batch 2` folder size will be 38.5GB.
+
+- `Training Batch 1` consist of 28 scans which are used for testing
+- `Training Batch 2` consist of 103 scans which are used for training
+
+Default directory structure looks like this
+
+    ├── data/
+    │   ├── Training Batch 1/
+            ├── segmentation-0.nii
+            ├── volume-0.nii
+            ├── ...
+            ├── volume-27.nii
+    │   ├── Training Batch 2/
+            ├── segmentation-28.nii
+            ├── volume-28.nii
+            ├── ...
+            ├── volume-130.nii
+
+For testing, you can have any number of files in Training Batch 1 and Training Batch 2. But make sure the naming
+convention is similar.
+
+To prepare LiTS dataset for training run
+
+```
+python data_preparation/preprocess_data.py
+```
+
+> Note: Because of the extensive preprocessing, it will take some time, so relax and wait.
+
+#### Final directory
+
+After completion, you will have a directories like this
+
+    ├── data/
+    │   ├── train/
+            ├── images
+                ├── image_28_0.png
+                ├── ...
+            ├── mask
+                ├── mask_28_0.png
+                ├── ...
+    │   ├── val/
+            ├── images
+                ├── image_0_0.png
+                ├── ...
+            ├── mask
+                ├── mask_0_0.png
+                ├── ...
+
+After processing the `train` folder size will be 5GB and `val` folder size will be 1.7GB.
+
+#### Free space (Optional)
+
+At this stage you can delete the intermediate scans files to free space, run below command
+
+```shell
+bash data_preparation/delete_extracted_scans_data.sh
+```
+
+You can also delete the data zip files using below command, but remember you cannot retrieve them back
+
+```shell
+bash data_preparation/delete_zip_data.sh
+```
+
+> Note: It is recommended to delete scan files but not zip data because you may need it again.
+
+## Train on custom data
+
+To train on custom dateset it's advised that you follow the same train and val directory structure like
+mentioned [above](#final-directory).
+
+In our case image file name can be mapped to it's corresponding mask file name by replacing `image` text with `mask`. If
+your data has different mapping then you need to update [image_to_mask_name](./../utils/images_utils.py#L63) function which
+is responsible for converting image name to it's corresponding file name.
+
+Each image should be a color image with 3 channels and `RGB` color format. Each mask is considered as a gray scale
+image, where each pixel value is the class on which each pixel belongs.
+
+Congratulations, now you can start training and testing on your new dataset!

TensorFlow2/Segmentation/Contrib/UNet3P/data_preparation/delete_extracted_scans_data.sh

@@ -0,0 +1,2 @@
+rm -r 'data/Training Batch 1/'
+rm -r 'data/Training Batch 2/'

TensorFlow2/Segmentation/Contrib/UNet3P/data_preparation/delete_zip_data.sh

@@ -0,0 +1,2 @@
+rm data/Training_Batch1.zip 
+rm data/Training_Batch2.zip

TensorFlow2/Segmentation/Contrib/UNet3P/data_preparation/extract_data.sh

@@ -0,0 +1,9 @@
+# extract testing data
+unzip data/Training_Batch1.zip -d data/
+mv  "data/media/nas/01_Datasets/CT/LITS/Training Batch 1/" "data/Training Batch 1/"
+rm -r data/media
+
+# extract training data
+unzip data/Training_Batch2.zip -d data/
+mv  "data/media/nas/01_Datasets/CT/LITS/Training Batch 2/" "data/Training Batch 2/"
+rm -r data/media

TensorFlow2/Segmentation/Contrib/UNet3P/data_preparation/preprocess_data.py

@@ -0,0 +1,242 @@
+"""
+Convert LiTS 2017 (Liver Tumor Segmentation) data into UNet3+ data format
+LiTS: https://competitions.codalab.org/competitions/17094
+"""
+import os
+import sys
+from glob import glob
+from pathlib import Path
+from tqdm import tqdm
+import numpy as np
+import multiprocessing as mp
+import cv2
+import nibabel as nib
+import hydra
+from omegaconf import DictConfig
+
+sys.path.append(os.path.abspath("./"))
+from utils.general_utils import create_directory, join_paths
+from utils.images_utils import resize_image
+
+
+def read_nii(filepath):
+    """
+    Reads .nii file and returns pixel array
+    """
+    ct_scan = nib.load(filepath).get_fdata()
+    # TODO: Verify images orientation
+    # in both train and test set, especially on train scan 130
+    ct_scan = np.rot90(np.array(ct_scan))
+    return ct_scan
+
+
+def crop_center(img, croph, cropw):
+    """
+    Center crop on given height and width
+    """
+    height, width = img.shape[:2]
+    starth = height // 2 - (croph // 2)
+    startw = width // 2 - (cropw // 2)
+    return img[starth:starth + croph, startw:startw + cropw, :]
+
+
+def linear_scale(img):
+    """
+    First convert image to range of 0-1 and them scale to 255
+    """
+    img = (img - img.min(axis=(0, 1))) / (img.max(axis=(0, 1)) - img.min(axis=(0, 1)))
+    return img * 255
+
+
+def clip_scan(img, min_value, max_value):
+    """
+    Clip scan to given range
+    """
+    return np.clip(img, min_value, max_value)
+
+
+def resize_scan(scan, new_height, new_width, scan_type):
+    """
+    Resize CT scan to given size
+    """
+    scan_shape = scan.shape
+    resized_scan = np.zeros((new_height, new_width, scan_shape[2]), dtype=scan.dtype)
+    resize_method = cv2.INTER_CUBIC if scan_type == "image" else cv2.INTER_NEAREST
+    for start in range(0, scan_shape[2], scan_shape[1]):
+        end = start + scan_shape[1]
+        if end >= scan_shape[2]: end = scan_shape[2]
+        resized_scan[:, :, start:end] = resize_image(
+            scan[:, :, start:end],
+            new_height, new_width,
+            resize_method
+        )
+
+    return resized_scan
+
+
+def save_images(scan, save_path, img_index):
+    """
+    Based on UNet3+ requirement "input image had three channels, including
+    the slice to be segmented and the upper and lower slices, which was
+    cropped to 320×320" save each scan as separate image with previous and
+    next scan concatenated.
+    """
+    scan_shape = scan.shape
+    for index in range(scan_shape[-1]):
+        before_index = index - 1 if (index - 1) > 0 else 0
+        after_index = index + 1 if (index + 1) < scan_shape[-1] else scan_shape[-1] - 1
+
+        new_img_path = join_paths(save_path, f"image_{img_index}_{index}.png")
+        new_image = np.stack(
+            (
+                scan[:, :, before_index],
+                scan[:, :, index],
+                scan[:, :, after_index]
+            )
+            , axis=-1)
+        new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)  # RGB to BGR
+        cv2.imwrite(new_img_path, new_image)  # save the images as .png
+
+
+def save_mask(scan, save_path, mask_index):
+    """
+    Save each scan as separate mask
+    """
+    for index in range(scan.shape[-1]):
+        new_mask_path = join_paths(save_path, f"mask_{mask_index}_{index}.png")
+        cv2.imwrite(new_mask_path, scan[:, :, index])  # save grey scale image
+
+
+def extract_image(cfg, image_path, save_path, scan_type="image", ):
+    """
+    Extract image from given scan path
+    """
+    _, index = str(Path(image_path).stem).split("-")
+
+    scan = read_nii(image_path)
+    scan = resize_scan(
+        scan,
+        cfg.DATA_PREPARATION.RESIZED_HEIGHT,
+        cfg.DATA_PREPARATION.RESIZED_WIDTH,
+        scan_type
+    )
+    if scan_type == "image":
+        scan = clip_scan(
+            scan,
+            cfg.DATA_PREPARATION.SCAN_MIN_VALUE,
+            cfg.DATA_PREPARATION.SCAN_MAX_VALUE
+        )
+        scan = linear_scale(scan)
+        scan = np.uint8(scan)
+        save_images(scan, save_path, index)
+    else:
+        # 0 for background/non-lesion, 1 for liver, 2 for lesion/tumor
+        # merging label 2 into label 1, because lesion/tumor is part of liver
+        scan = np.where(scan != 0, 1, scan)
+        # scan = np.where(scan==2, 1, scan)
+        scan = np.uint8(scan)
+        save_mask(scan, save_path, index)
+
+
+def extract_images(cfg, images_path, save_path, scan_type="image", ):
+    """
+    Extract images paths using multiprocessing and pass to
+    extract_image function for further processing .
+    """
+    # create pool
+    process_count = np.clip(mp.cpu_count() - 2, 1, 20)  # less than 20 workers
+    pool = mp.Pool(process_count)
+    for image_path in tqdm(images_path):
+        pool.apply_async(extract_image,
+                         args=(cfg, image_path, save_path, scan_type),
+                         )
+
+    # close pool
+    pool.close()
+    pool.join()
+
+
[email protected](version_base=None, config_path="../configs", config_name="config")
+def preprocess_lits_data(cfg: DictConfig):
+    """
+    Preprocess LiTS 2017 (Liver Tumor Segmentation) data by extractions
+    images and mask into UNet3+ data format
+    """
+    train_images_names = glob(
+        join_paths(
+            cfg.WORK_DIR,
+            cfg.DATA_PREPARATION.SCANS_TRAIN_DATA_PATH,
+            "volume-*.nii"
+        )
+    )
+    train_mask_names = glob(
+        join_paths(
+            cfg.WORK_DIR,
+            cfg.DATA_PREPARATION.SCANS_TRAIN_DATA_PATH,
+            "segmentation-*.nii"
+        )
+    )
+
+    assert len(train_images_names) == len(train_mask_names), \
+        "Train volumes and segmentations are not same in length"
+
+    val_images_names = glob(
+        join_paths(
+            cfg.WORK_DIR,
+            cfg.DATA_PREPARATION.SCANS_VAL_DATA_PATH,
+            "volume-*.nii"
+        )
+    )
+    val_mask_names = glob(
+        join_paths(
+            cfg.WORK_DIR,
+            cfg.DATA_PREPARATION.SCANS_VAL_DATA_PATH,
+            "segmentation-*.nii"
+        )
+    )
+    assert len(val_images_names) == len(val_mask_names), \
+        "Validation volumes and segmentations are not same in length"
+
+    train_images_names = sorted(train_images_names)
+    train_mask_names = sorted(train_mask_names)
+    val_images_names = sorted(val_images_names)
+    val_mask_names = sorted(val_mask_names)
+
+    train_images_path = join_paths(
+        cfg.WORK_DIR, cfg.DATASET.TRAIN.IMAGES_PATH
+    )
+    train_mask_path = join_paths(
+        cfg.WORK_DIR, cfg.DATASET.TRAIN.MASK_PATH
+    )
+    val_images_path = join_paths(
+        cfg.WORK_DIR, cfg.DATASET.VAL.IMAGES_PATH
+    )
+    val_mask_path = join_paths(
+        cfg.WORK_DIR, cfg.DATASET.VAL.MASK_PATH
+    )
+
+    create_directory(train_images_path)
+    create_directory(train_mask_path)
+    create_directory(val_images_path)
+    create_directory(val_mask_path)
+
+    print("\nExtracting train images")
+    extract_images(
+        cfg, train_images_names, train_images_path, scan_type="image"
+    )
+    print("\nExtracting train mask")
+    extract_images(
+        cfg, train_mask_names, train_mask_path, scan_type="mask"
+    )
+    print("\nExtracting val images")
+    extract_images(
+        cfg, val_images_names, val_images_path, scan_type="image"
+    )
+    print("\nExtracting val mask")
+    extract_images(
+        cfg, val_mask_names, val_mask_path, scan_type="mask"
+    )
+
+
+if __name__ == '__main__':
+    preprocess_lits_data()

TensorFlow2/Segmentation/Contrib/UNet3P/data_preparation/verify_data.py

@@ -0,0 +1,56 @@
+"""
+Verify for each image corresponding mask exist or not.
+Check against both train and val data
+"""
+import os
+import sys
+from omegaconf import DictConfig
+from tqdm import tqdm
+
+sys.path.append(os.path.abspath("./"))
+from utils.general_utils import join_paths
+from utils.images_utils import image_to_mask_name
+
+
+def check_image_and_mask(cfg, mode):
+    """
+    Check and print names of those images whose mask are not found.
+    """
+    images_path = join_paths(
+        cfg.WORK_DIR,
+        cfg.DATASET[mode].IMAGES_PATH
+    )
+    mask_path = join_paths(
+        cfg.WORK_DIR,
+        cfg.DATASET[mode].MASK_PATH
+    )
+
+    all_images = os.listdir(images_path)
+
+    both_found = True
+    for image in tqdm(all_images):
+        mask_name = image_to_mask_name(image)
+        if not (
+                os.path.exists(
+                    join_paths(images_path, image)
+                ) and
+                os.path.exists(
+                    join_paths(mask_path, mask_name)
+                )
+        ):
+            print(f"{mask_name} did not found against {image}")
+            both_found = False
+
+    return both_found
+
+
+def verify_data(cfg: DictConfig):
+    """
+    For both train and val data, check for each image its
+    corresponding mask exist or not. If not then stop the program.
+    """
+    assert check_image_and_mask(cfg, "TRAIN"), \
+        "Train images and mask should be same in length"
+
+    assert check_image_and_mask(cfg, "VAL"), \
+        "Validation images and mask should be same in length"

TensorFlow2/Segmentation/Contrib/UNet3P/evaluate.py

@@ -0,0 +1,125 @@
+"""
+Evaluation script used to calculate accuracy of trained model
+"""
+import os
+import hydra
+from omegaconf import DictConfig
+import tensorflow as tf
+from tensorflow.keras import mixed_precision
+
+from data_generators import data_generator
+from utils.general_utils import join_paths, set_gpus, suppress_warnings
+from models.model import prepare_model
+from losses.loss import DiceCoefficient
+from losses.unet_loss import unet3p_hybrid_loss
+
+
+def evaluate(cfg: DictConfig):
+    """
+    Evaluate or calculate accuracy of given model
+    """
+
+    # suppress TensorFlow and DALI warnings
+    suppress_warnings()
+
+    if cfg.USE_MULTI_GPUS.VALUE:
+        # change number of visible gpus for evaluation
+        set_gpus(cfg.USE_MULTI_GPUS.GPU_IDS)
+        # update batch size according to available gpus
+        data_generator.update_batch_size(cfg)
+
+    if cfg.OPTIMIZATION.AMP:
+        print("Enabling Automatic Mixed Precision(AMP) training")
+        policy = mixed_precision.Policy('mixed_float16')
+        mixed_precision.set_global_policy(policy)
+
+    if cfg.OPTIMIZATION.XLA:
+        print("Enabling Automatic Mixed Precision(XLA) training")
+        tf.config.optimizer.set_jit(True)
+
+    # create model
+    strategy = None
+    if cfg.USE_MULTI_GPUS.VALUE:
+        # multi gpu training using tensorflow mirrored strategy
+        strategy = tf.distribute.MirroredStrategy(
+            cross_device_ops=tf.distribute.HierarchicalCopyAllReduce()
+        )
+        print('Number of visible gpu devices: {}'.format(strategy.num_replicas_in_sync))
+        with strategy.scope():
+            optimizer = tf.keras.optimizers.Adam(
+                learning_rate=cfg.HYPER_PARAMETERS.LEARNING_RATE
+            )  # optimizer
+            if cfg.OPTIMIZATION.AMP:
+                optimizer = mixed_precision.LossScaleOptimizer(
+                    optimizer,
+                    dynamic=True
+                )
+            dice_coef = DiceCoefficient(post_processed=True, classes=cfg.OUTPUT.CLASSES)
+            dice_coef = tf.keras.metrics.MeanMetricWrapper(name="dice_coef", fn=dice_coef)
+            model = prepare_model(cfg, training=True)
+    else:
+        optimizer = tf.keras.optimizers.Adam(
+            learning_rate=cfg.HYPER_PARAMETERS.LEARNING_RATE
+        )  # optimizer
+        if cfg.OPTIMIZATION.AMP:
+            optimizer = mixed_precision.LossScaleOptimizer(
+                optimizer,
+                dynamic=True
+            )
+        dice_coef = DiceCoefficient(post_processed=True, classes=cfg.OUTPUT.CLASSES)
+        dice_coef = tf.keras.metrics.MeanMetricWrapper(name="dice_coef", fn=dice_coef)
+        model = prepare_model(cfg, training=True)
+
+    model.compile(
+        optimizer=optimizer,
+        loss=unet3p_hybrid_loss,
+        metrics=[dice_coef],
+    )
+
+    # weights model path
+    checkpoint_path = join_paths(
+        cfg.WORK_DIR,
+        cfg.CALLBACKS.MODEL_CHECKPOINT.PATH,
+        f"{cfg.MODEL.WEIGHTS_FILE_NAME}.hdf5"
+    )
+
+    assert os.path.exists(checkpoint_path), \
+        f"Model weight's file does not exist at \n{checkpoint_path}"
+
+    # TODO: verify without augment it produces same results
+    # load model weights
+    model.load_weights(checkpoint_path, by_name=True, skip_mismatch=True)
+    model.summary()
+
+    # data generators
+    val_generator = data_generator.get_data_generator(cfg, "VAL", strategy)
+    validation_steps = data_generator.get_iterations(cfg, mode="VAL")
+
+    # evaluation metric
+    evaluation_metric = "dice_coef"
+    if len(model.outputs) > 1:
+        evaluation_metric = f"{model.output_names[0]}_dice_coef"
+
+    result = model.evaluate(
+        x=val_generator,
+        steps=validation_steps,
+        workers=cfg.DATALOADER_WORKERS,
+        return_dict=True,
+    )
+
+    # return computed loss, validation accuracy, and it's metric name
+    return result, evaluation_metric
+
+
[email protected](version_base=None, config_path="configs", config_name="config")
+def main(cfg: DictConfig):
+    """
+    Read config file and pass to evaluate method
+    """
+    result, evaluation_metric = evaluate(cfg)
+    print(result)
+    print(f"Validation dice coefficient: {result[evaluation_metric]}")
+
+
+if __name__ == "__main__":
+    main()

TensorFlow2/Segmentation/Contrib/UNet3P/losses/loss.py

@@ -0,0 +1,114 @@
+"""
+Implementation of different loss functions
+"""
+import tensorflow as tf
+import tensorflow.keras.backend as K
+
+
+def iou(y_true, y_pred, smooth=1.e-9):
+    """
+    Calculate intersection over union (IoU) between images.
+    Input shape should be Batch x Height x Width x #Classes (BxHxWxN).
+    Using Mean as reduction type for batch values.
+    """
+    intersection = K.sum(K.abs(y_true * y_pred), axis=[1, 2, 3])
+    union = K.sum(y_true, [1, 2, 3]) + K.sum(y_pred, [1, 2, 3])
+    union = union - intersection
+    iou = K.mean((intersection + smooth) / (union + smooth), axis=0)
+    return iou
+
+
+def iou_loss(y_true, y_pred):
+    """
+    Jaccard / IoU loss
+    """
+    return 1 - iou(y_true, y_pred)
+
+
+def focal_loss(y_true, y_pred):
+    """
+    Focal loss
+    """
+    gamma = 2.
+    alpha = 4.
+    epsilon = 1.e-9
+
+    y_true_c = tf.convert_to_tensor(y_true, tf.float32)
+    y_pred_c = tf.convert_to_tensor(y_pred, tf.float32)
+
+    model_out = tf.add(y_pred_c, epsilon)
+    ce = tf.multiply(y_true_c, -tf.math.log(model_out))
+    weight = tf.multiply(y_true_c, tf.pow(
+        tf.subtract(1., model_out), gamma)
+                         )
+    fl = tf.multiply(alpha, tf.multiply(weight, ce))
+    reduced_fl = tf.reduce_max(fl, axis=-1)
+    return tf.reduce_mean(reduced_fl)
+
+
+def ssim_loss(y_true, y_pred, smooth=1.e-9):
+    """
+    Structural Similarity Index loss.
+    Input shape should be Batch x Height x Width x #Classes (BxHxWxN).
+    Using Mean as reduction type for batch values.
+    """
+    ssim_value = tf.image.ssim(y_true, y_pred, max_val=1)
+    return K.mean(1 - ssim_value + smooth, axis=0)
+
+
+class DiceCoefficient(tf.keras.metrics.Metric):
+    """
+    Dice coefficient metric. Can be used to calculate dice on probabilities
+    or on their respective classes
+    """
+
+    def __init__(self, post_processed: bool,
+                 classes: int,
+                 name='dice_coef',
+                 **kwargs):
+        """
+        Set post_processed=False if dice coefficient needs to be calculated
+        on probabilities. Set post_processed=True if probabilities needs to
+        be first converted/mapped into their respective class.
+        """
+        super(DiceCoefficient, self).__init__(name=name, **kwargs)
+        self.dice_value = self.add_weight(name='dice_value', initializer='zeros',
+                                          aggregation=tf.VariableAggregation.MEAN)  # SUM
+        self.post_processed = post_processed
+        self.classes = classes
+        if self.classes == 1:
+            self.axis = [1, 2, 3]
+        else:
+            self.axis = [1, 2, ]
+
+    def update_state(self, y_true, y_pred, sample_weight=None):
+        if self.post_processed:
+            if self.classes == 1:
+                y_true_ = y_true
+                y_pred_ = tf.where(y_pred > .5, 1.0, 0.0)
+            else:
+                y_true_ = tf.math.argmax(y_true, axis=-1, output_type=tf.int32)
+                y_pred_ = tf.math.argmax(y_pred, axis=-1, output_type=tf.int32)
+                y_true_ = tf.cast(y_true_, dtype=tf.float32)
+                y_pred_ = tf.cast(y_pred_, dtype=tf.float32)
+        else:
+            y_true_, y_pred_ = y_true, y_pred
+
+        self.dice_value.assign(self.dice_coef(y_true_, y_pred_))
+
+    def result(self):
+        return self.dice_value
+
+    def reset_state(self):
+        self.dice_value.assign(0.0)  # reset metric state
+
+    def dice_coef(self, y_true, y_pred, smooth=1.e-9):
+        """
+        Calculate dice coefficient.
+        Input shape could be either Batch x Height x Width x #Classes (BxHxWxN)
+        or Batch x Height x Width (BxHxW).
+        Using Mean as reduction type for batch values.
+        """
+        intersection = K.sum(y_true * y_pred, axis=self.axis)
+        union = K.sum(y_true, axis=self.axis) + K.sum(y_pred, axis=self.axis)
+        return K.mean((2. * intersection + smooth) / (union + smooth), axis=0)

TensorFlow2/Segmentation/Contrib/UNet3P/losses/unet_loss.py

@@ -0,0 +1,19 @@
+"""
+UNet 3+ Loss
+"""
+from .loss import focal_loss, ssim_loss, iou_loss
+
+
+def unet3p_hybrid_loss(y_true, y_pred):
+    """
+    Hybrid loss proposed in
+    UNET 3+ (https://arxiv.org/ftp/arxiv/papers/2004/2004.08790.pdf)
+    Hybrid loss for segmentation in three-level hierarchy – pixel,
+    patch and map-level, which is able to capture both large-scale
+    and fine structures with clear boundaries.
+    """
+    f_loss = focal_loss(y_true, y_pred)
+    ms_ssim_loss = ssim_loss(y_true, y_pred)
+    jacard_loss = iou_loss(y_true, y_pred)
+
+    return f_loss + ms_ssim_loss + jacard_loss

TensorFlow2/Segmentation/Contrib/UNet3P/models/backbones.py

@@ -0,0 +1,73 @@
+"""
+Unet3+ backbones
+"""
+import tensorflow as tf
+import tensorflow.keras as k
+from .unet3plus_utils import conv_block
+
+
+def vgg16_backbone(input_layer, ):
+    """ VGG-16 backbone as encoder for UNet3P """
+
+    base_model = tf.keras.applications.VGG16(
+        input_tensor=input_layer,
+        weights=None,
+        include_top=False
+    )
+
+    # block 1
+    e1 = base_model.get_layer("block1_conv2").output  # 320, 320, 64
+    # block 2
+    e2 = base_model.get_layer("block2_conv2").output  # 160, 160, 128
+    # block 3
+    e3 = base_model.get_layer("block3_conv3").output  # 80, 80, 256
+    # block 4
+    e4 = base_model.get_layer("block4_conv3").output  # 40, 40, 512
+    # block 5
+    e5 = base_model.get_layer("block5_conv3").output  # 20, 20, 512
+
+    return [e1, e2, e3, e4, e5]
+
+
+def vgg19_backbone(input_layer, ):
+    """ VGG-19 backbone as encoder for UNet3P """
+
+    base_model = tf.keras.applications.VGG19(
+        input_tensor=input_layer,
+        weights=None,
+        include_top=False
+    )
+
+    # block 1
+    e1 = base_model.get_layer("block1_conv2").output  # 320, 320, 64
+    # block 2
+    e2 = base_model.get_layer("block2_conv2").output  # 160, 160, 128
+    # block 3
+    e3 = base_model.get_layer("block3_conv4").output  # 80, 80, 256
+    # block 4
+    e4 = base_model.get_layer("block4_conv4").output  # 40, 40, 512
+    # block 5
+    e5 = base_model.get_layer("block5_conv4").output  # 20, 20, 512
+
+    return [e1, e2, e3, e4, e5]
+
+
+def unet3plus_backbone(input_layer, filters):
+    """ UNet3+ own backbone """
+    """ Encoder"""
+    # block 1
+    e1 = conv_block(input_layer, filters[0])  # 320*320*64
+    # block 2
+    e2 = k.layers.MaxPool2D(pool_size=(2, 2))(e1)  # 160*160*64
+    e2 = conv_block(e2, filters[1])  # 160*160*128
+    # block 3
+    e3 = k.layers.MaxPool2D(pool_size=(2, 2))(e2)  # 80*80*128
+    e3 = conv_block(e3, filters[2])  # 80*80*256
+    # block 4
+    e4 = k.layers.MaxPool2D(pool_size=(2, 2))(e3)  # 40*40*256
+    e4 = conv_block(e4, filters[3])  # 40*40*512
+    # block 5, bottleneck layer
+    e5 = k.layers.MaxPool2D(pool_size=(2, 2))(e4)  # 20*20*512
+    e5 = conv_block(e5, filters[4])  # 20*20*1024
+
+    return [e1, e2, e3, e4, e5]

TensorFlow2/Segmentation/Contrib/UNet3P/models/model.py

@@ -0,0 +1,100 @@
+"""
+Returns Unet3+ model
+"""
+import tensorflow as tf
+from omegaconf import DictConfig
+
+from .backbones import vgg16_backbone, vgg19_backbone, unet3plus_backbone
+from .unet3plus import unet3plus
+from .unet3plus_deep_supervision import unet3plus_deepsup
+from .unet3plus_deep_supervision_cgm import unet3plus_deepsup_cgm
+
+
+def prepare_model(cfg: DictConfig, training=False):
+    """
+    Creates and return model object based on given model type.
+    """
+
+    input_shape = [cfg.INPUT.HEIGHT, cfg.INPUT.WIDTH, cfg.INPUT.CHANNELS]
+    input_layer = tf.keras.layers.Input(
+        shape=input_shape,
+        name="input_layer"
+    )  # 320*320*3
+    filters = [64, 128, 256, 512, 1024]
+
+    #  create backbone
+    if cfg.MODEL.BACKBONE.TYPE == "unet3plus":
+        backbone_layers = unet3plus_backbone(
+            input_layer,
+            filters
+        )
+    elif cfg.MODEL.BACKBONE.TYPE == "vgg16":
+        backbone_layers = vgg16_backbone(input_layer, )
+    elif cfg.MODEL.BACKBONE.TYPE == "vgg19":
+        backbone_layers = vgg19_backbone(input_layer, )
+    else:
+        raise ValueError(
+            "Wrong backbone type passed."
+            "\nPlease check config file for possible options."
+        )
+    print(f"Using {cfg.MODEL.BACKBONE.TYPE} as a backbone.")
+
+    if cfg.MODEL.TYPE == "unet3plus":
+        #  training parameter does not matter in this case
+        outputs, model_name = unet3plus(
+            backbone_layers,
+            cfg.OUTPUT.CLASSES,
+            filters
+        )
+    elif cfg.MODEL.TYPE == "unet3plus_deepsup":
+        outputs, model_name = unet3plus_deepsup(
+            backbone_layers,
+            cfg.OUTPUT.CLASSES,
+            filters,
+            training
+        )
+    elif cfg.MODEL.TYPE == "unet3plus_deepsup_cgm":
+        if cfg.OUTPUT.CLASSES != 1:
+            raise ValueError(
+                "UNet3+ with Deep Supervision and Classification Guided Module"
+                "\nOnly works when model output classes are equal to 1"
+            )
+        outputs, model_name = unet3plus_deepsup_cgm(
+            backbone_layers,
+            cfg.OUTPUT.CLASSES,
+            filters,
+            training
+        )
+    else:
+        raise ValueError(
+            "Wrong model type passed."
+            "\nPlease check config file for possible options."
+        )
+
+    return tf.keras.Model(
+        inputs=input_layer,
+        outputs=outputs,
+        name=model_name
+    )
+
+
+if __name__ == "__main__":
+    """## Test model Compilation,"""
+    from omegaconf import OmegaConf
+
+    cfg = {
+        "WORK_DIR": "H:\\Projects\\UNet3P",
+        "INPUT": {"HEIGHT": 320, "WIDTH": 320, "CHANNELS": 3},
+        "OUTPUT": {"CLASSES": 1},
+        # available variants are unet3plus, unet3plus_deepsup, unet3plus_deepsup_cgm
+        "MODEL": {"TYPE": "unet3plus",
+                  # available variants are unet3plus, vgg16, vgg19
+                  "BACKBONE": {"TYPE": "vgg19", }
+                  }
+    }
+    unet_3P = prepare_model(OmegaConf.create(cfg), True)
+    unet_3P.summary()
+
+    # tf.keras.utils.plot_model(unet_3P, show_layer_names=True, show_shapes=True)
+
+    # unet_3P.save("unet_3P.hdf5")

TensorFlow2/Segmentation/Contrib/UNet3P/models/unet3plus.py

@@ -0,0 +1,104 @@
+"""
+UNet3+ base model
+"""
+import tensorflow as tf
+import tensorflow.keras as k
+from .unet3plus_utils import conv_block
+
+
+def unet3plus(encoder_layer, output_channels, filters):
+    """ UNet3+ base model """
+
+    """ Encoder """
+    e1 = encoder_layer[0]
+    e2 = encoder_layer[1]
+    e3 = encoder_layer[2]
+    e4 = encoder_layer[3]
+    e5 = encoder_layer[4]
+
+    """ Decoder """
+    cat_channels = filters[0]
+    cat_blocks = len(filters)
+    upsample_channels = cat_blocks * cat_channels
+
+    """ d4 """
+    e1_d4 = k.layers.MaxPool2D(pool_size=(8, 8))(e1)  # 320*320*64  --> 40*40*64
+    e1_d4 = conv_block(e1_d4, cat_channels, n=1)  # 320*320*64  --> 40*40*64
+
+    e2_d4 = k.layers.MaxPool2D(pool_size=(4, 4))(e2)  # 160*160*128 --> 40*40*128
+    e2_d4 = conv_block(e2_d4, cat_channels, n=1)  # 160*160*128 --> 40*40*64
+
+    e3_d4 = k.layers.MaxPool2D(pool_size=(2, 2))(e3)  # 80*80*256  --> 40*40*256
+    e3_d4 = conv_block(e3_d4, cat_channels, n=1)  # 80*80*256  --> 40*40*64
+
+    e4_d4 = conv_block(e4, cat_channels, n=1)  # 40*40*512  --> 40*40*64
+
+    e5_d4 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(e5)  # 80*80*256  --> 40*40*256
+    e5_d4 = conv_block(e5_d4, cat_channels, n=1)  # 20*20*1024  --> 20*20*64
+
+    d4 = k.layers.concatenate([e1_d4, e2_d4, e3_d4, e4_d4, e5_d4])
+    d4 = conv_block(d4, upsample_channels, n=1)  # 40*40*320  --> 40*40*320
+
+    """ d3 """
+    e1_d3 = k.layers.MaxPool2D(pool_size=(4, 4))(e1)  # 320*320*64 --> 80*80*64
+    e1_d3 = conv_block(e1_d3, cat_channels, n=1)  # 80*80*64 --> 80*80*64
+
+    e2_d3 = k.layers.MaxPool2D(pool_size=(2, 2))(e2)  # 160*160*256 --> 80*80*256
+    e2_d3 = conv_block(e2_d3, cat_channels, n=1)  # 80*80*256 --> 80*80*64
+
+    e3_d3 = conv_block(e3, cat_channels, n=1)  # 80*80*512 --> 80*80*64
+
+    e4_d3 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d4)  # 40*40*320 --> 80*80*320
+    e4_d3 = conv_block(e4_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64
+
+    e5_d3 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(e5)  # 20*20*320 --> 80*80*320
+    e5_d3 = conv_block(e5_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64
+
+    d3 = k.layers.concatenate([e1_d3, e2_d3, e3_d3, e4_d3, e5_d3])
+    d3 = conv_block(d3, upsample_channels, n=1)  # 80*80*320 --> 80*80*320
+
+    """ d2 """
+    e1_d2 = k.layers.MaxPool2D(pool_size=(2, 2))(e1)  # 320*320*64 --> 160*160*64
+    e1_d2 = conv_block(e1_d2, cat_channels, n=1)  # 160*160*64 --> 160*160*64
+
+    e2_d2 = conv_block(e2, cat_channels, n=1)  # 160*160*256 --> 160*160*64
+
+    d3_d2 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d3)  # 80*80*320 --> 160*160*320
+    d3_d2 = conv_block(d3_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    d4_d2 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d4)  # 40*40*320 --> 160*160*320
+    d4_d2 = conv_block(d4_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    e5_d2 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(e5)  # 20*20*320 --> 160*160*320
+    e5_d2 = conv_block(e5_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    d2 = k.layers.concatenate([e1_d2, e2_d2, d3_d2, d4_d2, e5_d2])
+    d2 = conv_block(d2, upsample_channels, n=1)  # 160*160*320 --> 160*160*320
+
+    """ d1 """
+    e1_d1 = conv_block(e1, cat_channels, n=1)  # 320*320*64 --> 320*320*64
+
+    d2_d1 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d2)  # 160*160*320 --> 320*320*320
+    d2_d1 = conv_block(d2_d1, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    d3_d1 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d3)  # 80*80*320 --> 320*320*320
+    d3_d1 = conv_block(d3_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64
+
+    d4_d1 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(d4)  # 40*40*320 --> 320*320*320
+    d4_d1 = conv_block(d4_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64
+
+    e5_d1 = k.layers.UpSampling2D(size=(16, 16), interpolation='bilinear')(e5)  # 20*20*320 --> 320*320*320
+    e5_d1 = conv_block(e5_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64
+
+    d1 = k.layers.concatenate([e1_d1, d2_d1, d3_d1, d4_d1, e5_d1, ])
+    d1 = conv_block(d1, upsample_channels, n=1)  # 320*320*320 --> 320*320*320
+
+    # last layer does not have batchnorm and relu
+    d = conv_block(d1, output_channels, n=1, is_bn=False, is_relu=False)
+
+    if output_channels == 1:
+        output = k.layers.Activation('sigmoid', dtype='float32')(d)
+    else:
+        output = k.layers.Activation('softmax', dtype='float32')(d)
+
+    return output, 'UNet_3Plus'

TensorFlow2/Segmentation/Contrib/UNet3P/models/unet3plus_deep_supervision.py

@@ -0,0 +1,132 @@
+"""
+UNet3+ with Deep Supervision
+"""
+import tensorflow as tf
+import tensorflow.keras as k
+from .unet3plus_utils import conv_block
+
+
+def unet3plus_deepsup(encoder_layer, output_channels, filters, training=False):
+    """ UNet_3Plus with Deep Supervision """
+
+    """ Encoder """
+    e1 = encoder_layer[0]
+    e2 = encoder_layer[1]
+    e3 = encoder_layer[2]
+    e4 = encoder_layer[3]
+    e5 = encoder_layer[4]
+
+    """ Decoder """
+    cat_channels = filters[0]
+    cat_blocks = len(filters)
+    upsample_channels = cat_blocks * cat_channels
+
+    """ d4 """
+    e1_d4 = k.layers.MaxPool2D(pool_size=(8, 8))(e1)  # 320*320*64  --> 40*40*64
+    e1_d4 = conv_block(e1_d4, cat_channels, n=1)  # 320*320*64  --> 40*40*64
+
+    e2_d4 = k.layers.MaxPool2D(pool_size=(4, 4))(e2)  # 160*160*128 --> 40*40*128
+    e2_d4 = conv_block(e2_d4, cat_channels, n=1)  # 160*160*128 --> 40*40*64
+
+    e3_d4 = k.layers.MaxPool2D(pool_size=(2, 2))(e3)  # 80*80*256  --> 40*40*256
+    e3_d4 = conv_block(e3_d4, cat_channels, n=1)  # 80*80*256  --> 40*40*64
+
+    e4_d4 = conv_block(e4, cat_channels, n=1)  # 40*40*512  --> 40*40*64
+
+    e5_d4 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(e5)  # 80*80*256  --> 40*40*256
+    e5_d4 = conv_block(e5_d4, cat_channels, n=1)  # 20*20*1024  --> 20*20*64
+
+    d4 = k.layers.concatenate([e1_d4, e2_d4, e3_d4, e4_d4, e5_d4])
+    d4 = conv_block(d4, upsample_channels, n=1)  # 40*40*320  --> 40*40*320
+
+    """ d3 """
+    e1_d3 = k.layers.MaxPool2D(pool_size=(4, 4))(e1)  # 320*320*64 --> 80*80*64
+    e1_d3 = conv_block(e1_d3, cat_channels, n=1)  # 80*80*64 --> 80*80*64
+
+    e2_d3 = k.layers.MaxPool2D(pool_size=(2, 2))(e2)  # 160*160*256 --> 80*80*256
+    e2_d3 = conv_block(e2_d3, cat_channels, n=1)  # 80*80*256 --> 80*80*64
+
+    e3_d3 = conv_block(e3, cat_channels, n=1)  # 80*80*512 --> 80*80*64
+
+    e4_d3 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d4)  # 40*40*320 --> 80*80*320
+    e4_d3 = conv_block(e4_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64
+
+    e5_d3 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(e5)  # 20*20*320 --> 80*80*320
+    e5_d3 = conv_block(e5_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64
+
+    d3 = k.layers.concatenate([e1_d3, e2_d3, e3_d3, e4_d3, e5_d3])
+    d3 = conv_block(d3, upsample_channels, n=1)  # 80*80*320 --> 80*80*320
+
+    """ d2 """
+    e1_d2 = k.layers.MaxPool2D(pool_size=(2, 2))(e1)  # 320*320*64 --> 160*160*64
+    e1_d2 = conv_block(e1_d2, cat_channels, n=1)  # 160*160*64 --> 160*160*64
+
+    e2_d2 = conv_block(e2, cat_channels, n=1)  # 160*160*256 --> 160*160*64
+
+    d3_d2 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d3)  # 80*80*320 --> 160*160*320
+    d3_d2 = conv_block(d3_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    d4_d2 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d4)  # 40*40*320 --> 160*160*320
+    d4_d2 = conv_block(d4_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    e5_d2 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(e5)  # 20*20*320 --> 160*160*320
+    e5_d2 = conv_block(e5_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    d2 = k.layers.concatenate([e1_d2, e2_d2, d3_d2, d4_d2, e5_d2])
+    d2 = conv_block(d2, upsample_channels, n=1)  # 160*160*320 --> 160*160*320
+
+    """ d1 """
+    e1_d1 = conv_block(e1, cat_channels, n=1)  # 320*320*64 --> 320*320*64
+
+    d2_d1 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d2)  # 160*160*320 --> 320*320*320
+    d2_d1 = conv_block(d2_d1, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    d3_d1 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d3)  # 80*80*320 --> 320*320*320
+    d3_d1 = conv_block(d3_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64
+
+    d4_d1 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(d4)  # 40*40*320 --> 320*320*320
+    d4_d1 = conv_block(d4_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64
+
+    e5_d1 = k.layers.UpSampling2D(size=(16, 16), interpolation='bilinear')(e5)  # 20*20*320 --> 320*320*320
+    e5_d1 = conv_block(e5_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64
+
+    d1 = k.layers.concatenate([e1_d1, d2_d1, d3_d1, d4_d1, e5_d1, ])
+    d1 = conv_block(d1, upsample_channels, n=1)  # 320*320*320 --> 320*320*320
+
+    # last layer does not have batch norm and relu
+    d1 = conv_block(d1, output_channels, n=1, is_bn=False, is_relu=False)
+
+    if output_channels == 1:
+        d1 = k.layers.Activation('sigmoid', dtype='float32')(d1)
+    else:
+        # d1 = k.activations.softmax(d1)
+        d1 = k.layers.Activation('softmax', dtype='float32')(d1)
+
+    """ Deep Supervision Part"""
+    if training:
+        d2 = conv_block(d2, output_channels, n=1, is_bn=False, is_relu=False)
+        d3 = conv_block(d3, output_channels, n=1, is_bn=False, is_relu=False)
+        d4 = conv_block(d4, output_channels, n=1, is_bn=False, is_relu=False)
+        e5 = conv_block(e5, output_channels, n=1, is_bn=False, is_relu=False)
+
+        # d1 = no need for up sampling
+        d2 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d2)
+        d3 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d3)
+        d4 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(d4)
+        e5 = k.layers.UpSampling2D(size=(16, 16), interpolation='bilinear')(e5)
+
+        if output_channels == 1:
+            d2 = k.layers.Activation('sigmoid', dtype='float32')(d2)
+            d3 = k.layers.Activation('sigmoid', dtype='float32')(d3)
+            d4 = k.layers.Activation('sigmoid', dtype='float32')(d4)
+            e5 = k.layers.Activation('sigmoid', dtype='float32')(e5)
+        else:
+            d2 = k.layers.Activation('softmax', dtype='float32')(d2)
+            d3 = k.layers.Activation('softmax', dtype='float32')(d3)
+            d4 = k.layers.Activation('softmax', dtype='float32')(d4)
+            e5 = k.layers.Activation('softmax', dtype='float32')(e5)
+
+    if training:
+        return [d1, d2, d3, d4, e5], 'UNet3Plus_DeepSup'
+    else:
+        return [d1, ], 'UNet3Plus_DeepSup'

TensorFlow2/Segmentation/Contrib/UNet3P/models/unet3plus_deep_supervision_cgm.py

@@ -0,0 +1,138 @@
+"""
+UNet_3Plus with Deep Supervision and Classification Guided Module
+"""
+import tensorflow as tf
+import tensorflow.keras as k
+from .unet3plus_utils import conv_block, dot_product
+
+
+def unet3plus_deepsup_cgm(encoder_layer, output_channels, filters, training=False):
+    """ UNet_3Plus with Deep Supervision and Classification Guided Module """
+
+    """ Encoder """
+    e1 = encoder_layer[0]
+    e2 = encoder_layer[1]
+    e3 = encoder_layer[2]
+    e4 = encoder_layer[3]
+    e5 = encoder_layer[4]
+
+    """ Classification Guided Module. Part 1"""
+    cls = k.layers.Dropout(rate=0.5)(e5)
+    cls = k.layers.Conv2D(2, kernel_size=(1, 1), padding="same", strides=(1, 1))(cls)
+    cls = k.layers.GlobalMaxPooling2D()(cls)
+    cls = k.layers.Activation('sigmoid', dtype='float32')(cls)
+    cls = tf.argmax(cls, axis=-1)
+    cls = cls[..., tf.newaxis]
+    cls = tf.cast(cls, dtype=tf.float32, )
+
+    """ Decoder """
+    cat_channels = filters[0]
+    cat_blocks = len(filters)
+    upsample_channels = cat_blocks * cat_channels
+
+    """ d4 """
+    e1_d4 = k.layers.MaxPool2D(pool_size=(8, 8))(e1)  # 320*320*64  --> 40*40*64
+    e1_d4 = conv_block(e1_d4, cat_channels, n=1)  # 320*320*64  --> 40*40*64
+
+    e2_d4 = k.layers.MaxPool2D(pool_size=(4, 4))(e2)  # 160*160*128 --> 40*40*128
+    e2_d4 = conv_block(e2_d4, cat_channels, n=1)  # 160*160*128 --> 40*40*64
+
+    e3_d4 = k.layers.MaxPool2D(pool_size=(2, 2))(e3)  # 80*80*256  --> 40*40*256
+    e3_d4 = conv_block(e3_d4, cat_channels, n=1)  # 80*80*256  --> 40*40*64
+
+    e4_d4 = conv_block(e4, cat_channels, n=1)  # 40*40*512  --> 40*40*64
+
+    e5_d4 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(e5)  # 80*80*256  --> 40*40*256
+    e5_d4 = conv_block(e5_d4, cat_channels, n=1)  # 20*20*1024  --> 20*20*64
+
+    d4 = k.layers.concatenate([e1_d4, e2_d4, e3_d4, e4_d4, e5_d4])
+    d4 = conv_block(d4, upsample_channels, n=1)  # 40*40*320  --> 40*40*320
+
+    """ d3 """
+    e1_d3 = k.layers.MaxPool2D(pool_size=(4, 4))(e1)  # 320*320*64 --> 80*80*64
+    e1_d3 = conv_block(e1_d3, cat_channels, n=1)  # 80*80*64 --> 80*80*64
+
+    e2_d3 = k.layers.MaxPool2D(pool_size=(2, 2))(e2)  # 160*160*256 --> 80*80*256
+    e2_d3 = conv_block(e2_d3, cat_channels, n=1)  # 80*80*256 --> 80*80*64
+
+    e3_d3 = conv_block(e3, cat_channels, n=1)  # 80*80*512 --> 80*80*64
+
+    e4_d3 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d4)  # 40*40*320 --> 80*80*320
+    e4_d3 = conv_block(e4_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64
+
+    e5_d3 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(e5)  # 20*20*320 --> 80*80*320
+    e5_d3 = conv_block(e5_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64
+
+    d3 = k.layers.concatenate([e1_d3, e2_d3, e3_d3, e4_d3, e5_d3])
+    d3 = conv_block(d3, upsample_channels, n=1)  # 80*80*320 --> 80*80*320
+
+    """ d2 """
+    e1_d2 = k.layers.MaxPool2D(pool_size=(2, 2))(e1)  # 320*320*64 --> 160*160*64
+    e1_d2 = conv_block(e1_d2, cat_channels, n=1)  # 160*160*64 --> 160*160*64
+
+    e2_d2 = conv_block(e2, cat_channels, n=1)  # 160*160*256 --> 160*160*64
+
+    d3_d2 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d3)  # 80*80*320 --> 160*160*320
+    d3_d2 = conv_block(d3_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    d4_d2 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d4)  # 40*40*320 --> 160*160*320
+    d4_d2 = conv_block(d4_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    e5_d2 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(e5)  # 20*20*320 --> 160*160*320
+    e5_d2 = conv_block(e5_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    d2 = k.layers.concatenate([e1_d2, e2_d2, d3_d2, d4_d2, e5_d2])
+    d2 = conv_block(d2, upsample_channels, n=1)  # 160*160*320 --> 160*160*320
+
+    """ d1 """
+    e1_d1 = conv_block(e1, cat_channels, n=1)  # 320*320*64 --> 320*320*64
+
+    d2_d1 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d2)  # 160*160*320 --> 320*320*320
+    d2_d1 = conv_block(d2_d1, cat_channels, n=1)  # 160*160*320 --> 160*160*64
+
+    d3_d1 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d3)  # 80*80*320 --> 320*320*320
+    d3_d1 = conv_block(d3_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64
+
+    d4_d1 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(d4)  # 40*40*320 --> 320*320*320
+    d4_d1 = conv_block(d4_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64
+
+    e5_d1 = k.layers.UpSampling2D(size=(16, 16), interpolation='bilinear')(e5)  # 20*20*320 --> 320*320*320
+    e5_d1 = conv_block(e5_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64
+
+    d1 = k.layers.concatenate([e1_d1, d2_d1, d3_d1, d4_d1, e5_d1, ])
+    d1 = conv_block(d1, upsample_channels, n=1)  # 320*320*320 --> 320*320*320
+
+    """ Deep Supervision Part"""
+    # last layer does not have batch norm and relu
+    d1 = conv_block(d1, output_channels, n=1, is_bn=False, is_relu=False)
+    if training:
+        d2 = conv_block(d2, output_channels, n=1, is_bn=False, is_relu=False)
+        d3 = conv_block(d3, output_channels, n=1, is_bn=False, is_relu=False)
+        d4 = conv_block(d4, output_channels, n=1, is_bn=False, is_relu=False)
+        e5 = conv_block(e5, output_channels, n=1, is_bn=False, is_relu=False)
+
+        # d1 = no need for up sampling
+        d2 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d2)
+        d3 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d3)
+        d4 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(d4)
+        e5 = k.layers.UpSampling2D(size=(16, 16), interpolation='bilinear')(e5)
+
+    """ Classification Guided Module. Part 2"""
+    d1 = dot_product(d1, cls)
+    d1 = k.layers.Activation('sigmoid', dtype='float32')(d1)
+
+    if training:
+        d2 = dot_product(d2, cls)
+        d3 = dot_product(d3, cls)
+        d4 = dot_product(d4, cls)
+        e5 = dot_product(e5, cls)
+
+        d2 = k.layers.Activation('sigmoid', dtype='float32')(d2)
+        d3 = k.layers.Activation('sigmoid', dtype='float32')(d3)
+        d4 = k.layers.Activation('sigmoid', dtype='float32')(d4)
+        e5 = k.layers.Activation('sigmoid', dtype='float32')(e5)
+
+    if training:
+        return [d1, d2, d3, d4, e5, cls], 'UNet3Plus_DeepSup_CGM'
+    else:
+        return [d1, ], 'UNet3Plus_DeepSup_CGM'

TensorFlow2/Segmentation/Contrib/UNet3P/models/unet3plus_utils.py

@@ -0,0 +1,31 @@
+"""
+Utility functions for Unet3+ models
+"""
+import tensorflow as tf
+import tensorflow.keras as k
+
+
+def conv_block(x, kernels, kernel_size=(3, 3), strides=(1, 1), padding='same',
+               is_bn=True, is_relu=True, n=2):
+    """ Custom function for conv2d:
+        Apply  3*3 convolutions with BN and relu.
+    """
+    for i in range(1, n + 1):
+        x = k.layers.Conv2D(filters=kernels, kernel_size=kernel_size,
+                            padding=padding, strides=strides,
+                            kernel_regularizer=tf.keras.regularizers.l2(1e-4),
+                            kernel_initializer=k.initializers.he_normal(seed=5))(x)
+        if is_bn:
+            x = k.layers.BatchNormalization()(x)
+        if is_relu:
+            x = k.activations.relu(x)
+
+    return x
+
+
+def dot_product(seg, cls):
+    b, h, w, n = k.backend.int_shape(seg)
+    seg = tf.reshape(seg, [-1, h * w, n])
+    final = tf.einsum("ijk,ik->ijk", seg, cls)
+    final = tf.reshape(final, [-1, h, w, n])
+    return final

TensorFlow2/Segmentation/Contrib/UNet3P/predict.py

@@ -0,0 +1,101 @@
+"""
+Prediction script used to visualize model output
+"""
+import os
+import hydra
+from omegaconf import DictConfig
+
+from data_generators import tf_data_generator
+from utils.general_utils import join_paths, suppress_warnings
+from utils.images_utils import display
+from utils.images_utils import postprocess_mask, denormalize_mask
+from models.model import prepare_model
+
+
+def predict(cfg: DictConfig):
+    """
+    Predict and visualize given data
+    """
+
+    # suppress TensorFlow and DALI warnings
+    suppress_warnings()
+
+    # set batch size to one
+    cfg.HYPER_PARAMETERS.BATCH_SIZE = 1
+
+    # data generator
+    val_generator = tf_data_generator.DataGenerator(cfg, mode="VAL")
+
+    # create model
+    model = prepare_model(cfg)
+
+    # weights model path
+    checkpoint_path = join_paths(
+        cfg.WORK_DIR,
+        cfg.CALLBACKS.MODEL_CHECKPOINT.PATH,
+        f"{cfg.MODEL.WEIGHTS_FILE_NAME}.hdf5"
+    )
+
+    assert os.path.exists(checkpoint_path), \
+        f"Model weight's file does not exist at \n{checkpoint_path}"
+
+    # load model weights
+    model.load_weights(checkpoint_path, by_name=True, skip_mismatch=True)
+    # model.summary()
+
+    # check mask are available or not
+    mask_available = True
+    if cfg.DATASET.VAL.MASK_PATH is None or \
+            str(cfg.DATASET.VAL.MASK_PATH).lower() == "none":
+        mask_available = False
+
+    showed_images = 0
+    for batch_data in val_generator:  # for each batch
+        batch_images = batch_data[0]
+        if mask_available:
+            batch_mask = batch_data[1]
+
+        # make prediction on batch
+        batch_predictions = model.predict_on_batch(batch_images)
+        if len(model.outputs) > 1:
+            batch_predictions = batch_predictions[0]
+
+        for index in range(len(batch_images)):
+
+            image = batch_images[index]  # for each image
+            if cfg.SHOW_CENTER_CHANNEL_IMAGE:
+                # for UNet3+ show only center channel as image
+                image = image[:, :, 1]
+
+            # do postprocessing on predicted mask
+            prediction = batch_predictions[index]
+            prediction = postprocess_mask(prediction, cfg.OUTPUT.CLASSES)
+            # denormalize mask for better visualization
+            prediction = denormalize_mask(prediction, cfg.OUTPUT.CLASSES)
+
+            if mask_available:
+                mask = batch_mask[index]
+                mask = postprocess_mask(mask, cfg.OUTPUT.CLASSES)
+                mask = denormalize_mask(mask, cfg.OUTPUT.CLASSES)
+
+            # if np.unique(mask).shape[0] == 2:
+            if mask_available:
+                display([image, mask, prediction], show_true_mask=True)
+            else:
+                display([image, prediction], show_true_mask=False)
+
+            showed_images += 1
+        # stop after displaying below number of images
+        # if showed_images >= 10: break
+
+
[email protected](version_base=None, config_path="configs", config_name="config")
+def main(cfg: DictConfig):
+    """
+    Read config file and pass to prediction method
+    """
+    predict(cfg)
+
+
+if __name__ == "__main__":
+    main()

TensorFlow2/Segmentation/Contrib/UNet3P/requirements.txt

@@ -0,0 +1,7 @@
+hydra-core
+opencv-python
+jupyter
+matplotlib
+tqdm
+nibabel
+numba

TensorFlow2/Segmentation/Contrib/UNet3P/train.py

@@ -0,0 +1,215 @@
+"""
+Training script
+"""
+import numpy as np
+from datetime import datetime, timedelta
+import hydra
+from omegaconf import DictConfig
+import tensorflow as tf
+from tensorflow.keras import mixed_precision
+from tensorflow.keras.callbacks import (
+    EarlyStopping,
+    ModelCheckpoint,
+    TensorBoard,
+    CSVLogger
+)
+
+from data_generators import data_generator
+from data_preparation.verify_data import verify_data
+from utils.general_utils import create_directory, join_paths, set_gpus, \
+    suppress_warnings
+from models.model import prepare_model
+from losses.loss import DiceCoefficient
+from losses.unet_loss import unet3p_hybrid_loss
+from callbacks.timing_callback import TimingCallback
+
+
+def create_training_folders(cfg: DictConfig):
+    """
+    Create directories to store Model CheckPoint and TensorBoard logs.
+    """
+    create_directory(
+        join_paths(
+            cfg.WORK_DIR,
+            cfg.CALLBACKS.MODEL_CHECKPOINT.PATH
+        )
+    )
+    create_directory(
+        join_paths(
+            cfg.WORK_DIR,
+            cfg.CALLBACKS.TENSORBOARD.PATH
+        )
+    )
+
+
+def train(cfg: DictConfig):
+    """
+    Training method
+    """
+
+    # suppress TensorFlow and DALI warnings
+    suppress_warnings()
+
+    print("Verifying data ...")
+    verify_data(cfg)
+
+    if cfg.MODEL.TYPE == "unet3plus_deepsup_cgm":
+        raise ValueError(
+            "UNet3+ with Deep Supervision and Classification Guided Module"
+            "\nModel exist but training script is not supported for this variant"
+            "please choose other variants from config file"
+        )
+
+    if cfg.USE_MULTI_GPUS.VALUE:
+        # change number of visible gpus for training
+        set_gpus(cfg.USE_MULTI_GPUS.GPU_IDS)
+        # update batch size according to available gpus
+        data_generator.update_batch_size(cfg)
+
+    # create folders to store training checkpoints and logs
+    create_training_folders(cfg)
+
+    if cfg.OPTIMIZATION.AMP:
+        print("Enabling Automatic Mixed Precision(AMP) training")
+        policy = mixed_precision.Policy('mixed_float16')
+        mixed_precision.set_global_policy(policy)
+
+    if cfg.OPTIMIZATION.XLA:
+        print("Enabling Accelerated Linear Algebra(XLA) training")
+        tf.config.optimizer.set_jit(True)
+
+    # create model
+    strategy = None
+    if cfg.USE_MULTI_GPUS.VALUE:
+        # multi gpu training using tensorflow mirrored strategy
+        strategy = tf.distribute.MirroredStrategy(
+            cross_device_ops=tf.distribute.HierarchicalCopyAllReduce()
+        )
+        print('Number of visible gpu devices: {}'.format(strategy.num_replicas_in_sync))
+        with strategy.scope():
+            optimizer = tf.keras.optimizers.Adam(
+                learning_rate=cfg.HYPER_PARAMETERS.LEARNING_RATE
+            )  # optimizer
+            if cfg.OPTIMIZATION.AMP:
+                optimizer = mixed_precision.LossScaleOptimizer(
+                    optimizer,
+                    dynamic=True
+                )
+            dice_coef = DiceCoefficient(post_processed=True, classes=cfg.OUTPUT.CLASSES)
+            dice_coef = tf.keras.metrics.MeanMetricWrapper(name="dice_coef", fn=dice_coef)
+            model = prepare_model(cfg, training=True)
+    else:
+        optimizer = tf.keras.optimizers.Adam(
+            learning_rate=cfg.HYPER_PARAMETERS.LEARNING_RATE
+        )  # optimizer
+        if cfg.OPTIMIZATION.AMP:
+            optimizer = mixed_precision.LossScaleOptimizer(
+                optimizer,
+                dynamic=True
+            )
+        dice_coef = DiceCoefficient(post_processed=True, classes=cfg.OUTPUT.CLASSES)
+        dice_coef = tf.keras.metrics.MeanMetricWrapper(name="dice_coef", fn=dice_coef)
+        model = prepare_model(cfg, training=True)
+
+    model.compile(
+        optimizer=optimizer,
+        loss=unet3p_hybrid_loss,
+        metrics=[dice_coef],
+    )
+    model.summary()
+
+    # data generators
+    train_generator = data_generator.get_data_generator(cfg, "TRAIN", strategy)
+    val_generator = data_generator.get_data_generator(cfg, "VAL", strategy)
+
+    # verify generator
+    # for i, (batch_images, batch_mask) in enumerate(val_generator):
+    #     print(len(batch_images))
+    #     if i >= 3: break
+
+    # the tensorboard log directory will be a unique subdirectory
+    # based on the start time for the run
+    tb_log_dir = join_paths(
+        cfg.WORK_DIR,
+        cfg.CALLBACKS.TENSORBOARD.PATH,
+        "{}".format(datetime.now().strftime("%Y.%m.%d.%H.%M.%S"))
+    )
+    print("TensorBoard directory\n" + tb_log_dir)
+
+    checkpoint_path = join_paths(
+        cfg.WORK_DIR,
+        cfg.CALLBACKS.MODEL_CHECKPOINT.PATH,
+        f"{cfg.MODEL.WEIGHTS_FILE_NAME}.hdf5"
+    )
+    print("Weights path\n" + checkpoint_path)
+
+    csv_log_path = join_paths(
+        cfg.WORK_DIR,
+        cfg.CALLBACKS.CSV_LOGGER.PATH,
+        f"training_logs_{cfg.MODEL.TYPE}.csv"
+    )
+    print("Logs path\n" + csv_log_path)
+
+    # evaluation metric
+    evaluation_metric = "val_dice_coef"
+    if len(model.outputs) > 1:
+        evaluation_metric = f"val_{model.output_names[0]}_dice_coef"
+
+    # Timing, TensorBoard, EarlyStopping, ModelCheckpoint, CSVLogger callbacks
+    timing_callback = TimingCallback()
+    callbacks = [
+        TensorBoard(log_dir=tb_log_dir, write_graph=False, profile_batch=0),
+        EarlyStopping(
+            patience=cfg.CALLBACKS.EARLY_STOPPING.PATIENCE,
+            verbose=cfg.VERBOSE
+        ),
+        ModelCheckpoint(
+            checkpoint_path,
+            verbose=cfg.VERBOSE,
+            save_weights_only=cfg.CALLBACKS.MODEL_CHECKPOINT.SAVE_WEIGHTS_ONLY,
+            save_best_only=cfg.CALLBACKS.MODEL_CHECKPOINT.SAVE_BEST_ONLY,
+            monitor=evaluation_metric,
+            mode="max"
+
+        ),
+        CSVLogger(
+            csv_log_path,
+            append=cfg.CALLBACKS.CSV_LOGGER.APPEND_LOGS
+        ),
+        timing_callback
+    ]
+
+    training_steps = data_generator.get_iterations(cfg, mode="TRAIN")
+    validation_steps = data_generator.get_iterations(cfg, mode="VAL")
+
+    # start training
+    model.fit(
+        x=train_generator,
+        steps_per_epoch=training_steps,
+        validation_data=val_generator,
+        validation_steps=validation_steps,
+        epochs=cfg.HYPER_PARAMETERS.EPOCHS,
+        callbacks=callbacks,
+        workers=cfg.DATALOADER_WORKERS,
+    )
+
+    training_time = timing_callback.train_end_time - timing_callback.train_start_time
+    training_time = timedelta(seconds=training_time)
+    print(f"Total training time {training_time}")
+
+    mean_time = np.mean(timing_callback.batch_time)
+    throughput = data_generator.get_batch_size(cfg) / mean_time
+    print(f"Training latency: {round(mean_time * 1e3, 2)} msec")
+    print(f"Training throughput/FPS: {round(throughput, 2)} samples/sec")
+
+
[email protected](version_base=None, config_path="configs", config_name="config")
+def main(cfg: DictConfig):
+    """
+    Read config file and pass to train method for training
+    """
+    train(cfg)
+
+
+if __name__ == "__main__":
+    main()

TensorFlow2/Segmentation/Contrib/UNet3P/utils/general_utils.py

@@ -0,0 +1,123 @@
+"""
+General Utility functions
+"""
+import os
+import tensorflow as tf
+from omegaconf import DictConfig
+from .images_utils import image_to_mask_name
+
+
+def create_directory(path):
+    """
+    Create Directory if it already does not exist.
+    """
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def join_paths(*paths):
+    """
+    Concatenate multiple paths.
+    """
+    return os.path.normpath(os.path.sep.join(path.rstrip(r"\/") for path in paths))
+
+
+def set_gpus(gpu_ids):
+    """
+    Change number of visible gpus for tensorflow.
+    gpu_ids: Could be integer or list of integers.
+    In case Integer: if integer value is -1 then use all available gpus.
+    otherwise if positive number, then use given number of gpus.
+    In case list of Integer: each integer will be considered as gpu id
+    """
+    all_gpus = tf.config.experimental.list_physical_devices('GPU')
+    all_gpus_length = len(all_gpus)
+    if isinstance(gpu_ids, int):
+        if gpu_ids == -1:
+            gpu_ids = range(all_gpus_length)
+        else:
+            gpu_ids = min(gpu_ids, all_gpus_length)
+            gpu_ids = range(gpu_ids)
+
+    selected_gpus = [all_gpus[gpu_id] for gpu_id in gpu_ids if gpu_id < all_gpus_length]
+
+    try:
+        tf.config.experimental.set_visible_devices(selected_gpus, 'GPU')
+    except RuntimeError as e:
+        # Visible devices must be set at program startup
+        print(e)
+
+
+def get_gpus_count():
+    """
+    Return length of available gpus.
+    """
+    return len(tf.config.experimental.list_logical_devices('GPU'))
+
+
+def get_data_paths(cfg: DictConfig, mode: str, mask_available: bool):
+    """
+    Return list of absolute images/mask paths.
+    There are two options you can either pass directory path or list.
+    In case of directory, it should contain relative path of images/mask
+    folder from project root path.
+    In case of list of images, every element should contain absolute path
+    for each image and mask.
+    For prediction, you can set mask path to None if mask are not
+    available for visualization.
+    """
+
+    # read images from directory
+    if isinstance(cfg.DATASET[mode].IMAGES_PATH, str):
+        # has only images name not full path
+        images_paths = os.listdir(
+            join_paths(
+                cfg.WORK_DIR,
+                cfg.DATASET[mode].IMAGES_PATH
+            )
+        )
+
+        if mask_available:
+            mask_paths = [
+                image_to_mask_name(image_name) for image_name in images_paths
+            ]
+            # create full mask paths from folder
+            mask_paths = [
+                join_paths(
+                    cfg.WORK_DIR,
+                    cfg.DATASET[mode].MASK_PATH,
+                    mask_name
+                ) for mask_name in mask_paths
+            ]
+
+        # create full images paths from folder
+        images_paths = [
+            join_paths(
+                cfg.WORK_DIR,
+                cfg.DATASET[mode].IMAGES_PATH,
+                image_name
+            ) for image_name in images_paths
+        ]
+    else:
+        # read images and mask from absolute paths given in list
+        images_paths = list(cfg.DATASET[mode].IMAGES_PATH)
+        if mask_available:
+            mask_paths = list(cfg.DATASET[mode].MASK_PATH)
+
+    if mask_available:
+        return images_paths, mask_paths
+    else:
+        return images_paths,
+
+
+def suppress_warnings():
+    """
+    Suppress TensorFlow warnings.
+    """
+    import logging
+    logging.getLogger('tensorflow').setLevel(logging.ERROR)
+    logging.getLogger('dali').setLevel(logging.ERROR)
+    os.environ["KMP_AFFINITY"] = "noverbose"
+    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+    import tensorflow as tf
+    tf.autograph.set_verbosity(3)

TensorFlow2/Segmentation/Contrib/UNet3P/utils/images_utils.py

@@ -0,0 +1,118 @@
+"""
+Utility functions for image processing
+"""
+import numpy as np
+import cv2
+from omegaconf import DictConfig
+import matplotlib.pyplot as plt
+
+
+def read_image(img_path, color_mode):
+    """
+    Read and return image as np array from given path.
+    In case of color image, it returns image in BGR mode.
+    """
+    return cv2.imread(img_path, color_mode)
+
+
+def resize_image(img, height, width, resize_method=cv2.INTER_CUBIC):
+    """
+    Resize image
+    """
+    return cv2.resize(img, dsize=(width, height), interpolation=resize_method)
+
+
+def prepare_image(path: str, resize: DictConfig, normalize_type: str):
+    """
+    Prepare image for model.
+    read image --> resize --> normalize --> return as float32
+    """
+    image = read_image(path, cv2.IMREAD_COLOR)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    if resize.VALUE:
+        # TODO verify image resizing method
+        image = resize_image(image, resize.HEIGHT, resize.WIDTH, cv2.INTER_AREA)
+
+    if normalize_type == "normalize":
+        image = image / 255.0
+
+    image = image.astype(np.float32)
+
+    return image
+
+
+def prepare_mask(path: str, resize: dict, normalize_mask: dict):
+    """
+        Prepare mask for model.
+        read mask --> resize --> normalize --> return as int32
+        """
+    mask = read_image(path, cv2.IMREAD_GRAYSCALE)
+
+    if resize.VALUE:
+        mask = resize_image(mask, resize.HEIGHT, resize.WIDTH, cv2.INTER_NEAREST)
+
+    if normalize_mask.VALUE:
+        mask = mask / normalize_mask.NORMALIZE_VALUE
+
+    mask = mask.astype(np.int32)
+
+    return mask
+
+
+def image_to_mask_name(image_name: str):
+    """
+    Convert image file name to it's corresponding mask file name e.g.
+    image name     -->     mask name
+    image_28_0.png         mask_28_0.png
+    replace image with mask
+    """
+
+    return image_name.replace('image', 'mask')
+
+
+def postprocess_mask(mask, classes, output_type=np.int32):
+    """
+    Post process model output.
+    Covert probabilities into indexes based on maximum value.
+    """
+    if classes == 1:
+        mask = np.where(mask > .5, 1.0, 0.0)
+    else:
+        mask = np.argmax(mask, axis=-1)
+    return mask.astype(output_type)
+
+
+def denormalize_mask(mask, classes):
+    """
+    Denormalize mask by multiplying each class with higher
+    integer (255 / classes) for better visualization.
+    """
+    mask = mask * (255 / classes)
+    return mask.astype(np.int32)
+
+
+def display(display_list, show_true_mask=False):
+    """
+    Show list of images. it could be
+    either [image, true_mask, predicted_mask] or [image, predicted_mask].
+    Set show_true_mask to True if true mask is available or vice versa
+    """
+    if show_true_mask:
+        title_list = ('Input Image', 'True Mask', 'Predicted Mask')
+        plt.figure(figsize=(12, 4))
+    else:
+        title_list = ('Input Image', 'Predicted Mask')
+        plt.figure(figsize=(8, 4))
+
+    for i in range(len(display_list)):
+        plt.subplot(1, len(display_list), i + 1)
+        if title_list is not None:
+            plt.title(title_list[i])
+        if len(np.squeeze(display_list[i]).shape) == 2:
+            plt.imshow(np.squeeze(display_list[i]), cmap='gray')
+            plt.axis('on')
+        else:
+            plt.imshow(np.squeeze(display_list[i]))
+            plt.axis('on')
+    plt.show()