# ---
# 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
# ---

# + [markdown] colab_type="text" id="5slcdn4iWSTq"
# # Kaggle Ship Segmentation - Create Model on Kaggle Environment
#
# Link to competition: https://www.kaggle.com/c/airbus-ship-detection
#
# This notebook was converted from my prior Kaggle notebook.  Migrated to TF 2.x and converted various methods to be more native TF.  This will create and save a trained model.  The model is built from scratch, not pre-trained.  I do have links to a pre-trained option, but it did not perform as well as starting from scratch.  (Pre-trained is smaller, so had a harder time learning features.)
#
# The image size is (224, 224, 3), you can change it in the Config settings.  If you want to get the best possible result, the size should be increased as some features are small and when downsized lose detail.
#
# I stopped training after 7 epocs, val was over 75%, but still slowly learning.
#
# This notebook was about applying learning to use TensorFlow 2.x and Datasets, not to create a final model for submission.  There is much more fine-tuning todo to obtain 80% plus scores.
#

# + colab={} colab_type="code" id="Tn_PWaH1WSTu"
# Change to True if using Kaggle environment....
USING_KAGGLE = True

# + colab={} colab_type="code" id="g3zIITg_WST0"
# Normal includes...
from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, random, warnings, time, copy, csv
import numpy as np 

import IPython.display as display
import pandas as pd
import matplotlib.pyplot as plt
# %matplotlib inline

import tensorflow as tf
print(tf.__version__)

from tensorflow.keras.models import load_model 

# This allows the runtime to decide how best to optimize CPU/GPU usage
AUTOTUNE = tf.data.experimental.AUTOTUNE


# + colab={} colab_type="code" id="XPUpxkxkPzGS"
# Config class that wraps global variable access, using personal libs is a pain in Kaggle, so copied in the class
# Not all of the global vars are used, easier to jsut copy class over from lib

