def setup_yolo_files(args):
# Train/valid split and write to train.txt and valid.txt
fnames = cf.get_filenames(args.f_yolo_images)
r = 0.2
fname_trains, fname_valids = cf.train_valid_split(fnames, ratio_valid=r)
cf.write_list(args.f_yolo_train, fname_trains)
cf.write_list(args.f_yolo_valid, fname_valids)
print("Split all all images into [valid/train={:.2f}/{:.2f}]:".format(
r, 1 - r))
print("\t{}".format(args.f_yolo_train))
print("\t{}".format(args.f_yolo_valid))
# Write yolo.cfg, which is yolo's network configurations file
n_labels = len(args.labels)
yolo_config = YoloConfig(NUM_CLASSES=n_labels)
yolo_config.write_to_file(args.f_yolo_config)
# Write yolo.data, which specifies the data path to train yolo
s = [""] * 4
s[0] = f"classes= {n_labels}"
s[1] = f"train= {args.f_yolo_train}"
s[2] = f"valid= {args.f_yolo_valid}"
s[3] = f"names= {args.f_yolo_classes}"
cf.write_list(filename=args.f_yolo_data, arr=s)
print("Write data config to: ", args.f_yolo_data)
def __init__(self, args,
preload_all_images=True,
crop_mask=True, # Crop out only a sub rectangular white region inside the mask
):
# Settings
img_folder = args.f_template_img
mask_folder = args.f_template_mask
self.preload_all_images = preload_all_images
self.crop_mask = crop_mask
# Read image filenames
fnames_img = cf.get_filenames(
img_folder, file_types=('*.jpg', '*.png'))
fnames_mask = cf.get_filenames(
mask_folder, file_types=('*.jpg', '*.png'))
# Check if files are matched
def get_basename(fname):
path, name, ext = cf.split_name(fname)
return name
fnames_img_basename = [get_basename(fname) for fname in fnames_img]
fnames_mask_basename = [get_basename(fname) for fname in fnames_mask]
if fnames_img_basename != fnames_mask_basename:
print(f"fnames_img_basename = {fnames_img_basename}")
print(f"fnames_mask_basename = {fnames_mask_basename}")
raise ValueError(
"Template images and masks have mis-matched names.")
# Save vars
self.num_templates = len(fnames_img)
self.fnames_img = fnames_img
self.fnames_mask = fnames_mask
# Load images
self.imgs = []
self.masks = []
if preload_all_images:
for i in range(self.num_templates):
img, mask = self.load_ith_image(i)
self.imgs.append(img)
self.masks.append(mask)
def setup_train_test_txt(args):
# Train/valid split and write to train.txt and valid.txt
fnames = cf.get_filenames(args.f_yolo_images)
rt = args.yolo["ratio_train"]
fname_trains, fname_valids = cf.train_valid_split(fnames, ratio_train=rt)
cf.write_list(args.f_yolo_train, fname_trains)
cf.write_list(args.f_yolo_valid, fname_valids)
print(
"Split all {} images into [train/valid={:.2f}/{:.2f}]:".format(len(fnames), rt, 1-rt))
print("\t{}".format(args.f_yolo_train))
print("\t{}".format(args.f_yolo_valid))
# Copy images in valid.txt to the "valid_images/" folder
if os.path.isdir(args.f_yolo_valid_images):
shutil.rmtree(args.f_yolo_valid_images)
cf.copy_files(fname_valids, args.f_yolo_valid_images)
print("\tSave valid images into {}".format(args.f_yolo_valid_images))
def __init__(
self,
args,
preload_all_images=True,
):
# Settings
img_folder = args.f_background
self.preload_all_images = preload_all_images
# Read image names
self.fnames_img = cf.get_filenames(img_folder,
file_types=('*.jpg', '*.png'))
N = len(self.fnames_img)
# Load images
if preload_all_images:
self.imgs = [self.load_ith_image(i) for i in range(N)]
else:
self.imgs = []
def __init__(self, args,
resize_to_rows=None,
preload_all_images=True,
):
# Settings
img_folder = args.f_background
self.preload_all_images = preload_all_images
self.resize_to_rows = resize_to_rows
# Read image names
self.fnames_img = cf.get_filenames(img_folder, file_types=('*.jpg', '*.png'))
N = len(self.fnames_img)
assert N>0, "Background images are empty. You must put them in 'data/xxxx/background/' folder."
# Load images
if preload_all_images:
self.imgs = [self.load_ith_image(i)
for i in range(N)]
else:
self.imgs = []
def write_object_labels_to_file(args):
'''
From "data/$data_name/template_img/", read in filenames to know the labels.
A filename should be like in this form "labelname_index.jpy", e.g.: "bottle_1.jpg"
'''
# Get object labels
if 0: # from "data/$data_name/template_img/"
fnames = cf.get_filenames(args.f_template_img,
file_types=('*.jpg', '*.png'))
labels = {get_label(fn)
for fn in fnames} # /folder/bottle_1.jpg --> bottle
labels = sorted(list(labels))
else: # from configurations
labels = args.data_labels
print(f"Object labels: {labels}")
# Write labels to txt for yolo training
cf.write_list(filename=args.f_yolo_classes, arr=labels)
print(f"Write labels to {args.f_yolo_classes}")
return labels
def __init__(
self,
args,
preload_all_images=True,
crop_mask=True, # Crop out only a sub rectangular white region inside the mask
):
# Settings
img_folder = args.f_template_img
mask_folder = args.f_template_mask
self.preload_all_images = preload_all_images
self.crop_mask = crop_mask
# Read image filenames
fnames_img = cf.get_filenames(img_folder,
file_types=('*.jpg', '*.png'))
fnames_mask = []
for fname in fnames_img:
path, name, ext = cf.split_name(fname)
mask_path = mask_folder + name + ext
if not os.path.isfile(mask_path):
print("mask_path = ", mask_path)
raise ValueError(
"The corresponding mask of an template image doesn't exist"
)
fnames_mask.append(mask_path)
self.num_templates = len(fnames_img)
self.fnames_img = fnames_img
self.fnames_mask = fnames_mask
# Load images
self.imgs = []
self.masks = []
if preload_all_images:
for i in range(self.num_templates):
img, mask = self.load_ith_image(i)
self.imgs.append(img)
self.masks.append(mask)
