Repository: nasa-nccs-hpda/tensorflow-caney (language: python)

Issue:
0.2.5



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/dockerhub-slbnn.yml b/.github/workflows/dockerhub-slbnn.yml
new file mode 100644
index 0000000..83eb29e
--- /dev/null
+++ b/.github/workflows/dockerhub-slbnn.yml
@@ -0,0 +1,43 @@
+name: CI to Dockerhub Dev
+
+on:
+  push:
+    branches:
+      - 'main'
+
+jobs:
+  docker:
+    runs-on: ubuntu-latest
+    steps:
+      -
+        name: Checkout
+        uses: actions/checkout@v3
+      -
+        name: Set up QEMU
+        uses: docker/setup-qemu-action@v2
+      -
+        name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      -
+        name: Login to Docker Hub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - 
+        name: Lower github-runner storage
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf "/usr/local/share/boost"
+          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+
+      -
+        name: Build and push
+        uses: docker/build-push-action@v4
+        with:
+          context: .
+          file: ./requirements/Dockerfile.SLBNN
+          push: true
+          tags: nasanccs/slbnn:latest
diff --git a/README.md b/README.md
index 1ee0af3..3b50fc5 100644
--- a/README.md
+++ b/README.md
@@ -32,9 +32,18 @@ directory so projects, examples, and notebooks can be run.
 Note: PIP installations do not include CUDA libraries for GPU support. Make sure NVIDIA libraries
 are installed locally in the system if not using conda/mamba.
 
+### Production Container
+
+```bash
+module load singularity
+singularity build --sandbox /lscratch/$USER/container/tensorflow-caney docker://nasanccs/tensorflow-caney:latest
+```
+
+## Development Container
+
 ```bash
 module load singularity
-singularity build --sandbox tensorflow-caney docker://nasanccs/tensorflow-caney:latest
+singularity build --sandbox /lscratch/$USER/container/tensorflow-caney docker://nasanccs/tensorflow-caney:dev
 ```
 
 ## Why Caney?
diff --git a/requirements/Dockerfile b/requirements/Dockerfile
index cc25845..475cdbd 100644
--- a/requirements/Dockerfile
+++ b/requirements/Dockerfile
@@ -8,15 +8,15 @@ FROM ${FROM_IMAGE}:${VERSION_DATE}-${TF_VER}-py3
 
 # Ubuntu needs noninteractive to be forced
 ENV DEBIAN_FRONTEND noninteractive
-ENV LD_LIBRARY_PATH="/mnt/data-store/FileGDB_API-64gcc51/lib/:/usr/local/lib"
 ENV CPLUS_INCLUDE_PATH="/usr/include/gdal"
 ENV C_INCLUDE_PATH="/usr/include/gdal"
+ENV PROJ_LIB="/usr/share/proj"
 ENV SM_FRAMEWORK="tf.keras"
 
 # System dependencies
 RUN apt-get update && \
     apt-get -y install software-properties-common && \
-    add-apt-repository ppa:ubuntugis/ppa && \
+    add-apt-repository ppa:ubuntugis/ubuntugis-unstable && \
     apt-get update && apt-get -y dist-upgrade && \
     apt-get -y install build-essential \
         libsm6 \
@@ -35,7 +35,9 @@ RUN apt-get update && \
         libgeos-dev \
         libsqlite3-dev \
         libx11-dev \
-        libproj-dev proj-data proj-bin \
+        libproj-dev \
+        proj-data \
+        proj-bin \
         libspatialindex-dev \
         wget \
         vim \
@@ -46,20 +48,15 @@ RUN apt-get update && \
         g++ \
         bzip2 \
         libssl-dev \