class GlobalParms(object):

    def __init__(self, **kwargs):
        self.keys_and_defaults = {
         "MODEL_NAME": "",  # if you leave .h5 off, puts into a subdirectory
         "ROOT_PATH": "",  # Location of the data for storing any data or files
         "TRAIN_DIR": "",  # Subdirectory in the Root for Training files
         "TEST_DIR": "",  # Optional subdirectory in  Root for Testing file
         "SUBMISSION_PATH": None,  # Optional subdirectory for Contest files
         "MODEL_PATH": None,  # Optional, subdirectory for saving/loading model
         "TRAIN_PATH": None,  # Subdirectory in the Root for Training files
         "TEST_PATH": None,  # Optional subdirectory in  Root for Testing file
         "SMALL_RUN": False,   # Optional, run size will be reduced
         "NUM_CLASSES": 0,  # Number of classes
         "CLASS_NAMES": [],  # list of class names
         "IMAGE_ROWS": 0,  # Row size of the image
         "IMAGE_COLS": 0,  # Col size of the image
         "IMAGE_CHANNELS": 0,  # Num of Channels, 1 for Greyscale, 3 for color
         "BATCH_SIZE": 0,  # Number of images in each batch
         "EPOCS": 0,  # Max number of training EPOCS
         "ROW_SCALE_FACTOR": 1,  # Optional, allows scaling of an image.
         "COL_SCALE_FACTOR": 1,  # Optional, allows scaling of an image.
         "IMAGE_EXT": ".jpg",  # Extent of the image file_ext
         # Optional, default is np.float64, reduce memory by using np.float32
         # or np.float16
         "IMAGE_DTYPE": np.float32,
         # Optional, change default if needed, can save memory space
         "Y_DTYPE": np.int,
         "LOAD_MODEL": False,  # Optional, If you want to load a saved model
         "SUBMISSION": "submission.csv",  # Optional, Mainly used for Kaggle
         "METRICS": ['accuracy'],  # ['categorical_accuracy'], ['accuracy']
         "FINAL_ACTIVATION": 'sigmoid',  # sigmoid, softmax
         "LOSS": ""  # 'binary_crossentropy', 'categorical_crossentropy'
        }

        self.__dict__.update(self.keys_and_defaults)
        self.__dict__.update((k, v) for k, v in kwargs.items()
                             if k in self.keys_and_defaults)

        # Automatically reduce the training parms, change as needed
        if self.__dict__["SMALL_RUN"]:
            self.__dict__["BATCH_SIZE"] = 1
            self.__dict__["EPOCS"] = 2
            self.__dict__["ROW_SCALE_FACTOR"] = 1
            self.__dict__["COL_SCALE_FACTOR"] = 1

        # Use configuration items to create real ones
        self.__dict__["SCALED_ROW_DIM"] = \
            np.int(self.__dict__["IMAGE_ROWS"] /
                   self.__dict__["ROW_SCALE_FACTOR"])

        self.__dict__["SCALED_COL_DIM"] =  \
            np.int(self.__dict__["IMAGE_COLS"] /
                   self.__dict__["COL_SCALE_FACTOR"])

        if self.__dict__["TRAIN_PATH"] is None:  # Not passed, so set it
            self.__dict__["TRAIN_PATH"] = \
                os.path.join(self.__dict__["ROOT_PATH"],
                             self.__dict__["TRAIN_DIR"])

        if self.__dict__["TEST_PATH"] is None:  # Not passed, so set it
            self.__dict__["TEST_PATH"] = \
                os.path.join(self.__dict__["ROOT_PATH"],
                             self.__dict__["TEST_DIR"])

        if self.__dict__["SUBMISSION_PATH"] is None:  # Not passed, so set
            self.__dict__["SUBMISSION_PATH"] = \
                os.path.join(self.__dict__["ROOT_PATH"],
                             self.__dict__["SUBMISSION"])
        else:
            self.__dict__["SUBMISSION_PATH"] = \
                os.path.join(self.__dict__["SUBMISSION_PATH"],
                             self.__dict__["SUBMISSION"])

        if self.__dict__["MODEL_PATH"] is None:  # Not passed, so set it
            self.__dict__["MODEL_PATH"] = \
                os.path.join(self.__dict__["ROOT_PATH"],
                             self.__dict__["MODEL_NAME"])
        else:
            self.__dict__["MODEL_PATH"] = \
                os.path.join(self.__dict__["MODEL_PATH"],
                             self.__dict__["MODEL_NAME"])

        self.__dict__["IMAGE_DIM"] = \
            (self.__dict__["SCALED_ROW_DIM"],
             self.__dict__["SCALED_COL_DIM"],
             self.__dict__["IMAGE_CHANNELS"])

        if self.__dict__["IMAGE_CHANNELS"] == 1:
            self.__dict__["COLOR_MODE"] = "grayscale"
        else:
            self.__dict__["COLOR_MODE"] = "rgb"

    def set_train_path(self, train_path):
        self.__dict__["TRAIN_PATH"] = train_path

    def set_class_names(self, class_name_list):
        self.__dict__["CLASS_NAMES"] = class_name_list

        if self.__dict__["NUM_CLASSES"] != \
           len(self.__dict__["CLASS_NAMES"]):
            raise ValueError("ERROR number of classses do not match, Classes: "
                             + str(self.__dict__["NUM_CLASSES"])
                             + " Class List: "
                             + str(self.__dict__["CLASS_NAMES"]))

    def print_contents(self):
        print(self.__dict__)

    def print_key_value(self):
        for key, value in self.__dict__.items():
            print(key, ":", value)



# + [markdown] colab_type="text" id="rMK-FCdiWST_"
# ## Various Methods

# + colab={} colab_type="code" id="QFpBD1NUWSUB"
# Set these to match your environment

if USING_KAGGLE:
    ROOT_PATH = "../input/airbus-ship-detection/"
else:
    ROOT_PATH = "/Users/john/Documents/ImageData/KaggleShip/"  ###### CHANGE FOR SPECIFIC ENVIRONMENT
        
# Establish global dictionary
parms = GlobalParms(ROOT_PATH=ROOT_PATH,
                    MODEL_NAME="airModel.h5",
                    MODEL_PATH=""
                    TRAIN_DIR="train_v2", 
                    NUM_CLASSES=1,
                    IMAGE_ROWS=224,
                    IMAGE_COLS=224,
                    IMAGE_CHANNELS=3,
                    BATCH_SIZE=16,
                    EPOCS=8,
                    FINAL_ACTIVATION="sigmoid",
                    IMAGE_EXT=".jpg")

