ai-training/preprocessing.py

import os.path
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.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

    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_name, os.path.basename(image_path))
        self.image = cv2.imread(image_path)

        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 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


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():
    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()