import os
import glob
import tensorflow as tf
import numpy as np
import cv2
from pylab import imshow
from IPython.display import Image
import matplotlib.pyplot as plt
import segmentation_models as sm
%matplotlib inline
#print(os.listdir("membrane/train"))
SIZE_X = 284 #Resize images (height = X, width = Y)
SIZE_Y = 284
MASK_SIZE_X = 196
MASK_SIZE_Y = 196
img_path_dir = "/Users/ginoprasad/ai_stuff/data/train"
mask_path_dir = "/Users/ginoprasad/ai_stuff/data/train_masks"
os.makedirs(f"models", exist_ok=True)
#Capture training image info as a list
train_images = []
train_image_names = []
for img_path in glob.glob(os.path.join(img_path_dir, "*.jpg"))[:1000]:
train_image_names.append(os.path.basename(img_path).rstrip(".jpg"))
img = cv2.imread(img_path, cv2.IMREAD_COLOR)
img = cv2.resize(img, (SIZE_Y, SIZE_X))
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
train_images.append(img)
train_images = np.array(train_images) / 255
train_images.shape
(1000, 284, 284, 3)
#Capture mask/label info as a list
train_masks = []
for name in train_image_names:
mask_path = os.path.join(mask_path_dir, name + "_mask.gif")
assert os.path.exists(mask_path)
gif = cv2.VideoCapture(mask_path)
valid, mask = gif.read()
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
# mask = cv2.imread(mask_path, 0)
mask = cv2.resize(mask, (MASK_SIZE_Y, MASK_SIZE_X))
mask = np.ma.make_mask(mask).astype(np.float32)
train_masks.append(mask)
train_masks = np.array(train_masks)
train_masks.shape
(1000, 196, 196)
#Use customary x_train and y_train variables
X = train_images
Y = train_masks
Y = np.expand_dims(Y, axis=3) #May not be necessary.. leftover from previous code
from sklearn.model_selection import train_test_split
x_train, x_val, y_train, y_val = train_test_split(X, Y, test_size=0.2, random_state=42)
x_train[0].shape
(284, 284, 3)
NUM_ENCODERS = 4
NUM_DECODERS = 3
NUM_CONVOLUTIONS = 2
ENCODER_FILTER_SIZES = [64, 128, 256, 512]
DECODER_FILTER_SIZES = [256, 128, 64]
LEARNING_RATE = 0.0001
EPOCHS = 30
BATCH_SIZE = 1
inputs = tf.keras.layers.Input(shape=x_train[0].shape)
layers = [inputs]
for i, filter_size in zip(range(NUM_ENCODERS), ENCODER_FILTER_SIZES):
sub_layers = []
if i:
# Max Pooling
sub_layers.append(tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=None))
for _ in range(NUM_CONVOLUTIONS):
# 2D Conv
sub_layers.append(tf.keras.layers.Conv2D(filters=filter_size, kernel_size=(3, 3), activation='relu'))
encoder = tf.keras.Sequential(sub_layers)(layers[-1])
layers.append(encoder)
for i, filter_size in zip(range(NUM_DECODERS), DECODER_FILTER_SIZES):
# deconvolve layer -1 -> size 2 * layer -1
deconv = tf.keras.Sequential([tf.keras.layers.UpSampling2D(size=(2, 2)), tf.keras.layers.Conv2D(filters=filter_size, kernel_size=(2, 2), activation='relu', padding='same')])(layers[-1])
# crop layer NUM_ENCODERS-i-1
crop_val = (layers[NUM_ENCODERS-i-1].shape[-3] - deconv.shape[-3]) // 2
cropped = tf.keras.layers.Cropping2D(cropping=(crop_val))(layers[NUM_ENCODERS-i-1])
decoder = fusion = tf.keras.layers.Concatenate()([cropped, deconv])
for _ in range(NUM_CONVOLUTIONS):
decoder = tf.keras.layers.Conv2D(filters=filter_size, kernel_size=(3, 3), activation='relu')(decoder)
layers.append(decoder)
outputs = tf.keras.layers.Conv2D(filters=1, kernel_size=(1, 1), activation='sigmoid')(layers[-1])
layers.append(outputs)
model = tf.keras.Model(inputs, outputs, name="test_unet")
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE), # Optimizer
# Loss function to minimize
loss=sm.losses.bce_jaccard_loss,
# List of metrics to monitor
metrics=[sm.metrics.iou_score],
)
tf.keras.utils.plot_model(model, 'unet_implementation.png', show_shapes=True)
callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=3)
history = model.fit(x_train, y_train, validation_data=(x_val, y_val), epochs=EPOCHS, batch_size=BATCH_SIZE, callbacks=[callback]).history
model.save(f"models/test_car_model_lr_{LEARNING_RATE}_epochs_{EPOCHS}.h5")
Epoch 1/30
2022-10-16 16:03:18.914037: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:114] Plugin optimizer for device_type GPU is enabled.
800/800 [==============================] - ETA: 0s - loss: 0.8779 - iou_score: 0.4541
2022-10-16 16:08:17.535113: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:114] Plugin optimizer for device_type GPU is enabled.
800/800 [==============================] - 321s 400ms/step - loss: 0.8779 - iou_score: 0.4541 - val_loss: 0.4666 - val_iou_score: 0.7286 Epoch 2/30 800/800 [==============================] - 315s 393ms/step - loss: 0.4553 - iou_score: 0.7110 - val_loss: 0.3662 - val_iou_score: 0.7586 Epoch 3/30 800/800 [==============================] - 328s 410ms/step - loss: 0.3572 - iou_score: 0.7693 - val_loss: 0.3508 - val_iou_score: 0.7880 Epoch 4/30 800/800 [==============================] - 313s 392ms/step - loss: 0.3219 - iou_score: 0.7920 - val_loss: 0.3893 - val_iou_score: 0.7513 Epoch 5/30 800/800 [==============================] - 318s 398ms/step - loss: 0.2944 - iou_score: 0.8088 - val_loss: 0.2822 - val_iou_score: 0.8173 Epoch 6/30 800/800 [==============================] - 311s 389ms/step - loss: 0.2773 - iou_score: 0.8207 - val_loss: 0.2875 - val_iou_score: 0.8233 Epoch 7/30 800/800 [==============================] - 318s 397ms/step - loss: 0.2502 - iou_score: 0.8375 - val_loss: 0.2680 - val_iou_score: 0.8407 Epoch 8/30 800/800 [==============================] - 218s 273ms/step - loss: 0.2280 - iou_score: 0.8518 - val_loss: 0.2485 - val_iou_score: 0.8414 Epoch 9/30 800/800 [==============================] - 166s 208ms/step - loss: 0.2092 - iou_score: 0.8639 - val_loss: 0.2367 - val_iou_score: 0.8577 Epoch 10/30 800/800 [==============================] - 175s 218ms/step - loss: 0.1977 - iou_score: 0.8719 - val_loss: 0.2116 - val_iou_score: 0.8656 Epoch 11/30 800/800 [==============================] - 175s 218ms/step - loss: 0.1763 - iou_score: 0.8852 - val_loss: 0.2095 - val_iou_score: 0.8721 Epoch 12/30 800/800 [==============================] - 198s 248ms/step - loss: 0.1729 - iou_score: 0.8879 - val_loss: 0.2115 - val_iou_score: 0.8712 Epoch 13/30 800/800 [==============================] - 207s 259ms/step - loss: 0.1569 - iou_score: 0.8980 - val_loss: 0.2097 - val_iou_score: 0.8737 Epoch 14/30 800/800 [==============================] - 212s 265ms/step - loss: 0.1417 - iou_score: 0.9078 - val_loss: 0.2002 - val_iou_score: 0.8799 Epoch 15/30 800/800 [==============================] - 191s 238ms/step - loss: 0.1270 - iou_score: 0.9173 - val_loss: 0.2023 - val_iou_score: 0.8837 Epoch 16/30 800/800 [==============================] - 194s 242ms/step - loss: 0.1195 - iou_score: 0.9220 - val_loss: 0.1925 - val_iou_score: 0.8915 Epoch 17/30 800/800 [==============================] - 192s 241ms/step - loss: 0.1069 - iou_score: 0.9303 - val_loss: 0.1918 - val_iou_score: 0.8895 Epoch 18/30 800/800 [==============================] - 197s 247ms/step - loss: 0.1016 - iou_score: 0.9338 - val_loss: 0.2069 - val_iou_score: 0.8926 Epoch 19/30 800/800 [==============================] - 202s 252ms/step - loss: 0.0978 - iou_score: 0.9364 - val_loss: 0.2099 - val_iou_score: 0.8853 Epoch 20/30 800/800 [==============================] - 205s 256ms/step - loss: 0.0851 - iou_score: 0.9446 - val_loss: 0.1974 - val_iou_score: 0.8969 Epoch 21/30 800/800 [==============================] - 225s 281ms/step - loss: 0.0821 - iou_score: 0.9465 - val_loss: 0.1853 - val_iou_score: 0.8988 Epoch 22/30 800/800 [==============================] - 250s 313ms/step - loss: 0.0746 - iou_score: 0.9513 - val_loss: 0.1949 - val_iou_score: 0.8977 Epoch 23/30 800/800 [==============================] - 254s 317ms/step - loss: 0.0733 - iou_score: 0.9523 - val_loss: 0.2038 - val_iou_score: 0.8974 Epoch 24/30 800/800 [==============================] - 248s 310ms/step - loss: 0.0692 - iou_score: 0.9549 - val_loss: 0.1912 - val_iou_score: 0.9007 Epoch 25/30 800/800 [==============================] - 247s 309ms/step - loss: 0.0685 - iou_score: 0.9555 - val_loss: 0.1920 - val_iou_score: 0.9025 Epoch 26/30 800/800 [==============================] - 161s 201ms/step - loss: 0.0646 - iou_score: 0.9580 - val_loss: 0.1941 - val_iou_score: 0.8967 Epoch 27/30 800/800 [==============================] - 155s 194ms/step - loss: 0.0626 - iou_score: 0.9594 - val_loss: 0.2051 - val_iou_score: 0.9030 Epoch 28/30 800/800 [==============================] - 160s 201ms/step - loss: 0.0561 - iou_score: 0.9635 - val_loss: 0.2006 - val_iou_score: 0.9039 Epoch 29/30 800/800 [==============================] - 165s 206ms/step - loss: 0.0561 - iou_score: 0.9635 - val_loss: 0.2000 - val_iou_score: 0.9007 Epoch 30/30 800/800 [==============================] - 166s 207ms/step - loss: 0.0542 - iou_score: 0.9648 - val_loss: 0.2025 - val_iou_score: 0.8935
plt.xlabel('Epoch #')
plt.plot(history['loss'], label='loss')
plt.plot(history['val_loss'], label='val_loss')
plt.plot(history['iou_score'], label='iou_score')
plt.plot(history['val_iou_score'], label='val_iou_score')
plt.legend(loc='upper right')
plt.savefig(f"loss_accuracy_plot.png", dpi='figure')
model = tf.keras.models.load_model(f"models/test_car_model_lr_{LEARNING_RATE}_epochs_{EPOCHS}.h5", custom_objects={'binary_crossentropy_plus_jaccard_loss': sm.losses.bce_jaccard_loss, 'iou_score': sm.metrics.iou_score})
val_index = 8
imshow(x_val[val_index])
plt.figure()
imshow(model.predict(x_val[val_index:val_index+1])[0])
2022-10-18 12:52:43.869092: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz
1/1 [==============================] - 1s 636ms/step
2022-10-18 12:52:44.302327: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:114] Plugin optimizer for device_type GPU is enabled.
<matplotlib.image.AxesImage at 0x2e29866d0>
