Made some changes

This commit is contained in:
Nazar Sturanec
2024-05-23 17:27:36 +03:00
parent c1d30d024d
commit 8529bd0c6f
2 changed files with 70 additions and 2 deletions
+1 -2
View File
@@ -42,7 +42,6 @@ def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]:
transform = A.Compose([ transform = A.Compose([
A.HorizontalFlip(p=0.5), A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.2), A.RandomBrightnessContrast(p=0.2),
A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, p=0.5),
], bbox_params=A.BboxParams(format='yolo', label_fields=['category_ids'])) ], bbox_params=A.BboxParams(format='yolo', label_fields=['category_ids']))
bboxes = bboxes bboxes = bboxes
@@ -56,7 +55,7 @@ def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]:
return transformed_image, transformed_bboxes, transformed_category_ids return transformed_image, transformed_bboxes, transformed_category_ids
def write_results(img_ann: ImageAnnotation): def write_results(img_ann: ImageAnnotation):
for i in range(10): for i in range(100):
transformed_image, transformed_bboxes, transformed_category_ids, = image_processing(img_ann) transformed_image, transformed_bboxes, transformed_category_ids, = image_processing(img_ann)
cv2.imwrite(os.path.join(current_dataset_dir, images_dir, str(i)+ImageAnnotation(img_ann).image_path + '.jpg'), transformed_image) cv2.imwrite(os.path.join(current_dataset_dir, images_dir, str(i)+ImageAnnotation(img_ann).image_path + '.jpg'), transformed_image)
with open(os.path.join(current_dataset_dir, labels_dir, str(i)+ImageAnnotation(img_ann).image_path + '.txt'), 'w') as f: with open(os.path.join(current_dataset_dir, labels_dir, str(i)+ImageAnnotation(img_ann).image_path + '.txt'), 'w') as f:
+69
View File
@@ -0,0 +1,69 @@
import math
import os
import datetime
import cv2
file_start_save = 'Zombobase-'+str(datetime.date.today())
file_txt = []
file_jpg = []
labels_dir = 'labels'
images_dir = 'images'
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:
for iii in range(len(annotations)):
f.write(annotations[iii])
annotations = []
def piercing_photo_file():
for foldername, subfolders, filenames in os.walk(current_dataset_dir):
for subfolder in subfolders:
cu = os.path.join('datasets', 'zombobase-current',subfolder)
for foldername, subfolders, filenames in os.walk(cu):
for filename in filenames:
f = filename.split('.')
if f[-1] == 'txt':
file_txt.append(filename)
elif f[-1] == 'jpg':
file_jpg.append(filename)
def file_validation(annotations, file_txt, file_jpg):
os.makedirs(os.path.join(file_start_save, current_dataset_dir,'validation'))
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:
lines = file.readlines()
for line in lines:
annotations.append(line)
save(annotations, image, i, 'validation')
def sort_file():
os.makedirs(os.path.join(file_start_save, current_dataset_dir))
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))
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)
file_validation(annotations, file_txt, file_jpg)
def main():
piercing_photo_file()
sort_file()
main()