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')