# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # 3D Instance Segmentation with Discriminative Instance Loss
# ---
# Implemntation of paper:
#
# [Semantic Instance Segmentation with a Discriminative Loss Function](https://arxiv.org/abs/1708.02551)

# +
import os
import errno
import datetime

import numpy as np

import deepcell
# -

# ## Load the data
#
# ### Download the data from `deepcell.datasets`
#
# `deepcell.datasets` provides access to a set of annotated live-cell imaging datasets which can be used for training cell segmentation and tracking models.
# All dataset objects share the `load_data()` method, which allows the user to specify the name of the file (`path`), the fraction of data reserved for testing (`test_size`) and a `seed` which is used to generate the random train-test split.
# Metadata associated with the dataset can be accessed through the `metadata` attribute.

# +
# Download the data (saves to ~/.keras/datasets)
filename = 'mousebrain.npz'
test_size = 0.2 # % of data saved as test
seed = 0 # seed for random train-test split

(X_train, y_train), (X_test, y_test) = deepcell.datasets.mousebrain.load_data(
    filename, test_size=test_size, seed=seed)

print('X.shape: {}\ny.shape: {}'.format(X_train.shape, y_train.shape))
# -

# ### Set up filepath constants

# +
# the path to the data file is currently required for `train_model_()` functions

# change DATA_DIR if you are not using `deepcell.datasets`
DATA_DIR = os.path.expanduser(os.path.join('~', '.keras', 'datasets'))

# DATA_FILE should be a npz file, preferably from `make_training_data`
DATA_FILE = os.path.join(DATA_DIR, filename)

# confirm the data file is available
assert os.path.isfile(DATA_FILE)

# +
# Set up other required filepaths

# If the data file is in a subdirectory, mirror it in MODEL_DIR and LOG_DIR
PREFIX = os.path.relpath(os.path.dirname(DATA_FILE), DATA_DIR)

ROOT_DIR = '/data'  # TODO: Change this! Usually a mounted volume
MODEL_DIR = os.path.abspath(os.path.join(ROOT_DIR, 'models', PREFIX))
LOG_DIR = os.path.abspath(os.path.join(ROOT_DIR, 'logs', PREFIX))

# create directories if they do not exist
for d in (MODEL_DIR, LOG_DIR):
    try:
        os.makedirs(d)
    except OSError as exc:  # Guard against race condition
        if exc.errno != errno.EEXIST:
            raise
# -

# ## Create the Foreground/Background FeatureNet Model
#
# Here we instantiate two `FeatureNet` models from `deepcell.model_zoo` for foreground/background separation as well as the interior/edge segmentation.

norm_method = 'whole_image'  # data normalization
receptive_field = 61  # should be adjusted for the scale of the data
n_skips = 1  # number of skip-connections (only for FC training)
frames_per_batch = 3
embedding_dim = 3

# +
from deepcell import model_zoo

fgbg_model = model_zoo.bn_feature_net_skip_3D(
    receptive_field=receptive_field,
    n_features=2,  # segmentation mask (is_cell, is_not_cell)
    n_frames=frames_per_batch,
    n_skips=n_skips,
    n_conv_filters=32,
    n_dense_filters=128,
    input_shape=tuple([frames_per_batch] + list(X_train.shape[2:])),
    multires=False,
    last_only=False,
    norm_method=norm_method)
# -

# ## Prepare for training
#
# ### Set up training parameters.
#
# There are a number of tunable hyper parameters necessary for training deep learning models:
#
# **model_name**: Incorporated into any files generated during the training process.
#
# **n_epoch**: The number of complete passes through the training dataset.
#
# **lr**: The learning rate determines the speed at which the model learns. Specifically it controls the relative size of the updates to model values after each batch.
#
# **optimizer**: The TensorFlow module [tf.keras.optimizers](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers) offers optimizers with a variety of algorithm implementations. DeepCell typically uses the Adam or the SGD optimizers.
#
# **lr_sched**: A learning rate scheduler allows the learning rate to adapt over the course of model training. Typically a larger learning rate is preferred during the start of the training process, while a small learning rate allows for fine-tuning during the end of training.
#
# **batch_size**: The batch size determines the number of samples that are processed before the model is updated. The value must be greater than one and less than or equal to the number of samples in the training dataset.