parms.print_contents()


# + colab={} colab_type="code" id="KWpSlnmVWSUG"
# Encodes a mask, only used to verify training
def multi_rle_encode(img):
    labels = label(img[:, :, 0])
    return [rle_encode(labels==k) for k in np.unique(labels[labels>0])]

# ref: https://www.kaggle.com/paulorzp/run-length-encode-and-decode
def rle_encode(img):
    '''
    img: numpy array, 1 - mask, 0 - background
    Returns run length as string formated
    '''
    pixels = img.T.flatten()
    pixels = np.concatenate([[0], pixels, [0]])
    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
    runs[1::2] -= runs[::2]
    return ' '.join(str(x) for x in runs)

def rle_decode(mask_rle, shape=(768, 768)):
    '''
    mask_rle: run-length as string formated (start length)
    shape: (height,width) of array to return 
    Returns numpy array, 1 - mask, 0 - background
    '''
    
    s = mask_rle.split()
    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]
    starts -= 1
    ends = starts + lengths
    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)
    for lo, hi in zip(starts, ends):
        img[lo:hi] = 1
    return img.reshape(shape).T  # Needed to align to RLE direction

def masks_as_image(in_mask_list):
    # Take the individual ship masks and create a single mask array for all ships
    all_masks = np.zeros((768, 768), dtype = np.int16)
    #if isinstance(in_mask_list, list):
    for mask in in_mask_list:
        if isinstance(mask, str):
            all_masks += rle_decode(mask)
    return np.expand_dims(all_masks, -1)



# + colab={} colab_type="code" id="VT2xKt1nWSUL"
def show_batch_mask(display_list):
    plt.figure(figsize=(15, 15))

    title = ['Input Image', 'True Mask', 'Predicted Mask']

    for i in range(len(display_list)):
        plt.subplot(1, len(display_list), i+1)
        plt.title(title[i])
        plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
        plt.axis('off')
    plt.show()


# + [markdown] colab_type="text" id="ES3pwdKpWSUR"
# ## Create training and validation files
#
# The number of ships is not balanced and the size of some of the images are very small.  The approach can be changed if you want.  

# + _cell_guid="" _uuid="" colab={} colab_type="code" id="YLfIJShCWSUT"
# https://www.kaggle.com/hmendonca/u-net-model-with-submission
all_df = pd.read_csv(os.path.join(parms.ROOT_PATH, "train_ship_segmentations_v2.csv"))
not_empty = pd.notna(all_df.EncodedPixels)
print(not_empty.sum(), 'masks in', all_df[not_empty].ImageId.nunique(), 'images')
print((~not_empty).sum(), 'empty images in', all_df.ImageId.nunique(), 'total images')
all_df.head()

# + colab={} colab_type="code" id="iSgWYkcYWSUZ"
# Add columns and built a unique list of image_ids.

all_df['ships'] = all_df['EncodedPixels'].map(lambda c_row: 1 if isinstance(c_row, str) else 0)
unique_img_ids = all_df.groupby('ImageId').agg({'ships': 'sum'}).reset_index()
unique_img_ids['has_ship'] = unique_img_ids['ships'].map(lambda x: 1.0 if x>0 else 0.0)
unique_img_ids['has_ship_vec'] = unique_img_ids['has_ship'].map(lambda x: [x])
# some files are too small/corrupt
unique_img_ids['file_size_kb'] = unique_img_ids['ImageId'].map(lambda c_img_id: 
                                                               os.stat(os.path.join(parms.TRAIN_PATH, 
                                                                                    c_img_id)).st_size/1024)
unique_img_ids = unique_img_ids[unique_img_ids['file_size_kb'] > 50] # keep only +50kb files
unique_img_ids['file_size_kb'].hist()
all_df.drop(['ships'], axis=1, inplace=True)
unique_img_ids.sample(7)

# + colab={} colab_type="code" id="1rK6dMVSWSUe"
# Shows the unblanced, most have no ships, so need to change training set to have more with ships and 
# less without ships
unique_img_ids['ships'].hist(bins=unique_img_ids['ships'].max())

