Add check.py

This commit is contained in:
Nazar Sturanec
2024-05-25 03:13:19 +03:00
parent 8529bd0c6f
commit 6adc4c37e9
21 changed files with 95 additions and 31 deletions
+30 -12
View File
@@ -12,9 +12,10 @@ current_dataset_dir = os.path.join('datasets', 'zombobase-current')
def save(annotations, image, i, f_n):
cv2.imwrite(os.path.join(file_start_save, current_dataset_dir, f_n, file_jpg[i]), image)
with open(os.path.join(file_start_save, current_dataset_dir, f_n, file_txt[i]), 'w') as f:
cv2.imwrite(os.path.join(file_start_save, current_dataset_dir, f_n, images_dir, file_jpg[i]), image)
with open(os.path.join(file_start_save, current_dataset_dir, f_n, labels_dir, file_txt[i]), 'w') as f:
for iii in range(len(annotations)):
f.write(annotations[iii])
annotations = []
@@ -33,7 +34,12 @@ def piercing_photo_file():
def file_validation(annotations, file_txt, file_jpg):
os.makedirs(os.path.join(file_start_save, current_dataset_dir,'validation'))
try:
os.makedirs(os.path.join(file_start_save, current_dataset_dir,'validation'))
os.makedirs(os.path.join(file_start_save, current_dataset_dir, 'validation', images_dir))
os.makedirs(os.path.join(file_start_save, current_dataset_dir, 'validation', labels_dir))
except FileExistsError:
pass
for i in range(len(file_txt)):
image = cv2.imread(os.path.join(current_dataset_dir, images_dir, file_jpg[i]))
with open(os.path.join(current_dataset_dir, labels_dir, file_txt[i]), 'r') as file:
@@ -41,29 +47,41 @@ def file_validation(annotations, file_txt, file_jpg):
for line in lines:
annotations.append(line)
save(annotations, image, i, 'validation')
annotations = []
def sort_file():
os.makedirs(os.path.join(file_start_save, current_dataset_dir))
try:
os.makedirs(os.path.join(file_start_save, current_dataset_dir))
except FileExistsError:
pass
annotations = []
folder_name = ['test', 'train']
percent_file = [0.20, 0.10]
for f_n, p_f in zip(folder_name, percent_file):
os.makedirs(os.path.join(file_start_save, current_dataset_dir, f_n))
try:
os.makedirs(os.path.join(file_start_save, current_dataset_dir, f_n))
os.makedirs(os.path.join(file_start_save, current_dataset_dir, f_n, images_dir))
os.makedirs(os.path.join(file_start_save, current_dataset_dir, f_n, labels_dir))
except FileExistsError:
pass
for i in range(math.ceil(len(file_txt) * p_f)):
image = cv2.imread(os.path.join(current_dataset_dir, images_dir, file_jpg[i]))
with open(os.path.join(current_dataset_dir, labels_dir, file_txt[i]), 'r') as file:
lines = file.readlines()
for line in lines:
annotations.append(line)
save(annotations, image, i, f_n)
file_txt.pop(i)
file_jpg.pop(i)
save(annotations, image, i, f_n)
annotations = []
print(annotations)
file_txt.pop(i)
file_jpg.pop(i)
file_validation(annotations, file_txt, file_jpg)
def main():
piercing_photo_file()
sort_file()
main()
sort_file()