def merge_single_classes(input_dirs, merge_dir):
i = 0
for input_dir in input_dirs:
for fn in file_helper.enumerate_files(input_dir):
try:
i += 1
print("{} - {}".format(i, fn), end=end_txt)
dir_name, name, extension = file_helper.get_file_name_extension(
fn)
if extension != ".txt":
fn_input_txt = file_helper.path_join(
input_dir, name + ".txt")
if not os.path.isfile(fn_input_txt):
raise Exception("no label file!")
fn_output_txt = file_helper.path_join(
merge_dir, name + ".txt")
fn_input_img = fn
fn_output_img = file_helper.path_join(
merge_dir, name + ".jpg")
if not os.path.isfile(fn_output_img):
mat = cv2.imread(fn_input_img)
cv2.imwrite(fn_output_img, mat)
if not os.path.isfile(fn_output_txt):
file_helper.copy_file(fn_input_txt, fn_output_txt)
else:
file_helper.append_lines(
fn_output_txt,
file_helper.read_lines(fn_input_txt))
print('Merged: {}'.format(fn_output_txt), end=end_txt)
except Exception as e:
print('Error - merge_single_classes - file: {} msg: {}'.format(
fn, str(e)))
def save_images_with_cv(source_dir_name,
target_dir_name,
max_dim=None,
recursive=False):
if not os.path.isdir(target_dir_name):
file_helper.create_dir(target_dir_name)
i = 0
for fn in file_helper.enumerate_files(source_dir_name,
recursive=recursive):
try:
i += 1
print("{} - {}".format(i, fn), end=end_txt)
dir_name, name, extension = file_helper.get_file_name_extension(fn)
if string_helper.equals_case_insensitive(extension, ".txt"):
new_fn = file_helper.path_join(target_dir_name,
name + extension)
file_helper.copy_file(fn, new_fn)
else:
mat = cv2.imread(fn)
if max_dim is not None:
mat = image_helper.resize_if_larger(mat, max_dim)
new_fn = file_helper.path_join(target_dir_name, name + ".jpg")
cv2.imwrite(new_fn, mat)
except Exception as e:
print("error - save_images_with_cv: {}".format(fn))
print("save_images_with_cv finished")
def test_copy_file(self):
src_folder = os.path.join('tests','unit_tests','resources')
tgt_folder = os.path.join(src_folder,'tmp')
file_helper.copy_file(
os.path.join(src_folder,'2.png'),
os.path.join(tgt_folder,'2.png'))
assert(len(os.listdir(tgt_folder)) == 1)
def coco_separate_classes(input_dir, output_dir):
classes_fn = file_helper.path_join(input_dir, "_classes.txt")
classes = []
for line in file_helper.read_lines(classes_fn):
classes.append(line)
i = 0
ann_fn = file_helper.path_join(input_dir, "_annotations.txt")
img_fn = None
for line in file_helper.read_lines(ann_fn):
try:
items = line.split(" ")
img_fn = file_helper.path_join(input_dir, items[0])
dir_name, name, extension = file_helper.get_file_name_extension(
img_fn)
i += 1
print("{} - {}".format(i, img_fn), end=end_txt)
mat = cv2.imread(img_fn)
h, w = image_helper.image_h_w(mat)
for j in range(1, len(items)):
item = items[j]
values_str = item.split(",")
values = []
for v in values_str:
values.append(int(v))
class_name = classes[values[4]]
write_dir = file_helper.path_join(output_dir, class_name)
if not os.path.isdir(write_dir):
file_helper.create_dir(write_dir)
write_img_fn = file_helper.path_join(write_dir,
name + extension)
if not os.path.isfile(write_img_fn):
file_helper.copy_file(img_fn, write_img_fn)
x1, y1, x2, y2 = list(values[:4])
w_ = (x2 - x1)
h_ = (y2 - y1)
cx = (x1 + w_ * 0.5) / w
cy = (y1 + h_ * 0.5) / h
line = "0 {} {} {} {}".format(cx, cy, w_ / w, h_ / h)
write_lbl_fn = file_helper.path_join(write_dir, name + ".txt")
if os.path.isfile(write_lbl_fn):
file_helper.append_line(write_lbl_fn, line)
else:
file_helper.write_lines(write_lbl_fn, [line])
except Exception as e:
print('Error - file:{} msg:{}'.format(img_fn, str(e)))
def test_collect_folder_as(self):
file_helper.copy_file(
os.path.join(self.resources_folder,'1.jpg'),
os.path.join(self.tmp_folder,'1.jpg')
)
file_helper.copy_file(
os.path.join(self.resources_folder,'2.png'),
os.path.join(self.tmp_folder,'2.png'))
self.image_collector.collect_files_from(
folder=self.tmp_folder,
destination=self.test_destination)
assert(file_helper.does_file_exist(self.test_destination))
def _oidv6_to_yolo(class_id, class_name, images_dir, labels_dir, yolo_dir):
for label_fn in file_helper.enumerate_files(labels_dir, recursive=False):
try:
print(label_fn, end="\r")
name = os.path.basename(label_fn)
image_fn = file_helper.path_join(images_dir,
name.replace(".txt", ".jpg"))
if os.path.isfile(image_fn):
write_image_fn = file_helper.path_join(
yolo_dir, name.replace(".txt", ".jpg"))
write_fn_txt = file_helper.path_join(yolo_dir, name)
if os.path.isfile(write_image_fn) and os.path.isfile(
write_fn_txt):
continue
mat = cv2.imread(image_fn)
h, w = image_helper.image_h_w(mat)
lines = []
for line in file_helper.read_lines(label_fn):
line = line.replace("\t", " ").replace(" ", " ").replace(
" ", " ").replace(" ", " ").replace(" ", " ")
arr0 = line.split(" ")
arr = [class_id]
x1 = float(arr0[1])
y1 = float(arr0[2])
x2 = float(arr0[3])
y2 = float(arr0[4])
arr.append((x1 + x2) * 0.5 / w)
arr.append((y1 + y2) * 0.5 / h)
arr.append((x2 - x1) / w)
arr.append((y2 - y1) / h)
line = ' '.join(str(e) for e in arr)
lines.append(line)
if len(lines) > 0:
# write_image_fn = file_helper.path_join(yolo_dir, name.replace(".txt", ".jpg"))
# cv2.imwrite(write_image_fn, mat)
file_helper.copy_file(image_fn, write_image_fn)
with contextlib.suppress(FileNotFoundError):
os.remove(write_fn_txt)
file_helper.write_lines(write_fn_txt, lines)
else:
print("nok: " + label_fn)
else:
print("no image: " + image_fn)
except Exception as e:
print('Error - file:{} msg:{}'.format(label_fn, str(e)))
print("finished: " + class_name)
def oidv6_to_yolo2(input_images_dir, input_labels_dir,
output_labels_and_images_dir, class_id):
i = 0
if not os.path.isdir(output_labels_and_images_dir):
file_helper.create_dir(output_labels_and_images_dir)
for label_fn in file_helper.enumerate_files(input_labels_dir,
recursive=False,
wildcard_pattern="*.txt"):
try:
i += 1
print("{} - {}".format(i, label_fn), end=end_txt)
dir_name, name, extension = file_helper.get_file_name_extension(
label_fn)
image_fn = file_helper.path_join(input_images_dir, name + ".jpg")
if os.path.isfile(image_fn):
mat = cv2.imread(image_fn)
h, w = image_helper.image_h_w(mat)
lines = []
for line in file_helper.read_lines(label_fn):
line = line.replace("\t", " ").replace(" ", " ").replace(
" ", " ").replace(" ", " ").replace(" ", " ")
arr0 = line.split(" ")
arr = [class_id]
x1 = float(arr0[1])
y1 = float(arr0[2])
x2 = float(arr0[3])
y2 = float(arr0[4])
arr.append((x1 + x2) * 0.5 / w)
arr.append((y1 + y2) * 0.5 / h)
arr.append((x2 - x1) / w)
arr.append((y2 - y1) / h)
line = ' '.join(str(e) for e in arr)
lines.append(line)
out_img_fn = file_helper.path_join(
output_labels_and_images_dir, name + ".jpg")
out_lbl_fn = file_helper.path_join(
output_labels_and_images_dir, name + ".txt")
if os.path.isfile(out_img_fn):
os.remove(out_img_fn)
if os.path.isfile(out_lbl_fn):
os.remove(out_lbl_fn)
file_helper.write_lines(out_lbl_fn, lines)
file_helper.copy_file(image_fn, out_img_fn)
else:
print("no image: " + image_fn)
except Exception as e:
print('Error - file:{} msg:{}'.format(label_fn, str(e)))
def mirror_images(source_dir_name, target_dir_name):
if not os.path.isdir(target_dir_name):
file_helper.create_dir(target_dir_name)
i = 0
for fn in file_helper.enumerate_files(source_dir_name):
try:
i += 1
print("{} - {}".format(i, fn), end=end_txt)
dir_name, name, extension = file_helper.get_file_name_extension(fn)
if string_helper.equals_case_insensitive(extension, ".txt"):
if name == "classes" and extension == ".txt":
new_fn = file_helper.path_join(target_dir_name,
name + extension)
if not os.path.isfile(new_fn):
file_helper.copy_file(fn, new_fn)
continue
new_fn = file_helper.path_join(target_dir_name,
name + "_mirror" + extension)
if os.path.isfile(new_fn):
raise Exception()
for line in file_helper.read_lines(fn):
lst = line.split(" ")
lst[1] = str(1 - float(lst[1]))
new_line = string_helper.join(lst, " ")
file_helper.append_line(new_fn, new_line)
else:
mat = cv2.imread(fn)
new_fn = file_helper.path_join(target_dir_name,
name + "_mirror" + ".jpg")
mat = cv2.flip(mat, 1)
cv2.imwrite(new_fn, mat)
except Exception as e:
print("error - mirror_images: {}".format(fn))
print("mirror_images finished")
def generate_train_txt(output_dir,
model_name,
class_names,
images_dir,
ratio_train=0.7,
ratio_val=0.2,
ratio_test=0.1):
yaml = file_helper.path_join(output_dir, model_name + ".yaml")
if os.path.isfile(yaml):
os.remove(yaml)
with open(yaml, "a") as file:
train_fn = file_helper.path_join(output_dir, model_name + "_train.txt")
file.write(f"train: {train_fn}\n")
val_fn = file_helper.path_join(output_dir, model_name + "_val.txt")
file.write(f"val: {val_fn}\n")
if ratio_test > 0:
test_fn = file_helper.path_join(output_dir,
model_name + "_test.txt")
file.write(f"test: {test_fn}\n")
file.write(f"nc: {str(len(class_names))}" + "\n")
file.write("names: " + str(class_names))
all_data = file_helper.path_join(output_dir, model_name + "_all_data.txt")
train = file_helper.path_join(output_dir, model_name + "_train.txt")
val = file_helper.path_join(output_dir, model_name + "_val.txt")
test = file_helper.path_join(output_dir, model_name, "_test.txt")
if os.path.isfile(all_data):
os.remove(all_data)
if os.path.isfile(train):
os.remove(train)
if os.path.isfile(val):
os.remove(val)
if os.path.isfile(test):
os.remove(test)
labels_dir = os.path.join(
os.path.abspath(os.path.join(images_dir, os.pardir)), "labels")
if not os.path.isdir(labels_dir):
file_helper.create_dir(labels_dir)
for fn in file_helper.enumerate_files(labels_dir, recursive=False):
if str.endswith(fn, ".txt"):
os.remove(fn)
with open(all_data, "a") as file:
for file_full_name in file_helper.enumerate_files(images_dir):
dir_name, name, extension = file_helper.get_file_name_extension(
file_full_name)
if extension != ".txt":
file.write(file_full_name + "\n")
else:
new_file_name = file_helper.path_join(labels_dir,
name + extension)
file_helper.copy_file(file_full_name, new_file_name)
with open(all_data) as f:
content = f.readlines()
content = [x.strip() for x in content]
from random import shuffle
shuffle(content)
train_count = int(len(content) * ratio_train)
val_count = train_count + int(len(content) * ratio_val)
# test %10
if ratio_test > 0:
i = 0
for line in content:
i += 1
if i < train_count:
fn = train
elif i < val_count:
fn = val
else:
fn = test
with open(fn, "a") as file:
file.write(line + "\n")
# print(line)
else:
i = 0
for line in content:
i += 1
if i < train_count:
fn = train
else:
fn = val
with open(fn, "a") as file:
file.write(line + "\n")
# print(line)
print("generate_train_txt finished: " + output_dir)
def mirror_images_of_classes(input_images_dir,
input_labels_dir,
output_dir,
class_ids,
copy_other_classes=True):
out_images_dir = file_helper.path_join(output_dir, "images")
out_labels_dir = file_helper.path_join(output_dir, "labels")
for dir_name in [out_images_dir, out_labels_dir]:
if not os.path.isdir(dir_name):
file_helper.create_dir(dir_name)
i = 0
for fn_label in file_helper.enumerate_files(input_labels_dir,
recursive=False,
wildcard_pattern="*.txt"):
try:
i += 1
print("{} - {}".format(i, fn_label), end=end_txt)
dir_name, name, extension = file_helper.get_file_name_extension(
fn_label)
mirror = False
for line in file_helper.read_lines(fn_label):
class_id = int(line.split(" ")[0])
if class_id in class_ids:
mirror = True
break
if mirror:
fn_image = file_helper.path_join(input_images_dir,
name + ".jpg")
if not os.path.isfile(fn_image):
print("No image: {}".format(fn_image))
else:
new_image_fn = file_helper.path_join(
out_images_dir, name + "_mirror" + ".jpg")
cv2.imwrite(new_image_fn,
image_helper.mirror(cv2.imread(fn_image)))
new_label_fn1 = file_helper.path_join(
out_labels_dir, name + "_mirror" + ".txt")
new_label_fn2 = file_helper.path_join(
out_images_dir, name + "_mirror" + ".txt")
new_label_file_names = [new_label_fn1, new_label_fn2]
for new_label_fn in new_label_file_names:
if os.path.isfile(new_label_fn):
raise Exception()
for line in file_helper.read_lines(fn_label):
lst = line.split(" ")
lst[1] = str(1 - float(lst[1]))
new_line = string_helper.join(lst, " ")
file_helper.append_line(new_label_fn, new_line)
if mirror or copy_other_classes:
fn_image = file_helper.path_join(input_images_dir,
name + ".jpg")
if not os.path.isfile(fn_image):
print("No image: {}".format(fn_image))
else:
new_image_fn = file_helper.path_join(
out_images_dir, name + ".jpg")
cv2.imwrite(new_image_fn, cv2.imread(fn_image))
new_label_fn1 = file_helper.path_join(
out_labels_dir, name + ".txt")
new_label_fn2 = file_helper.path_join(
out_images_dir, name + ".txt")
new_label_file_names = [new_label_fn1, new_label_fn2]
for new_label_fn in new_label_file_names:
if os.path.isfile(new_label_fn):
raise Exception()
file_helper.copy_file(fn_label, new_label_fn)
except Exception as e:
print("error - mirror_images_of_classes: {}".format(fn_label))
print("mirror_images_of_classes {} finished".format(class_ids))
