train.py

import tensorflow as tf
import numpy as np

from tensorflow.keras.preprocessing.image import ImageDataGenerator

from model import *
from utils import *

train_image_generator = ImageDataGenerator(rescale=1./255)
validation_image_generator = ImageDataGenerator(rescale=1./255) 

train_data_gen = train_image_generator.flow_from_directory(batch_size=BATCH_SIZE,
                                                           directory=TRAIN_DATASET_PATH,
                                                           shuffle=True,
                                                           target_size=(IMAGE_SIZE, IMAGE_SIZE),
                                                           class_mode='binary')

val_data_gen = validation_image_generator.flow_from_directory(batch_size=BATCH_SIZE,
                                                              directory=VAL_DATASET_PATH,
                                                              target_size=(IMAGE_SIZE, IMAGE_SIZE),
                                                              class_mode='binary')

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

model.summary()

history = model.fit_generator(
    train_data_gen,
    steps_per_epoch=9980 // BATCH_SIZE,
    epochs=N_EPOCHS,
    validation_data=val_data_gen
)
model.save_weights('weights/', save_format='tf')