forked from s444426/AIProjekt
36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
import matplotlib.pyplot as plt
|
|
from tensorflow.keras.preprocessing.image import ImageDataGenerator
|
|
from tensorflow.keras.applications import MobileNetV2
|
|
from tensorflow.keras.applications.mobilenet import preprocess_input, decode_predictions
|
|
from tensorflow.keras.layers import Dense
|
|
from tensorflow.keras.models import Model
|
|
import numpy as np
|
|
import os, random
|
|
|
|
#podstawa modelu
|
|
base_model = MobileNetV2(include_top=False, weights="imagenet", pooling='avg')
|
|
|
|
#model
|
|
x=base_model.output
|
|
preds=Dense(4,activation='softmax')(x)
|
|
model=Model(inputs=base_model.input,outputs=preds)
|
|
|
|
#tylko ostatnie 20 warstw uczymy
|
|
for layer in model.layers[:20]:
|
|
layer.trainable=False
|
|
for layer in model.layers[20:]:
|
|
layer.trainable=True
|
|
|
|
#generator obrazkow
|
|
train_datagen=ImageDataGenerator(preprocessing_function=preprocess_input)
|
|
|
|
train_generator=train_datagen.flow_from_directory('./dataset', target_size=(224,224), color_mode='rgb', batch_size=32, class_mode='categorical', shuffle=True)
|
|
|
|
|
|
#uczenie //to dzielenie i podloga
|
|
model.compile(optimizer='Adam',loss='categorical_crossentropy',metrics=['accuracy'])
|
|
step_size_train=train_generator.n//train_generator.batch_size
|
|
model.fit_generator(generator=train_generator, steps_per_epoch=step_size_train, epochs=10)
|
|
|
|
#zapis
|
|
model.save('moj_model.h5') |