Блог

Я построил модель CNN для обнаружения двух видов дефектов на изображении. Эти классы «большие» и «маленькие», и точность действительно хорошая. Архитектура моей модели представлена ниже:

 inputs = tf.keras.Input(shape=(120, 120, 3))
x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu')(inputs)
x = tf.keras.layers.MaxPool2D(pool_size=(2, 2))(x)
x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu')(x)
x = tf.keras.layers.MaxPool2D(pool_size=(2, 2))(x)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)

model = tf.keras.Model(inputs=inputs, outputs=outputs)

model.compile(
    optimizer='adam',
    loss='binary_crossentropy',
    metrics=['accuracy']
)

print(model.summary())

history = model.fit(
    train_data,
    validation_data=val_data,
    epochs=100,
    callbacks=[
        tf.keras.callbacks.EarlyStopping(
            monitor='val_loss',
            patience=3,
            restore_best_weights=True
        )
    ]
)
 

Теперь я хочу использовать эту модель CNN для нескольких классов, и классы будут «большими», «маленькими», «другими». У меня есть набор данных, но я не знаю, как изменить модель для трех классов. Кроме того, в конце я хочу протестировать одно изображение в своей модели CNN и получить метку, если вставленное изображение большое, маленькое или другое, но я не знаю, как это сделать.

Попробуйте это:

 model = tf.keras.Sequential()
model.add(tf.keras.layers.Conv2D(input_shape = (120, 120, 3), filters=16, kernel_size=(3, 3), activation='relu'))
model.add(tf.keras.layers.MaxPool2D(pool_size=(2, 2)))
model.add(tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu'))
model.add(tf.keras.layers.MaxPool2D(pool_size=(2, 2)))
model.add(tf.keras.layers.GlobalAveragePooling2D())
model.add(tf.keras.layers.Dense(3, activation='softmax'))

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
 

для прогнозирования вы можете использовать этот код:

 from PIL import Image
import numpy as np
from skimage import transform

def load(filename):
    np_image = Image.open(filename)
    np_image = np.array(np_image).astype('float32')/255
    np_image = transform.resize(np_image, (120, 120, 3))
    np_image = np.expand_dims(np_image, axis=0)
    return np_image


folder_path = 'Dataset/test/4.jpg'
image = load(folder_path)
pred = model.predict_classes(image)
pred.tolist()[0]