137 lines
3.5 KiB
Python
137 lines
3.5 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import tensorflow as tf
|
|
import os
|
|
|
|
DIR = "data/binary-classification"
|
|
|
|
#Import Data
|
|
PATH = os.path.join(os.getcwd(), DIR)
|
|
|
|
training_data = os.path.join(PATH, "train")
|
|
validation_data = os.path.join(PATH, "valid")
|
|
test_data = os.path.join(PATH, "test")
|
|
|
|
BATCH_SIZE = 64
|
|
IMG_SIZE = (224,224)
|
|
|
|
#TODO: Import data from both directories, then resplit into test, train, and validation
|
|
|
|
print(f"Train: {len(training_data)}\nValid: {len(validation_data)}\nTest: {len(test_data)}")
|
|
|
|
|
|
#View Data
|
|
plt.figure(figsize=(10,10))
|
|
for images, labels in training_data.take(1):
|
|
for i in range(9):
|
|
ax = plt.subplot(3, 3, i+1)
|
|
plt.imshow(images[i].numpy().astype("uint8"))
|
|
plt.title(class_names[labels[i]])
|
|
plt.axis("off")
|
|
plt.show()
|
|
|
|
|
|
#Init Prefetching
|
|
AUTOTUNE = tf.data.AUTOTUNE
|
|
training_data = training_data.prefetch(buffer_size=AUTOTUNE)
|
|
validation_data = validation_data.prefetch(buffer_size=AUTOTUNE)
|
|
test_data = test_data.prefetch(buffer_size=AUTOTUNE)
|
|
|
|
|
|
#Data Augmentation
|
|
data_augmentation = tf.keras.Sequential([
|
|
tf.keras.layers.RandomFlip('horizontal'),
|
|
tf.keras.layers.RandomRotation(0.2)
|
|
])
|
|
|
|
#Create Base Model From MobileNetV3
|
|
IMG_SHAPE = IMG_SIZE + (3,)
|
|
base_model = tf.keras.applications.MobileNetV3Large(
|
|
input_shape=IMG_SHAPE,
|
|
include_top=False,
|
|
weights="imagenet"
|
|
)
|
|
|
|
image_batch, label_batch = next(iter(training_data))
|
|
feature_batch = base_model(image_batch)
|
|
|
|
base_model.trainable = False
|
|
|
|
|
|
#View Base Model
|
|
base_model.summary()
|
|
|
|
#Add Classification Header
|
|
|
|
|
|
global_avg_layer = tf.keras.layers.GlobalAveragePooling2D()
|
|
feature_batch_avg = global_avg_layer(feature_batch)
|
|
|
|
prediction_layer = tf.keras.layers.Dense(38, activation="softmax")
|
|
predication_batch = prediction_layer(feature_batch_avg)
|
|
|
|
inputs = tf.keras.Input(shape=(160,160,3))
|
|
x = data_augmentation(inputs)
|
|
x = base_model(x, training=False)
|
|
x = global_avg_layer(x)
|
|
x = tf.keras.layers.Dropout(0.2)(x)
|
|
outputs = prediction_layer(x)
|
|
|
|
model = tf.keras.Model(inputs, outputs)
|
|
|
|
#View Model with Classification Head
|
|
model.summary()
|
|
|
|
|
|
#Compile the Model
|
|
base_learning_rate = 0.0001
|
|
|
|
training_data = training_data.map(lambda x,y: (x, tf.one_hot(y,38)))
|
|
validation_data = validation_data.map(lambda x,y: (x, tf.one_hot(y,38)))
|
|
test_data = test_data.map(lambda x,y: (x, tf.one_hot(y,38)))
|
|
|
|
optimizer = tf.keras.optimizers.Adam(learning_rate=base_learning_rate)
|
|
loss = tf.keras.losses.CategoricalCrossentropy()
|
|
metrics = [tf.keras.metrics.CategoricalAccuracy()]
|
|
|
|
|
|
|
|
model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
|
|
|
|
#Train the Model
|
|
initial_epochs = 50
|
|
|
|
loss0, accuracy0 = model.evaluate(validation_data)
|
|
|
|
print(f"initial loss: {loss0}")
|
|
print(f"initial accuracy: {accuracy0}")
|
|
|
|
lr_schedule = tf.keras.callbacks.ReduceLROnPlateau(
|
|
monitor="val_loss",
|
|
factor=0.1,
|
|
patience=5,
|
|
min_lr=1e-6
|
|
)
|
|
|
|
early_stopping = tf.keras.callbacks.EarlyStopping(
|
|
monitor="val_loss",
|
|
patience=10,
|
|
restore_best_weights=True
|
|
)
|
|
|
|
history = model.fit(training_data,
|
|
epochs=initial_epochs,
|
|
validation_data=validation_data,
|
|
callbacks=[lr_schedule, early_stopping])
|
|
|
|
model.save("crop-classifier-better-test.keras")
|
|
|
|
#Evaluate Model
|
|
results = model.evaluate(validation_data)
|
|
print(f"Validation Loss: {results[0]}")
|
|
print(f"Validation Accuracy: {results[1]}")
|
|
|
|
results = model.evaluate(test_data)
|
|
print(f"Test Loss: {results[0]}")
|
|
print(f"Test Accuracy: {results[1]}")
|