import os
import random
from dataclasses import dataclass
from typing import List, Optional, Sequence, Tuple
import cv2
import numpy as np
from kitcar_ml.utils.data_generation import augmentation_utils, utils
from kitcar_ml.utils.data_generation.generation_config import GenerationConfiguration
from kitcar_ml.utils.data_generation.sample_generator import SampleGenerator
[docs]def calculate_overlay(bb_old, bb_new):
bb_area_old = abs(bb_old[2] - bb_old[0]) * abs(bb_old[3] - bb_old[1])
bb_area_new = abs(bb_new[2] - bb_new[0]) * abs(bb_new[3] - bb_new[1])
x1_IR = max(bb_new[0], bb_old[0])
y1_IR = max(bb_new[1], bb_old[1])
x2_IR = min(bb_new[2], bb_old[2])
y2_IR = min(bb_new[3], bb_old[3])
IR_Area = max((x2_IR - x1_IR), 0) * max((y2_IR - y1_IR), 0)
overlay = float(IR_Area) / min(float(bb_area_old), float(bb_area_new))
return overlay
[docs]def clean_bbox(
x_coords: Sequence[int],
y_coords: Sequence[int],
crop_box: Tuple[int, int, int, int],
min_visibility,
min_area,
max_overlay,
other_bboxes,
) -> Optional[Tuple[int, int, int, int]]:
"""Sort/Clip coordinates and filter invalid/overlayed bboxes."""
x1, *_, x2 = sorted(x_coords)
y1, *_, y2 = sorted(y_coords)
original_size = (y2 - y1) * (x2 - x1)
def shift_and_clip(*coords, offset, min_, max_):
return tuple(int(np.clip(ord + offset, min_, max_).item()) for ord in coords)
x1, x2 = shift_and_clip(x1, x2, offset=-crop_box[1], min_=0, max_=crop_box[3])
y1, y2 = shift_and_clip(y1, y2, offset=-crop_box[0], min_=0, max_=crop_box[2])
clipped_size = (y2 - y1) * (x2 - x1)
if clipped_size / original_size < min_visibility or clipped_size < min_area:
return None
for bb_old in other_bboxes:
overlay = calculate_overlay(bb_old, (x1, y1, x2, y2))
if overlay >= max_overlay:
return None
return x1, y1, x2, y2
[docs]@dataclass
class GeneratedObject:
img_coords: Sequence[int]
aug_idx: int
label: str
[docs]def sample_objects(
number_of_objects: int,
name_list: List[str],
sample_generator: SampleGenerator,
angle_x_max: float,
angle_z_max: float,
overlay_trshld: float,
crop_box: float,
bbox_min_visibility: float,
bbox_min_area: int,
) -> List[GeneratedObject]:
"""Try to create randomized objects that have valid bboxes."""
objects = []
existing_boxes = []
attempts = 0
max_attempts = 2 * number_of_objects
while len(objects) < number_of_objects and attempts < max_attempts:
attempts += 1
# Sample position of object
position = sample_generator.sample_world_vec()
# Sample type of object
rand_sign_number = random.randint(0, len(name_list) - 1)
aug_name, _ = os.path.splitext(os.path.basename(name_list[rand_sign_number]))
angle_x = np.clip(
random.normalvariate(0, angle_x_max / 2), a_min=-angle_x_max, a_max=angle_x_max
)
angle_z = np.clip(
random.normalvariate(0, angle_z_max / 2), a_min=-angle_z_max, a_max=angle_z_max
)
# Sign coordinates for position
sign_positions = [
p for p in sample_generator.sign_coords(aug_name, angle_x, angle_z)
]
# Transform into img coordinates
# Points in columns
img_coords = np.array(
[
sample_generator.vehicle_point_to_img(position + p)[:2]
for p in sign_positions
]
).T
# Check if bbox is valid and modify if outside of image
bb_new = clean_bbox(
x_coords=img_coords[0],
y_coords=img_coords[1],
crop_box=crop_box,
min_visibility=bbox_min_visibility,
min_area=bbox_min_area,
max_overlay=overlay_trshld,
other_bboxes=existing_boxes,
)
if bb_new:
img_position = sample_generator.vehicle_point_to_img(position)
img_position = int(img_position[0] - crop_box[1]), int(
img_position[1] - crop_box[0]
)
sample = (aug_name, rand_sign_number, bb_new, img_position)
objects.append(GeneratedObject(img_coords, rand_sign_number, sample))
existing_boxes.append(bb_new)
return objects
[docs]def add_objects_to_image(
max_number_of_objects: int,
name_list: List[str],
aug_images: List[np.ndarray],
input_image: np.ndarray,
sample_generator: SampleGenerator,
supersampling: float,
angle_x_max: float,
angle_z_max: float,
noise_sigma: float,
noise_mean: float,
gauss_blur_sign: float,
overlay_trshld: float,
crop_box: float,
bbox_min_visibility: float,
bbox_min_area: int,
hor_motion_blur_size: int,
) -> Tuple[np.ndarray, List[str]]:
"""Add multiple objects to image.
Args:
max_number_of_objects: # of inserted objects.
name_list: Names of available objects.
aug_images: Available generation images.
input_image: Image that should be altered.
sample_generator: Generator of sample points for generations.
supersampling: Supersampling factor.
angle_x_max: Maximum rotation of generation around x axis.
angle_z_max: Maximum rotation of generation in z axis (after x rotation).
noise_sigma: Stddev of noise applied to generation image.
noise_mean: Mean of noise applied to generation image.
gauss_blur_sign: Filter kernel size of gaussian blur.
overlay_trshld: Threshold of max. overlay between two generations.
crop_box: ROI.
min_pxl_size: Minimal size of generation in pixels.
bbox_min_visibility: Minimal portion of generated bbox that is visible.
bbox_min_area: Minimal area of generated bbox.
hor_motion_blur_size: Kernel size of motion blur applied in hor direction.
Return:
Augmented image and list of labels.
"""
number_of_generations = random.randint(0, max_number_of_objects)
if number_of_generations == 0:
return input_image, []
# Create objects
generated_objects = sample_objects(
number_of_generations,
name_list,
sample_generator,
angle_x_max,
angle_z_max,
overlay_trshld,
crop_box,
bbox_min_visibility,
bbox_min_area,
)
# Supersample background img to improve resulting quality
input_image = augmentation_utils.supersample(input_image, supersampling)
bckg_img_stddev = np.std(input_image)
# Add objects to image
for obj in generated_objects:
input_image = augmentation_utils.add_object_to_background(
background_img=input_image,
aug_img=aug_images[obj.aug_idx],
img_coords=supersampling * obj.img_coords,
noise_sigma=noise_sigma,
noise_mean=noise_mean,
gauss_blur_sign=gauss_blur_sign,
bckg_img_stddev=bckg_img_stddev,
hor_motion_blur_size=hor_motion_blur_size,
)
input_image = augmentation_utils.revert_supersampling(input_image, supersampling)
labels = [obj.label for obj in generated_objects]
return input_image, labels
[docs]def create_artificial_objects_in_image(
image: List[np.ndarray],
config: GenerationConfiguration,
name_list,
aug_images,
sample_generator,
) -> Tuple[np.ndarray, List[str]]:
"""Add artificial objects to the given image.
The objects are sampled from aug_images and added as labels to the dataset.
Args:
image: Image.
config_data: Configurations; e.g. max number of objects in one image.
name_list: Class names of objects.
aug_images: Images of the objects (e.g. image of a stop sign)
sample_generator: Generator of positions in vehicle coordinates.
dataset: Dataset
naming_convention: Callable that returns a unique name for each index.
Return:
Altered image and labels.
"""
image, sample_data = add_objects_to_image(
config.max_number_of_objects,
name_list,
aug_images,
image,
sample_generator,
config.supersampling,
config.angle_x_max,
config.angle_z_max,
config.noise_sigma,
config.noise_mean,
config.gauss_blur_sign,
config.overlay_treshold,
config.crop_box,
config.bbox_min_visibility,
config.bbox_min_area,
config.hor_motion_blur_size,
)
# Add sample independent noise
image = augmentation_utils.add_random_intensity_and_contrast(
image,
config.sigma_all,
config.mean_all,
config.gauss_blur_all,
)
image = image.astype(np.uint8)
# Crop image to region of interest
crop_box = config.crop_box
image = image[crop_box[0] : crop_box[2], crop_box[1] : crop_box[3]]
image = image.reshape(image.shape[0], image.shape[1], 1) # 3 dim 1 channel
return image, sample_data
[docs]def load_and_create_artificial_images(
img_names: List[str],
dataset,
config,
name_list,
aug_images,
sample_generator,
):
"""Load provided images and generate synthetic images out of them.
Args:
img_names: Names(paths) to all images that should be used.
dataset: The dataset.
config: Configuration of the image generation.
name_list: Names of available generations.
aug_images: Available generation images.
sample_generator: Generator of sample points for generations.
"""
for img_path in img_names:
img = cv2.imread(img_path, -1).astype(np.float32)
generated_img, labels = create_artificial_objects_in_image(
img,
config,
name_list,
aug_images,
sample_generator,
)
# Add labels to dataset and save image
img_name = os.path.basename(img_path)
utils.write_annotations_to_dataset(dataset, img_name, labels)
cv2.imwrite(dataset._base_path + "/" + img_name, generated_img)