Source code for kitcar_ml.utils.data_generation.utils

import glob
import os
from typing import List, Tuple

import cv2
import numpy as np

from kitcar_ml.utils.data.labeled_dataset import LabeledDataset


def read_augmentation_data(
    path_to_augmentation_images: str,
) -> Tuple[List[np.ndarray], List[str]]:
    """Read augmentation images and file names."""
    print("load augmentation data")
    name_list = glob.glob(path_to_augmentation_images + "/*.png")
    name_list = sorted(name_list)

    aug_img_list = [cv2.imread(x, -1).astype(np.float32) for x in name_list]

    print("augmentation data loaded")
    return aug_img_list, name_list


def create_dataset(path_dataset_generated: str, class_names: List[str]) -> LabeledDataset:
    """Create dataset.

    Args:
        path_dataset_generated: Path to dataset directory.
        class_names: Classes added to dataset.
    """
    os.makedirs(path_dataset_generated, exist_ok=True)

    dataset = LabeledDataset(_base_path=path_dataset_generated)
    dataset.classes = {
        i: os.path.splitext(os.path.basename(name))[0] for i, name in enumerate(class_names)
    }
    dataset.attributes = [
        "img_name",
        "x1",
        "y1",
        "x2",
        "y2",
        "class_id",
        "x_ground",
        "y_ground",
    ]
    return dataset


def write_annotations_to_dataset(dataset: LabeledDataset, name, sample_data):
    print(sample_data)
    for label in sample_data:
        class_id = label[1]
        x1, y1, x2, y2 = label[2]
        x_ground, y_ground = label[3]
        dataset.append_label(name, [name, x1, y1, x2, y2, class_id, x_ground, y_ground])