# + colab={} colab_type="code" id="sLTB1uWrWSUj"
# Bqlance rows
SAMPLES_PER_GROUP = 1500
balanced_train_df = unique_img_ids.groupby('ships').apply(lambda x: x.sample(SAMPLES_PER_GROUP) if len(x) > SAMPLES_PER_GROUP else x)
balanced_train_df['ships'].hist(bins=balanced_train_df['ships'].max()+1)
print(balanced_train_df.shape[0], 'masks')


# + colab={} colab_type="code" id="LEqMBiMRWSUo"
# Create training and validation lists from the balanced df

from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle

train_ids, valid_ids = train_test_split(balanced_train_df, 
                 test_size = 0.2, 
                 stratify = balanced_train_df['ships'])  #Try and make sure nice distribution between train and val
train_df = pd.merge(all_df, train_ids)
valid_df = pd.merge(all_df, valid_ids)
train_df = shuffle(train_df) # Shuffle since same image would be grouped
print(train_df.shape[0], 'training masks')
print(valid_df.shape[0], 'validation masks')


# + colab={} colab_type="code" id="P_S8I3ViWSUt"
# Can double check....
#valid_df['ships'].hist(bins=train_df['ships'].max()+1)
#train_df['ships'].hist()

# + colab={} colab_type="code" id="8e8VIlUFWSUw"
# set lengths and steps
train_len = len(train_df)
val_len = len(valid_df)
images_list_len = train_len + val_len

