-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
86 lines (74 loc) · 2.38 KB
/
model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import tensorflow as tf
import os
import shutil
import random
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.layers import BatchNormalization, Dense, Dropout
from tensorflow.keras import Model, Input
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.optimizers import Adam
# Define the training and test data directories
data_dir='asl_dataset'
# Define the data generators for training and validation
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,
fill_mode = 'nearest',
validation_split = 0.2)
height = 224
width = 224
channels = 3
batch_size = 32
img_shape = (height, width, channels)
img_size = (height, width)
train_data = datagen.flow_from_directory(
data_dir,
target_size = img_size,
batch_size = batch_size,
class_mode = 'categorical',
subset = 'training')
val_data = datagen.flow_from_directory(
data_dir,
target_size = img_size,
batch_size = batch_size,
class_mode='categorical',
subset = 'validation')
# Define the MobileNetV2 model as the base model
base_model = MobileNetV2(
input_shape=(224, 224, 3),
include_top=False,
weights='imagenet',
pooling='avg'
)
base_model.trainable = False
# Freeze the base model layers
for layer in base_model.layers:
layer.trainable = False
# Combine the base model with the new layers
x = base_model.output
x = BatchNormalization()(x)
x = Dense(1024, activation='relu')(x)
x = Dropout(rate=0.2)(x)
outputs = Dense(units=36, activation="softmax")(x)
model = Model(inputs=base_model.input, outputs=outputs)
# Compile the model
LEARNING_RATE = 0.001
model.compile(optimizer=Adam(learning_rate=LEARNING_RATE),
loss='categorical_crossentropy',
metrics=['accuracy'])
# Train the model
STEP_SIZE_TRAIN = train_data.n // train_data.batch_size
STEP_SIZE_VALID = val_data.n // val_data.batch_size
history = model.fit_generator(train_data,
steps_per_epoch = STEP_SIZE_TRAIN,
validation_data = val_data,
validation_steps = STEP_SIZE_VALID,
epochs = 30,
verbose = 1)
# Saving the model
model.save('asl_image_classification.h5')