train.py

# USAGE
# python train.py --dataset dataset/category --model category.model --labelbin lb.pickle
# python train.py --dataset dataset/texture --model texture.model --labelbin lb.pickle

# *** Should comment cropped augmentation-data when you train category

# set the matplotlib backend so figures can be saved in the background
import matplotlib

matplotlib.use("Agg")

# import the necessary packages
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import img_to_array
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.utils import class_weight
from pyimagesearch.smallervggnet import SmallerVGGNet
import matplotlib.pyplot as plt
from imutils import paths
import tensorflow as tf
import numpy as np
import argparse
import random
import pickle
import cv2
import os
from PIL import Image

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Memory growth must be set before GPUs have been initialized
    print(e)

# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-d", "--dataset", required=True,
	help="path to input dataset (i.e., directory of images)")
ap.add_argument("-m", "--model", required=True,
	help="path to output model")
ap.add_argument("-l", "--labelbin", required=True,
	help="path to output label binarizer")
ap.add_argument("-p", "--plot", type=str, default="plot.png",
	help="path to output accuracy/loss plot")
args = vars(ap.parse_args())

# initialize the number of epochs to train for, initial learning rate,
# batch size, and image dimensions
EPOCHS = 1000
INIT_LR = 1e-3
BS = 32
IMAGE_DIMS = (96, 96, 3) # grayScale인경우 IMAGE_DIMS = (96, 96, 1) 아닌경우 (96, 96, 3)

# initialize the data and labels
data = []
labels = []

# grab the image paths and randomly shuffle them
print("[INFO] loading images...")
imagePaths = sorted(list(paths.list_images(args["dataset"])))
random.seed(42)
random.shuffle(imagePaths)

# loop over the input images
for imagePath in imagePaths:
	# load the image, pre-process it, and store it in the data list
	image = cv2.imread(imagePath)
	origin_img = image
	image = cv2.resize(image, (IMAGE_DIMS[1], IMAGE_DIMS[0]))
	# grayscale 추가함
	# image = tf.image.rgb_to_grayscale(image)
	# image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # gray scale (모델 train 시 전처리에 사용)
	image = img_to_array(image)
	data.append(image)
	# extract the class label from the image path and update the
	# labels list
	label = imagePath.split(os.path.sep)[-2]
	labels.append(label)
	# data augmentation
	aug_images = []
	# img_cropped = tf.image.central_crop(image, central_fraction=0.5)
	# img_flipped = tf.image.flip_left_right(image)
	img_brightness = tf.image.adjust_brightness(image, -0.1)
	img_saturated = tf.image.adjust_saturation(image, 3)
	# aug_images.append(img_cropped) # Should comment this line when you train category
	# aug_images.append(img_flipped)
	aug_images.append(img_brightness)
	aug_images.append(img_saturated)
	for aug_img in aug_images:
		aug_img = tf.image.resize(aug_img, (IMAGE_DIMS[1], IMAGE_DIMS[0]))
		aug_img = img_to_array(aug_img)
		data.append(aug_img)
		labels.append(label)

# scale the raw pixel intensities to the range [0, 1]
data = np.array(data, dtype="float") / 255.0
labels = np.array(labels)
print("[INFO] data matrix: {:.2f}MB".format(
	data.nbytes / (1024 * 1000.0)))

# class weight -> 데이터 불균형 해결 (texture 학습시 추가함)
class_weights = class_weight.compute_class_weight('balanced', np.unique(labels), labels)

# binarize the labels
lb = LabelBinarizer()
labels = lb.fit_transform(labels)

# partition the data into training and testing splits using 80% of
# the data for training and the remaining 20% for testing
(trainX, testX, trainY, testY) = train_test_split(data,
	labels, test_size=0.2, random_state=42)

# construct the image generator for data augmentation
aug = ImageDataGenerator(rotation_range=25, width_shift_range=0.2, shear_range=0.2,
	horizontal_flip=True, fill_mode="nearest")

# texture train
# aug = ImageDataGenerator(rotation_range=25, width_shift_range=0.2,
# 	height_shift_range=0.3, shear_range=0.2, zoom_range=0.3,
# 	horizontal_flip=True, vertical_flip=True, fill_mode="nearest")



# initialize the model
print("[INFO] compiling model...")
model = SmallerVGGNet.build(width=IMAGE_DIMS[1], height=IMAGE_DIMS[0],
	depth=IMAGE_DIMS[2], classes=len(lb.classes_))
opt = Adam(lr=INIT_LR, decay=INIT_LR / EPOCHS)
model.compile(loss="categorical_crossentropy", optimizer=opt, metrics=["accuracy"])
# model.compile(loss="binary_crossentropy", optimizer=opt, metrics=["accuracy"])
# model.compile(loss="sparse_categorical_crossentropy", optimizer=opt, metrics=["accuracy"])

# train the network
print("[INFO] training network...")
H = model.fit(
	x=aug.flow(trainX, trainY, batch_size=BS),
	validation_data=(testX, testY),
	steps_per_epoch=len(trainX) // BS,
	epochs=EPOCHS, verbose=1, class_weight=class_weights) # class weight 추가

# save the model to disk
print("[INFO] serializing network...")
model.save(args["model"], save_format="h5")

# save the label binarizer to disk
print("[INFO] serializing label binarizer...")
f = open(args["labelbin"], "wb")
f.write(pickle.dumps(lb))
f.close()

# plot the training loss and accuracy
plt.style.use("ggplot")
plt.figure()
N = EPOCHS
plt.plot(np.arange(0, N), H.history["loss"], label="train_loss")
plt.plot(np.arange(0, N), H.history["val_loss"], label="val_loss")
plt.plot(np.arange(0, N), H.history["accuracy"], label="train_acc")
plt.plot(np.arange(0, N), H.history["val_accuracy"], label="val_acc")
plt.title("Training Loss and Accuracy")
plt.xlabel("Epoch #")
plt.ylabel("Loss/Accuracy")
plt.legend(loc="upper left")
plt.savefig(args["plot"])