# +
from tensorflow.keras.optimizers import SGD
from deepcell.utils.train_utils import rate_scheduler

fgbg_model_name = 'disc_fgbg'
disc_3d_model_name = 'disc_3d'

n_epoch = 5  # Number of training epochs

lr = 0.01

optimizer = SGD(learning_rate=lr, decay=1e-6, momentum=0.9, nesterov=True)

lr_sched = rate_scheduler(lr=lr, decay=0.99)

batch_size = 1  # FC training uses 1 image per batch
# -

# ### Create the DataGenerators
#
# The `MovieDataGenerator` outputs a raw image (`X`) with it's labeled annotation mask (`y`). Additionally, it can apply a transform to `y` to change the task the model learns. Below we generate 2 training and validation data sets for both the foreground/background model and the pixelwise model.

# +
from deepcell.image_generators import MovieDataGenerator

datagen = MovieDataGenerator(
    rotation_range=180,
    zoom_range=(.8, 1.2),
    horizontal_flip=True,
    vertical_flip=True)

datagen_val = MovieDataGenerator()

# +
# Create the foreground/background data iterators

fgbg_train_data = datagen.flow(
    {'X': X_train, 'y': y_train},
    seed=seed,
    skip=n_skips,
    transform='fgbg',
    batch_size=batch_size)

fgbg_val_data = datagen_val.flow(
    {'X': X_test, 'y': y_test},
    seed=seed,
    skip=n_skips,
    transform='fgbg',
    batch_size=batch_size)

# +
# Create the pixelwise data iterators

disc_train_data = datagen.flow(
    {'X': X_train, 'y': y_train},
    seed=seed,
    skip=n_skips,
    transform='disc',
    batch_size=batch_size)

disc_val_data = datagen_val.flow(
    {'X': X_test, 'y': y_test},
    seed=seed,
    skip=n_skips,
    transform='disc',
    batch_size=batch_size)
# -

# ### Compile the model with a loss function
#
# Each model is trained with it's own loss function. `weighted_categorical_crossentropy` is often used for classification models, and `disc_loss` is used for the discriminative instance models.

# +
from deepcell import losses

def loss_function(y_true, y_pred):
    return losses.weighted_categorical_crossentropy(
        y_true, y_pred,
        n_classes=2,
        from_logits=False)

fgbg_model.compile(
    loss=loss_function,
    optimizer=SGD(learning_rate=lr, decay=1e-6, momentum=0.9, nesterov=True),
    metrics=['accuracy'])
# -

# ## Train the foreground/background model
#
# Call `fit()` on the compiled model, along with a default set of callbacks.

# +
from deepcell.utils.train_utils import get_callbacks
from deepcell.utils.train_utils import count_gpus

model_path = os.path.join(MODEL_DIR, '{}.h5'.format(fgbg_model_name))
loss_path = os.path.join(MODEL_DIR, '{}.npz'.format(fgbg_model_name))

num_gpus = count_gpus()

print('Training on', num_gpus, 'GPUs.')

train_callbacks = get_callbacks(
    model_path,
    lr_sched=lr_sched,
    save_weights_only=num_gpus >= 2,
    monitor='val_loss',
    verbose=1)

loss_history = fgbg_model.fit(
    fgbg_train_data,
    steps_per_epoch=fgbg_train_data.y.shape[0] // batch_size,
    epochs=n_epoch,
    validation_data=fgbg_val_data,
    validation_steps=fgbg_val_data.y.shape[0] // batch_size,
    callbacks=train_callbacks)
# -

# ## Create the 3D vector embedding FeatureNet Model
#
# Instatiate a `FeatureNet` for discriminative instance loss, which maps each pixel to a N-dimensional vector.

# +
from deepcell import model_zoo

disc_3d_model = model_zoo.bn_feature_net_skip_3D(
    fgbg_model=fgbg_model,
    receptive_field=receptive_field,
    n_skips=n_skips,
    n_features=embedding_dim,
    n_dense_filters=128,
    n_conv_filters=32,
    input_shape=tuple([frames_per_batch] + list(X_train.shape[2:])),
    norm_method=norm_method)
