diff --git a/breeding_date_set.py b/breeding_date_set.py index d193b95..2f39933 100644 --- a/breeding_date_set.py +++ b/breeding_date_set.py @@ -1,78 +1,119 @@ +import math +import datetime import cv2 import albumentations as A import numpy as np import os -file_txt = [] -file_jpg = [] +def file_magnification(folder_path): + file_txt = [] + file_jpg = [] + for foldername, subfolders, filenames in os.walk(folder_path): + for filename in filenames: + f = filename.split('.') -folder_path = ('D:\\train') + if f[-1] == 'txt': + file_txt.append(filename) + elif f[-1] == 'jpg': + file_jpg.append(filename) + + for k in range(len(file_jpg)): + image = cv2.imread(f'{folder_path}\\{file_jpg[k]}') + annotations = [] + with open(f'{folder_path}\\{file_txt[k]}', 'r') as file: + lines = file.readlines() + for line in lines: + annotations.append(line) + + + main_fillet_yolo_conversion = [] + fillet_yolo_bboxes = [] + fillet_yolo_class = [] + print(annotations) + for ii in range(len(annotations)): + a = annotations[ii].split(' ') + for i in range(len(a)): + try: + main_fillet_yolo_conversion.append(int(a[i])) + except ValueError: + main_fillet_yolo_conversion.append(float(a[i])) + print(main_fillet_yolo_conversion) + fillet_yolo_class.append(main_fillet_yolo_conversion[0]) + del main_fillet_yolo_conversion[0] + fillet_yolo_bboxes.append(main_fillet_yolo_conversion) + main_fillet_yolo_conversion = [] + + for o in range(10): + bboxes = fillet_yolo_bboxes + category_ids = fillet_yolo_class + + transform = A.Compose([ + A.HorizontalFlip(p=0.5), + 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'])) + + transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids) + transformed_image = transformed['image'] + transformed_bboxes = transformed['bboxes'] + transformed_category_ids = transformed['category_ids'] + + cv2.imwrite(f'{folder_path}\\{o}_{file_jpg[k]}', transformed_image) + with open(f'{folder_path}\\{o}_{file_txt[k]}', 'w') as f: + for bbox, category_id in zip(transformed_bboxes, transformed_category_ids): + x_center, y_center, width, height = bbox + cla = category_id + f.write(f"{cla} {x_center} {y_center} {width} {height}\n") +file_txt_1 = [] +file_jpg_1 = [] +file_start = 'Zombobase-'+str(datetime.date.today()) +folder_path = ('train') for foldername, subfolders, filenames in os.walk(folder_path): for subfolder in subfolders: - folder_path = (f'D:\\train\\{subfolder}') + folder_path = (f'train\\{subfolder}') for foldername, subfolders, filenames in os.walk(folder_path): for filename in filenames: f = filename.split('.') if f[-1] == 'txt': - file_txt.append(filename) + file_txt_1.append(filename) elif f[-1] == 'jpg': - file_jpg.append(filename) + file_jpg_1.append(filename) -for k in range(len(file_jpg)): - image = cv2.imread(f'D:\\train\\images\\{file_jpg[k]}') - annotations = [] - with open(f'D:\\train\\labels\\{file_txt[k]}', 'r') as file: +annotations = [] +os.makedirs(file_start) +file = ['test', 'train'] +percent_fille = [0.20,0.10] +for fi, p_f in zip(file, percent_fille): + os.makedirs(f'{file_start}\\{fi}') + + for i in range(math.ceil(len(file_txt_1)* p_f)): + image = cv2.imread(f'D:\\train\\images\\{file_jpg_1[i]}') + with open(f'D:\\train\\labels\\{file_txt_1[i]}', 'r') as file: + lines = file.readlines() + for line in lines: + annotations.append(line) + cv2.imwrite(f'{file_start}\\{fi}\\{file_jpg_1[i]}', image) + with open(f'{file_start}\\{fi}\\{file_txt_1[i]}', 'w') as f: + for iii in range(len(annotations)): + f.write(annotations[iii]) + + annotations = [] + del file_txt_1[i] + del file_jpg_1[i] + +os.makedirs(f'{file_start}\\validation') +for a, j in zip(file_txt_1, file_jpg_1): + image = cv2.imread(f'D:\\train\\images\\{j}') + with open(f'D:\\train\\labels\\{a}', 'r') as file: lines = file.readlines() for line in lines: - print(line) annotations.append(line) - - - main_fillet_yolo_conversion = [] - fillet_yolo_bboxes = [] - fillet_yolo_class = [] - print(annotations) - for ii in range(len(annotations)): - a = annotations[ii].split(' ') - - for i in range(len(a)): - - try: - main_fillet_yolo_conversion.append(int(a[i])) - except ValueError: - main_fillet_yolo_conversion.append(float(a[i])) - print(main_fillet_yolo_conversion) - fillet_yolo_class.append(main_fillet_yolo_conversion[0]) - del main_fillet_yolo_conversion[0] - fillet_yolo_bboxes.append(main_fillet_yolo_conversion) - print(fillet_yolo_bboxes) - print(fillet_yolo_class) - main_fillet_yolo_conversion = [] - - for o in range(10): - if image is None: - raise ValueError("Image not found or the path is incorrect") - if not isinstance(image, np.ndarray): - raise TypeError("Image must be a numpy array") - - bboxes = fillet_yolo_bboxes - category_ids = fillet_yolo_class - - transform = A.Compose([ - A.HorizontalFlip(p=0.5), - 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'])) - - transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids) - transformed_image = transformed['image'] - transformed_bboxes = transformed['bboxes'] - transformed_category_ids = transformed['category_ids'] - - cv2.imwrite(f'D:\\python\\prodgect_2\\pythonProject1\\op\\1\\{o}_{file_jpg[k]}', transformed_image) - with open(f'D:\\python\\prodgect_2\\pythonProject1\\op\\1\\{o}_{file_txt[k]}', 'w') as f: - for bbox, category_id in zip(transformed_bboxes, transformed_category_ids): - x_center, y_center, width, height = bbox - cla = category_id - f.write(f"{cla} {x_center} {y_center} {width} {height}\n") \ No newline at end of file + cv2.imwrite(f'{file_start}\\Validation\\{j}', image) + with open(f'{file_start}\\Validation\\{a}', 'w') as f: + for iii in range(len(annotations)): + f.write(annotations[iii]) + annotations = [] +file = ['test', 'train','validation'] +for i in file: + file_magnification(f'{file_start}\\{i}') \ No newline at end of file