+ color model and train script

This commit is contained in:
countingthedots 2023-06-02 11:24:21 +02:00
parent 563fa7da6a
commit a0c48a91ca
2 changed files with 69 additions and 7 deletions

View File

@ -7,14 +7,18 @@ import random
class VacuumRecognizer:
model = keras.models.load_model("D:/Image_dataset/model.h5")
model = keras.models.load_model("D:/Image_dataset/model_color.h5")
#model = keras.models.load_model("D:/Image_dataset/model.h5")
def recognize(self, image_path) -> str:
class_names = ['Banana', 'Cat', 'Earings', 'Plant']
img = cv2.imread(image_path, flags=cv2.IMREAD_GRAYSCALE)
img = cv2.imread(image_path)
#img = cv2.imread(image_path, flags=cv2.IMREAD_GRAYSCALE)
# print(img.shape)
cv2.imshow("lala", img)
cv2.imshow("img", img)
cv2.waitKey(0)
img = (np.expand_dims(img, 0))
@ -29,10 +33,15 @@ class VacuumRecognizer:
image_paths = []
image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Banana/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Cat/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Earings/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Plant/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Banana/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Cat/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Earings/')
image_paths.append('D:/Image_dataset/Image_datasetJPGnew/Image_datasetJPGnew/test/Plant/')
# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/train/Banana/')
# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Cat/')
# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Earings/')
# image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Plant/')
uio = VacuumRecognizer()

53
AI_brain/train.py Normal file
View File

@ -0,0 +1,53 @@
import numpy as np
import tensorflow as tf
from tensorflow import keras
def norm_pix(image, label):
return image / 255, label
directoryTRAIN = "D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/train"
directoryVALIDATION = "D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/validation"
train_ds = tf.keras.utils.image_dataset_from_directory(directoryTRAIN,
seed=123, batch_size=16,
image_size=(200, 200), color_mode='rgb')#color_mode='grayscale')
val_ds = tf.keras.utils.image_dataset_from_directory(directoryVALIDATION,
seed=123, batch_size=16,
image_size=(200, 200), color_mode='rgb')#color_mode='grayscale')
train_ds = train_ds.map(norm_pix)
val_ds = val_ds.map(norm_pix)
model = keras.Sequential([
#keras.layers.Conv2D(64, (3, 3), activation='relu', input_shape=(200, 200, 1)),
keras.layers.Conv2D(64, (3, 3), activation='relu', input_shape=(200, 200, 3)),
keras.layers.MaxPool2D((2, 2)),
keras.layers.Conv2D(128, (3, 3), activation='relu'),
keras.layers.MaxPool2D((2, 2)),
keras.layers.Conv2D(256, (3, 3), activation='relu'),
keras.layers.MaxPool2D((2, 2)),
keras.layers.Flatten(),
keras.layers.Dense(1024, activation='relu'),
keras.layers.Dense(4, activation='softmax')
])
print(model.summary())
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
trainHistory = model.fit(train_ds, epochs=8, validation_data=val_ds)
#model.save("model.h5")
model.save("model_color.h5")
(loss, accuracy) = model.evaluate(val_ds)
print(loss)
print(accuracy)