Update ML setup guide and add model scripts
This commit is contained in:
parent
073d395ae6
commit
ff1561a704
4 changed files with 121 additions and 3 deletions
82
train_model.py
Normal file
82
train_model.py
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
import os
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras import layers, models, optimizers
|
||||
from tensorflow.keras.preprocessing.image import ImageDataGenerator
|
||||
|
||||
def create_model(num_classes):
|
||||
# Use MobileNetV2 as the base model
|
||||
base_model = tf.keras.applications.MobileNetV2(
|
||||
input_shape=(64, 64, 3),
|
||||
include_top=False,
|
||||
weights='imagenet'
|
||||
)
|
||||
base_model.trainable = False # Freeze base model for transfer learning
|
||||
|
||||
model = models.Sequential([
|
||||
base_model,
|
||||
layers.GlobalAveragePooling2D(),
|
||||
layers.Dropout(0.2),
|
||||
layers.Dense(num_classes, activation='softmax')
|
||||
])
|
||||
|
||||
model.compile(
|
||||
optimizer=optimizers.Adam(),
|
||||
loss='categorical_crossentropy',
|
||||
metrics=['accuracy']
|
||||
)
|
||||
return model
|
||||
|
||||
def train(model_type, dataset_root, output_path):
|
||||
# Set parameters based on model type
|
||||
if model_type == 'suit':
|
||||
num_classes = 4
|
||||
elif model_type == 'value':
|
||||
num_classes = 9
|
||||
else:
|
||||
raise ValueError("model_type must be 'suit' or 'value'")
|
||||
|
||||
# Data augmentation as per ML_SETUP_GUIDE.md
|
||||
datagen = ImageDataGenerator(
|
||||
rescale=1./255,
|
||||
rotation_range=15,
|
||||
brightness_range=[0.8, 1.2],
|
||||
validation_split=0.2
|
||||
)
|
||||
|
||||
train_generator = datagen.flow_from_directory(
|
||||
os.path.join(dataset_root, f'{model_type}_model'),
|
||||
target_size=(64, 64),
|
||||
batch_size=32,
|
||||
class_mode='categorical',
|
||||
subset='training'
|
||||
)
|
||||
|
||||
validation_generator = datagen.flow_from_directory(
|
||||
os.path.join(dataset_root, f'{model_type}_model'),
|
||||
target_size=(64, 64),
|
||||
batch_size=32,
|
||||
class_mode='categorical',
|
||||
subset='validation'
|
||||
)
|
||||
|
||||
model = create_model(num_classes)
|
||||
|
||||
print(f"Training {model_type} model...")
|
||||
model.fit(
|
||||
train_generator,
|
||||
epochs=20,
|
||||
validation_data=validation_generator
|
||||
)
|
||||
|
||||
model.save(output_path)
|
||||
print(f"Model saved to {output_path}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--type', choices=['suit', 'value'], required=True)
|
||||
parser.add_argument('--dataset', default='dataset')
|
||||
parser.add_argument('--output', required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
train(args.type, args.dataset, args.output)
|
||||
Loading…
Add table
Add a link
Reference in a new issue