# -

# ### Compile the model with a loss function
#
# Just like the foreground/background model, the `disc` model is compiled with the discriminative instance loss (`disc_loss`) function.

# +
import tensorflow as tf
from deepcell import losses

disc_3d_model.compile(
    loss=losses.discriminative_instance_loss,
    optimizer=SGD(learning_rate=lr, decay=1e-6, momentum=0.9, nesterov=True),
    metrics=[tf.keras.metrics.Accuracy()])
# -

# ## Train the 3D vector embedding `disc` model
#
# Call `fit()` on the compiled model, along with a default set of callbacks.

# +
from deepcell.utils.train_utils import get_callbacks
from deepcell.utils.train_utils import count_gpus

model_path = os.path.join(MODEL_DIR, '{}.h5'.format(disc_3d_model_name))
loss_path = os.path.join(MODEL_DIR, '{}.npz'.format(disc_3d_model_name))

num_gpus = count_gpus()

print('Training on', num_gpus, 'GPUs.')

train_callbacks = get_callbacks(
    model_path,
    lr_sched=lr_sched,
    save_weights_only=num_gpus >= 2,
    monitor='val_loss',
    verbose=1)

loss_history = disc_3d_model.fit(
    disc_train_data,
    steps_per_epoch=disc_train_data.y.shape[0] // batch_size,
    epochs=n_epoch,
    validation_data=disc_val_data,
    validation_steps=disc_val_data.y.shape[0] // batch_size,
    callbacks=train_callbacks)
# -

# ### Run the model

# +
# With the trained model, make predictions on testing data
test_images = disc_3d_model.predict(X_test[:1, :frames_per_batch])
test_images_fgbg = fgbg_model.predict(X_test[:1, :frames_per_batch])[-1]

print(test_images.shape)
print(test_images_fgbg.shape)
# -

# ## Predict on test data
#
# Use the trained model to predict on new data and post-process the results into a label mask.

# #### Threshold the foreground/background

# threshold the foreground/background
# and remove background from vector embedding
fg_thresh = test_images_fgbg[..., 1] > 0.9
fg_thresh = np.expand_dims(fg_thresh, axis=-1)
test_images_post_fgbg = test_images * fg_thresh

# #### Mean Shift

# +
from itertools import cycle
from sklearn.cluster import MeanShift, estimate_bandwidth

def fit_mean_shift(images, index, frame, ndim, n_samples=500):
    test_plot = images[index, frame].reshape(-1, ndim)
    bandwidth = estimate_bandwidth(test_plot, n_samples=n_samples)

    ms = MeanShift(bandwidth=bandwidth, bin_seeding=True)
    ms.fit(test_plot)
    return ms


# -

# #### DB SCAN

# +
from sklearn.cluster import DBSCAN
from sklearn import metrics

def fit_dbscan(images, index, frame, ndim, eps=.15):
    test_plot = images[index, frame].reshape(-1, ndim)
    db = DBSCAN(eps=eps, min_samples=5, algorithm='kd_tree').fit(test_plot)
    return db


# -

# ### Plot the results

# #### Scatter plots of the embedded vectors

# +
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

index = 0#np.random.randint(low=0, high=test_images_post_fgbg.shape[0])
frame = 2#np.random.randint(low=0, high=test_images_post_fgbg.shape[1])

fig = plt.figure(figsize=(15, 15))

# 3D Vector Embedding
ax = fig.add_subplot(221, projection='3d')
test_plot = test_images[index, frame].reshape(-1, embedding_dim)
ax.scatter(test_plot[:, 0], test_plot[:, 1], test_plot[:, 2], c='b', marker='o', s=5)
ax.set_title('3D Vector Embedding')

# 3D Vector Embedding sans Background
ax = fig.add_subplot(222, projection='3d')
test_plot = test_images_post_fgbg[index, frame].reshape(-1, 3)
ax.scatter(test_plot[:, 0], test_plot[:, 1], test_plot[:, 2], c='b', marker='o', s=5)
ax.set_title('3D Vector Embedding sans Background')

