diff --git a/.gitignore b/.gitignore index 9f11b75..fb9edd7 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ .idea/ +*labels/ +*images/ \ No newline at end of file diff --git a/breeding_date_set.py b/breeding_date_set.py deleted file mode 100644 index 2f39933..0000000 --- a/breeding_date_set.py +++ /dev/null @@ -1,119 +0,0 @@ -import math -import datetime -import cv2 -import albumentations as A -import numpy as np -import os -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('.') - - 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'train\\{subfolder}') - for foldername, subfolders, filenames in os.walk(folder_path): - for filename in filenames: - - f = filename.split('.') - - if f[-1] == 'txt': - file_txt_1.append(filename) - elif f[-1] == 'jpg': - file_jpg_1.append(filename) - -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: - annotations.append(line) - 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 diff --git a/datasets/zombobase-current/images/Himars_vs_Tulpan_000091.jpg b/datasets/zombobase-current/images/Himars_vs_Tulpan_000091.jpg deleted file mode 100644 index afc1e60..0000000 Binary files a/datasets/zombobase-current/images/Himars_vs_Tulpan_000091.jpg and /dev/null differ diff --git a/datasets/zombobase-current/labels/Himars_vs_Tulpan_000091.txt b/datasets/zombobase-current/labels/Himars_vs_Tulpan_000091.txt deleted file mode 100644 index 50d774d..0000000 --- a/datasets/zombobase-current/labels/Himars_vs_Tulpan_000091.txt +++ /dev/null @@ -1,3 +0,0 @@ -3 0.41877 0.64332 0.06107 0.07926 -7 0.40442 0.78827 0.06779 0.05212 -7 0.19860 0.90662 0.05314 0.05429 \ No newline at end of file diff --git a/images/Himars_vs_Tulpan_000091.jpg b/images/Himars_vs_Tulpan_000091.jpg deleted file mode 100644 index afc1e60..0000000 Binary files a/images/Himars_vs_Tulpan_000091.jpg and /dev/null differ diff --git a/labels/Himars_vs_Tulpan_000091.txt b/labels/Himars_vs_Tulpan_000091.txt deleted file mode 100644 index 50d774d..0000000 --- a/labels/Himars_vs_Tulpan_000091.txt +++ /dev/null @@ -1,3 +0,0 @@ -3 0.41877 0.64332 0.06107 0.07926 -7 0.40442 0.78827 0.06779 0.05212 -7 0.19860 0.90662 0.05314 0.05429 \ No newline at end of file diff --git a/preprocessing.py b/preprocessing.py index 13538d7..7e796a4 100644 --- a/preprocessing.py +++ b/preprocessing.py @@ -1,60 +1,83 @@ import os.path import time - -import cv2 -import albumentations as alb -from os import listdir -from os.path import isfile, join +from array import * from pathlib import Path +import albumentations as A +import cv2 + labels_dir = 'labels' images_dir = 'images' current_dataset_dir = os.path.join('datasets', 'zombobase-current') +current_images_dir = os.path.join(current_dataset_dir, 'images') +current_labels_dir = os.path.join(current_dataset_dir, 'labels') -class ImageAnnotation: - def read_annotations(self) -> [[]]: - with open(self.annotation_path, 'r') as f: - rows = f.readlines() - arr = [] - for row in rows: - str_coordinates = row.split(' ') - class_num = str_coordinates.pop(0) - coordinates = [float(n) for n in str_coordinates] - coordinates.append(class_num) - arr.append(coordinates) - return arr +class ImageLabel: - def __init__(self, image_path): + def __init__(self, image_path, image, labels_path, labels): self.image_path = image_path - self.image_name = Path(image_path).stem - self.dataset_image_path = os.path.join(current_dataset_dir, images_dir, self.image_name, os.path.basename(image_path)) - self.image = cv2.imread(image_path) + self.image = image + self.labels_path = labels_path + self.labels = labels - self.annotation_path = os.path.join(labels_dir, self.image_name, '.txt') - self.dataset_annotation_path = os.path.join(current_dataset_dir, labels_dir, self.image_name, '.txt') - self.annotations = self.read_annotations() -def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]: - # return structure example: - # utilize transform albumentations here - return [ImageAnnotation(f'{img_ann.image_name}1', image1, bboxes1 ), - ImageAnnotation(f'{img_ann.image_name}2', image2, bboxes2), - ... - ] +def image_processing(img_ann: ImageLabel) -> [ImageLabel]: + transforms = [ + A.Compose([A.HorizontalFlip(always_apply=True)], bbox_params=A.BboxParams(format='yolo')), + A.Compose([A.RandomBrightnessContrast(always_apply=True)], bbox_params=A.BboxParams(format='yolo')), + A.Compose([A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, always_apply=True)], + bbox_params=A.BboxParams(format='yolo')) + ] -def write_results(img_ann: ImageAnnotation): - # write image cv2.imwrite(, image) dataset_image_path - # write img_ann.annotations into new file with name dataset_annotation_path + results = [] + for i, transform in enumerate(transforms): + res = transform(image=img_ann.image, bboxes=img_ann.labels) + path = Path(img_ann.image_path) + name = f'{path.stem}_{i+1}' + results.append(ImageLabel( + image=res['image'], + labels=res['bboxes'], + image_path=os.path.join(current_images_dir, f'{name}{path.suffix}'), + labels_path=os.path.join(current_labels_dir, f'{name}.txt') + )) + return results + + +def write_result(img_ann: ImageLabel): + cv2.imwrite(img_ann.image_path, img_ann.image) + with open(img_ann.labels_path, 'w') as f: + lines = [f'{ann[4]} {round(ann[0], 5)} {round(ann[1], 5)} {round(ann[2], 5)} {round(ann[3], 5)}\n' for ann in img_ann.labels] + f.writelines(lines) + f.close() + + +def read_labels(labels_path) -> [[]]: + with open(labels_path, 'r') as f: + rows = f.readlines() + arr = [] + for row in rows: + str_coordinates = row.split(' ') + class_num = str_coordinates.pop(0) + coordinates = [float(n) for n in str_coordinates] + coordinates.append(class_num) + arr.append(coordinates) + return arr def process_image(img_ann): results = image_processing(img_ann) for res_ann in results: - write_results(res_ann) - write_results(img_ann) + write_result(res_ann) + + write_result(ImageLabel( + image=img_ann.image, + labels=img_ann.labels, + image_path=os.path.join(current_images_dir, Path(img_ann.image_path).name), + labels_path=os.path.join(current_labels_dir, Path(img_ann.labels_path).name) + )) os.remove(img_ann.image_path) - os.remove(img_ann.annotation_path) + os.remove(img_ann.labels_path) def main(): @@ -66,7 +89,14 @@ def main(): for image in images: image_path = os.path.join(images_dir, image) - process_image(ImageAnnotation(image_path)) + labels_path = os.path.join(labels_dir, f'{Path(image_path).stem}.txt') + process_image(ImageLabel( + image_path=image_path, + image=cv2.imread(image_path), + labels_path=labels_path, + labels=read_labels(labels_path) + )) + if __name__ == '__main__': main() \ No newline at end of file