Tensorflow CNN with binary categories

Related tutorials:

Tensorflow Convolutional Neural Network for Computer Vision
Tensorflow Convolutional Neural Network for Computer Vision (Part 2)

Code:

!wget https://cdn....zip #load your dataset
!unzip ....zip #unzip your dataset
 
PATH = 'your_path_to_the_dataset_root'
 
train_dir = os.path.join(PATH, 'train')
validation_dir = os.path.join(PATH, 'validation')
test_dir = os.path.join(PATH, 'test')
 
# Get number of files in each directory. The train and validation directories
# each have the subdirecories for the categories
total_train = sum([len(files) for r, d, files in os.walk(train_dir)])
total_val = sum([len(files) for r, d, files in os.walk(validation_dir)])
total_test = len(os.listdir(test_dir))
 
# Variables for pre-processing and training.
batch_size = 128
epochs = 15
IMG_HEIGHT = 150
IMG_WIDTH = 150
 
# We created training, validation, and testimages and shuffle them
 
train_image_generator = ImageDataGenerator(rescale=1./255)
validation_image_generator = ImageDataGenerator(rescale=1./255)
test_image_generator = ImageDataGenerator(rescale=1./255)
 
train_data_gen = train_image_generator.flow_from_directory(train_dir, target_size=(IMG_HEIGHT, IMG_WIDTH), batch_size=batch_size, class_mode='binary', shuffle=True)
val_data_gen = validation_image_generator.flow_from_directory(validation_dir, target_size=(IMG_HEIGHT, IMG_WIDTH), batch_size=batch_size, class_mode='binary', shuffle=True)
test_data_gen = test_image_generator.flow_from_directory(directory=PATH, target_size=(IMG_HEIGHT, IMG_WIDTH), batch_size=batch_size, class_mode='binary', shuffle=False, classes=['test'])
 
# Now we resclae the images
train_image_generator = ImageDataGenerator(height_shift_range=0.4, width_shift_range=0.4, rotation_range=50, horizontal_flip=True, vertical_flip=True,shear_range=0.2,zoom_range=0.2, rescale=1./255)
 
# put the rescaled images to the model we want to train
train_data_gen = train_image_generator.flow_from_directory(batch_size=batch_size,
                                                     directory=train_dir,
                                                     target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                     class_mode='binary')
 
#train it
model = Sequential()
 
model.add(Conv2D(16, (3, 3), input_shape=(IMG_HEIGHT, IMG_WIDTH, 3), activation='relu'))
model.add(MaxPooling2D((2,2)))
model.add(Conv2D(32, (3, 3), activation='relu'))
model.add(MaxPooling2D((2,2)))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(MaxPooling2D((2,2)))
model.add(Conv2D(168, (3, 3), activation='relu'))
 
model.add(Flatten())
model.add(Dense(64, activation='relu'))
model.add(Dense(32, activation='relu'))
model.add(Dense(2))
 
model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer='adam', metrics=['accuracy'])
 
model.fit(train_data_gen, epochs=epochs, steps_per_epoch=(train_data_gen.n / train_data_gen.batch_size), validation_steps=(val_data_gen.n / val_data_gen.batch_size), validation_data=val_data_gen)

Imports:

import tensorflow as tf
 
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPooling2D
from tensorflow.keras.preprocessing.image import ImageDataGenerator
 
import os
import numpy as np