# Scatter plot after MeanShift
ms = fit_mean_shift(test_images_post_fgbg, index=index, frame=frame, ndim=embedding_dim, n_samples=2000)
n_clusters_ms = np.unique(ms.labels_).size
ax = fig.add_subplot(223, projection='3d')
for k, col in zip(range(n_clusters_ms), cycle('bgrcmyk')):
    my_members = ms.labels_ == k
    ax.scatter(test_plot[my_members, 0],
               test_plot[my_members, 1],
               test_plot[my_members, 2],
               c=col, s=5)
ax.set_title('MeanShift: %d estimated clusters' % n_clusters_ms)

# Scatter plot after DBSCAN
db = fit_dbscan(test_images_post_fgbg, index=index, frame=frame, ndim=embedding_dim, eps=0.2)
# Number of clusters in labels, ignoring noise if present.
n_clusters_db = len(set(db.labels_)) - (1 if -1 in db.labels_ else 0)

ax = fig.add_subplot(224, projection='3d')

core_samples_mask = np.zeros_like(db.labels_, dtype=bool)
core_samples_mask[db.core_sample_indices_] = True
# Scatter plot after DBScan
unique_labels = set(db.labels_)
colors = [plt.cm.Spectral(each) for each in np.linspace(0, 1, len(unique_labels))]

for k, col in zip(unique_labels, colors):
    if k == -1:
        col = [0, 0, 0, 1]  # Black used for noise.

    class_member_mask = (db.labels_ == k)

    xy = test_plot[class_member_mask & core_samples_mask]
    ax.scatter(xy[:, 0], xy[:, 1], xy[:, 2], 'o', c=col, s=5)

    xy = test_plot[class_member_mask & ~core_samples_mask]
    ax.scatter(xy[:, 0], xy[:, 1], xy[:, 2], 'o', c=col, s=5)

ax.set_title('DBSCAN: %d estimated clusters' % n_clusters_db)

plt.show()
# -

# #### Plot segmented images

# +
import matplotlib.pyplot as plt

index = 0#np.random.randint(low=0, high=test_images_post_fgbg.shape[0])
frame = 0#np.random.randint(low=0, high=test_images_post_fgbg.shape[1])

fig, axes = plt.subplots(ncols=2, nrows=4, figsize=(15, 30), sharex=True, sharey=True)
ax = axes.ravel()

ax[0].imshow(X_test[index, frame, ..., 0])
ax[0].set_title('Source Image')

ax[1].imshow(test_images_fgbg[index, frame, ..., 1])
ax[1].set_title('Segmentation Prediction')

ax[2].imshow(test_images[index, frame, ..., 0], cmap='jet')
ax[2].set_title('Vector Embedding: Component 0')

ax[3].imshow(test_images[index, frame, ..., 1], cmap='jet')
ax[3].set_title('Vector Embedding: Component 1')

ax[4].imshow(test_images_post_fgbg[index, frame, ..., 0], cmap='jet')
ax[4].set_title('Vector Embedding: Component 0 sans Background')

ax[5].imshow(test_images_post_fgbg[index, frame, ..., 0], cmap='jet')
ax[5].set_title('Vector Embedding: Component 1 sans Background')

ms = fit_mean_shift(test_images_post_fgbg, index=index, frame=frame, ndim=3, n_samples=2000)
ax[6].imshow(ms.labels_.reshape(test_images_post_fgbg.shape[2:-1]), cmap='jet')
ax[6].set_title('MeanShift: %d Clusters' % np.unique(ms.labels_).size)

db = fit_dbscan(test_images_post_fgbg, index=index, frame=frame, ndim=3, eps=0.2)
ax[7].imshow(db.labels_.reshape(test_images_post_fgbg.shape[2:-1]), cmap='jet')
ax[7].set_title('DBSCAN: %d Clusters' % (len(set(db.labels_)) - int(-1 in db.labels_)))

fig.tight_layout()
plt.show()
# -



StarcoderdataJupyterScriptsDedupFiltered