GuzmanDo/image_classification

1import os2import glob3import shutil4import tensorflow as tf5 6from tensorflow.keras.preprocessing.image import ImageDataGenerator7 8# init necessary variables9BATCH_SIZE = 3210IMG_SIZE = (224, 224)11IMG_SHAPE = IMG_SIZE + (3,)12num_classes = 513 14# 1. Download dataset15_URL = "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"16path_to_zip = tf.keras.utils.get_file(origin=_URL, fname="flower_photos.tgz", extract=True)17base_dir = os.path.join(os.path.dirname(path_to_zip), 'flower_photos')18classes = ['roses', 'daisy', 'dandelion', 'sunflowers', 'tulips']19 20# create train and validation set21for cl in classes:22    img_path = os.path.join(base_dir, cl)23    images = glob.glob(img_path + '/*.jpg')24    # print("{}: {} Images".format(cl, len(images)))25    num_train = int(round(len(images)*0.8))26    train, val = images[:num_train], images[num_train:]27 28    for t in train:29        if not os.path.exists(os.path.join(base_dir, 'train', cl)):30            os.makedirs(os.path.join(base_dir, 'train', cl))31        shutil.move(t, os.path.join(base_dir, 'train', cl))32 33    for v in val:34        if not os.path.exists(os.path.join(base_dir, 'val', cl)):35            os.makedirs(os.path.join(base_dir, 'val', cl))36        shutil.move(v, os.path.join(base_dir, 'val', cl))37 38train_dir = os.path.join(base_dir, 'train')39val_dir = os.path.join(base_dir, 'val')40 41# 2. Data augmentation42image_gen = ImageDataGenerator(tf.keras.applications.mobilenet_v2.preprocess_input,43                               rescale=1./255, rotation_range=45,44                               width_shift_range=0.2, height_shift_range=0.2,45                               shear_range=0.3, zoom_range=0.5, horizontal_flip=True)46 47train_data_gen = image_gen.flow_from_directory(batch_size=BATCH_SIZE, directory=train_dir,48                                               shuffle=True, target_size=IMG_SIZE, class_mode='sparse')49 50val_data_gen = image_gen.flow_from_directory(batch_size=BATCH_SIZE, directory=val_dir,51                                             shuffle=True, target_size=IMG_SIZE, class_mode='sparse')52 53# 3. Create the base model from the pre-trained model MobileNet V254base_model = tf.keras.applications.mobilenet_v2.MobileNetV2(input_shape=IMG_SHAPE,55                                                            weights='imagenet', include_top=False)56# Freeze base model57base_model.trainable = False58 59# Connect new predict output to base model60x = base_model.output61x = tf.keras.layers.GlobalAveragePooling2D()(x)62x = tf.keras.layers.Dense(1024, activation='relu')(x)63x = tf.keras.layers.Dropout(0.2)(x)64x = tf.keras.layers.Dense(512, activation='relu')(x)65outputs = tf.keras.layers.Dense(num_classes)(x)66model = tf.keras.Model(base_model.inputs, outputs)67 68# Set up the learning process69model.compile(optimizer='adam',70              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),71              metrics=['accuracy'])72 73# 4. Train74epochs = 1075# save best model76checkpoint = tf.keras.callbacks.ModelCheckpoint('flowers/best.h5', monitor='val_loss',77                                                save_best_only=True, mode='auto')78callback_list = [checkpoint]79 80history = model.fit(train_data_gen, epochs=10, validation_data=val_data_gen, callbacks=callback_list)81