def get_annotation_data(image_base_name: str, path_to_images: str,
ann_format: str, label_mapping: pd.DataFrame,
class_mapping: pd.DataFrame) -> pd.DataFrame:
"""
Calculates the annotation data related to the labels of the given image.
:param image_base_name: the base name of the image (without extension)
:param path_to_images: the path where the images relative to the dataset are stored
:param ann_format: the annotation format, which can be either YOLOv2 or darkflow
:param label_mapping: the image-labels mapping
:param class_mapping: the class string to class number mapping
:return: the annotation of the current image as a dataframe
"""
# Get all the labels of the image as a string
try:
labels = label_mapping.loc[image_base_name, 'labels']
except KeyError:
labels = ''
# Convert the string of labels to list
labels = [line[:-1] for line in re.findall(r"(?:\S*\s){5}", str(labels))]
# Get the width and height of the image
img_width, img_height = Image.open(
os.path.join(path_to_images, to_file_name(image_base_name))).size
convert_to = {
'YOLOv2': convert_to_yolov2,
'darkflow': convert_to_darkflow,
'frcnn': convert_to_frcnn,
}
# Create a list of lists to store the annotation data
img_path = os.path.join(path_to_images, image_base_name)
annotation_data = [
convert_to[ann_format](label=label,
class_mapping=class_mapping,
image_width=img_width,
image_height=img_height,
image_path=img_path) for label in labels
]
# If the image has no labels, insert a default row. For frcnn 'bg' is special background class
if not annotation_data:
annotation_data = [['', '', '', '', '', img_width, img_height]] \
if ann_format != 'frcnn' else [[to_file_name(img_path), 0, 0, img_width, img_height, 'bg']]
data_format = {
'YOLOv2': ['class', 'x_c', 'y_c', 'bb_width', 'bb_height'],
'darkflow':
['class', 'xmin', 'ymin', 'xmax', 'ymax', 'img_width', 'img_height'],
'frcnn': ['filepath', 'xmin', 'ymin', 'xmax', 'ymax', 'class_name']
}
# Create a dataframe to store the whole annotation
annotation = pd.DataFrame(annotation_data, columns=data_format[ann_format])
return annotation
def write_as_darkflow(annotation: pd.DataFrame, path_to_annotations: str,
image_id: str):
"""
Writes an annotation on file according to the darkflow format.
:param annotation: the annotation data
:param path_to_annotations: the path to the annotations files
:param image_id: the base name of the image (without file extension)
"""
file_name = to_file_name(image_id, 'jpg').replace('-', '_')
# Create a writer object for the image
writer = Writer(path=os.path.join(path_to_annotations, file_name),
filename=os.path.join('..', 'images', file_name),
width=annotation['img_width'][0],
height=annotation['img_height'][0],
database='training')
# Add the data related to each row (label) of the dataframe
for _, row in annotation.iterrows():
writer.addObject(name=row['class'],
xmin=row['xmin'],
ymin=row['ymin'],
xmax=row['xmax'],
ymax=row['ymax'])
# Set the path to the XML annotation
path_to_xml = os.path.join(path_to_annotations,
to_file_name(image_id, 'xml'))
# Write the data to an XML file
writer.save(path_to_xml)
# If the XML file is empty because the image has no objects
if not annotation.loc[0]['class']:
# Get the XML tree of the annotation
tree = ElementTree.parse(path_to_xml)
# Get the root of the annotation
root = tree.getroot()
# Delete the object tag
obj = root.getchildren()[-1]
root.remove(obj)
ElementTree.dump(root)
tree.write(path_to_xml)
def visualize_training_data(path_to_train_images: str, image_id: str,
mapping: pd.DataFrame):
"""
Takes an image_id and return an image with bounding boxes around each character
:param path_to_train_images: the path to the images of the training set
:param mapping: the image-labels mapping
:param image_id: an image base name (i.e. without the file extension)
"""
print(
'\nDisplaying the image {} from the training set with bounding boxes...'
.format(image_id))
# Get all the characters and the position of the bounding boxes for the image
mapping.set_index('image_id', inplace=True)
labels = mapping.loc[image_id, 'labels']
labels = np.array(labels.split(" ")).reshape(-1, 5)
plt = display_image(os.path.join(path_to_train_images,
to_file_name(image_id)),
show=False)
ax = plt.gca()
for label in labels:
ax = draw_box_and_text(ax, label)
plt.show()
def visualize_null_images(path_to_train_images: str, mapping: pd.DataFrame):
"""
Visualizes some of the images which have no characters.
:param path_to_train_images: the path to the images of the training set
:param mapping: the image-labels mapping
"""
# Configure the plot
rows, columns = 2, 2
fig = plt.figure(figsize=(20, 20))
# Get all the images with no characters
images_nan_labels = mapping[mapping.isna().labels]['image_id'].tolist()
# Plot some of the images with no characters
for i in range(1, rows * columns + 1):
rnd_img_nan_nbr = random.randint(0, len(images_nan_labels) - 1)
img_nan = Image.open(
os.path.join(path_to_train_images,
to_file_name(images_nan_labels[rnd_img_nan_nbr])))
fig.add_subplot(rows, columns, i)
plt.imshow(img_nan, aspect='equal')
print('\nDisplaying some images with no characters...')
plt.show()
def resize_dataset(size: int,
path_to_images: str,
path_to_annotations: str = None):
"""
Resizes the dataset to the given size deleting all the exceeding items
and corresponding annotations.
:param size: the new size of the dataset
:param path_to_images: the path to the images of the dataset
:param path_to_annotations: the path to the annotations of the dataset
"""
print('Resizing the dataset from {old_n} to {new_n} items...'.format(
old_n=len(os.listdir(path_to_images)), new_n=size))
# Iterate over all the dataset items (i.e. names of images)
for i, img_file_name in enumerate(sorted(os.listdir(path_to_images))):
# Get the base name of the image (without extension)
img_id = to_id(img_file_name)
# Remove all the exceeding items
if i >= size:
# Delete the image
os.remove(os.path.join(path_to_images, to_file_name(img_id,
'jpg')))
# If annotations have been specified
if path_to_annotations:
path_to_annotation = os.path.join(path_to_annotations,
to_file_name(img_id, 'xml'))
# Delete the annotation
if os.path.isfile(path_to_annotation):
os.remove(path_to_annotation)
print('Resizing Complete!')
print('The dataset now has:')
print('* {} images'.format(len(os.listdir(path_to_images))))
if path_to_annotations:
print('* {} annotations'.format(len(os.listdir(path_to_annotations))))
def convert_to_frcnn(label: [], image_path: str, **_) -> []:
"""
Converts the label data to the darkflow data format (i.e. PASCAL VOC XML).
The dataset is provided with labels in the format: <unicode> <x_bl> <y_bl> <abs_bb_width> <abs_bb_height>
Where:
- <unicode> : is the class label as a unicode
- <x_bl> : is the bottom-left corner x coordinate of the bounding box
- <y_bl> : is the bottom-left corner y coordinate of the bounding box
- <abs_bb_width> : is the width of the bounding box
- <abs_bb_height> : is the height of the bounding box
The darkflow data format is: <class_name> <xmin> <ymin> <xmax> <ymax> <img_width> <img_height>
Where:
- <class_name> : is the name of the class object as string
- <img_width> : is the width of the image
- <img_height> : is the height of the image
- <xmin> : is the bottom-left x coordinate of the bounding box
- <ymin> : is the bottom-left y coordinate of the bounding box
- <xmax> : is the top-right x coordinate of the bounding box
- <ymax> : is the top-right y coordinate of the bounding box
:param image_path: the path to the image, including image extension
:param label: the label for one character in the image
:return: a list with the converted data
"""
# Get the label data in the dataset format
unicode, xmin, ymin, abs_bb_width, abs_bb_height = label.split()
# Make sure the class is a string
class_name = str(unicode)
# Cast each data to int
xmin = int(xmin)
ymin = int(ymin)
abs_bb_width = int(abs_bb_width)
abs_bb_height = int(abs_bb_height)
# Calculate the top-right coordinates of the bounding box
xmax = xmin + abs_bb_width
ymax = ymin + abs_bb_height
return [to_file_name(image_path, 'jpg'),
xmin,
ymin,
xmax,
ymax,
class_name]
def write_as_yolov2(annotation: pd.DataFrame, path_to_annotations: str,
image_id: str):
"""
Writes an annotation on file according to YOLOv2 format.
:param annotation: the annotation data
:param path_to_annotations: the path to the annotations files
:param image_id: the base name of the image (without file extension)
"""
annotation.to_csv(os.path.join(path_to_annotations,
to_file_name(image_id, 'txt')),
header=None,
index=None,
sep=' ',
mode='a')
def visualize_bounding_boxes(path_to_images: str, image_id: str, labels: []):
"""
Takes an image_id and return an image with bounding boxes around each character.
:param path_to_images: the path to the images of the training set
:param labels: the labels of the image to be tested
:param image_id: an image base name (i.e. without the file extension)
"""
plt = display_image(os.path.join(path_to_images, to_file_name(image_id)),
show=False)
ax = plt.gca()
for label in labels:
ax = draw_box_and_text(ax, label)
plt.show()
def visualize_random_image(path_to_train_images: str):
"""
Visualizes a random image from the training set
:param path_to_train_images: the path to the images of the training set
"""
# Get a random index for the selection of the image
rnd_nbr = random.randint(0, len(os.listdir(path_to_train_images)))
# Get the random image base name
rnd_img = os.listdir(path_to_train_images)[rnd_nbr]
# Get the full path to the random image
path_to_rnd_image = os.path.join(path_to_train_images,
to_file_name(rnd_img))
print('\nDisplaying the random image {}...'.format(rnd_img))
# Display the random image
display_image(path_to_rnd_image)
def restore_dataset_from_backup(path_to_images: str, path_to_annotations: str,
path_to_backup: str):
"""
Copies all the images from the backup to the main folder and deletes the old files.
:param path_to_images: the path to the old images
:param path_to_annotations: the path to the old annotations
:param path_to_backup: the path to the backup images
"""
print('Restoring the dataset from backup...\n'
'The dataset now counts:\n'
'* {n_img} images\n'
'* {n_ann} annotations'.format(n_img=len(os.listdir(path_to_images)),
n_ann=len(
os.listdir(path_to_annotations))))
# Delete the current images and annotations
for img_file_name in os.listdir(path_to_images):
# Delete the image
os.remove(os.path.join(path_to_images, img_file_name))
# Delete the corresponding annotation
path_to_annotation = os.path.join(
path_to_annotations, to_file_name(to_id(img_file_name), 'xml'))
if os.path.isfile(path_to_annotation):
os.remove(path_to_annotation)
# Copy the backup images into the main folder
[
copy(os.path.join(path_to_backup, image), path_to_images)
for image in os.listdir(path_to_backup)
]
print('Restoring performed successfully!\n'
'The dataset now counts:\n'
'* {n_img} images\n'
'* {n_ann} annotations'.format(n_img=len(os.listdir(path_to_images)),
n_ann=len(
os.listdir(path_to_annotations))))
def visualize_from_training(path_to_images: str, path_to_annotations: str,
image_id: str):
"""
Visualizes the image from the current training set.
:param path_to_images: the path to the backup images
:param path_to_annotations: the path to the annotations
:param image_id: the id of the image (without file extension)
"""
print('\nRetrieving generated annotation...')
xml_doc = minidom.parse(
os.path.join(path_to_annotations, to_file_name(image_id, 'xml')))
objects = xml_doc.getElementsByTagName('object')
labels = []
for obj in objects:
obj_class = obj.getElementsByTagName('name')[0].firstChild.nodeValue
bounding_box = obj.getElementsByTagName('bndbox')[0]
xmin = bounding_box.getElementsByTagName(
'xmin')[0].firstChild.nodeValue
ymin = bounding_box.getElementsByTagName(
'ymin')[0].firstChild.nodeValue
xmax = bounding_box.getElementsByTagName(
'xmax')[0].firstChild.nodeValue
ymax = bounding_box.getElementsByTagName(
'ymax')[0].firstChild.nodeValue
label_data = [
obj_class, xmin, ymin,
str(int(xmax) - int(xmin)),
str(int(ymax) - int(ymin))
]
labels.append(label_data)
visualize_bounding_boxes(path_to_images, image_id, labels)