def create_masked_template_to_verify(args):
'''
Read in template mask drawn by the user, verify it, and then write it to dst folder.
Output format: gray image with white/black values
'''
# Settings
folder_dst = args.f_data_dst + "masked_template/"
# Check
cf.create_folder(folder_dst)
# Vars
dataset_tp = TemplatesDataset(args) # template dataset
# Start
for i in range(len(dataset_tp)):
# Read in template image and corresponding mask
img, mask = dataset_tp.load_ith_image(i)
# Create masked image
img = pi.add_mask(img, mask) # put mask onto the image
res_img = np.hstack((img, pi.cvt_binary2color(mask)))
# Write to file
fname = dataset_tp.get_ith_filenames(i, base_name_only=True)[0]
cv2.imwrite(folder_dst + fname, res_img)
def set_args():
args = cf.SimpleNamespace()
configs = read_all_args(ROOT + "config/config.yaml")
# 1. Path to model definition file
# e.g.: "data/custom1_generated/yolo.cfg"
args.f_yolo_config = configs.f_yolo_config
# 2. Path to data config file
# e.g.: "data/custom1_generated/yolo.data"
args.f_yolo_data = configs.f_yolo_data
# 3. Other training settings
training_args = set_training_args()
args.__dict__.update(training_args.__dict__)
# 4. Folder to save model
save_model_to = "checkpoints/" + cf.get_readable_time(no_blank=True) + "/"
cf.create_folder(save_model_to)
args.save_model_to = save_model_to
cf.write_dict(save_model_to + "configs.yaml", args.__dict__)
# return
print("\nArgs:\n", args)
return args
def augment_images(args):
# Settings
pass
# Check
cf.create_folder(args.f_yolo_images)
cf.create_folder(args.f_yolo_labels)
cf.create_folder(args.f_yolo_images_with_bbox)
# Vars
dataset_tp = TemplatesDataset(args) # template dataset
dataset_bg = BackgroundDataset(args) # background dataset
yolo_labels = YoloLabels(args) # label parser
aug = ImgAugmenter(args.template_aug_effects)
# Functions
def get_random_template():
i = np.random.randint(len(dataset_tp))
image, mask = dataset_tp[i]
filename = dataset_tp.get_ith_filenames(i)[0]
label, label_idx = yolo_labels.parse_label(filename)
return image, mask, label, label_idx
def get_random_background():
i = np.random.randint(len(dataset_bg))
return dataset_bg[i] # return image
def random_object_number():
l, r = args.img_aug_nums["objects_per_image"]
return np.random.randint(l, r + 1)
# Start
for ith_background in range(args.img_aug_nums["num_new_images"]):
print("Generating {}th augmented image ...".format(ith_background))
# Read background image
bg_img = get_random_background()
# Vars to store
masks = []
labels = []
# 1st augment (apply small affine to background image)
bg_img = aug.augment_by_transform(bg_img)
# Add many templates onto background image
for ith_object in range(random_object_number()):
# Read template
tp_img, tp_mask, label, label_idx = get_random_template()
# put template onto the background image
new_bg, new_mask = aug.put_object_onto_background(
tp_img, tp_mask, bg_img)
# Store vars
bg_img = new_bg
masks.append(new_mask)
labels.append([label_idx])
# Last augment (add noise to the new background image)
bg_img = aug.augment_by_noises(bg_img)
# Get/Save/Plot bounding boxt of background image
bg_img_with_bbox = bg_img.copy()
for i, mask in enumerate(masks):
# Get bbox
x, y, w, h = pi.getBbox(mask, norm=True)
# Store and draw
labels[i].extend([x, y, w, h])
draw_bbox(bg_img_with_bbox, bbox=(x, y, w, h))
# Display the new background image
if 0:
# show((tp_img, tp_mask), figsize=(10, 5))
show((new_mask, bg_img, bg_img_with_bbox),
figsize=(15, 6),
layout=(1, 3))
# Write the image and its labels
filename = "{:06d}".format(ith_background) + get_time()
f_image = args.f_yolo_images + filename + ".jpg"
f_labels = args.f_yolo_labels + filename + ".txt"
f_image_with_bbox = args.f_yolo_images_with_bbox + filename + ".jpg"
cv2.imwrite(f_image, bg_img) # new image
cf.write_listlist(f_labels, labels) # its labels file
cv2.imwrite(f_image_with_bbox,
bg_img_with_bbox) # new image with annotated bbox on it
continue
# End
print("\n" + "-" * 80)
print("Image augmentation completes.")
print("Generated {} images. See folder = {}".format(
ith_background + 1, args.f_data_dst))
return