steps_per_epoch = np.ceil(train_len // parms.BATCH_SIZE) # set step sizes based on train & batch
validation_steps = np.ceil(val_len // parms.BATCH_SIZE) # set step sizes based on val & batch

print("Total number: ", images_list_len, "  Train number: ", train_len, "  Val number: ", val_len)
print("Steps/EPOC: ", steps_per_epoch, "  Steps/Validation: ", validation_steps)


# + [markdown] colab_type="text" id="rj8gJz3AWSU1"
# ## Build, load and augment TensorFlow Datasets

# + colab={} colab_type="code" id="Wo5SrgekWSU2"
# Decode the image, convert to float, normalize by 255 and resize
def decode_img(image: tf.Tensor) -> tf.Tensor:
    # convert the compressed string to a 3D uint8 tensor
    image = tf.image.decode_jpeg(image, channels=parms.IMAGE_CHANNELS)
    # Use `convert_image_dtype` to convert to floats in the [0,1] range.
    image = tf.image.convert_image_dtype(image, parms.IMAGE_DTYPE)
    return image

def resize_image_mask(image: tf.Tensor, mask: tf.Tensor) -> tf.Tensor:
    image = tf.image.resize(image, [parms.IMAGE_ROWS, parms.IMAGE_COLS])
    mask = tf.image.resize(mask, [parms.IMAGE_ROWS, parms.IMAGE_COLS])
    return image, mask

def image_mask_aug(image, mask):
    
    if tf.random.uniform(()) > 0.5:    
        k = tf.random.uniform(shape=[], minval=1, maxval=3, dtype=tf.int32)
        image = tf.image.rot90(image, k) #0-4, 0/360, 90/180/270
        mask = tf.image.rot90(mask, k) #0-4, 0/360, 90/180/270

    if tf.random.uniform(()) > 0.5:
        image = tf.image.flip_left_right(image)
        mask = tf.image.flip_left_right(mask)
        
    if tf.random.uniform(()) > 0.5:
        image = tf.image.flip_up_down(image)
        mask = tf.image.flip_up_down(mask)

    return image, mask

def mask_wrapper(image_id_in):
    image_id = image_id_in.numpy().decode("utf-8")
    #print(type(image_id), image_id)
    encoded_pixels = all_df.query('ImageId==@image_id')['EncodedPixels']
    #print("EP ", encoded_pixels)
    mask = masks_as_image(encoded_pixels)
    return tf.convert_to_tensor(mask, dtype=tf.int16)


# method mapped to load image and mask
def process_train_image_id(image_id: tf.Tensor) -> tf.Tensor:

    [mask,] = tf.py_function(mask_wrapper, [image_id], [tf.int16])  #parms must be tensors
    mask.set_shape((768,768, 1))

    file_path = parms.TRAIN_PATH + "/" + image_id
    # load the raw data from the file as a string
    image = tf.io.read_file(file_path)
    image = decode_img(image)
    
    image, mask = resize_image_mask(image, mask)
    image, mask = image_mask_aug(image, mask)
    
    return image, mask

def process_val_image_id(image_id: tf.Tensor) -> tf.Tensor:

    [mask,] = tf.py_function(mask_wrapper, [image_id], [tf.int16])  #parms must be tensors
    mask.set_shape((768,768, 1))

    file_path = parms.TRAIN_PATH + "/" + image_id
    # load the raw data from the file as a string
    image = tf.io.read_file(file_path)
    image = decode_img(image)
    
    image, mask = resize_image_mask(image, mask)
       
    return image, mask


# + colab={} colab_type="code" id="yj6RWCqJWSU8"
# Create Dataset from pf
train_dataset = tf.data.Dataset.from_tensor_slices(train_df["ImageId"].values)

# Verify image paths were loaded
for image_id in train_dataset.take(2):
    print("Image id: ", image_id.numpy().decode("utf-8"))

# map training images to processing, includes any augmentation
train_dataset = train_dataset.map(process_train_image_id, num_parallel_calls=AUTOTUNE)

# Verify the mapping worked
for image, mask in train_dataset.take(1):
    print("Image shape: {}  Max: {}  Min: {}".format(image.numpy().shape, np.max(image.numpy()), np.min(image.numpy())))
    print("Encoded Pixels shape: ", mask.numpy().shape)
    some_image = image.numpy()
    some_mask = mask.numpy()

#show_batch_mask([some_image, some_mask])

train_dataset = train_dataset.batch(parms.BATCH_SIZE).repeat()

# + colab={} colab_type="code" id="c-HBARzmWSVp"
# Show the images, execute this cell multiple times to see the images
for image, mask in train_dataset.take(1):
    sample_image, sample_mask = image[0], mask[0]
show_batch_mask([sample_image, sample_mask])

# + colab={} colab_type="code" id="dGNwIor4WSVt"
# Create Dataset from pd
val_dataset = tf.data.Dataset.from_tensor_slices(valid_df["ImageId"].values)

# Verify image paths were loaded
for image_id in val_dataset.take(2):
    print("Image id: ", image_id.numpy().decode("utf-8"))

    # map training images to processing, includes any augmentation
val_dataset = val_dataset.map(process_val_image_id, num_parallel_calls=AUTOTUNE)

# Verify the mapping worked
for image, mask in val_dataset.take(1):
    print("Image shape: {}  Max: {}  Min: {}".format(image.numpy().shape, np.max(image.numpy()), np.min(image.numpy())))
    print("Encoded Pixels shape: ", mask.numpy().shape)
    some_image = image.numpy()
    some_mask = mask.numpy()

#show_batch_mask([some_image, some_mask])

val_dataset = val_dataset.batch(parms.BATCH_SIZE).repeat()

# + colab={} colab_type="code" id="V2_PmP3CWSVw"
# Final check before model training.  I added a string of the mask non-zero counts - need to make sure the masks 
# were created ok.  (got bit by this one after a small change....)

# Test Validation or Train by changing the dataset

mask_cnt_str = ""
sample_image = None
sample_mask = None
#for image, mask in train_dataset.take(1):
for image, mask in val_dataset.take(1):
    image_np = image.numpy()
    mask_np = mask.numpy()
    for i in range(len(image_np)):
        #show_batch_mask([image[i], mask[i]])  # Will show all of the batch
        mask_cnt_str = mask_cnt_str + str(np.count_nonzero(mask_np[i])) + "  "
        #print("Mask shape: {}  Max: {}  Min: {}".format(mask.numpy().shape, np.max(mask.numpy()), np.min(mask.numpy())))

        if np.count_nonzero(mask_np[i]) > 0:
            sample_image, sample_mask = image[i], mask[i]
            
print("Mask counts: ", mask_cnt_str)
show_batch_mask([sample_image, sample_mask])  # Will show the sample masks, if errors, then no mask was found content


# + colab={} colab_type="code" id="eaBs1MOaWSVy"
# Pre-trained model from this article.  Could not build it on Kaggle, built it locally, then loaded as personal data set.
# Did not train as well as building a new model.  Move cell down if you want to play with it....
# The article is VERY good!!
#https://www.tensorflow.org/tutorials/images/segmentation

#https://tensorlayer.com
#def iou_coe(output, target, threshold=0.5, axis=(1, 2, 3), smooth=1e-5):
#    pre = tf.cast(output > threshold, dtype=tf.float32)
#    truth = tf.cast(target > threshold, dtype=tf.float32)
#    inse = tf.reduce_sum(tf.multiply(pre, truth), axis=axis)  # AND
#    union = tf.reduce_sum(tf.cast(tf.add(pre, truth) >= 1, dtype=tf.float32), axis=axis)  # OR
#    batch_iou = (inse + smooth) / (union + smooth)
#    iou = tf.reduce_mean(batch_iou, name='iou_coe')
#    return iou  # , pre, truth, inse, union

#def compile_model_pre_trained(parms, model):
#    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3, decay=1e-6),
#                loss=combo_loss,
#                metrics=[iou_coe])
#    return model

#model = load_model("../input/unetmodel/baseModel.h5")
#model = compile_model_pre_trained(parms, model)

# + colab={} colab_type="code" id="vzVfXGUrWSV1"
# If you want to see the improvements after each EPOC, add to the callback. Helps to make sure show_predictions works...
class DisplayCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        clear_output(wait=True)
        show_predictions()
        print ('\nSample Prediction after epoch {}\n'.format(epoch+1))

# Normal callbacks
# I monitor val_loss, just seemed to work better....
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, CSVLogger
reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.33, patience=1, verbose=1, mode='min', min_delta=0.0001, 
                              cooldown=0, min_lr=1e-8)
earlystopper = EarlyStopping(monitor="val_loss", mode="min", verbose=2, patience=10)
checkpointer = ModelCheckpoint(parms.MODEL_PATH, monitor='val_loss', verbose=1, mode="auto", save_best_only=True)
 
# Methods to support training verification
def create_mask(pred_mask):
    pred_mask = np.where(pred_mask > 0.5, 1, 0)
    return pred_mask[0]

# Shows the image, original mask and predicted mask
def show_predictions(dataset=None, num=1):
    if dataset:
        for image, mask in dataset.take(num):
            pred_mask = model.predict(image)
            show_batch_mask([image[0], mask[0], create_mask(pred_mask)])
    else:
        show_batch_mask([sample_image, sample_mask,
             create_mask(model.predict(sample_image[tf.newaxis, ...]))])
       
# https://lars76.github.io/neural-networks/object-detection/losses-for-segmentation/
def combo_loss(y_true, y_pred):
    def dice_loss(y_true, y_pred):
        numerator = 2 * tf.reduce_sum(y_true * y_pred, axis=(1,2,3))
        denominator = tf.reduce_sum(y_true + y_pred, axis=(1,2,3))
        return tf.reshape(1 - numerator / denominator, (-1, 1, 1))
    return tf.keras.losses.binary_crossentropy(y_true, y_pred, from_logits=True) + dice_loss(y_true, y_pred)


def compile_model_unet(parms, model):
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3, decay=1e-6),
                loss=combo_loss,
                metrics=[tf.keras.metrics.MeanIoU(num_classes=2)])
    return model



