Spaces:
Sleeping
Sleeping
File size: 965 Bytes
b39cc0b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
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')
|