from keras.models import Model
import keras
from keras.utils import plot_model
from keras.layers import Activation, Dropout, UpSampling2D, concatenate, Input
from keras.layers import Conv2DTranspose, Conv2D, MaxPooling2D
from PIL import Image
import numpy as np
import pickle
import cv2
import os
import tensorflow as tf
from keras import layers
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
from keras.models import load_model

Import necessary items from Keras

from keras.layers.normalization import BatchNormalization
from keras.preprocessing.image import ImageDataGenerator

from keras import regularizers
from keras import optimizers
import matplotlib.pyplot as plt
import keras.backend as K

from keras.optimizers import Adam
from keras.optimizers import SGD
from keras.layers import Dense,AveragePooling2D,ZeroPadding2D
from keras.layers import add,Flatten
from tensorflow.keras.callbacks import TensorBoard
import time
img_rows=80
img_cols=160
########################
def training_vis(hist):
loss = hist.history[‘loss’]
val_loss = hist.history[‘val_loss’]
acc = hist.history[‘acc’]
val_acc = hist.history[‘val_acc’]

# make a figure
fig = plt.figure(figsize=(8,4))
# subplot loss
ax1 = fig.add_subplot(121)
ax1.plot(loss,label='train_loss')
ax1.plot(val_loss,label='val_loss')
ax1.set_xlabel('Epochs')
ax1.set_ylabel('Loss')
ax1.set_title('Loss on Training and Validation Data')
ax1.legend()
# subplot acc
ax2 = fig.add_subplot(122)
ax2.plot(acc,label='train_acc')
ax2.plot(val_acc,label='val_acc')
ax2.set_xlabel('Epochs')
ax2.set_ylabel('Accuracy')
ax2.set_title('Accuracy  on Training and Validation Data')
ax2.legend()
plt.tight_layout()

log_filepath=’’#tensorboard路径文件
model_name = “resunet-cnn-64x2-{}”.format(int(time.time()))
tensorboard = TensorBoard(log_dir=log_filepath.format(model_name))

train_images=[]
labels=[]
#loading data
path=#训练图片路径
for imag in os.listdir(path):
str=path+imag
lab=cv2.imread(str)
y = cv2.cvtColor(lab, cv2.COLOR_RGB2GRAY).reshape(80,160,1)

labels.append(y)

for imag in os.listdir('path):
str=path+imag
img=cv2.imread(str).reshape(80,160,3)

train_images.append(img)

if len(labels)==len(train_images):
print(‘data load success’)
else:
print(‘data load filed’)

Make into arrays as the neural network wants these

train_images = np.array(train_images)
labels = np.array(labels)
print(‘label:’,labels[0].shape)
print(‘image:’,train_images[0].shape)

Normalize labels - training images get normalized to start in the network

labels = labels /50.
#255

Shuffle images along with their labels, then split into training/validation sets

train_images, labels = shuffle(train_images, labels)

Test size may be 10% or 20%

X_train, X_val, y_train, y_val = train_test_split(train_images, labels, test_size=0.2)

Batch size, epochs and pool size below are all paramaters to fiddle with for optimization

batch_size =16# 16
epochs =200

Here is the actual neural network

Normalizes incoming inputs. First layer needs the input shape to work

#BatchNormalization(input_shape=input_shape)
inputs = Input(batch_shape=(None, 80, 160, 3))

#################
def get_unet(inputs):

conv1 = Conv2D(32, (3, 3), activation='relu', padding='same')(inputs)
conv1 = Conv2D(32, (3, 3), activation='relu', padding='same')(conv1)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
# pool1 = Dropout(0.25)(pool1)
pool1 = BatchNormalization()(pool1)


conv2 = Conv2D(64, (3, 3), activation='relu', padding='same')(pool1)
conv2 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv2)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
# pool2 = Dropout(0.5)(pool2)
pool2 = BatchNormalization()(pool2)


conv3 = Conv2D(128, (3, 3), activation='relu', padding='same')(pool2)
conv3 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv3)
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
# pool3 = Dropout(0.5)(pool3)
pool3 = BatchNormalization()(pool3)


conv4 = Conv2D(256, (3, 3), activation='relu', padding='same')(pool3)
conv4 = Conv2D(256, (3, 3), activation='relu', padding='same')(conv4)
pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)
# pool4 = Dropout(0.5)(pool4)
pool4 = BatchNormalization()(pool4)


conv5 = Conv2D(512, (3, 3), activation='relu', padding='same')(pool4)
conv5 = Conv2D(512, (3, 3), activation='relu', padding='same')(conv5)


up6 = concatenate([Conv2DTranspose(256, (2, 2), strides=(
    2, 2), padding='same')(conv5), conv4], axis=3)
# up6 = Dropout(0.5)(up6)
up6 = BatchNormalization()(up6)
conv6 = Conv2D(256, (3, 3), activation='relu', padding='same')(up6)
conv6 = Conv2D(256, (3, 3), activation='relu', padding='same')(conv6)


up7 = concatenate([Conv2DTranspose(128, (2, 2), strides=(
    2, 2), padding='same')(conv6), conv3], axis=3)
# up7 = Dropout(0.5)(up7)
up7 = BatchNormalization()(up7)
conv7 = Conv2D(128, (3, 3), activation='relu', padding='same')(up7)
conv7 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv7)

up8 = concatenate([Conv2DTranspose(64, (2, 2), strides=(
    2, 2), padding='same')(conv7), conv2], axis=3)
# up8 = Dropout(0.5)(up8)
up8 = BatchNormalization()(up8)
conv8 = Conv2D(64, (3, 3), activation='relu', padding='same')(up8)
conv8 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv8)

up9 = concatenate([Conv2DTranspose(32, (2, 2), strides=(
    2, 2), padding='same')(conv8), conv1], axis=3)
# up9 = Dropout(0.5)(up9)
up9 = BatchNormalization()(up9)
conv9 = Conv2D(32, (3, 3), activation='relu', padding='same')(up9)
conv9 = Conv2D(32, (3, 3), activation='relu', padding='same')(conv9)

# conv9 = Dropout(0.5)(conv9)

conv10 = Conv2D(1, (1, 1), activation='sigmoid')(conv9)

model = Model(inputs=[inputs], outputs=[conv10])

#model.compile(optimizer=Adam(lr=1e-5),
#              loss=dice_coef_loss, metrics=[dice_coef])
model.compile(optimizer='Adam', loss='mean_squared_error',metrics=['accuracy'])
model.summary()
#plot_model(model, to_file='unet.png')
return model

inputs = Input(batch_shape=(None, 80, 160, 3))
datagen = ImageDataGenerator(channel_shift_range=0.2)
datagen.fit(X_train)

Freeze layers since training is done

model=get_unet(inputs)
#weight_file = # 保存/加载地址
#model = load_model(weight_file) # 载入已保存的模型文件

hist=model.fit_generator(datagen.flow(X_train, y_train, batch_size=batch_size), steps_per_epoch=len(X_train)/batch_size,
epochs=epochs, verbose=1, validation_data=(X_val, y_val), callbacks=[tensorboard])

#model.compile(loss=‘mean_squared_error’, optimizer=sgd)
model.compile(optimizer=‘Adam’, loss=‘mean_squared_error’, metrics=[‘accuracy’])

Save model architecture and weights

model.save(‘.h5’)#模型保存路径
training_vis(hist)