# + colab={} colab_type="code" id="3q6fKKGGWSV4"
#UNet Builder: https://www.kaggle.com/yushas/imageprocessingtips
# I use this whenever I need a u-net, very easy to shrink or grow size/levels as needed.

def conv_block_mod(m, dim, acti, bn, res, do=0):
    n = tf.keras.layers.Conv2D(dim, 3, activation=acti, padding='same')(m)
    n = tf.keras.layers.BatchNormalization()(n) if bn else n
    n = tf.keras.layers.Dropout(do)(n) if do else n
    n = tf.keras.layers.Conv2D(dim, 3, activation=acti, padding='same')(n)
    n = tf.keras.layers.BatchNormalization()(n) if bn else n
    return tf.keras.layers.Concatenate()([m, n]) if res else n

def level_block_mod(m, dim, depth, inc, acti, do, bn, mp, up, res):
    if depth > 0:
        n = conv_block_mod(m, dim, acti, bn, res)
        m = tf.keras.layers.MaxPooling2D()(n) if mp else Conv2D(dim, 3, strides=2, padding='same')(n)
        m = level_block_mod(m, int(inc*dim), depth-1, inc, acti, do, bn, mp, up, res)#再帰
        if up:
            m = tf.keras.layers.UpSampling2D()(m)
            m = tf.keras.layers.Conv2D(dim, 2, activation=acti, padding='same')(m)
        else:
            m = tf.keras.layers.Conv2DTranspose(dim, 3, strides=2, activation=acti, padding='same')(m)
        n = tf.keras.layers.Concatenate()([n, m])
        m = conv_block_mod(n, dim, acti, bn, res)
    else:
        m = conv_block_mod(m, dim, acti, bn, res, do)
    return m

