Scribd
Upload
35 views3 pages

Original Code

Uploaded by

AlexandruVlad
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd
35 views3 pages

Original Code

Uploaded by

AlexandruVlad
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd
You are on page 1/ 3

import tensorflow as tf

import numpy as np
import pickle
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dropout, UpSampling2D
from tensorflow.keras.layers import Conv2DTranspose, Conv2D, MaxPooling2D
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.python.keras.layers import BatchNormalization
import matplotlib.pyplot as plt

tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

def create_model(input_shape, pool_size):


# Create the actual neural network here
model = Sequential()
# Normalizes incoming inputs. First layer needs the input shape to work
model.add(BatchNormalization(input_shape=input_shape))

# Below layers were re-named for easier reading of model summary; this not
necessary
# Conv Layer 1
model.add(Conv2D(8, (3, 3), strides=(1, 1), activation='relu', name='Conv1'))

# Conv Layer 2
model.add(Conv2D(16, (3, 3), strides=(1, 1), activation='relu', name='Conv2'))

# Pooling 1
model.add(MaxPooling2D(pool_size=pool_size))

# Conv Layer 3
model.add(Conv2D(16, (3, 3), strides=(1, 1), activation='relu', name='Conv3'))
model.add(Dropout(0.2))

# Conv Layer 4
model.add(Conv2D(32, (3, 3), strides=(1, 1), activation='relu', name='Conv4'))
model.add(Dropout(0.2))

# Conv Layer 5
model.add(Conv2D(32, (3, 3), strides=(1, 1), activation='relu', name='Conv5'))
model.add(Dropout(0.2))

# Pooling 2
model.add(MaxPooling2D(pool_size=pool_size))

# Conv Layer 6
model.add(Conv2D(64, (3, 3), strides=(1, 1), activation='relu', name='Conv6'))
model.add(Dropout(0.2))

# Conv Layer 7
model.add(Conv2D(64, (3, 3), strides=(1, 1), activation='relu', name='Conv7'))
model.add(Dropout(0.2))

# Pooling 3
model.add(MaxPooling2D(pool_size=pool_size))

# Upsample 1
model.add(UpSampling2D(size=pool_size))

# Deconv 1
model.add(Conv2DTranspose(64, (3, 3), strides=(1, 1), activation='relu',
name='Deconv1'))
model.add(Dropout(0.2))

# Deconv 2
model.add(Conv2DTranspose(64, (3, 3), strides=(1, 1), activation='relu',
name='Deconv2'))
model.add(Dropout(0.2))

# Upsample 2
model.add(UpSampling2D(size=pool_size))

# Deconv 3
model.add(Conv2DTranspose(32, (3, 3), strides=(1, 1), activation='relu',
name='Deconv3'))
model.add(Dropout(0.2))

# Deconv 4
model.add(Conv2DTranspose(32, (3, 3), strides=(1, 1), activation='relu',
name='Deconv4'))
model.add(Dropout(0.2))

# Deconv 5
model.add(Conv2DTranspose(16, (3, 3), strides=(1, 1), activation='relu',
name='Deconv5'))
model.add(Dropout(0.2))

# Upsample 3
model.add(UpSampling2D(size=pool_size))

# Deconv 6
model.add(Conv2DTranspose(16, (3, 3), strides=(1, 1), activation='relu',
name='Deconv6'))

# Final layer - only including one channel so 1 filter


model.add(Conv2DTranspose(1, (3, 3), strides=(1, 1), activation='relu',
name='Final'))

return model

def main():
# Load training images
train_images = pickle.load(open("full_CNN_train.p", "rb"))

# Load image labels


labels = pickle.load(open("full_CNN_labels.p", "rb"))

# Make into arrays as the neural network wants these


train_images = np.array(train_images)
labels = np.array(labels)

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


labels = labels / 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.1)

# Batch size, epochs and pool size below are all paramaters to fiddle with for
optimization
batch_size = 128
epochs = 10
pool_size = (2, 2)
input_shape = X_train.shape[1:]

# Create the neural network


model = create_model(input_shape, pool_size) # type: Sequential

# Using a generator to help the model use less data


# Channel shifts help with shadows slightly
datagen = ImageDataGenerator(channel_shift_range=0.2)
datagen.fit(X_train)

# Compiling and training the model


model.compile(optimizer='Adam', loss='mean_squared_error',
metrics=['accuracy'])
model.fit_generator(datagen.flow(X_train, y_train, batch_size=batch_size),
steps_per_epoch=len(X_train) / batch_size, epochs=epochs,
validation_data=(X_val, y_val))

# Freeze layers since training is done


model.trainable = False
model.compile(optimizer='Adam', loss='mean_squared_error')

# Save model architecture and weights


model.save('full_CNN_model_test.h5')

# Show summary of model


model.summary()

if __name__ == '__main__':
main()

You might also like