76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
import os
|
|
|
|
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
|
|
|
|
import numpy as np
|
|
|
|
import tensorflow as tf
|
|
|
|
|
|
class EfficientNetWrapper:
|
|
|
|
def __init__(self, classes, base_weights_path=None, weights_path=None):
|
|
self.__classes = classes
|
|
self.__input_shape = (224, 224, 3)
|
|
self.model = self.__build(base_weights_path)
|
|
self.model.load_weights(weights_path)
|
|
|
|
@property
|
|
def input_shape(self):
|
|
return self.__input_shape
|
|
|
|
@property
|
|
def classes(self):
|
|
return self.__classes
|
|
|
|
@staticmethod
|
|
def __preprocess_tensor(tensor):
|
|
return tf.keras.applications.efficientnet.preprocess_input(tensor)
|
|
|
|
@staticmethod
|
|
def __images_to_tensor(images):
|
|
return np.array(list(map(tf.keras.preprocessing.image.img_to_array, images)))
|
|
|
|
def __resize_and_convert(self, image):
|
|
return image.resize(self.input_shape[:-1]).convert("RGB")
|
|
|
|
def prep_images(self, images):
|
|
images = map(self.__resize_and_convert, images)
|
|
tensor = self.__images_to_tensor(images)
|
|
tensor = self.__preprocess_tensor(tensor)
|
|
|
|
return tensor
|
|
|
|
def __build(self, base_weights=None) -> tf.keras.models.Model:
|
|
input_img = tf.keras.layers.Input(shape=self.input_shape)
|
|
|
|
pretrained = tf.keras.applications.efficientnet.EfficientNetB0(
|
|
include_top=False, input_tensor=tf.keras.layers.Input(shape=self.input_shape), weights=base_weights
|
|
)
|
|
|
|
pretrained.trainable = False
|
|
|
|
for layer in pretrained.layers:
|
|
layer.trainable = False
|
|
|
|
pretrained = pretrained(input_img)
|
|
|
|
finetuned = tf.keras.layers.Flatten()(pretrained)
|
|
finetuned = tf.keras.layers.Dense(512, activation="relu")(finetuned)
|
|
finetuned = tf.keras.layers.Dropout(0.2)(finetuned)
|
|
finetuned = tf.keras.layers.Dense(128, activation="relu")(finetuned)
|
|
finetuned = tf.keras.layers.Dropout(0.2)(finetuned)
|
|
finetuned = tf.keras.layers.Dense(32, activation="relu")(finetuned)
|
|
finetuned = tf.keras.layers.Dropout(0.2)(finetuned)
|
|
finetuned = tf.keras.layers.Dense(len(self.classes), activation="softmax")(finetuned)
|
|
|
|
model = tf.keras.models.Model(inputs=input_img, outputs=finetuned)
|
|
|
|
model.compile(
|
|
loss="categorical_crossentropy",
|
|
optimizer="adam",
|
|
metrics=[tf.keras.metrics.Recall(), tf.keras.metrics.Precision()],
|
|
)
|
|
|
|
return model
|