def histogram_labels(dataset_name: str,
double_bins: bool = False,
show_labels: bool = False):
""" Plot an histogram of the labels. double_bins allow to have have more space between bins in the plot """
annotations_df = pd.read_csv(
constants.ANNOTATIONS_TRAIN_CSV_PATH.format(dataset_name))
# if use_label_name:
labels_values = list(helpers.get_labels_dict(dataset_name).values())
labels_ids = annotations_df['label'].unique()
n_bins = len(labels_ids)
if double_bins:
n_bins *= 2
fig, ax = plt.subplots()
bins = np.arange(n_bins + 1) - 0.5
counts, bins, patches = ax.hist(annotations_df['label'], bins=bins)
plt.xlim([-1, n_bins])
if show_labels:
plt.xticks(range(n_bins))
# Label below the x-axis...
bin_centers = bins + 0.5
i = 0
for count, x in zip(counts, bin_centers):
ax.annotate(labels_values[i], (x, 0),
xycoords=('data', 'axes fraction'),
xytext=(0, -18),
textcoords='offset points',
va='top',
ha='center',
rotation=-90)
i += 1
plt.subplots_adjust(bottom=0.15)
plt.show()
def detect(model,
dataset_name: str,
image_path: str,
title: str = "",
output_path: str = None):
labels_map_dict = helpers.get_labels_dict(dataset_name)
img = tf.image.decode_image(open(image_path, 'rb').read(), channels=3)
img = tf.expand_dims(img, 0)
img_shape = np.array(
[img.shape[2], img.shape[1], img.shape[2], img.shape[1]])
img = helpers.transform_images(img, model.image_res)
boxes, scores, classes, nums = model.inference_model(img)
boxes, scores, classes = [
np.array(elem[0]) for elem in [boxes, scores, classes]
]
annotations = []
for object_index in range(nums[0]):
bbox = [int(elem) for elem in list(boxes[object_index] * img_shape)]
annotations.append(bbox + [int(classes[object_index])])
print(constants.C_OKGREEN, 'Object',
labels_map_dict[str(int(classes[object_index]))],
'with probability', scores[object_index], 'is in',
annotations[object_index][:-1], constants.C_ENDC)
visualize.plot_image_annotations(image_path, annotations, labels_map_dict,
scores, title, output_path)
if len(annotations) == 0:
print(constants.C_WARNING, "No object found.", title, constants.C_ENDC)
def create_label_map_pbtxt(dataset_name: str):
""" This kind of label map is used in the google models research repository """
labels_dict = helpers.get_labels_dict(dataset_name)
with open(constants.LABEL_MAP_PBTXT_PATH.format(dataset_name), 'a') as t:
for key, value in labels_dict.items():
id = str(
int(key) + 1
) # pbtxt files are used in google models and they require id 0 to be background
t.write("item {\n name: \"" + value + "\"\n id: " + id +
"\n display_name: \"" + value + "\"\n}\n")
def plot_image_annotations_simple(image_path: str, title: str = ""):
""" the image file should finish with: /dataset_name/train_val_test/image"""
info = image_path.split('/')
dataset_name, train_val_test, image = info[-3:]
annotations = helpers.get_annotations_dict(dataset_name,
train_val_test)[image]
labels_map_dict = helpers.get_labels_dict(dataset_name)
plot_image_annotations(image_path,
annotations,
labels_map_dict,
title=title)
def __init__(self,
dataset_name: str,
img_res: int,
train: bool,
model_version: str,
train_cluster: bool = False):
self.model_name = "Recognizer_" + model_version
if dataset_name == '':
self.dataset_name = 'FAKE'
self.n_classes = 10
else:
self.dataset_name = dataset_name
self.label_map_json_path = constants.LABEL_MAP_JSON_PATH.format(
dataset_name)
self.n_classes = helpers.get_n_classes(self.label_map_json_path)
self.labels_map_dict = helpers.get_labels_dict(dataset_name)
if train:
self.checkpoints_path, self.logs_path, self.figs_path = helpers.model_directories(
constants.RECOGNIZER, self.dataset_name, train_cluster)
self.model_description = ''
if model_version == 'mini':
self.train_model = self.model_mini(img_res)
elif model_version == 'mini_2':
self.train_model = self.model_mini_2(img_res)
elif model_version == 'kaggle':
self.train_model = self.kaggle_model(img_res)
elif model_version == 'mod':
self.train_model = self.mod(img_res)
elif model_version == 'mod_bn':
self.train_model = self.mod_bn(img_res)
elif model_version == 'mod_bn_reg':
self.train_model = self.mod_bn_reg(img_res)
elif model_version == 'v0':
self.train_model = self.model_v0(img_res)
elif model_version == 'v0.1':
self.train_model = self.model_v0_1(img_res)
elif model_version == 'v0.2':
self.train_model = self.model_v0_2(img_res)
elif model_version == 'v0.3':
self.train_model = self.model_v0_3(img_res)
elif model_version == 'v0.4':
self.train_model = self.model_v0_4(img_res)
elif model_version == 'v1':
self.train_model = self.model_v1(img_res)
elif model_version == 'v2':
self.train_model = self.model_v2(img_res)
elif model_version == 'v3':
self.train_model = self.model_v3(img_res)
self.img_res = img_res
self.model_description += 'The model has {:,} parameters.'.format(
self.train_model.count_params())
self.inference_model = self.train_model
def plot_images_and_boxes(img,
boxes,
switch=False,
multi=True,
i=None,
title="",
dataset_name=None,
colors=None):
if type(boxes) != np.ndarray:
boxes = np.array(boxes)
fig, ax = plt.subplots(1)
ax.imshow(img)
img_shape = np.array(
[img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
if dataset_name:
labels_map = helpers.get_labels_dict(dataset_name)
labels = []
colors = get_colors(boxes, labels)
if switch:
boxes[:, [0, 1, 2, 3]] = boxes[:, [1, 0, 3, 2]]
for bbox in boxes:
original = bbox
if multi:
bbox = [int(elem) for elem in list(bbox[:4] * img_shape)
] + ([int(bbox[-1])] if len(bbox) == 5 else [])
if np.sum(bbox[:4]) == 0:
continue
# print("Box added:", bbox, "original", original)
rect = patches.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
linewidth=1,
edgecolor=colors[labels.index(bbox[-1])],
facecolor='none')
if dataset_name and len(bbox) == 5:
add_label_to_plot(ax, colors[labels.index(bbox[-1])],
labels_map[str(bbox[-1])], bbox[0], bbox[1])
plt.scatter(bbox[0], bbox[1], c='r', marker='x')
plt.scatter(bbox[2], bbox[3], c='g', marker='x')
ax.add_patch(rect)
plt.title(title)
if i is not None:
plt.savefig('{:03d}.png'.format(i))
plt.show()
if i is not None:
return i + 1
else:
return i
def __init__(self,
dataset_name: str,
img_res: int,
train: bool,
model_name: str = "Recognizer"):
self.model_name = model_name
if dataset_name == '' or dataset_name == 'FAKE':
self.dataset_name = 'FAKE'
self.n_classes = 10
else:
self.dataset_name = dataset_name
self.label_map_json_path = constants.LABEL_MAP_JSON_PATH.format(
dataset_name)
self.n_classes = helpers.get_n_classes(self.label_map_json_path)
self.labels_map_dict = helpers.get_labels_dict(dataset_name)
if train:
self.checkpoints_path, self.logs_path, self.figs_path = helpers.model_directories(
model_name, self.dataset_name)
self.img_res = img_res
leaky_relu = lambda x: tf.keras.activations.relu(x, alpha=0.1)
self.train_model = tf.keras.models.Sequential([
Conv2D(filters=32,
kernel_size=(5, 5),
kernel_regularizer=l2(0.0005),
activation=leaky_relu,
input_shape=(img_res, img_res, 3)),
BatchNormalization(),
MaxPool2D(pool_size=(2, 2)),
Conv2D(filters=64,
kernel_size=(3, 3),
activation=leaky_relu,
kernel_regularizer=l2(0.0005)),
MaxPool2D(pool_size=(2, 2)),
Dropout(rate=0.25),
Conv2D(filters=64,
kernel_size=(3, 3),
activation=leaky_relu,
kernel_regularizer=l2(0.0005)),
MaxPool2D(pool_size=(2, 2)),
Dropout(rate=0.25),
Conv2D(filters=128,
kernel_size=(1, 1),
activation=leaky_relu,
kernel_regularizer=l2(0.0005)),
Dropout(rate=0.5),
Flatten(),
Dense(self.n_classes, 'softmax')
])
self.inference_model = self.train_model
self.model_description = 'The model has {:,} parameters.'.format(
self.train_model.count_params())
def __init__(self, dataset_name: str, tiny: bool = True):
if dataset_name == '':
self.dataset_name = 'FAKE'
self.n_classes = 80
self.tiny = tiny
print(constants.C_OKBLUE, "FAKE dataset created", constants.C_ENDC)
return
self.dataset_name = dataset_name
self.label_map_json_path = constants.LABEL_MAP_JSON_PATH.format(dataset_name)
self.tf_paths = constants.PROCESSED_PROJECT_FOLDER_PATH + constants.TFRECORDS_PATH
self.n_classes = helpers.get_n_classes(self.label_map_json_path)
self.labels_map_dict = helpers.get_labels_dict(dataset_name)
self.train_data, self.validation_data, self.test_data = None, None, None
self.image_res, self.loading_train = None, None # To be defined when called load_datasets
print(constants.C_OKBLUE, "Dataset", dataset_name, "created", constants.C_ENDC)
def load_tensorboard_image(image, groundtruth_ann: list, annotations: list,
probs: list, dataset_name: str):
labels_map_dict = helpers.get_labels_dict(dataset_name)
# Create figure and axes
fig, ax = plt.subplots(1)
labels = []
# Get the unique labels to assign a color to each
colors = get_colors(groundtruth_ann + annotations, labels)
separation = 10
# Concatenate the same image with a white space in between
new_image = np.concatenate((image, 255 * np.ones(
(image.shape[0], separation, 3)).astype(np.int32), image),
axis=1)
# Load the original image with its groundtruth annotations
ax.imshow(new_image)
for i, bbox in enumerate(groundtruth_ann):
if sum(bbox) == 0:
continue
plot_image_annotations_help(ax=ax,
bbox=bbox,
colors=colors,
labels=labels,
labels_map=labels_map_dict,
show_axis=False)
# Load the predicted labels
for i, bbox in enumerate(annotations):
if sum(bbox) == 0:
continue
# Move the bounding boxes to the left to align with the second image
bbox[:3] = (np.array(bbox[:3]) +
(separation + image.shape[1], 0,
separation + image.shape[1])).tolist()
plot_image_annotations_help(ax=ax,
bbox=bbox,
colors=colors,
labels=labels,
labels_map=labels_map_dict,
show_axis=False,
probability=probs[i])
# Workaround for returning the image with the annotations as a tensor, save and load it.
plt.axis('off')
plt.savefig('/tmp/tb_img.png', dpi=500, bbox_inches='tight', pad_inches=0)
plt.close()
img = tf.image.decode_jpeg(open('/tmp/tb_img.png', 'rb').read(),
channels=3)
img = tf.cast(tf.image.resize(img, (400, 1200)), tf.uint8)
return img
def train_recognition(_argv):
"""
Train an image classifier. The images should be saved each in a folder with name the class labels.
It will save a dictionary that maps each output neuron to the corresponding label as a pickle
file named 'neuron_to_class_dict.p'. Divides de training set into train and validation sets and uses
the test set to analyze the accuracy of the model.
It saves the parameters passed to train the model in a file named 'train_info.txt' in the checkpoints
directory.
For the moment it has been checked only with the GTSR dataset.
"""
img_res = FLAGS.img_res
gpu_aval = tf.test.is_gpu_available(cuda_only=True)
model = Recognizer(FLAGS.dataset_name, img_res=img_res, train=True, model_name=FLAGS.model_name)
gpus = 0
if gpu_aval:
for x in device_lib.list_local_devices():
if x.device_type == "GPU":
gpus += 1
print(constants.C_WARNING, "Are CUDA gpus available? \n\t-",
(constants.C_OKBLUE + ('Yes, ' + str(gpus)) if gpu_aval
else constants.C_FAIL + 'No.'), constants.C_ENDC)
batch_size = FLAGS.batch_size
optimizer = tf.keras.optimizers.Adam(lr=FLAGS.lr)
if gpus > 1:
strategy = tf.distribute.MirroredStrategy()
# Here the batch size scales up by number of workers since
# `tf.data.Dataset.batch` expects the global batch size. Previously we used 'batch_size',
# and now this is multiplied by the number of workers.
batch_size *= gpus
with strategy.scope():
model.train_model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])
else:
model.train_model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])
model.train_model.summary()
train_info = FLAGS.extra.replace('\\n', '\n') + '\n'
train_info += "Train model: {} For {} epochs and {} as Database. \nParameters used:\n " \
"- Frequency = {}\n " \
"- Batch size = {}\n " \
"- Image resolution = {}\n".format(model.model_description, FLAGS.epochs, FLAGS.dataset_name,
FLAGS.save_freq, FLAGS.batch_size, FLAGS.img_res)
print(constants.C_WARNING, FLAGS.extra.replace('\\n', '\n'), constants.C_ENDC)
if model.dataset_name != 'FAKE':
ds_directory = constants.DATASET_PATH.format(FLAGS.dataset_name)
train_directory = ds_directory + "train"
if FLAGS.simple_aug:
generator = tf.keras.preprocessing.image.ImageDataGenerator(
rescale=1. / 255.0,
validation_split=0.1,
rotation_range=20,
height_shift_range=0.2,
width_shift_range=0.2,
zoom_range=0.3,
shear_range=0.3,
brightness_range=(0.2, 0.8)
)
train_generator = generator.flow_from_directory(
train_directory,
target_size=(img_res, img_res),
batch_size=batch_size,
class_mode='categorical',
subset='training'
)
val_generator = generator.flow_from_directory(
train_directory,
target_size=(img_res, img_res),
batch_size=batch_size,
class_mode='categorical',
subset='validation'
)
class_to_neuron_dict = train_generator.class_indices
else:
generator = Generator(directory=train_directory,
batch_size=FLAGS.batch_size,
image_dimensions=(img_res, img_res),
validation_split=0.1)
train_generator = generator.train_generator
val_generator = generator.val_generator
class_to_neuron_dict = generator.class_indices
neuron_to_class_dict = {}
labels_dict = helpers.get_labels_dict(model.dataset_name)
neuron_labels = []
for class_id, neuron in class_to_neuron_dict.items():
neuron_to_class_dict[str(neuron)] = str(class_id)
neuron_labels.append(labels_dict[str(neuron)])
print('class_to_neuron_dict', class_to_neuron_dict)
start = time.time()
history, history_callback = train.train(model=model,
epochs=FLAGS.epochs,
train_data=train_generator,
val_data=val_generator,
save_freq=FLAGS.save_freq,
initial_lr=FLAGS.lr,
train_info=train_info,
use_fit_generator=gpus <= 1,
use_cosine_lr=FLAGS.use_cosine_lr)
model.train_model.save_weights(model.checkpoints_path + 'weights.ckpt')
test_acc = 100 * test_model(model, neuron_to_class_dict)
history_callback.test_acc = test_acc
print(constants.C_OKGREEN, "Test accuracy {:0.2f} %".format(test_acc),
constants.C_ENDC)
with open(model.checkpoints_path + 'neuron_to_class_dict.json', 'w') as fp:
json.dump(neuron_to_class_dict, fp, indent=4)
with open(model.checkpoints_path + 'train_info.txt', 'a') as t:
t.write("Test accuracy {:0.2f} %".format(test_acc))
with open(model.checkpoints_path + 'neuron_labels.txt', 'a') as t:
for label in neuron_labels:
t.write(label + '\n')
else:
print("Train with fake data")
train_data, val_data = helpers.load_fake_dataset_recognition(img_res)
start = time.time()
history, history_callback = train.train(model, FLAGS.epochs, train_data, val_data, FLAGS.save_freq,
FLAGS.lr, 'Use FAKE DS\n', False, FLAGS.use_cosine_lr)
helpers.save_history(FLAGS, model.model_name, model.dataset_name, history, start, 'recognition')
if FLAGS.tflite:
converter = tf.lite.TFLiteConverter.from_keras_model(model.inference_model)
tflite_model = converter.convert()
open(model.checkpoints_path + "model.tflite", "wb").write(tflite_model)
return model, history, history_callback