Note
As there are a lot of related modules in the project, I recently posted several similar posts(because all content cannot fit due to character limit) and someone indicated that this might be against the website policy, so I edited and included only the features of what my code does, I got some votes for closing the question, therefore, i'll include a few modules [trainer.py, evaluator.py] here and you can check the rest on github and review whatever parts you prefer.
All modules:
- detector.py
- evaluator.py
- models.py
- trainer.py
- anchors.py
- annotation_parsers.py
- augmentor.py
- dataset_handlers.py
- utils.py
- visual_tools.py
Description
yolov3-keras-tf2 is an implementation of yolov3 (you only look once) which is is a state-of-the-art, real-time object detection system that is extremely fast and accurate. There are many implementations that support tensorflow, only a few that support tensorflow v2 and as I did not find versions that suit my needs so, I decided to create this version which is very flexible and customizable.
Features
- tensorflow-2.X--keras-functional-api.
- cpu-gpu support.
- Random weights and DarkNet weights support.
- csv-xml annotation parsers.
- Anchor generator.
matplotlibvisualization of all stages.tf.datainput pipeline.pandas&numpydata handling.imgaugaugmentation pipeline(customizable).loggingcoverage.- All-in-1 custom trainer.
- Stop and resume training support.
- Fully vectorized mAP evaluation.
- Photo & video detection.
Directory structure
yolov3-keras-tf2
├── Config
│ ├── __pycache__
│ │ └── augmentation_options.cpython-37.pyc
│ ├── augmentation_options.py
│ ├── beverly_hills.txt
│ ├── coco.names
│ ├── set_annotation_conf.py
│ └── voc_conf.json
├── Data
│ ├── Photos
│ ├── TFRecords
│ ├── XML\ Labels
│ └── bh_labels.csv
├── Docs
│ ├── Augmentor.md
│ ├── Evaluator.md
│ ├── Predictor.md
│ └── Trainer.md
├── Helpers
│ ├── __pycache__
│ │ ├── anchors.cpython-37.pyc
│ │ ├── annotation_parsers.cpython-37.pyc
│ │ ├── dataset_handlers.cpython-37.pyc
│ │ ├── utils.cpython-37.pyc
│ │ └── visual_tools.cpython-37.pyc
│ ├── anchors.py
│ ├── annotation_parsers.py
│ ├── augmentor.py
│ ├── dataset_handlers.py
│ ├── scratch
│ │ └── label_coordinates.csv
│ ├── utils.py
│ └── visual_tools.py
├── LICENSE
├── Logs
│ └── session.log
├── Main
│ ├── __pycache__
│ │ ├── evaluator.cpython-37.pyc
│ │ └── models.cpython-37.pyc
│ ├── detector.py
│ ├── evaluator.py
│ ├── models.py
│ └── trainer.py
├── Models
├── Output
│ ├── Data
│ ├── Detections
│ ├── Evaluation
│ └── Plots
├── README.md
├── Samples
│ ├── anchors.png
│ ├── anchors_sample.png
│ ├── aug1.png
│ ├── data.png
│ ├── detections.png
│ ├── map.png
│ ├── pr.png
│ ├── sample_image.png
│ └── true_false.png
├── requirements.txt
└── test.py
Features
tensorflow 2.2 & keras functional api
This program leverages features that were introduced in tensorflow 2.0 including:
- Eager execution: an imperative programming environment that evaluates operations immediately, without building graphs check here
tf.function: A JIT compilation decorator that speeds up some components of the program check heretf.data: API for input pipelines check here
CPU & GPU support
The program detects and uses available GPUs at runtime(training/detection) if no GPUs available, the CPU will be used(slow).
Random weights and DarkNet weights support
Both options are available, and NOTE in case of using DarkNet yolov3 weights you must maintain the same number of COCO classes (80 classes) as transfer learning to models with different classes will be supported in future versions of this program.
csv-xml annotation parsers
There are 2 currently supported formats that the program is able to read and translate to input.
- XML VOC format which looks like the following example:
<annotation>
<folder>/path/to/image/folder</folder>
<filename>image_filename.png</filename>
<path>/path/to/image/folder/image_filename.png</path>
<size>
<width>image_width</width>
<height>image_height</height>
<depth>image_depth</depth>
</size>
<object>
<name>obj1_name</name>
<bndbox>
<xmin>382.99999987200005</xmin>
<ymin>447.000000174</ymin>
<xmax>400.00000051200004</xmax>
<ymax>469.000000098</ymax>
</bndbox>
</annotation>
- CSV with relative labels that looks like the following example:
Anchor generator
A k-means algorithm finds the optimal sizes and generates anchors with process visualization.
matplotlib visualization of all stages
Including:
- k-means visualization:
- Generated anchors:
- Precision and recall curves:
- Evaluation bar charts:
- Actual vs. detections:
You can always visualize different stages of the program using my other repo labelpix which is tool for drawing bounding boxes, but can also be used to visualize bounding boxes over images using csv files in the format mentioned above
tf.data input pipeline
TFRecords a simple format for storing a sequence of binary records. Protocol buffers are a cross-platform, cross-language library for efficient serialization of structured data and are used as input pipeline to store and read data efficiently the program takes as input images and their respective annotations and builds training and validation(optional) TFRecords to be further used for all operations and TFRecords are also used in the evaluation(mid/post) training, so it's valid to say you can delete images to free space after conversion to TFRecords.
pandas & numpy data handling
Most of the operations are using numpy and pandas for efficiency and vectorization.
imgaug augmentation pipeline(customizable)
Special thanks to the amazing imgaug creators, an augmentation pipeline(optional) is available and NOTE that the augmentation is conducted before the training not during the training due to technical complications to integrate tensorflow and imgaug. If you have a small dataset, augmentation is an option and it can be preconfigured before the training
logging
Different operations are recorded using logging module.
All-in-1 custom Trainer class
For custom training, Trainer class accepts configurations for augmentation,
new anchor generation, new dataset(TFRecord(s)) creation, mAP evaluation
mid-training and post training. So all you have to do is place images
in Data > Photos, provide the configuration that suits you and start the training
process, all operations are managed from the same place for convenience.
For detailed instructions check
Stop and resume training support
by default the trainer checkpoints to Models > checkpoint_name.tf at the end of each training epoch which enables the training to be resumed at any given point by loading the checkpoint which would be the most recent.
Fully vectorized mAP evaluation
Evaluation is optional during the training every n epochs(not recommended for large datasets as it predicts every image in the dataset) and one evaluation at the end which is optional as well. Training and validation datasets can be evaluated separately and calculate mAP(mean average precision) as well as precision and recall curves for every class in the model.
trainer.py
import tensorflow as tf
import os
import numpy as np
import pandas as pd
from pathlib import Path
import sys
sys.path.append('..')
from tensorflow.keras.callbacks import (
ReduceLROnPlateau,
TensorBoard,
ModelCheckpoint,
Callback,
EarlyStopping,
)
import shutil
from Helpers.dataset_handlers import read_tfr, save_tfr, get_feature_map
from Helpers.annotation_parsers import parse_voc_folder
from Helpers.anchors import k_means, generate_anchors
from Helpers.augmentor import DataAugment
from Config.augmentation_options import augmentations
from Main.models import V3Model
from Helpers.utils import transform_images, transform_targets
from Helpers.annotation_parsers import adjust_non_voc_csv
from Helpers.utils import calculate_loss, timer, default_logger, activate_gpu
from Main.evaluator import Evaluator
class Trainer(V3Model):
"""
Create a training instance.
"""
def __init__(
self,
input_shape,
classes_file,
image_width,
image_height,
train_tf_record=None,
valid_tf_record=None,
anchors=None,
masks=None,
max_boxes=100,
iou_threshold=0.5,
score_threshold=0.5,
):
"""
Initialize training.
Args:
input_shape: tuple, (n, n, c)
classes_file: File containing class names \n delimited.
image_width: Width of the original image.
image_height: Height of the original image.
train_tf_record: TFRecord file.
valid_tf_record: TFRecord file.
anchors: numpy array of (w, h) pairs.
masks: numpy array of masks.
max_boxes: Maximum boxes of the TFRecords provided(if any) or
maximum boxes setting.
iou_threshold: float, values less than the threshold are ignored.
score_threshold: float, values less than the threshold are ignored.
"""
self.classes_file = classes_file
self.class_names = [
item.strip() for item in open(classes_file).readlines()
]
super().__init__(
input_shape,
len(self.class_names),
anchors,
masks,
max_boxes,
iou_threshold,
score_threshold,
)
self.train_tf_record = train_tf_record
self.valid_tf_record = valid_tf_record
self.image_folder = (
Path(os.path.join('..', 'Data', 'Photos')).absolute().resolve()
)
self.image_width = image_width
self.image_height = image_height
def get_adjusted_labels(self, configuration):
"""
Adjust labels according to given configuration.
Args:
configuration: A dictionary containing any of the following keys:
- relative_labels
- from_xml
- adjusted_frame
Returns:
pandas DataFrame with adjusted labels.
"""
labels_frame = None
check = 0
if configuration.get('relative_labels'):
labels_frame = adjust_non_voc_csv(
configuration['relative_labels'],
self.image_folder,
self.image_width,
self.image_height,
)
check += 1
if configuration.get('from_xml'):
if check:
raise ValueError(f'Got more than one configuration')
labels_frame = parse_voc_folder(
os.path.join('..', 'Data', 'XML Labels'),
os.path.join('..', 'Config', 'voc_conf.json'),
)
labels_frame.to_csv(
os.path.join('..', 'Output', 'Data', 'parsed_from_xml.csv'),
index=False,
)
check += 1
if configuration.get('adjusted_frame'):
if check:
raise ValueError(f'Got more than one configuration')
labels_frame = pd.read_csv(configuration['adjusted_frame'])
check += 1
return labels_frame
def generate_new_anchors(self, new_anchors_conf):
"""
Create new anchors according to given configuration.
Args:
new_anchors_conf: A dictionary containing the following keys:
- anchors_no
and one of the following:
- relative_labels
- from_xml
- adjusted_frame
Returns:
None
"""
anchor_no = new_anchors_conf.get('anchor_no')
if not anchor_no:
raise ValueError(f'No "anchor_no" found in new_anchors_conf')
labels_frame = self.get_adjusted_labels(new_anchors_conf)
relative_dims = np.array(
list(
zip(
labels_frame['Relative Width'],
labels_frame['Relative Height'],
)
)
)
centroids, _ = k_means(relative_dims, anchor_no, frame=labels_frame)
self.anchors = (
generate_anchors(self.image_width, self.image_height, centroids)
/ self.input_shape[0]
)
default_logger.info('Changed default anchors to generated ones')
def generate_new_frame(self, new_dataset_conf):
"""
Create new labels frame according to given configuration.
Args:
new_dataset_conf: A dictionary containing the following keys:
- dataset_name
and one of the following:
- relative_labels
- from_xml
- adjusted_frame
- coordinate_labels(optional in case of augmentation)
- augmentation(optional)
and this implies the following:
- sequences
- workers(optional, defaults to 32)
- batch_size(optional, defaults to 64)
- new_size(optional, defaults to None)
Returns:
pandas DataFrame adjusted for building the dataset containing
labels or labels and augmented labels combined
"""
if not new_dataset_conf.get('dataset_name'):
raise ValueError('dataset_name not found in new_dataset_conf')
labels_frame = self.get_adjusted_labels(new_dataset_conf)
if new_dataset_conf.get('augmentation'):
labels_frame = self.augment_photos(new_dataset_conf)
return labels_frame
def initialize_dataset(self, tf_record, batch_size, shuffle_buffer=512):
"""
Initialize and prepare TFRecord dataset for training.
Args:
tf_record: TFRecord file.
batch_size: int, training batch size
shuffle_buffer: Buffer size for shuffling dataset.
Returns:
dataset.
"""
dataset = read_tfr(
tf_record, self.classes_file, get_feature_map(), self.max_boxes
)
dataset = dataset.shuffle(shuffle_buffer)
dataset = dataset.batch(batch_size)
dataset = dataset.map(
lambda x, y: (
transform_images(x, self.input_shape[0]),
transform_targets(
y, self.anchors, self.masks, self.input_shape[0]
),
)
)
dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
return dataset
@staticmethod
def augment_photos(new_dataset_conf):
"""
Augment photos in self.image_paths
Args:
new_dataset_conf: A dictionary containing the following keys:
one of the following:
- relative_labels
- from_xml
- adjusted_frame
- coordinate_labels(optional)
and:
- sequences
- workers(optional, defaults to 32)
- batch_size(optional, defaults to 64)
- new_size(optional, defaults to None)
Returns:
pandas DataFrame with both original and augmented data.
"""
sequences = new_dataset_conf.get('sequences')
relative_labels = new_dataset_conf.get('relative_labels')
coordinate_labels = new_dataset_conf.get('coordinate_labels')
workers = new_dataset_conf.get('workers')
batch_size = new_dataset_conf.get('batch_size')
new_augmentation_size = new_dataset_conf.get('new_size')
if not sequences:
raise ValueError(f'"sequences" not found in new_dataset_conf')
if not relative_labels:
raise ValueError(f'No "relative_labels" found in new_dataset_conf')
augment = DataAugment(
relative_labels, augmentations, workers or 32, coordinate_labels
)
augment.create_sequences(sequences)
return augment.augment_photos_folder(
batch_size or 64, new_augmentation_size
)
@timer(default_logger)
def evaluate(
self,
weights_file,
merge,
workers,
shuffle_buffer,
min_overlaps,
display_stats=True,
plot_stats=True,
save_figs=True,
):
"""
Evaluate on training and validation datasets.
Args:
weights_file: Path to trained .tf file.
merge: If False, training and validation datasets will be evaluated separately.
workers: Parallel predictions.
shuffle_buffer: Buffer size for shuffling datasets.
min_overlaps: a float value between 0 and 1, or a dictionary
containing each class in self.class_names mapped to its
minimum overlap
display_stats: If True evaluation statistics will be printed.
plot_stats: If True, evaluation statistics will be plotted including
precision and recall curves and mAP
save_figs: If True, resulting plots will be save to Output folder.
Returns:
stats, map_score.
"""
default_logger.info('Starting evaluation ...')
evaluator = Evaluator(
self.input_shape,
self.train_tf_record,
self.valid_tf_record,
self.classes_file,
self.anchors,
self.masks,
self.max_boxes,
self.iou_threshold,
self.score_threshold,
)
predictions = evaluator.make_predictions(
weights_file, merge, workers, shuffle_buffer
)
if isinstance(predictions, tuple):
training_predictions, valid_predictions = predictions
if any([training_predictions.empty, valid_predictions.empty]):
default_logger.info(
'Aborting evaluations, no detections found'
)
return
training_actual = pd.read_csv(
os.path.join('..', 'Data', 'TFRecords', 'training_data.csv')
)
valid_actual = pd.read_csv(
os.path.join('..', 'Data', 'TFRecords', 'test_data.csv')
)
training_stats, training_map = evaluator.calculate_map(
training_predictions,
training_actual,
min_overlaps,
display_stats,
'Train',
save_figs,
plot_stats,
)
valid_stats, valid_map = evaluator.calculate_map(
valid_predictions,
valid_actual,
min_overlaps,
display_stats,
'Valid',
save_figs,
plot_stats,
)
return training_stats, training_map, valid_stats, valid_map
actual_data = pd.read_csv(
os.path.join('..', 'Data', 'TFRecords', 'full_data.csv')
)
if predictions.empty:
default_logger.info('Aborting evaluations, no detections found')
return
stats, map_score = evaluator.calculate_map(
predictions,
actual_data,
min_overlaps,
display_stats,
save_figs=save_figs,
plot_results=plot_stats,
)
return stats, map_score
@staticmethod
def clear_outputs():
"""
Clear Output folder.
Returns:
None
"""
for file_name in os.listdir(os.path.join('..', 'Output')):
if not file_name.startswith('.'):
full_path = (
Path(os.path.join('..', 'Output', file_name))
.absolute()
.resolve()
)
if os.path.isdir(full_path):
shutil.rmtree(full_path)
else:
os.remove(full_path)
default_logger.info(f'Deleted old output: {full_path}')
def create_new_dataset(self, new_dataset_conf):
"""
Build new dataset and respective TFRecord(s).
Args:
new_dataset_conf: A dictionary containing the following keys:
one of the following:
- relative_labels
- from_xml
- adjusted_frame
- coordinate_labels(optional)
and:
- sequences
- workers(optional, defaults to 32)
- batch_size(optional, defaults to 64)
- new_size(optional, defaults to None)
Returns:
None
"""
default_logger.info(f'Generating new dataset ...')
test_size = new_dataset_conf.get('test_size')
labels_frame = self.generate_new_frame(new_dataset_conf)
save_tfr(
labels_frame,
os.path.join('..', 'Data', 'TFRecords'),
new_dataset_conf['dataset_name'],
test_size,
self,
)
def check_tf_records(self):
"""
Ensure TFRecords are specified to start training.
Returns:
None
"""
if not self.train_tf_record:
issue = 'No training TFRecord specified'
default_logger.error(issue)
raise ValueError(issue)
if not self.valid_tf_record:
issue = 'No validation TFRecord specified'
default_logger.error(issue)
raise ValueError(issue)
@staticmethod
def create_callbacks(checkpoint_name):
"""
Create a list of tf.keras.callbacks.
Args:
checkpoint_name: Name under which the checkpoint is saved.
Returns:
callbacks.
"""
return [
ReduceLROnPlateau(verbose=3),
ModelCheckpoint(
os.path.join(checkpoint_name),
verbose=1,
save_weights_only=True,
),
TensorBoard(log_dir=os.path.join('..', 'Logs')),
EarlyStopping(monitor='val_loss', patience=6, verbose=1),
]
@timer(default_logger)
def train(
self,
epochs,
batch_size,
learning_rate,
new_anchors_conf=None,
new_dataset_conf=None,
dataset_name=None,
weights=None,
evaluate=True,
merge_evaluation=True,
evaluation_workers=8,
shuffle_buffer=512,
min_overlaps=None,
display_stats=True,
plot_stats=True,
save_figs=True,
clear_outputs=False,
n_epoch_eval=None,
):
"""
Train on the dataset.
Args:
epochs: Number of training epochs.
batch_size: Training batch size.
learning_rate: non-negative value.
new_anchors_conf: A dictionary containing anchor generation configuration.
new_dataset_conf: A dictionary containing dataset generation configuration.
dataset_name: Name of the dataset for model checkpoints.
weights: .tf or .weights file
evaluate: If False, the trained model will not be evaluated after training.
merge_evaluation: If False, training and validation maps will
be calculated separately.
evaluation_workers: Parallel predictions.
shuffle_buffer: Buffer size for shuffling datasets.
min_overlaps: a float value between 0 and 1, or a dictionary
containing each class in self.class_names mapped to its
minimum overlap
display_stats: If True and evaluate=True, evaluation statistics will be displayed.
plot_stats: If True, Precision and recall curves as well as
comparative bar charts will be plotted
save_figs: If True and plot_stats=True, figures will be saved
clear_outputs: If True, old outputs will be cleared
n_epoch_eval: Conduct evaluation every n epoch.
Returns:
history object, pandas DataFrame with statistics, mAP score.
"""
min_overlaps = min_overlaps or 0.5
if clear_outputs:
self.clear_outputs()
activate_gpu()
default_logger.info(f'Starting training ...')
if new_anchors_conf:
default_logger.info(f'Generating new anchors ...')
self.generate_new_anchors(new_anchors_conf)
self.create_models()
if weights:
self.load_weights(weights)
if new_dataset_conf:
self.create_new_dataset(new_dataset_conf)
self.check_tf_records()
training_dataset = self.initialize_dataset(
self.train_tf_record, batch_size, shuffle_buffer
)
valid_dataset = self.initialize_dataset(
self.valid_tf_record, batch_size, shuffle_buffer
)
optimizer = tf.keras.optimizers.Adam(learning_rate)
loss = [
calculate_loss(
self.anchors[mask], self.classes, self.iou_threshold
)
for mask in self.masks
]
self.training_model.compile(optimizer=optimizer, loss=loss)
checkpoint_name = os.path.join(
'..', 'Models', f'{dataset_name or "trained"}_model.tf'
)
callbacks = self.create_callbacks(checkpoint_name)
if n_epoch_eval:
mid_train_eval = MidTrainingEvaluator(
self.input_shape,
self.classes_file,
self.image_width,
self.image_height,
self.train_tf_record,
self.valid_tf_record,
self.anchors,
self.masks,
self.max_boxes,
self.iou_threshold,
self.score_threshold,
n_epoch_eval,
merge_evaluation,
evaluation_workers,
shuffle_buffer,
min_overlaps,
display_stats,
plot_stats,
save_figs,
checkpoint_name,
)
callbacks.append(mid_train_eval)
history = self.training_model.fit(
training_dataset,
epochs=epochs,
callbacks=callbacks,
validation_data=valid_dataset,
)
default_logger.info('Training complete')
if evaluate:
evaluations = self.evaluate(
checkpoint_name,
merge_evaluation,
evaluation_workers,
shuffle_buffer,
min_overlaps,
display_stats,
plot_stats,
save_figs,
)
return evaluations, history
return history
class MidTrainingEvaluator(Callback, Trainer):
"""
Tool to evaluate trained model on the go(during the training, every n epochs).
"""
def __init__(
self,
input_shape,
classes_file,
image_width,
image_height,
train_tf_record,
valid_tf_record,
anchors,
masks,
max_boxes,
iou_threshold,
score_threshold,
n_epochs,
merge,
workers,
shuffle_buffer,
min_overlaps,
display_stats,
plot_stats,
save_figs,
weights_file,
):
"""
Initialize mid-training evaluation settings.
Args:
input_shape: tuple, (n, n, c)
classes_file: File containing class names \n delimited.
image_width: Width of the original image.
image_height: Height of the original image.
train_tf_record: TFRecord file.
valid_tf_record: TFRecord file.
anchors: numpy array of (w, h) pairs.
masks: numpy array of masks.
max_boxes: Maximum boxes of the TFRecords provided(if any) or
maximum boxes setting.
iou_threshold: float, values less than the threshold are ignored.
score_threshold: float, values less than the threshold are ignored.
n_epochs: int, perform evaluation every n epochs
merge: If True, The whole dataset(train + valid) will be evaluated
workers: Parallel predictions
shuffle_buffer: Buffer size for shuffling datasets
min_overlaps: a float value between 0 and 1, or a dictionary
containing each class in self.class_names mapped to its
minimum overlap
display_stats: If True, statistics will be displayed at the end.
plot_stats: If True, precision and recall curves as well as
comparison bar charts will be plotted.
save_figs: If True and display_stats, plots will be save to Output folder
weights_file: .tf file(most recent checkpoint)
"""
Trainer.__init__(
self,
input_shape,
classes_file,
image_width,
image_height,
train_tf_record,
valid_tf_record,
anchors,
masks,
max_boxes,
iou_threshold,
score_threshold,
)
self.n_epochs = n_epochs
self.evaluation_args = [
weights_file,
merge,
workers,
shuffle_buffer,
min_overlaps,
display_stats,
plot_stats,
save_figs,
]
def on_epoch_end(self, epoch, logs=None):
"""
Start evaluation in valid epochs.
Args:
epoch: int, epoch number.
logs: dict, Tensorboard log.
Returns:
None
"""
if not (epoch + 1) % self.n_epochs == 0:
return
self.evaluate(*self.evaluation_args)
os.mkdir(
os.path.join(
'..', 'Output', 'Evaluation', f'epoch-{epoch}-evaluation'
)
)
for file_name in os.listdir(
os.path.join('..', 'Output', 'Evaluation')
):
if not os.path.isdir(file_name) and (
file_name.endswith('.png') or 'prediction' in file_name
):
full_path = str(
Path(os.path.join('..', 'Output', 'Evaluation', file_name))
.absolute()
.resolve()
)
new_path = str(
Path(
os.path.join(
'..',
'Output',
'Evaluation',
f'epoch-{epoch}-evaluation',
file_name,
)
)
.absolute()
.resolve()
)
shutil.move(full_path, new_path)
evaluator.py
import cv2
import pandas as pd
import numpy as np
import tensorflow as tf
import os
import sys
sys.path.append('..')
from concurrent.futures import ThreadPoolExecutor, as_completed
from Main.models import V3Model
from Helpers.dataset_handlers import read_tfr, get_feature_map
from Helpers.utils import (
transform_images,
get_detection_data,
default_logger,
timer,
)
from Helpers.visual_tools import visualize_pr, visualize_evaluation_stats
class Evaluator(V3Model):
def __init__(
self,
input_shape,
train_tf_record,
valid_tf_record,
classes_file,
anchors=None,
masks=None,
max_boxes=100,
iou_threshold=0.5,
score_threshold=0.5,
):
"""
Evaluate a trained model.
Args:
input_shape: input_shape: tuple, (n, n, c)
train_tf_record: Path to training TFRecord file.
valid_tf_record: Path to validation TFRecord file.
classes_file: File containing class names \n delimited.
anchors: numpy array of (w, h) pairs.
masks: numpy array of masks.
max_boxes: Maximum boxes of the TFRecords provided.
iou_threshold: Minimum overlap value.
score_threshold: Minimum confidence for detection to count
as true positive.
"""
self.classes_file = classes_file
self.class_names = [
item.strip() for item in open(classes_file).readlines()
]
super().__init__(
input_shape,
len(self.class_names),
anchors,
masks,
max_boxes,
iou_threshold,
score_threshold,
)
self.train_tf_record = train_tf_record
self.valid_tf_record = valid_tf_record
self.train_dataset_size = sum(
1 for _ in tf.data.TFRecordDataset(train_tf_record)
)
self.valid_dataset_size = sum(
1 for _ in tf.data.TFRecordDataset(valid_tf_record)
)
self.dataset_size = self.train_dataset_size + self.valid_dataset_size
self.predicted = 1
def predict_image(self, image_data, features):
"""
Make predictions on a single image from the TFRecord.
Args:
image_data: image as numpy array
features: features of the TFRecord.
Returns:
pandas DataFrame with detection data.
"""
image_path = bytes.decode(features['image_path'].numpy())
image_name = os.path.basename(image_path)
image = tf.expand_dims(image_data, 0)
resized = transform_images(image, self.input_shape[0])
outs = self.inference_model(resized)
adjusted = cv2.cvtColor(image_data.numpy(), cv2.COLOR_RGB2BGR)
result = (
get_detection_data(adjusted, image_name, outs, self.class_names),
image_name,
)
return result
@staticmethod
def get_dataset_next(dataset):
try:
return next(dataset)
except tf.errors.UnknownError as e: # sometimes encountered when reading from google drive
default_logger.error(
f'Error occurred during reading from dataset\n{e}'
)
def predict_dataset(
self, dataset, workers=16, split='train', batch_size=64
):
"""
Predict entire dataset.
Args:
dataset: MapDataset object.
workers: Parallel predictions.
split: str representation of the dataset 'train' or 'valid'
batch_size: Prediction batch size.
Returns:
pandas DataFrame with entire dataset predictions.
"""
predictions = []
sizes = {
'train': self.train_dataset_size,
'valid': self.valid_dataset_size,
}
size = sizes[split]
current_prediction = 0
with ThreadPoolExecutor(max_workers=workers) as executor:
while current_prediction < size:
current_batch = []
for _ in range(min(batch_size, size - current_prediction)):
item = self.get_dataset_next(dataset)
if item is not None:
current_batch.append(item)
future_predictions = {
executor.submit(
self.predict_image, img_data, features
): features['image_path']
for img_data, labels, features in current_batch
}
for future_prediction in as_completed(future_predictions):
result, completed_image = future_prediction.result()
predictions.append(result)
completed = f'{self.predicted}/{self.dataset_size}'
percent = (self.predicted / self.dataset_size) * 100
print(
f'\rpredicting {completed_image} {completed}\t{percent}% completed',
end='',
)
self.predicted += 1
current_prediction += 1
return pd.concat(predictions)
@timer(default_logger)
def make_predictions(
self,
trained_weights,
merge=False,
workers=16,
shuffle_buffer=512,
batch_size=64,
):
"""
Make predictions on both training and validation data sets
and save results as csv in Output folder.
Args:
trained_weights: Trained .tf weights or .weights file(in case self.classes = 80).
merge: If True a single file will be saved for training
and validation sets predictions combined.
workers: Parallel predictions.
shuffle_buffer: int, shuffle dataset buffer size.
batch_size: Prediction batch size.
Returns:
1 combined pandas DataFrame for entire dataset predictions
or 2 pandas DataFrame(s) for training and validation
data sets respectively.
"""
self.create_models()
self.load_weights(trained_weights)
features = get_feature_map()
train_dataset = read_tfr(
self.train_tf_record,
self.classes_file,
features,
self.max_boxes,
get_features=True,
)
valid_dataset = read_tfr(
self.valid_tf_record,
self.classes_file,
features,
self.max_boxes,
get_features=True,
)
train_dataset.shuffle(shuffle_buffer)
valid_dataset.shuffle(shuffle_buffer)
train_dataset = iter(train_dataset)
valid_dataset = iter(valid_dataset)
train_predictions = self.predict_dataset(
train_dataset, workers, 'train', batch_size
)
valid_predictions = self.predict_dataset(
valid_dataset, workers, 'valid', batch_size
)
if merge:
predictions = pd.concat([train_predictions, valid_predictions])
save_path = os.path.join(
'..', 'Output', 'Data', 'full_dataset_predictions.csv'
)
predictions.to_csv(save_path, index=False)
return predictions
train_path = os.path.join(
'..', 'Output', 'Data', 'train_dataset_predictions.csv'
)
valid_path = os.path.join(
'..', 'Output', 'Data', 'valid_dataset_predictions.csv'
)
train_predictions.to_csv(train_path, index=False)
valid_predictions.to_csv(valid_path, index=False)
return train_predictions, valid_predictions
@staticmethod
def get_area(frame, columns):
"""
Calculate bounding boxes areas.
Args:
frame: pandas DataFrame that contains prediction data.
columns: column names that represent x1, y1, x2, y2.
Returns:
pandas Series(area column)
"""
x1, y1, x2, y2 = [frame[column] for column in columns]
return (x2 - x1) * (y2 - y1)
def get_true_positives(self, detections, actual, min_overlaps):
"""
Filter True positive detections out of all detections.
Args:
detections: pandas DataFrame with all detections.
actual: pandas DataFrame with real data.
min_overlaps: a float value between 0 and 1, or a dictionary
containing each class in self.class_names mapped to its
minimum overlap
Returns:
pandas DataFrame that contains detections that satisfy
True positive constraints.
"""
if detections.empty:
raise ValueError(f'Empty predictions frame')
if isinstance(min_overlaps, float):
assert 0 <= min_overlaps < 1, (
f'min_overlaps should be '
f'between 0 and 1, {min_overlaps} is given'
)
if isinstance(min_overlaps, dict):
assert all(
[0 < min_overlap < 1 for min_overlap in min_overlaps.values()]
)
assert all([obj in min_overlaps for obj in self.class_names]), (
f'{[item for item in self.class_names if item not in min_overlaps]} '
f'are missing in min_overlaps'
)
actual = actual.rename(
columns={'Image Path': 'image', 'Object Name': 'object_name'}
)
actual['image'] = actual['image'].apply(lambda x: os.path.split(x)[-1])
random_gen = np.random.default_rng()
if 'detection_key' not in detections.columns:
detection_keys = random_gen.choice(
len(detections), size=len(detections), replace=False
)
detections['detection_key'] = detection_keys
total_frame = actual.merge(detections, on=['image', 'object_name'])
assert (
not total_frame.empty
), 'No common image names found between actual and detections'
total_frame['x_max_common'] = total_frame[['X_max', 'x2']].min(1)
total_frame['x_min_common'] = total_frame[['X_min', 'x1']].max(1)
total_frame['y_max_common'] = total_frame[['Y_max', 'y2']].min(1)
total_frame['y_min_common'] = total_frame[['Y_min', 'y1']].max(1)
true_intersect = (
total_frame['x_max_common'] > total_frame['x_min_common']
) & (total_frame['y_max_common'] > total_frame['y_min_common'])
total_frame = total_frame[true_intersect]
actual_areas = self.get_area(
total_frame, ['X_min', 'Y_min', 'X_max', 'Y_max']
)
predicted_areas = self.get_area(total_frame, ['x1', 'y1', 'x2', 'y2'])
intersect_areas = self.get_area(
total_frame,
['x_min_common', 'y_min_common', 'x_max_common', 'y_max_common'],
)
iou_areas = intersect_areas / (
actual_areas + predicted_areas - intersect_areas
)
total_frame['iou'] = iou_areas
if isinstance(min_overlaps, float):
return total_frame[total_frame['iou'] >= min_overlaps]
if isinstance(min_overlaps, dict):
class_data = [
(name, total_frame[total_frame['object_name'] == name])
for name in self.class_names
]
thresholds = [min_overlaps[item[0]] for item in class_data]
frames = [
item[1][item[1]['iou'] >= threshold]
for (item, threshold) in zip(class_data, thresholds)
if not item[1].empty
]
return pd.concat(frames)
@staticmethod
def get_false_positives(detections, true_positive):
"""
Filter out False positives in all detections.
Args:
detections: pandas DataFrame with detection data.
true_positive: pandas DataFrame with True positive data.
Returns:
pandas DataFrame with False positives.
"""
keys_before = detections['detection_key'].values
keys_after = true_positive['detection_key'].values
false_keys = np.where(np.isin(keys_before, keys_after, invert=True))
false_keys = keys_before[false_keys]
false_positives = detections.set_index('detection_key').loc[false_keys]
return false_positives.reset_index()
@staticmethod
def combine_results(true_positive, false_positive):
"""
Combine True positives and False positives.
Args:
true_positive: pandas DataFrame with True positive data.
false_positive: pandas DataFrame with False positive data.
Returns:
pandas DataFrame with all detections combined.
"""
true_positive['true_positive'] = 1
true_positive['false_positive'] = 0
true_positive = true_positive[
[
'image',
'object_name',
'score',
'x_min_common',
'y_min_common',
'x_max_common',
'y_max_common',
'iou',
'image_width',
'image_height',
'true_positive',
'false_positive',
'detection_key',
]
]
true_positive = true_positive.rename(
columns={
'x_min_common': 'x1',
'y_min_common': 'y1',
'x_max_common': 'x2',
'y_max_common': 'y2',
}
)
false_positive['iou'] = 0
false_positive['true_positive'] = 0
false_positive['false_positive'] = 1
false_positive = false_positive[
[
'image',
'object_name',
'score',
'x1',
'y1',
'x2',
'y2',
'iou',
'image_width',
'image_height',
'true_positive',
'false_positive',
'detection_key',
]
]
return pd.concat([true_positive, false_positive])
def calculate_stats(
self,
actual_data,
detection_data,
true_positives,
false_positives,
combined,
):
"""
Calculate prediction statistics for every class in self.class_names.
Args:
actual_data: pandas DataFrame with real data.
detection_data: pandas DataFrame with all detection data before filtration.
true_positives: pandas DataFrame with True positives.
false_positives: pandas DataFrame with False positives.
combined: pandas DataFrame with True and False positives combined.
Returns:
pandas DataFrame with statistics for all classes.
"""
class_stats = []
for class_name in self.class_names:
stats = dict()
stats['Class Name'] = class_name
stats['Average Precision'] = (
combined[combined['object_name'] == class_name][
'average_precision'
].sum()
* 100
)
stats['Actual'] = len(
actual_data[actual_data["Object Name"] == class_name]
)
stats['Detections'] = len(
detection_data[detection_data["object_name"] == class_name]
)
stats['True Positives'] = len(
true_positives[true_positives["object_name"] == class_name]
)
stats['False Positives'] = len(
false_positives[false_positives["object_name"] == class_name]
)
stats['Combined'] = len(
combined[combined["object_name"] == class_name]
)
class_stats.append(stats)
total_stats = pd.DataFrame(class_stats).sort_values(
by='Average Precision', ascending=False
)
return total_stats
@staticmethod
def calculate_ap(combined, total_actual):
"""
Calculate average precision for a single object class.
Args:
combined: pandas DataFrame with True and False positives combined.
total_actual: Total number of actual object class boxes.
Returns:
pandas DataFrame with average precisions calculated.
"""
combined = combined.sort_values(
by='score', ascending=False
).reset_index(drop=True)
combined['acc_tp'] = combined['true_positive'].cumsum()
combined['acc_fp'] = combined['false_positive'].cumsum()
combined['precision'] = combined['acc_tp'] / (
combined['acc_tp'] + combined['acc_fp']
)
combined['recall'] = combined['acc_tp'] / total_actual
combined['m_pre1'] = combined['precision'].shift(1, fill_value=0)
combined['m_pre'] = combined[['m_pre1', 'precision']].max(axis=1)
combined['m_rec1'] = combined['recall'].shift(1, fill_value=0)
combined.loc[
combined['m_rec1'] != combined['recall'], 'valid_m_rec'
] = 1
combined['average_precision'] = (
combined['recall'] - combined['m_rec1']
) * combined['m_pre']
return combined
@timer(default_logger)
def calculate_map(
self,
prediction_data,
actual_data,
min_overlaps,
display_stats=False,
fig_prefix='',
save_figs=True,
plot_results=True,
):
"""
Calculate mAP(mean average precision) for the trained model.
Args:
prediction_data: pandas DataFrame containing predictions.
actual_data: pandas DataFrame containing actual data.
min_overlaps: a float value between 0 and 1, or a dictionary
containing each class in self.class_names mapped to its
minimum overlap
display_stats: If True, statistics will be displayed.
fig_prefix: Prefix for plot titles.
save_figs: If True, figures will be saved.
plot_results: If True, results will be calculated.
Returns:
pandas DataFrame with statistics, mAP score.
"""
actual_data['Object Name'] = actual_data['Object Name'].apply(
lambda x: x.replace("b'", '').replace("'", '')
)
class_counts = actual_data['Object Name'].value_counts().to_dict()
true_positives = self.get_true_positives(
prediction_data, actual_data, min_overlaps
)
false_positives = self.get_false_positives(
prediction_data, true_positives
)
combined = self.combine_results(true_positives, false_positives)
class_groups = combined.groupby('object_name')
calculated = pd.concat(
[
self.calculate_ap(group, class_counts.get(object_name))
for object_name, group in class_groups
]
)
stats = self.calculate_stats(
actual_data,
prediction_data,
true_positives,
false_positives,
calculated,
)
map_score = stats['Average Precision'].mean()
if display_stats:
pd.set_option(
'display.max_rows',
None,
'display.max_columns',
None,
'display.width',
None,
)
print(stats.sort_values(by='Average Precision', ascending=False))
print(f'mAP score: {map_score}%')
pd.reset_option('display.[max_rows, max_columns, width]')
if plot_results:
visualize_pr(calculated, save_figs, fig_prefix)
visualize_evaluation_stats(stats, fig_prefix)
return stats, map_score
