diff --git a/preprocessing.py b/preprocessing.py index 25fa4ad..13538d7 100644 --- a/preprocessing.py +++ b/preprocessing.py @@ -1,15 +1,20 @@ import os.path -import albumentations as A +import time + import cv2 +import albumentations as alb +from os import listdir +from os.path import isfile, join from pathlib import Path labels_dir = 'labels' images_dir = 'images' current_dataset_dir = os.path.join('datasets', 'zombobase-current') + class ImageAnnotation: def read_annotations(self) -> [[]]: - with open(self.dataset_annotation_path, 'r') as f: + with open(self.annotation_path, 'r') as f: rows = f.readlines() arr = [] for row in rows: @@ -18,73 +23,50 @@ class ImageAnnotation: coordinates = [float(n) for n in str_coordinates] coordinates.append(class_num) arr.append(coordinates) - return arr def __init__(self, image_path): 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_path + '.jpg') - self.image = cv2.imread(self.dataset_image_path) + 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.annotation_path = os.path.join(labels_dir, self.image_path + '.txt') - self.dataset_annotation_path = os.path.join(current_dataset_dir, labels_dir, self.image_path + '.txt') + 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]: - category_ids = [] - bboxes = ImageAnnotation(img_ann).read_annotations() - for i in range(len(bboxes)): - category_ids.append(bboxes[i][4]) - bboxes[i].pop(4) - - 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'])) - - bboxes = bboxes - imag = ImageAnnotation(img_ann).image - - transformed = transform(image=imag, bboxes=bboxes, category_ids=category_ids) - transformed_image = transformed['image'] - transformed_bboxes = transformed['bboxes'] - transformed_category_ids = transformed['category_ids'] - - return transformed_image, transformed_bboxes, transformed_category_ids + # 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 write_results(img_ann: ImageAnnotation): - for i in range(10): - 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) - with open(os.path.join(current_dataset_dir, labels_dir, str(i)+ImageAnnotation(img_ann).image_path + '.txt'), '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") -# -# -def process_image(): - file_annotation = [] - file_annotation_finished =[] - for foldername, subfolders, filenames in os.walk(os.path.join(current_dataset_dir,images_dir)): - file_annotation.append(filenames) - print(file_annotation) + # write image cv2.imwrite(, image) dataset_image_path + # write img_ann.annotations into new file with name dataset_annotation_path - for i in range(len(file_annotation[0])): - ff = file_annotation[0][i].split('.') - ff.pop(-1) - ff = '.'.join(ff) - file_annotation_finished.append(ff) - return file_annotation_finished +def process_image(img_ann): + results = image_processing(img_ann) + for res_ann in results: + write_results(res_ann) + write_results(img_ann) + os.remove(img_ann.image_path) + os.remove(img_ann.annotation_path) + def main(): - process_image() - for i in process_image(): - write_results(i) + while True: + images = os.listdir(images_dir) + if len(images) == 0: + time.sleep(5) + continue + + for image in images: + image_path = os.path.join(images_dir, image) + process_image(ImageAnnotation(image_path)) if __name__ == '__main__': main() \ No newline at end of file