import os
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from models.model import build_model

train_dir = os.path.join('data', 'train')
validation_dir = os.path.join('data', 'validation')

train_datagen = ImageDataGenerator(rescale=1./255, rotation_range=40, width_shift_range=0.2,
                                   height_shift_range=0.2, shear_range=0.2, zoom_range=0.2, horizontal_flip=True)
validation_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(train_dir, target_size=(150, 150), batch_size=20, class_mode='categorical')
validation_generator = validation_datagen.flow_from_directory(validation_dir, target_size=(150, 150), batch_size=20, class_mode='categorical')

model = build_model()

history = model.fit(train_generator, steps_per_epoch=100, epochs=30,
                    validation_data=validation_generator, validation_steps=50)

model.save('bodybuilding_pose_classifier.h5')