ai-training/preprocessing.py

import os.path
import time
from pathlib import Path
import albumentations as A
import cv2
from matplotlib import pyplot as plt
from AnnotationClass import AnnotationClass

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')
annotation_classes = AnnotationClass.read_json()


class ImageLabel:
    def __init__(self, image_path, image, labels_path, labels):
        self.image_path = image_path
        self.image = image
        self.labels_path = labels_path
        self.labels = labels

    def visualize(self):
        img = cv2.cvtColor(self.image.copy(), cv2.COLOR_BGR2RGB)
        height, width, channels = img.shape
        for label in self.labels:
            class_num = int(label[-1])
            x_c = float(label[0])
            y_c = float(label[1])
            w = float(label[2])
            h = float(label[3])
            x_min = x_c - w / 2
            y_min = y_c - h / 2
            x_max = x_min + w
            y_max = y_min + h
            color = annotation_classes[class_num].color_tuple

            cv2.rectangle(img, (int(x_min * width), int(y_min * height)), (int(x_max * width), int(y_max * height)),
                          color=color, thickness=3)
        plt.figure(figsize=(12, 12))
        plt.axis('off')
        plt.imshow(img)
        plt.show()


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.SafeRotate(limit=90, always_apply=True)],
                  bbox_params=A.BboxParams(format='yolo')),
        A.Compose([A.SafeRotate(limit=90, always_apply=True),
                   A.RandomBrightnessContrast(always_apply=True)],
                  bbox_params=A.BboxParams(format='yolo')),
        A.Compose([A.ShiftScaleRotate(scale_limit=0.2, always_apply=True),
                   A.VerticalFlip(always_apply=True),],
                  bbox_params=A.BboxParams(format='yolo')),
        A.Compose([A.ShiftScaleRotate(scale_limit=0.2, always_apply=True)],
                  bbox_params=A.BboxParams(format='yolo')),
        A.Compose([A.SafeRotate(limit=90, always_apply=True),
                   A.RandomBrightnessContrast(always_apply=True)],
                  bbox_params=A.BboxParams(format='yolo'))
    ]

    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}'
        img = 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')
        )
        results.append(img)
    return results


def write_result(img_ann: ImageLabel, show_image=False):
    os.makedirs(os.path.dirname(img_ann.image_path), exist_ok=True)
    os.makedirs(os.path.dirname(img_ann.labels_path), exist_ok=True)

    if show_image:
        img_ann.visualize()

    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_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.labels_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)
            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()