def main(args):
dataset = args.dataset
data_type = args.data_type
# data dir, image size and batch size
DATA_DIR = 'data'
TRAIN_DIR = os.path.join(DATA_DIR, 'train')
VALID_DIR = os.path.join(DATA_DIR, 'valid')
SIZE = (224, 224)
BATCH_SIZE = 4
# remove files
try:
shutil.rmtree(DATA_DIR)
except:
pass
train_ratio = 0.8 # 80% data for training, the rest for testing
# get the list of filenames and corresponding list of labels for training et validation
train_filenames = []
train_labels = []
val_filenames = []
val_labels = []
# read files into label and frame lists
with open('../data/labels/' + dataset + '_' + data_type +
'_label.csv') as f:
frames_labels = [(line.strip().split(',')[0],
line.strip().split(',')[1]) for line in f]
# re-organize data by labels
file_dir = '../data/' + dataset + '/' + data_type + '/jpg/'
file_format = '.jpeg'
dict_frame = {
} # key: label value, value: a list of indice in the original file
for fr_lb in frames_labels:
fr, lb = fr_lb
if (lb not in dict_frame):
dict_frame[lb] = []
dict_frame[lb].append(file_dir + fr + file_format)
random.seed() # using current time as the seed
# generate filenames and labels for training and validation dataset
for lb in dict_frame:
# pick random indices for training data for lb in dict_frame
train_index = random.sample(range(0, len(dict_frame[lb])),
int(train_ratio * len(dict_frame[lb])))
for index in range(len(dict_frame[lb])):
# training data
if (index in train_index):
train_filenames.append(dict_frame[lb][index])
train_labels.append(int(lb) - 1)
# validation data
else:
val_filenames.append(dict_frame[lb][index])
val_labels.append(int(lb) - 1)
assert set(train_labels) == set(
val_labels), "Train and val labels don't correspond:\n{}\n{}".format(
set(train_labels), set(val_labels))
# create new dir data/train/label_x and data/valid/label_x
for label in set(train_labels):
os.makedirs(os.path.join(TRAIN_DIR, str(label)), exist_ok=True)
os.makedirs(os.path.join(VALID_DIR, str(label)), exist_ok=True)
# copy files
for tr_file, label in zip(train_filenames, train_labels):
shutil.copy2(tr_file, os.path.join(TRAIN_DIR, str(label)))
for val_file, label in zip(val_filenames, val_labels):
shutil.copy2(val_file, os.path.join(VALID_DIR, str(label)))
# train models
num_train_samples = sum([len(files) for r, d, files in os.walk(TRAIN_DIR)])
num_valid_samples = sum([len(files) for r, d, files in os.walk(VALID_DIR)])
num_train_steps = math.floor(num_train_samples / BATCH_SIZE)
num_valid_steps = math.floor(num_valid_samples / BATCH_SIZE)
gen = image.ImageDataGenerator()
val_gen = image.ImageDataGenerator(horizontal_flip=True,
vertical_flip=True)
batches = gen.flow_from_directory(TRAIN_DIR,
target_size=SIZE,
class_mode='categorical',
shuffle=True,
batch_size=BATCH_SIZE)
val_batches = val_gen.flow_from_directory(VALID_DIR,
target_size=SIZE,
class_mode='categorical',
shuffle=True,
batch_size=BATCH_SIZE)
model = DenseNet121()
classes = list(iter(batches.class_indices))
model.layers.pop()
for layer in model.layers:
layer.trainable = False
last = model.layers[-1].output
x = Dense(len(classes), activation="softmax")(last)
finetuned_model = Model(model.input, x)
finetuned_model.compile(optimizer=Adam(lr=0.0001),
loss='categorical_crossentropy',
metrics=['accuracy'])
for c in batches.class_indices:
classes[batches.class_indices[c]] = c
finetuned_model.classes = classes
early_stopping = EarlyStopping(patience=450)
checkpointer = ModelCheckpoint('./densenet_model/densenet_121_best.h5',
verbose=1,
save_best_only=True)
finetuned_model.fit_generator(batches,
steps_per_epoch=num_train_steps,
epochs=450,
callbacks=[early_stopping, checkpointer],
validation_data=val_batches,
validation_steps=num_valid_steps)
finetuned_model.save('./densenet_model/densenet_121_final.h5')
def efficientnet_model(channels,
init_block_channels,
final_block_channels,
kernel_sizes,
strides_per_stage,
expansion_factors,
dropout_rate=0.2,
tf_mode=False,
bn_epsilon=1e-5,
in_channels=3,
in_size=(224, 224),
classes=1000):
"""
EfficientNet(-B0) model from 'EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks,'
https://arxiv.org/abs/1905.11946.
Parameters:
----------
channels : list of list of int
Number of output channels for each unit.
init_block_channels : list of 2 int
Numbers of output channels for the initial unit.
final_block_channels : int
Number of output channels for the final block of the feature extractor.
kernel_sizes : list of list of int
Number of kernel sizes for each unit.
strides_per_stage : list int
Stride value for the first unit of each stage.
expansion_factors : list of list of int
Number of expansion factors for each unit.
dropout_rate : float, default 0.2
Fraction of the input units to drop. Must be a number between 0 and 1.
tf_mode : bool, default False
Whether to use TF-like mode.
bn_epsilon : float, default 1e-5
Small float added to variance in Batch norm.
in_channels : int, default 3
Number of input channels.
in_size : tuple of two ints, default (224, 224)
Spatial size of the expected input image.
classes : int, default 1000
Number of classification classes.
"""
input_shape = (in_channels, in_size[0], in_size[1]) if is_channels_first() else\
(in_size[0], in_size[1], in_channels)
input = nn.Input(shape=input_shape)
activation = "swish"
x = effi_init_block(x=input,
in_channels=in_channels,
out_channels=init_block_channels,
bn_epsilon=bn_epsilon,
activation=activation,
tf_mode=tf_mode,
name="features/init_block")
in_channels = init_block_channels
for i, channels_per_stage in enumerate(channels):
kernel_sizes_per_stage = kernel_sizes[i]
expansion_factors_per_stage = expansion_factors[i]
for j, out_channels in enumerate(channels_per_stage):
kernel_size = kernel_sizes_per_stage[j]
expansion_factor = expansion_factors_per_stage[j]
strides = strides_per_stage[i] if (j == 0) else 1
if i == 0:
x = effi_dws_conv_unit(x=x,
in_channels=in_channels,
out_channels=out_channels,
strides=strides,
bn_epsilon=bn_epsilon,
activation=activation,
tf_mode=tf_mode,
name="features/stage{}/unit{}".format(
i + 1, j + 1))
else:
x = effi_inv_res_unit(x=x,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
strides=strides,
expansion_factor=expansion_factor,
bn_epsilon=bn_epsilon,
activation=activation,
tf_mode=tf_mode,
name="features/stage{}/unit{}".format(
i + 1, j + 1))
in_channels = out_channels
x = conv1x1_block(x=x,
in_channels=in_channels,
out_channels=final_block_channels,
bn_epsilon=bn_epsilon,
activation=activation,
name="features/final_block")
in_channels = final_block_channels
x = nn.GlobalAveragePooling2D(name="features/final_pool")(x)
if dropout_rate > 0.0:
x = nn.Dropout(rate=dropout_rate, name="output/dropout")(x)
x = nn.Dense(units=classes,
input_dim=in_channels,
name="output/fc",
activation='softmax')(x)
model = Model(inputs=input, outputs=x)
model.in_size = in_size
model.classes = classes
return model
