def augment_img(aug, suffix, img_path, count=1):
"""Iteratively augment a single image with a given augmentation function.
Generates transformed labels for each augmentation
Parameters:
aug Augmentation function from albumentations
suffix String suffix appended to each augmentation file and label
img_path Filesystem path to the image to augment
count Number of augmentations to perform
"""
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
base_name = img_path[:-4]
label_path = get_label_path(img_path)
boxes, field_ids = parse_label(label_path)
i = 0
while i < count:
aug_path = f"{base_name.replace('images', 'aug-images')}_{suffix}-{i}.png"
new_txt_path = get_label_path(aug_path)
os.makedirs(os.path.dirname(new_txt_path), exist_ok=True)
os.makedirs(os.path.dirname(aug_path), exist_ok=True)
if os.path.exists(aug_path) and os.path.exists(new_txt_path):
i += 1
continue
try:
result = aug(image=img, bboxes=boxes, classes=field_ids)
except IndexError:
continue
aug_img = result["image"]
new_bboxes = [" ".join(map(str, bbox)) for bbox in result["bboxes"]]
# Check if bounding boxes have been removed due to visibility criteria
if len(new_bboxes) != len(boxes):
continue
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(aug_path, aug_img)
with open(new_txt_path, "w+") as out:
for box_i, bbox_str in enumerate(new_bboxes):
out.write(f"{field_ids[box_i]} {bbox_str}\n")
i += 1
def augment(self, imgs_per_class, major_aug, min_visibility=0.75):
"""Augment all images in the image folder, adding the augmentations to the folder.
Parameters:
imgs_per_class Target number of images. If there are more samples in the folder
than this number for a class, no augmentation will be performed.
major_aug A boolean variable determining if 'major' transformations will be used.
min_visibility Minimum visibility of the resultant bounding boxes after augmentation.
This is a value in (0.0, 1.0] relative to the area of the bounding box.
"""
incr_factors = self.get_incr_factors(imgs_per_class)
bbox_params = alb.BboxParams("yolo",
min_visibility=min_visibility,
label_fields=["classes"])
aug = multi_aug(get_augmentations(), major_aug, bbox_params)
pbar = tqdm(desc="Augmenting training images",
total=sum(incr_factors.values()))
for img, count in incr_factors.items():
augment_img(aug, "compose", img, count=count)
new_imgs = {
f"{img[:-4].replace('images', 'aug-images')}_compose-{i}.png"
for i in range(count)
}
self.img_folder.imgs.update(new_imgs)
self.img_folder.labels.update(
{get_label_path(img)
for img in new_imgs})
pbar.update(count)
pbar.close()
def make_img_dict(self):
"""Get a dictionary of image paths and their corresponding classes."""
img_dict = dict()
for img in self.imgs:
classes = self.get_classes(get_label_path(img))
if len(classes) != 0:
img_dict[img] = classes
return img_dict
def filter_images(self):
"""Remove non-labeled images from the folder."""
labeled_imgs = set()
for img in self.imgs:
label_path = get_label_path(img)
if os.path.exists(label_path):
labeled_imgs.add(img)
self.imgs = labeled_imgs
def __init__(
self,
img_list,
img_size=416,
multiscale=True,
normalized_labels=True,
):
self.img_files = img_list
self.label_files = [get_label_path(path) for path in self.img_files]
self.img_size = img_size
self.multiscale = multiscale
self.normalized_labels = normalized_labels
self.min_size = self.img_size - 3 * 32
self.max_size = self.img_size + 3 * 32
self.batch_count = 0
def label(self, classes, ground_truth_func):
"""Label images in the folder in accordance to a provided ground truth function.
Ground truth function should return a list of labels, each in a tuple format with
class name and normalized bounding box dimensions in Darknet format.
The input class list should be ordered and consistent with the ground truth labels.
"""
for img in self.imgs:
labels = ground_truth_func(img, classes)
if len(labels) == 0:
return
text_label = open(get_label_path(img), "w+")
for i, (class_display_name, x_cent, y_cent, w,
h) in enumerate(labels):
class_num = classes.index(class_display_name)
if i > 0:
text_label.write("\n")
text_label.write(f"{class_num} {x_cent} {y_cent} {w} {h}")
text_label.close()
def get_labels(self):
"""Get a set of labels corresponding to images in the folder."""
return {get_label_path(img) for img in self.imgs}
def make_results_df(config, img_folder, detections_by_img, total_epochs):
metrics = [
"file",
"actual",
"detected",
"conf",
"conf_std",
"hit",
]
results = pd.DataFrame(columns=metrics)
classes = utils.load_classes(config["class_list"])
for path, detections in detections_by_img.items():
ground_truths = img_folder.get_classes(utils.get_label_path(path))
detection_pairs = list()
if detections is not None:
region_detections, regions_std = yoloutils.group_average_bb(
detections, total_epochs, config["iou_thres"])
# evaluate.save_image(region_detections, path, config, classes)
if len(region_detections) == 1:
detected_class = int(region_detections.numpy()[0][-1])
if detected_class in ground_truths:
label = detected_class
elif len(ground_truths) == 1:
label = ground_truths[0]
else:
label = None
detection_pairs = [(label, region_detections[0])]
else:
test_img = LabeledSet([path], len(classes))
detection_pairs = evaluate.match_detections(
test_img, region_detections.unsqueeze(0), config)
for (truth, box) in detection_pairs:
if box is None:
continue
obj_conf, class_conf, pred_class = box.numpy()[4:]
obj_std, class_std = regions_std[round(float(class_conf), 3)]
row = {
"file": path,
"detected": classes[int(pred_class)],
"actual": classes[int(truth)] if truth is not None else "",
"conf": obj_conf * class_conf,
"conf_std": math.sqrt(obj_std**2 + class_std**2),
}
row["hit"] = row["actual"] == row["detected"]
results = results.append(row, ignore_index=True)
if truth is not None:
ground_truths.remove(int(truth))
# Add rows for those missing detections
for truth in ground_truths:
row = {
"file": path,
"detected": "",
"actual": classes[int(truth)],
"conf": 0.0,
"hit": False,
"conf_std": 0.0,
}
results = results.append(row, ignore_index=True)
return results
def simple_benchmark_avg(img_folder,
prefix,
start,
end,
total_epochs,
config,
roll=False):
"""Deprecated version of benchmark averaging, meant for single object
detection within an image. Used for a fair comparison baseline on old models
"""
loader = DataLoader(
img_folder,
batch_size=1,
shuffle=False,
num_workers=config["n_cpu"],
)
results = pd.DataFrame(
columns=["file", "confs", "actual", "detected", "conf", "hit"])
results.set_index("file")
classes = utils.load_classes(config["class_list"])
if roll:
checkpoints_i = list(range(max(1, end - total_epochs + 1), end + 1))
else:
checkpoints_i = list(
sorted(
set(
np.linspace(start,
end,
total_epochs,
dtype=np.dtype(np.int16)))))
single = total_epochs == 1
if not single:
print("Benchmarking on epochs", checkpoints_i)
for n in tqdm(checkpoints_i, "Benchmarking epochs", disable=single):
ckpt = get_checkpoint(config["checkpoints"], prefix, n)
model_def = yoloutils.parse_model_config(config["model_config"])
model = models.get_eval_model(model_def, config["img_size"], ckpt)
for (img_paths, input_imgs) in loader:
path = img_paths[0]
if path not in results.file:
actual_class = classes[img_folder.get_classes(
utils.get_label_path(path))[0]]
results.loc[path] = [
path, dict(), actual_class, None, None, None
]
detections = evaluate.detect(input_imgs, config["conf_thres"],
model)
confs = results.loc[path]["confs"]
for detection in detections:
if detection is None:
continue
(_, _, _, _, _, cls_conf, cls_pred) = detection.numpy()[0]
if cls_pred not in confs.keys():
confs[cls_pred] = [cls_conf]
else:
confs[cls_pred].append(cls_conf)
for _, row in results.iterrows():
best_class = None
best_conf = float("-inf")
for class_name, confs in row["confs"].items():
avg_conf = sum(confs) / len(checkpoints_i)
if avg_conf > best_conf:
best_conf = avg_conf
best_class = class_name
if best_class is not None:
row["detected"] = classes[int(best_class)]
row["conf"] = best_conf
row["hit"] = row["actual"] == row["detected"]
else:
row["detected"] = ""
row["conf"] = 0.0
row["hit"] = False
return results