+        libzmq3-dev \
+        libpng-dev \
+        libfreetype6-dev \ 
         locales && \
     apt-get -y install gdal-bin libgdal-dev && \
     apt-get -y autoremove && \
     rm -rf /var/cache/apt /var/lib/apt/lists/*
 
 
-        libfreetype6-dev \
-        libpng-dev \
-        libzmq3-dev \
-        libspatialindex-dev \
-        gdal-bin \
-        libgdal-dev \
-        python3-gdal \
-
 # Install shiftc
 WORKDIR /app
 RUN git clone --single-branch --branch master https://github.com/pkolano/shift.git && \
@@ -78,12 +75,16 @@ RUN git clone --single-branch --branch master https://github.com/pkolano/shift.g
     rm -rf /app
 
 # Pip
-RUN pip --no-cache-dir install omegaconf \
+RUN pip install \
+    scipy \
+    omegaconf \
     rasterio \
     rioxarray \
     xarray \
     geopandas \
     tensorflow_probability \
+    keras-cv \
+    keras-unet-collection \
     opencv-python \
     opencv-python-headless \
     opencv-contrib-python \
@@ -101,6 +102,7 @@ RUN pip --no-cache-dir install omegaconf \
     rasterstats \
     sphinx \
     sphinx_rtd_theme \
+    geedim \
     geemap \
     GDAL==`ogrinfo --version | grep -Eo '[0-9]\.[0-9]\.[0-9]+'` \
     git+https://github.com/dshean/pygeotools.git@master \
diff --git a/requirements/Dockerfile.SLBNN b/requirements/Dockerfile.SLBNN
new file mode 100644
index 0000000..8450393
--- /dev/null
+++ b/requirements/Dockerfile.SLBNN
@@ -0,0 +1,56 @@
+# Arguments to pass to the image
+ARG TF_VER=tf2
+ARG VERSION_DATE=23.01
+ARG FROM_IMAGE=nvcr.io/nvidia/tensorflow
+
+# Import RAPIDS container as the BASE Image (cuda base image)
+FROM ${FROM_IMAGE}:${VERSION_DATE}-${TF_VER}-py3
+
+# Ubuntu needs noninteractive to be forced
+ENV DEBIAN_FRONTEND noninteractive
+ENV SM_FRAMEWORK="tf.keras"
+
+# System dependencies
+RUN apt-get update && \
+    apt-get -y install software-properties-common && \
+    add-apt-repository ppa:ubuntugis/ppa && \
+    apt-get update && apt-get -y dist-upgrade && \
+    apt-get -y install build-essential \
+        wget \
+        vim \
+        curl \
+        git \
+        procps \
+        gcc \
+        g++ \
+        bzip2 \
+        libssl-dev \
+        locales && \
+    apt-get -y autoremove && \
+    rm -rf /var/cache/apt /var/lib/apt/lists/*
+
+# Pip
+RUN pip --no-cache-dir install omegaconf \
+    xarray \
+    tensorflow_probability \
+    opencv-python \
+    opencv-python-headless \
+    opencv-contrib-python \
+    opencv-contrib-python-headless \
+    tifffile \
+    webcolors \
+    Pillow \
+    seaborn \
+    xgboost \
+    tiler \
+    segmentation-models \
+    pytest \
+    coveralls \
+    rtree \
+    sphinx \
+    sphinx_rtd_theme \
+    astropy
+
+HEALTHCHECK NONE
+ENTRYPOINT []
+CMD ["/bin/bash"]
diff --git a/requirements/Dockerfile.dev b/requirements/Dockerfile.dev
index 8b3d7a7..4d4bb95 100644
--- a/requirements/Dockerfile.dev
+++ b/requirements/Dockerfile.dev
@@ -84,6 +84,7 @@ RUN pip install \
     geopandas \
     tensorflow_probability \
     keras-cv \
+    keras-unet-collection \
     opencv-python \
     opencv-python-headless \
     opencv-contrib-python \
diff --git a/tensorflow_caney/config/cnn_config.py b/tensorflow_caney/config/cnn_config.py
index 66db3d2..01c5c22 100644
--- a/tensorflow_caney/config/cnn_config.py
+++ b/tensorflow_caney/config/cnn_config.py
@@ -181,3 +181,6 @@ class Config:
 
     # Regex to find model metadata
     metadata_regex: Optional[str] = None
+
+    # Validation database for validation of predictions
+    validation_database: Optional[str] = None
diff --git a/tensorflow_caney/inference/inference.py b/tensorflow_caney/inference/inference.py
index b8f22bf..caffce5 100644
--- a/tensorflow_caney/inference/inference.py
+++ b/tensorflow_caney/inference/inference.py
@@ -269,6 +269,12 @@ def sliding_window_tiler_multiclass(
     tile_channels = model.layers[0].input_shape[0][-1]
     # n_classes = out of the output layer, output_shape
 
+    if tile_size is None:
+        tile_size = 256
+    
+    if tile_channels is None:
+        tile_channels = xraster.shape[-1]
+
     tiler_image = Tiler(
         data_shape=xraster.shape,
         tile_shape=(tile_size, tile_size, tile_channels),
@@ -291,8 +297,8 @@ def sliding_window_tiler_multiclass(
     merger = Merger(tiler=tiler_mask, window=window)
     xraster = normalize_image(xraster, normalize)
 
-    if rescale is not None:
-        xraster = rescale_image(xraster, rescale)
+    # if rescale is not None:
+    #    xraster = rescale_image(xraster, rescale)
 
     # Iterate over the data in batches
     for batch_id, batch_i in tiler_image(xraster, batch_size=batch_size):
@@ -300,11 +306,21 @@ def sliding_window_tiler_multiclass(
         # Standardize
         batch = batch_i.copy()
 
-        if standardization is not None:
+        if standardization != 'None' and standardization is not None:
+            # print("STANDARDIZING")
             for item in range(batch.shape[0]):
                 batch[item, :, :, :] = standardize_image(
                     batch[item, :, :, :], standardization, mean, std)
 
+        if rescale != 'None' and rescale is not None:
+            # print("RESCALING")
+            for item in range(batch.shape[0]):
+                means = batch[item, :, :, :].mean(axis=(0, 1))
+                stds = batch[item, :, :, :].std(axis=(0, 1))
+                newMin = means - (2 * stds)
+                newMax = means + (2 * stds)
+                batch[item, :, :, :] = (batch[item, :, :, :] - newMin) / (newMax - newMin)
+
         # Predict
         batch = model.predict(batch, batch_size=batch_size, verbose=0)
 
@@ -313,16 +329,34 @@ def sliding_window_tiler_multiclass(
 
     prediction = merger.merge(unpad=True)
 
+    # print("BEFORE ARGMAX", probability_map, prediction.shape)
+
+    # return prediction map
     if not probability_map:
+
+        # print("ENTERING SINGLE ARGMAX")
+
         if prediction.shape[-1] > 1:
             prediction = np.argmax(prediction, axis=-1)
+            # print("DOING NUMPY ARGMAX", prediction.shape)
         else:
             prediction = np.squeeze(
-                np.where(prediction > threshold, 1, 0).astype(np.int16)
-            )
+                np.where(prediction > threshold, 1, 0).astype(np.int16))
+
+        # print("BEFORE RETURN", prediction.shape)
+        return prediction, None
+
+    # return prediction and probability map
     else:
-        prediction = np.squeeze(prediction)
-    return prediction
+        # print("ENTERING PROBABILITY")
+        if prediction.shape[-1] > 1:
+            probability = np.amax(prediction, axis=-1)
+            prediction = np.argmax(prediction, axis=-1)
+        else:
+            probability = np.squeeze(prediction)
+            prediction = np.squeeze(
+                np.where(prediction > threshold, 1, 0).astype(np.int16))
+        return prediction, probability
 
 
 def get_extract_pred_scatter(
diff --git a/tensorflow_caney/inference/regression_inference.py b/tensorflow_caney/inference/regression_inference.py
index 0df1785..a75a2ae 100644
--- a/tensorflow_caney/inference/regression_inference.py
+++ b/tensorflow_caney/inference/regression_inference.py
@@ -1,4 +1,5 @@
 import numpy as np
+import tensorflow as tf
 from tiler import Tiler, Merger
 from ..utils.data import standardize_image
 
@@ -59,7 +60,10 @@ def sliding_window_tiler(
                 batch[item, :, :, :] = standardize_image(
                     batch[item, :, :, :], standardization, mean, std)
 
-        # Predict
+        # Predict - working on TensorRT
+        # if isinstance(model, tf.python.saved_model.load._WrapperFunction):
+        # batch = model(batch)['probs'].numpy()
+        # else:
         batch = model.predict(batch, batch_size=batch_size, verbose=0)
 
         # Merge the updated data in the array
diff --git a/tensorflow_caney/model/config/cnn_config.py b/tensorflow_caney/model/config/cnn_config.py
index 4da63d6..90c61d6 100644
--- a/tensorflow_caney/model/config/cnn_config.py
+++ b/tensorflow_caney/model/config/cnn_config.py
@@ -163,7 +163,7 @@ class Config:
     inference_treshold: Optional[float] = 0.5
 
     # Inference padding value to replace no-data boundaries
-    inference_pad_value: Optional[int] = 2000
+    inference_pad_value: Optional[int] = 1000
 
     # Window algorithm for smoothing prediction
     window_algorithm: Optional[str] = 'triang'
@@ -188,3 +188,6 @@ class Config:
 
     # Regex to find model metadata
     metadata_regex: Optional[str] = None
+
+    # Validation database for validation of predictions
+    validation_database: Optional[str] = None
diff --git a/tensorflow_caney/model/dataloaders/segmentation.py b/tensorflow_caney/model/dataloaders/segmentation.py
index ce7f596..0a1b92c 100644
--- a/tensorflow_caney/model/dataloaders/segmentation.py
+++ b/tensorflow_caney/model/dataloaders/segmentation.py
@@ -113,6 +113,36 @@ def __init__(
             self.metadata = read_metadata(
                 conf.metadata_regex, conf.input_bands, conf.output_bands)
 
+        """
+        # sometimes = lambda aug, prob=0.5: iaa.Sometimes(prob, aug)
+        self.seq = iaa.Sequential([
+            # Basic aug without changing any values
+            iaa.Fliplr(0.5),  # horizontally flip 50% of all images
+            iaa.Flipud(0.5),  # vertically flip 20% of all images
+            #sometimes(iaa.Crop(percent=(0, 0.1))),  # random crops
+
+            # Gaussian blur and gamma contrast
+            sometimes(iaa.GaussianBlur(sigma=(0, 0.3)), 0.3),
+            # sometimes(iaa.GammaContrast(gamma=0.5, per_channel=True), 0.3),
+
+            # Sharpen
+            #sometimes(iaa.Sharpen(alpha=(0.0, 1.0), lightness=(0.75, 2.0))),
+            # AddToHueAndSaturation
+            # sometimes(iaa.AddToHueAndSaturation((-50, 50), per_channel=True)),
+
+            # sometimes(iaa.Affine(scale={"x": (0.5, 1.5), "y": (0.5, 1.5)})),
+            #sometimes(iaa.Affine(rotate=(-90, 90))),
+
+            # iaa.CoarseDropout((0.03, 0.25), size_percent=(0.02, 0.05), per_channel=True)
+            # sometimes(iaa.Multiply((0.75, 1.25), per_channel=True), 0.3),
+            sometimes(iaa.LinearContrast((0.3, 1.2)), 0.3),
+            # iaa.Add(value=(-0.5,0.5),per_channel=True),
+            #sometimes(iaa.PiecewiseAffine(0.05), 0.3),
+            #sometimes(iaa.PerspectiveTransform(0.01), 0.1)
+        ],
+            random_order=True)
+        """
+
     def tf_dataset(
                 self, x: list, y: list,
                 read_func: Any,
@@ -180,6 +210,12 @@ def load_data(self, x, y):
         # Normalization should be done at preprocessing step
         # x = normalize_image(x, self.conf.normalize)
 
+        # seq_det = self.seq.to_deterministic()
+        # x = seq_det.augment_image(x)
+        # y would have two channels, i.e. annotations and weights.
+        # We need to augment y for operations such as crop and transform
+        # y = seq_det.augment_image(y)
+
         # Standardize
         if self.conf.metadata_regex is not None and \
                 self.conf.standardization is not None:
@@ -187,15 +223,27 @@ def load_data(self, x, y):
                 x, self.metadata[timestamp], subtract='median'
             )
 
-        elif self.conf.standardization is not None:
+        elif self.conf.standardization is not None and \
+                self.conf.standardization != 'None':
             x = standardize_image(
                 x, self.conf.standardization, self.mean, self.std)
 
+        # Rescale
+        if self.conf.rescale is not None and \
+                self.conf.rescale != 'None':
+            # print("RESCALING")
+            means = x.mean(axis=(0, 1))
+            stds = x.std(axis=(0, 1))
+            newMin = means - (2 * stds)
+            newMax = means + (2 * stds)
+            x = (x - newMin) / (newMax - newMin)
+
         # Crop
         if self.conf.center_crop:
             x = center_crop(x, (self.conf.tile_size, self.conf.tile_size))
             y = center_crop(y, (self.conf.tile_size, self.conf.tile_size))
 
+       
         # Augment
         if self.conf.augment:
 
diff --git a/tensorflow_caney/model/networks/vision/segmentation_unet.py b/tensorflow_caney/model/networks/vision/segmentation_unet.py
index 1f0fecb..58b3e7f 100644
--- a/tensorflow_caney/model/networks/vision/segmentation_unet.py
+++ b/tensorflow_caney/model/networks/vision/segmentation_unet.py
@@ -309,7 +309,10 @@ def Attention_UNet(input_shape, NUM_CLASSES=1, dropout_rate=0.0, batch_norm=True
     conv_final = layers.Conv2D(NUM_CLASSES, kernel_size=(1,1))(up_conv_128)
     conv_final = layers.BatchNormalization(axis=3)(conv_final)
     # Change to softmax for multichannel
-    conv_final = layers.Activation('sigmoid')(conv_final)
+    if NUM_CLASSES == 1:
+        conv_final = layers.Activation('sigmoid')(conv_final)
+    else:
+        conv_final = layers.Activation('softmax')(conv_final)
 
     # Model integration
     model = models.Model(inputs, conv_final, name="Attention_UNet")
@@ -388,7 +391,10 @@ def Attention_ResUNet(input_shape, NUM_CLASSES=1, dropout_rate=0.0, batch_norm=T
     conv_final = layers.Conv2D(NUM_CLASSES, kernel_size=(1,1))(up_conv_128)
     conv_final = layers.BatchNormalization(axis=axis)(conv_final)
     # Change to softmax for multichannel
-    conv_final = layers.Activation('sigmoid')(conv_final)
+    if NUM_CLASSES == 1:
+        conv_final = layers.Activation('sigmoid')(conv_final)
+    else:
+        conv_final = layers.Activation('softmax')(conv_final)
 
     # Model integration
     model = models.Model(inputs, conv_final, name="AttentionResUNet")
diff --git a/tensorflow_caney/model/pipelines/cnn_regression.py b/tensorflow_caney/model/pipelines/cnn_regression.py
index 5f10e45..d89d691 100644
--- a/tensorflow_caney/model/pipelines/cnn_regression.py
+++ b/tensorflow_caney/model/pipelines/cnn_regression.py
@@ -435,7 +435,8 @@ def predict(self) -> None:
                 image = image.transpose("y", "x", "band")
 
                 # Remove no-data values to account for edge effects
-                temporary_tif = xr.where(image > -100, image, 600)
+                temporary_tif = xr.where(
+                    image > -100, image, self.conf.inference_pad_value)
 
                 # Sliding window prediction
                 prediction = regression_inference.sliding_window_tiler(
@@ -450,6 +451,8 @@ def predict(self) -> None:
                     normalize=self.conf.normalize,
                     window=self.conf.window_algorithm
                 ) * self.conf.normalize_label
+
+                # Predictions lower than 0 are 0'ed out
                 prediction[prediction < 0] = 0
 
                 # Drop image band to allow for a merge of mask
@@ -474,6 +477,8 @@ def predict(self) -> None:
 
                 # Set nodata values on mask
                 nodata = prediction.rio.nodata
+                # print("THIS IS MY NODATA FROM PREDICTION", nodata)
+
                 prediction = prediction.where(image != nodata)
                 prediction.rio.write_nodata(
                     self.conf.prediction_nodata, encoded=True, inplace=True)
diff --git a/tensorflow_caney/model/pipelines/cnn_segmentation.py b/tensorflow_caney/model/pipelines/cnn_segmentation.py
index fcc7f7f..d50d179 100644
--- a/tensorflow_caney/model/pipelines/cnn_segmentation.py
+++ b/tensorflow_caney/model/pipelines/cnn_segmentation.py
@@ -330,7 +330,7 @@ def train(self) -> None:
 
                 # get full model for training
                 model = get_model(self.conf.model)
-                model.trainable = False
+                model.trainable = True
                 pretrained = load_model(
                     model_filename=self.conf.transfer_learning_weights,
                     model_dir=self.model_dir
@@ -470,6 +470,12 @@ def predict(self) -> None:
                     image = rxr.open_rasterio(filename)
                     logging.info(f'Prediction shape: {image.shape}')
 
+                    # check bands in imagery, do not proceed if one band
+                    if image.shape[0] == 1:
+                        logging.info(
+                            'Skipping file because of non sufficient bands')
+                        continue
+
                 except rasterio.errors.RasterioIOError:
                     logging.info(f'Skipped {filename}, probably corrupted.')
                     continue
@@ -495,10 +501,10 @@ def predict(self) -> None:
                 # Remove no-data values to account for edge effects
                 temporary_tif = xr.where(image > -100, image, 600)
 
-                print("ENTERING PREDICTION")
+                # print("ENTERING PREDICTION")
 
                 # Sliding window prediction
-                prediction = \
+                prediction, probability = \
                     inference.sliding_window_tiler_multiclass(
                         xraster=temporary_tif,
                         model=model,
@@ -515,8 +521,8 @@ def predict(self) -> None:
                         probability_map=self.conf.probability_map
                     )
 
-                if isinstance(prediction, tuple):
-                    prediction, probability = prediction
+                # if isinstance(prediction, tuple):
+                #     prediction, probability = prediction
 
                 # Drop image band to allow for a merge of mask
                 image = image.drop(
@@ -524,7 +530,7 @@ def predict(self) -> None:
                     labels=image.coords["band"].values[1:],
                 )
 
-                # Get metadata to save raster
+                # Get metadata to save raster prediction
                 prediction = xr.DataArray(
                     np.expand_dims(prediction, axis=-1),
                     name=self.conf.experiment_type,
@@ -554,14 +560,13 @@ def predict(self) -> None:
                 )
                 del prediction
 
-                print(Path(output_filename).with_suffix('.tif'))
+                # print(Path(output_filename).with_suffix('.tif'))
 
-                """
                 # save probability map
-                if self.conf.probability_map:
+                if probability is not None:
 
                     probability = xr.DataArray(
-                        np.expand_dims(probability[0], axis=-1),
+                        np.expand_dims(probability, axis=-1), # probability[0]
                         name=self.conf.experiment_type,
                         coords=image.coords,
                         dims=image.dims,
@@ -585,14 +590,14 @@ def predict(self) -> None:
 
                     # Save output raster file to disk
                     probability.rio.to_raster(
-                        output_filename,
+                        Path(output_filename).with_suffix('.prob.tif'),
                         BIGTIFF="IF_SAFER",
                         compress=self.conf.prediction_compress,
                         driver=self.conf.prediction_driver,
                         dtype='float32'
                     )
                     del probability
-                """
+
                 # delete lock file
                 try:
                     os.remove(lock_filename)
diff --git a/tensorflow_caney/networks/regression/unet_regression.py b/tensorflow_caney/networks/regression/unet_regression.py
index dec05b5..6ed0291 100644
--- a/tensorflow_caney/networks/regression/unet_regression.py
+++ b/tensorflow_caney/networks/regression/unet_regression.py
@@ -6,7 +6,8 @@
 from tensorflow.keras.regularizers import l2
 from tensorflow.keras import layers
 from tensorflow.keras import models
-
+from tensorflow.keras import regularizers, mixed_precision
+from tensorflow.keras import backend as K
 
 def unet_batchnorm_regression(
             nclass=1,
@@ -133,6 +134,147 @@ def get_model_new():
     return model
 
 
+def attention_unet_regression(
+        input_shape,
+        input_label_channel_count:int,
+        layer_count=64,
+        regularizers=regularizers.l2(0.0001),
+        weight_file=None,
+         summary=False,
+         countbranch=False
+    ):
+    """ Method to declare the UNet model.
+    Args:
+        input_shape: tuple(int, int, int, int)
+            Shape of the input in the format (batch, height, width, channels).
+        input_label_channel_count: int
+            index count of label channels, used for calculating the number of channels in model output.
+        layer_count: (int, optional)
+            Count of kernels in first layer. Number of kernels in other layers grows with a fixed factor.
+        regularizers: keras.regularizers
+            regularizers to use in each layer.
+        weight_file: str
+            path to the weight file.
+        summary: bool
+            Whether to print the model summary
+    """
+
+    # NOTE(Jesse): Use "mixed_bfloat16" string for TPUs
+    mixed_precision.set_global_policy("mixed_bfloat16")
+
+    input_img = layers.Input(input_shape[1:], dtype="float32", name='Input')
+    pp_in_layer = input_img
+
+    c1 = layers.Conv2D(
+        1 * layer_count, (3, 3), activation='relu', padding='same')(pp_in_layer)
+    c1 = layers.Conv2D(
+        1 * layer_count, (3, 3), activation='relu', padding='same')(c1)
+    n1 = layers.BatchNormalization()(c1)
+    p1 = layers.MaxPooling2D((2, 2))(n1)
+
+    c2 = layers.Conv2D(
+        2 * layer_count, (3, 3), activation='relu', padding='same')(p1)
+    c2 = layers.Conv2D(
+        2 * layer_count, (3, 3), activation='relu', padding='same')(c2)
+    n2 = layers.BatchNormalization()(c2)
+    p2 = layers.MaxPooling2D((2, 2))(n2)
+
+    c3 = layers.Conv2D(
+        4 * layer_count, (3, 3), activation='relu', padding='same')(p2)
+    c3 = layers.Conv2D(
+        4 * layer_count, (3, 3), activation='relu', padding='same')(c3)
+    n3 = layers.BatchNormalization()(c3)
+    p3 = layers.MaxPooling2D((2, 2))(n3)
+
+    c4 = layers.Conv2D(
+        8 * layer_count, (3, 3), activation='relu', padding='same')(p3)
+    c4 = layers.Conv2D(
+        8 * layer_count, (3, 3), activation='relu', padding='same')(c4)
+    n4 = layers.BatchNormalization()(c4)
+    p4 = layers.MaxPooling2D(pool_size=(2, 2))(n4)
+
+    c5 = layers.Conv2D(
+        16 * layer_count, (3, 3), activation='relu', padding='same')(p4)
+    c5 = layers.Conv2D(
+        16 * layer_count, (3, 3), activation='relu', padding='same')(c5)
+    n5 = layers.BatchNormalization()(c5)
+
+    u6 = attention_up_and_concat(n5, n4)
+    c6 = layers.Conv2D(
+        8 * layer_count, (3, 3), activation='relu', padding='same')(u6)
+    c6 = layers.Conv2D(
+        8 * layer_count, (3, 3), activation='relu', padding='same')(c6)
+    n6 = layers.BatchNormalization()(c6)
+
+    u7 = attention_up_and_concat(n6, n3)
+    c7 = layers.Conv2D(
+        4 * layer_count, (3, 3), activation='relu', padding='same')(u7)
+    c7 = layers.Conv2D(
+        4 * layer_count, (3, 3), activation='relu', padding='same')(c7)
+    n7 = layers.BatchNormalization()(c7)
+
+    u8 = attention_up_and_concat(n7, n2)
+    c8 = layers.Conv2D(
+        2 * layer_count, (3, 3), activation='relu', padding='same')(u8)
+    c8 = layers.Conv2D(
+        2 * layer_count, (3, 3), activation='relu', padding='same')(c8)
+    n8 = layers.BatchNormalization()(c8)
+
+    u9 = attention_up_and_concat(n8, n1)
+    c9 = layers.Conv2D(
+        1 * layer_count, (3, 3), activation='relu', padding='same')(u9)
+    c9 = layers.Conv2D(
+        1 * layer_count, (3, 3), activation='relu', padding='same')(c9)
+    n9 = layers.BatchNormalization()(c9)
+
+# =============================================================================
+    # density map
+    if countbranch:
+        d = layers.Conv2D(
+            input_label_channel_count, (1, 1), activation='sigmoid',
+            kernel_regularizer=regularizers, name='output_seg')(n9)
+
+        d2 = layers.Conv2D(
+            input_label_channel_count, (1, 1), activation='linear',
+            dtype='float32', kernel_regularizer= regularizers, name = 'output_dens')(n9)
+
+        seg_model = models.Model(inputs=[input_img], outputs=[d, d2])
+# =============================================================================
+    else:
+        d = layers.Conv2D(
+            input_label_channel_count, (1, 1), activation='sigmoid',
+            dtype='float32', kernel_regularizer=regularizers)(n9)
+        seg_model = models.Model(inputs=[input_img], outputs=[d])
+    if weight_file:
+        seg_model.load_weights(weight_file)
+    if summary:
+        seg_model.summary()
+
+    return seg_model
+
+
+def attention_up_and_concat(down_layer, layer):
+    in_channel = down_layer.get_shape().as_list()[3]
+    up = layers.UpSampling2D(size=(2, 2), interpolation="bilinear")(down_layer)
+    layer = attention_block_2d(
+        x=layer, g=up, inter_channel=in_channel // 4)
+    my_concat = layers.Lambda(lambda x: K.concatenate([x[0], x[1]], axis=3))
+    concat = my_concat([up, layer])
+
+    return concat
+
+
+def attention_block_2d(x, g, inter_channel):
+    theta_x = layers.Conv2D(inter_channel, [1, 1], strides=[1, 1])(x)
+    phi_g = layers.Conv2D(inter_channel, [1, 1], strides=[1, 1])(g)
+    f = layers.Activation('relu')(layers.add([theta_x, phi_g]))
+    psi_f = layers.Conv2D(1, [1, 1], strides=[1, 1])(f)
+    rate = layers.Activation('sigmoid')(psi_f)
+    att_x = layers.multiply([x, rate])
+
+    return att_x
+
+
 if __name__ == "__main__":
 
     # Can add different images sizes - for now: (256,256,6)
diff --git a/tensorflow_caney/networks/unet.py b/tensorflow_caney/networks/unet.py
index ac1fe08..70ca33f 100644
--- a/tensorflow_caney/networks/unet.py
+++ b/tensorflow_caney/networks/unet.py
@@ -148,8 +148,10 @@ def unet_batchnorm(nclass=19, input_size=(256, 256, 8), weight_file=None,
     if nclass == 1:
         actv = 'sigmoid'
 
-    c10 = Conv2D(nclass, (1, 1), activation=actv, kernel_regularizer=kr)(c9)
-    model = Model(inputs=inputs, outputs=c10, name="UNetBatchNorm")
+    c10 = Conv2D(nclass, (1, 1), kernel_regularizer=kr)(c9)
+    outputs = Activation(actv, dtype='float32', name='predictions')(c10)
+
+    model = Model(inputs=inputs, outputs=outputs, name="UNetBatchNorm")
 
     if weight_file:
         model.load_weights(weight_file).expect_partial()
diff --git a/tensorflow_caney/utils/data.py b/tensorflow_caney/utils/data.py
index a063a27..1ed4414 100644
--- a/tensorflow_caney/utils/data.py
+++ b/tensorflow_caney/utils/data.py
@@ -62,6 +62,7 @@ def gen_random_tiles(
                 num_classes,
                 tile_size,
                 augment,
+                xp,
                 expand_dims,
                 out_image_dir,
                 out_label_dir
@@ -84,8 +85,13 @@ def gen_random_tiles(
                 image_tile.shape[1] != tile_size:
             continue
 
-        # first condition, time must have valid classes
-        if (image_tile.min() < -1 or label_tile.min() < 0) \
+        # make sure new tile has the same size
+        if label_tile.shape[0] != tile_size or \
+                label_tile.shape[1] != tile_size:
+            continue
+
+        # first condition, tile must have valid classes
+        if (image_tile.min() < -1000 or label_tile.min() < 0) \
                 and use_case == 'segmentation':
             continue
 
@@ -190,6 +196,7 @@ def gen_random_tiles_from_json(
             num_classes: int,
             tile_size: int = 128,
             augment: bool = True,
+            xp: object = np,
             expand_dims: bool = True,
             out_image_dir: str = 'image',
             out_label_dir: str = 'label',
@@ -207,6 +214,8 @@ def gen_random_tiles_from_json(
 
     for tile_filename in tiles_metadata:
 
+        # print("MY JSON FILENAME", tile_filename)
+
         # Tile indices
         x = tiles_metadata[tile_filename]['x']
         y = tiles_metadata[tile_filename]['y']
@@ -239,14 +248,15 @@ def gen_random_tiles_from_json(
                 label_tile = cp.rot90(label_tile, 3)
 
         if num_classes >= 2:
-            label_tile = cp.eye(num_classes, dtype='uint8')[label_tile]
+            label_tile = xp.eye(num_classes, dtype='uint8')[label_tile]
         else:
             if expand_dims:
-                label_tile = cp.expand_dims(label_tile, axis=-1)
+                label_tile = xp.expand_dims(label_tile, axis=-1)
 
         # save tiles to disk
-        cp.save(os.path.join(out_image_dir, tile_filename), image_tile)
-        cp.save(os.path.join(out_label_dir, tile_filename), label_tile)
+        xp.save(os.path.join(out_image_dir, tile_filename), image_tile)
+        xp.save(os.path.join(out_label_dir, tile_filename), label_tile)
+
     return
 
 
diff --git a/tensorflow_caney/utils/indices.py b/tensorflow_caney/utils/indices.py
index 2b1d442..cbee087 100644
--- a/tensorflow_caney/utils/indices.py
+++ b/tensorflow_caney/utils/indices.py
@@ -305,7 +305,7 @@ def sr(raster):
     'ndvi': ndvi,
     'ndwi': ndwi,
     'si': si,
-    'sr': sr
+    'sr': sr,
 }
 
 
diff --git a/tensorflow_caney/utils/losses.py b/tensorflow_caney/utils/losses.py
index d0ab51c..9c92296 100644
--- a/tensorflow_caney/utils/losses.py
+++ b/tensorflow_caney/utils/losses.py
@@ -6,6 +6,7 @@
 import tensorflow_caney as tfc
 
 from typing import Any
+import keras.backend as K
 from tensorflow.keras.losses import Loss, Reduction
 
 
@@ -15,6 +16,9 @@
     "JaccardDistanceLoss", "TanimotoDistanceLoss"
 ]
 
+epsilon = 1e-5
+smooth = 1
+
 
 def get_loss(loss: str) -> Any:
     """
@@ -324,3 +328,73 @@ def call(self, y_true, y_pred):
         loss = tf.reduce_sum(loss, axis=-1)
 
         return loss
+
+
+@tf.function
+def tversky(y_true, y_pred, alpha=0.40, beta=0.60):
+    """
+    Function to calculate the Tversky loss for imbalanced data
+    :param prediction: the logits
+    :param ground_truth: the segmentation ground_truth
+    :param alpha: weight of false positives
+    :param beta: weight of false negatives
+    :param weight_map:
+    :return: the loss
+    """
+
+    y_t = y_true[..., 0]
+    y_t = y_t[..., np.newaxis]
+    # weights
+    y_weights = y_true[..., 1]
+    y_weights = y_weights[..., np.newaxis]
+
+    ones = 1
+    p0 = y_pred  # proba that voxels are class i
+    p1 = ones - y_pred  # proba that voxels are not class i
+    g0 = y_t
+    g1 = ones - y_t
+
+    tp = tf.reduce_sum(y_weights * p0 * g0)
+    fp = alpha * tf.reduce_sum(y_weights * p0 * g1)
+    fn = beta * tf.reduce_sum(y_weights * p1 * g0)
+
+    EPSILON = 0.00001
+    numerator = tp
+    denominator = tp + fp + fn + EPSILON
+    score = numerator / denominator
+    return 1.0 - tf.reduce_mean(score)
+
+
+def class_tversky(y_true, y_pred):
+    smooth = 1
+
+    y_true = K.permute_dimensions(y_true, (3,1,2,0))
+    y_pred = K.permute_dimensions(y_pred, (3,1,2,0))
+
+    y_true_pos = K.batch_flatten(y_true)
+    y_pred_pos = K.batch_flatten(y_pred)
+    true_pos = K.sum(y_true_pos * y_pred_pos, 1)
+    false_neg = K.sum(y_true_pos * (1-y_pred_pos), 1)
+    false_pos = K.sum((1-y_true_pos)*y_pred_pos, 1)
+    alpha = 0.7
+    return (true_pos + smooth)/(true_pos + alpha*false_neg + (1-alpha)*false_pos + smooth)
+
+
+def focal_tversky_loss(y_true,y_pred):
+    pt_1 = class_tversky(y_true, y_pred)
+    gamma = 0.75
+    return K.sum(K.pow((1-pt_1), gamma))
+
+
+def binary_tversky(y_true, y_pred):
+    y_true_pos = K.flatten(y_true)
+    y_pred_pos = K.flatten(y_pred)
+    true_pos = K.sum(y_true_pos * y_pred_pos)
+    false_neg = K.sum(y_true_pos * (1-y_pred_pos))
+    false_pos = K.sum((1-y_true_pos)*y_pred_pos)
+    alpha = 0.7
+    return (true_pos + smooth)/(true_pos + alpha*false_neg + (1-alpha)*false_pos + smooth)
+
+
+def binary_tversky_loss(y_true, y_pred):
+    return 1 - binary_tversky(y_true, y_pred)
diff --git a/tensorflow_caney/utils/metrics.py b/tensorflow_caney/utils/metrics.py
index f5b7c35..d67df9d 100644
--- a/tensorflow_caney/utils/metrics.py
+++ b/tensorflow_caney/utils/metrics.py
@@ -1,12 +1,30 @@
 import sys
+import logging
+import numpy as np
 import tensorflow as tf
+from typing import Any, List
 import segmentation_models as sm
 import tensorflow_addons as tfa
+from tensorflow.keras import backend as K
+from sklearn.metrics import accuracy_score
+from sklearn.metrics import precision_score
+from sklearn.metrics import recall_score
 
-from typing import Any, List
-
+__author__ = "Jordan A Caraballo-Vega, Science Data Processing Branch"
+__email__ = "jordan.a.caraballo-vega@nasa.gov"
+__status__ = "Production"
 __all__ = ["get_metrics"]
 
+# ---------------------------------------------------------------------------
+# module metrics
+#
+# General functions to compute custom metrics.
+# ---------------------------------------------------------------------------
+
+# ---------------------------------------------------------------------------
+# Module Methods
+# ---------------------------------------------------------------------------
+
 
 def get_metrics(metrics: List[str]) -> Any:
     """
@@ -23,3 +41,100 @@ def get_metrics(metrics: List[str]) -> Any:
         except (NameError, AttributeError) as err:
             sys.exit(f'{err}. Accepted metrics from {tf}, {sm}, {tfa}')
     return metric_functions
+
+
+def iou_coef(y_true, y_pred, smooth=1):
+    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])-intersection
+    iou = K.mean((intersection + smooth) / (union + smooth), axis=0)
+    return iou
+
+
+def dice_coef(y_true, y_pred, smooth=1):
+    intersection = K.sum(y_true * y_pred, axis=[1, 2, 3])
+    union = K.sum(y_true, axis=[1, 2, 3]) + K.sum(y_pred, axis=[1, 2, 3])
+    dice = K.mean((2. * intersection + smooth)/(union + smooth), axis=0)
+    return dice
+
+
+def iou_val(y_true, y_pred):
+    intersection = np.logical_and(y_true, y_pred)
+    union = np.logical_or(y_true, y_pred)
+    iou_score = np.sum(intersection) / np.sum(union)
+    return iou_score
+
+
+def acc_val(y_true, y_pred):
+    return accuracy_score(y_true, y_pred)
+
+
+def prec_val(y_true, y_pred):
+    return precision_score(y_true, y_pred, average='macro'), \
+        precision_score(y_true, y_pred, average=None)
+
+
+def recall_val(y_true, y_pred):
+    return recall_score(y_true, y_pred, average='macro'), \
+        recall_score(y_true, y_pred, average=None)
+
+
+def compute_imf_weights(ground_truth, n_classes=None,
+                        ignored_classes=[]
+                        ) -> np.array:
+    """
+    Compute inverse median frequency weights for class balancing.
+
+    For each class i, it computes its frequency f_i, i.e the ratio between
+    the number of pixels from class i and the total number of pixels.
+    Then, it computes the median m of all frequencies. For each class the
+    associated weight is m/f_i.
+
+    :param ground_truth: the annotations array
+    :param nclasses: number of classes (defaults to max(ground_truth))
+    :param ignored_classes: id of classes to ignore (optional)
+    :return: numpy array with the IMF coefficients
+    """
+    n_classes = np.max(ground_truth) if n_classes is None else n_classes
+    weights = np.zeros(n_classes)
+    frequencies = np.zeros(n_classes)
+
+    for c in range(0, n_classes):
+        if c in ignored_classes:
+            continue
+        frequencies[c] = np.count_nonzero(ground_truth == c)
+
+    # Normalize the pixel counts to obtain frequencies
+    frequencies /= np.sum(frequencies)
+    # Obtain the median on non-zero frequencies
+    idx = np.nonzero(frequencies)
+    median = np.median(frequencies[idx])
+    weights[idx] = median / frequencies[idx]
+    weights[frequencies == 0] = 0.
+    return weights
+
+
+def getOccurrences(labels=[], fname='occurences.csv', nclasses=7):
+    """
+    Return pixel occurences per class.
+    :param labels: numpy array with labels in int format
+    :param fname: filename to save output
+    :param nclasses: number of classes to look for
+    :return: CSV file with class and occurrence per class
+    """
+    f = open(fname, "w+")
+    f.write("class,occurence\n")
+    for classes in range(nclasses):
+        occ = np.count_nonzero(labels == classes)
+        f.write(f'{classes},{occ}\n')
+    f.close()
+
+
+# -------------------------------------------------------------------------------
+# module metrics Unit Tests
+# -------------------------------------------------------------------------------
+
+if __name__ == "__main__":
+
+    logging.basicConfig(level=logging.INFO)
+
+    # Add unit tests here
diff --git a/tensorflow_caney/utils/model.py b/tensorflow_caney/utils/model.py
index 59f1e4e..8782c8e 100644
--- a/tensorflow_caney/utils/model.py
+++ b/tensorflow_caney/utils/model.py
@@ -5,6 +5,7 @@
 import segmentation_models as sm
 import tensorflow_addons as tfa
 import tensorflow_caney as tfc
+from keras_unet_collection import models as kuc
 import tensorflow_caney
 
 from typing import Any
@@ -17,6 +18,12 @@
 
 __all__ = ["get_model", "load_model"]
 
+CUSTOM_OBJECTS = {
+    'iou_score': sm.metrics.iou_score,
+    'focal_tversky_loss': tfc.utils.losses.focal_tversky_loss,
+    'binary_tversky_loss': tfc.utils.losses.binary_tversky_loss,
+}
+
 
 def get_model(model: str) -> Any:
     """
@@ -29,14 +36,14 @@ def get_model(model: str) -> Any:
     try:
         model_function = eval(model)
     except (NameError, AttributeError) as err:
-        sys.exit(f'{err}. Accepted models from {tf}, {sm}, {tfa}, {tfc}')
+        sys.exit(f'{err}. Accepted models from {tf}, {sm}, {tfa}, {tfc}, {kuc}')
     return model_function
 
 
 def load_model(
             model_filename: str = None,
             model_dir: str = None,
-            custom_objects: dict = {'iou_score': sm.metrics.iou_score},
+            custom_objects: dict = CUSTOM_OBJECTS,
             model_extension: str = '*.hdf5',
             conf: OmegaConf = None
         ) -> Any:


Test patch:
diff --git a/tests/data/land_cover_test.csv b/tests/data/land_cover_test.csv
index 34711cb..d1048c7 100644
--- a/tests/data/land_cover_test.csv
+++ b/tests/data/land_cover_test.csv
@@ -1,2 +1,2 @@
 data,label,ntiles
-/explore/nobackup/projects/3sl/data/Tappan/Tappan02_WV02_20120218_M1BS_103001001077BE00_data.tif,/explore/nobackup/projects/3sl/labels/landcover/2m/CAS/Tappan02_WV02_20120218_M1BS_103001001077BE00_mask_segs_reclassified.tif,50
\ No newline at end of file
+/explore/nobackup/projects/3sl/data/Tappan/Tappan02_WV02_20120218_M1BS_103001001077BE00_data.tif,/explore/nobackup/projects/3sl/labels/landcover/2m/eCAS/Tappan02_WV02_20120218_M1BS_103001001077BE00_mask_segs_reclassified.tif,50
\ No newline at end of file