def UNet_mod(img_shape, out_ch=1, start_ch=64, depth=4, inc_rate=2., activation='relu', 
         dropout=False, batchnorm=True, maxpool=True, upconv=True, residual=False):
    i = tf.keras.layers.Input(shape=img_shape)
#    s = Lambda(lambda x: x / 255) (i)
    o = level_block_mod(i, start_ch, depth, inc_rate, activation, dropout, batchnorm, maxpool, upconv, residual)#Unet
    o = tf.keras.layers.Conv2D(out_ch, 1, activation=parms.FINAL_ACTIVATION)(o)
    return tf.keras.Model(inputs=i, outputs=o)



# + colab={} colab_type="code" id="Pv_g7bWGWSV8"
# Build the model
model=UNet_mod(parms.IMAGE_DIM,
               out_ch=parms.NUM_CLASSES,
               start_ch=16,
               depth=4,
               batchnorm=True,
               dropout=0.5)
model = compile_model_unet(parms, model)

# + colab={} colab_type="code" id="5RNwmgTSWSV_"
# Reload the model from prior runs
#model = load_model(parms.MODEL_PATH, custom_objects={'combo_loss': combo_loss})

# + colab={} colab_type="code" id="_AiBKnfWWSWC"
# Train from scratch
# You need to download the saved model and/or move to your personal dataset
# Once session ends, temp workspace is lost

history = model.fit(train_dataset,
                    validation_data=val_dataset,
                    epochs=parms.EPOCS, 
                    steps_per_epoch=steps_per_epoch,
                    validation_steps=validation_steps,
                    callbacks=[reduce_lr, earlystopper, checkpointer])


# + colab={} colab_type="code" id="PFoZt2xBWSWF"

history_df = pd.DataFrame(history.history)
plt.figure()
history_df[['loss', 'val_loss']].plot(title="Loss")
plt.xlabel('Epocs')
plt.ylabel('Loss')
history_df[['mean_io_u', 'val_mean_io_u']].plot(title="Mean IOU")
plt.xlabel('Epocs')
plt.ylabel('Accuracy')
#history_df[['accuracy', 'val_accuracy']].plot(title="Accuracy")
#plt.xlabel('Epocs')
#plt.ylabel('Accuracy')

plt.show()

# + colab={} colab_type="code" id="O5bBr0MtWSWI"
#history_df

# + [markdown] colab_type="text" id="HpP4qbg2WSWL"
# ### Validate the training...

# + colab={} colab_type="code" id="NPgQbjuwWSWM"
# Create Dataset from pd
test_df = shuffle(valid_df)
test_dataset = tf.data.Dataset.from_tensor_slices(test_df["ImageId"].values)

# Verify image paths were loaded
for image_id in test_dataset.take(2):
    print(image_id.numpy().decode("utf-8"))

    # map training images to processing, includes any augmentation
test_dataset = test_dataset.map(process_val_image_id, num_parallel_calls=AUTOTUNE)

# Verify the mapping worked
for image, mask in test_dataset.take(1):
    print("Image shape: {}  Max: {}  Min: {}".format(image.numpy().shape, np.max(image.numpy()), np.min(image.numpy())))
    print("Encoded Pixels shape: ", mask.numpy().shape)
    some_image = image.numpy()
    some_mask = mask.numpy()

#show_batch_mask([some_image, some_mask])

test_dataset = test_dataset.batch(1).repeat()

# + colab={} colab_type="code" id="3Ps11g7PWSWP"
show_predictions(test_dataset, 16)

StarcoderdataJupyterScriptsDedupFiltered