101 lines
2.8 KiB
Python
101 lines
2.8 KiB
Python
import numpy as np
|
|
from keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
|
|
from keras.models import Sequential
|
|
from keras.layers import Conv2D, MaxPooling2D
|
|
from keras.layers import Activation, Dropout, Flatten, Dense, BatchNormalization
|
|
from keras import backend as K
|
|
img_width, img_height = 299, 299
|
|
|
|
train_data_dir = 'resources/zbior_uczacy'
|
|
validation_data_dir = 'resources/smieci'
|
|
nb_train_samples = 1599
|
|
nb_validation_samples = 1574
|
|
epochs = 1
|
|
batch_size = 16
|
|
|
|
def stworzModel():
|
|
|
|
|
|
if K.image_data_format() == 'channels_first':
|
|
input_shape = (3, img_width, img_height)
|
|
else:
|
|
input_shape = (img_width, img_height, 3)
|
|
|
|
model = Sequential()
|
|
model.add(Conv2D(32, (2, 2), input_shape=input_shape))
|
|
model.add(Activation('relu'))
|
|
model.add(MaxPooling2D(pool_size=(2, 2)))
|
|
|
|
model.add(Conv2D(32, (2, 2)))
|
|
model.add(Activation('relu'))
|
|
model.add(MaxPooling2D(pool_size=(2, 2)))
|
|
|
|
model.add(Conv2D(64, (2, 2)))
|
|
model.add(Activation('relu'))
|
|
model.add(MaxPooling2D(pool_size=(2, 2)))
|
|
|
|
model.add(Flatten())
|
|
model.add(Dense(64))
|
|
model.add(Activation('relu'))
|
|
model.add(Dropout(0.5))
|
|
model.add(Dense(4))
|
|
model.add(Activation('softmax'))
|
|
return model
|
|
|
|
def trainModel():
|
|
model = stworzModel()
|
|
model.compile(loss='categorical_crossentropy',
|
|
optimizer='rmsprop',
|
|
metrics=['accuracy'])
|
|
|
|
train_datagen = ImageDataGenerator(
|
|
rescale=1. / 255,
|
|
shear_range=0.2,
|
|
zoom_range=0.2,
|
|
horizontal_flip=True)
|
|
|
|
test_datagen = ImageDataGenerator(rescale=1. / 255)
|
|
|
|
train_generator = train_datagen.flow_from_directory(
|
|
train_data_dir,
|
|
target_size=(img_width, img_height),
|
|
batch_size=batch_size,
|
|
class_mode='categorical',
|
|
shuffle=True)
|
|
|
|
validation_generator = test_datagen.flow_from_directory(
|
|
validation_data_dir,
|
|
target_size=(img_width, img_height),
|
|
batch_size=batch_size,
|
|
class_mode='categorical',
|
|
shuffle=True)
|
|
|
|
model.fit_generator(
|
|
train_generator,
|
|
steps_per_epoch=nb_train_samples // batch_size,
|
|
epochs=epochs,
|
|
validation_data=validation_generator,
|
|
validation_steps=nb_validation_samples // batch_size,
|
|
shuffle=True)
|
|
|
|
model.save_weights('nowy_wytrenowany.h5')
|
|
|
|
def przewidz(path):
|
|
model = stworzModel()
|
|
model.load_weights('wytrenowany.h5')
|
|
img = load_img(path, target_size=(299, 299))
|
|
img_array = img_to_array(img)
|
|
img_array = np.expand_dims(img_array, axis=0)
|
|
prediction = model.predict(img_array)
|
|
np.argmax(prediction[0])
|
|
kategoria = np.argmax(prediction[0])
|
|
if kategoria == 0:
|
|
return "glass"
|
|
elif kategoria == 1:
|
|
return "metal"
|
|
elif kategoria == 2:
|
|
return "paper"
|
|
elif kategoria == 3:
|
|
return "plastic"
|
|
|