from .random_annotation import RandomAnnotation
from .ground_truth import GroundTruth
from .image import Image
from .gt_sample_info import GTSampleInfo
from iaa_od.utils import sample_global_avg_bboxes_per_image, sample_global_avg_bbox_size
from dataclasses import dataclass, field, InitVar
from datetime import datetime
from random import gauss, randint, uniform
import copy
[docs]
@dataclass(slots=True, kw_only=True)
class RandomGroundTruth:
"""
A class to represent a randomly generated ground truth dataset, based on statistics sampled from existing ground truths.
Attributes:
name (str): The name of the randomly generated ground truth dataset.
images (dict[str, Image]): A dictionary mapping image filenames to Image objects.
annotations (dict[str, list[RandomAnnotation]]): A dictionary mapping image filenames to lists of RandomAnnotation objects.
categories_dict (dict[int, str]): A dictionary mapping category IDs to category names.
deviation (float): The standard deviation to use for the Gaussian distribution when generating random bounding box sizes. Default is 0.5.
"""
# Properties
name: str = field(init=False)
images: dict[str, Image] = field(init=False)
annotations: dict[str, list[RandomAnnotation]] = field(init=False)
categories_dict: dict[int, str] = field(init=False)
deviation: float = field(default=0.5)
# Init-only fields
gts: InitVar[list[GroundTruth]]
def __post_init__(self, gts: list[GroundTruth]):
# Generate a random name for the randomly-generated ground truth
self.name = self._generate_random_name()
# Get the images dictionary from the first GT passed to the class
self.images = copy.deepcopy(gts[0].images)
# Do the same with the categories dictionary
self.categories_dict = gts[0].categories_dict.copy()
# Generate the random annotations
self.annotations = self._generate_annotations(gts)
def _generate_random_name(self) -> str:
"""
Generate a random name for the randomly generated ground truth dataset, using the current timestamp to ensure uniqueness.
Returns:
str: A randomly generated name for the ground truth dataset.
"""
timestamp: str = datetime.now().strftime("%Y%m%d_%H%M%S")
return f"random_gt_{timestamp}"
def _generate_annotations(self, gts: list[GroundTruth]) -> dict[str, list[RandomAnnotation]]:
"""
Generate random annotations for each image in the ground truth dataset, based on statistics sampled from the provided ground truths.
Parameters:
gts (list[GroundTruth]): A list of GroundTruth objects to sample statistics from for generating the random annotations.
Returns:
dict[str, list[RandomAnnotation]]: A dictionary mapping image filenames to lists of randomly generated annotations for each image.
"""
# Sample ground truths to get statistics for random generation
sample_info: GTSampleInfo = GTSampleInfo(
avg_bbox_area=sample_global_avg_bbox_size(gts),
avg_bboxes_per_image=sample_global_avg_bboxes_per_image(gts)
)
# Initialize the annotations dictionary
annotations: dict[str, list[RandomAnnotation]] = {image_id: [] for image_id in self.images.keys()}
# Define the means for the distributions generating bounding box sizes and counts per image
mean_size: float = sample_info.avg_bbox_area
mean_count: float = sample_info.avg_bboxes_per_image
# Define the upper bound for random category ID number generation
max_category_id: int = len(self.categories_dict)
# Define a counter for annotation IDs
annotation_id: int = 0
for image_id, image in self.images.items():
# Generate the actual number of bounding boxes to create per image
actual_count: int = round(uniform(0, round(mean_count * 1.5)))
for _ in range(actual_count):
# Generate a random category ID to assign to the annotation
random_category_id: int = randint(1, max_category_id)
# Generate the random size and starting coordinates for the annotation
actual_size: float = gauss(mean_size, self.deviation)
skew_parameter: float = uniform(0.2, 0.8)
bbox_width: int = int(actual_size ** 0.5 * skew_parameter)
bbox_height: int = int(actual_size ** 0.5 * (1 - skew_parameter))
bbox_x: int = randint(0, max(0, image.width - bbox_width))
bbox_y: int = randint(0, max(0, image.height - bbox_height))
bbox_coords: list[int] = [bbox_x, bbox_y, bbox_width, bbox_height]
# Create and store the annotation
annotations[image_id].append(RandomAnnotation(
gt_name=self.name,
id=str(annotation_id),
category_id=random_category_id,
image_id=image_id,
coordinates=bbox_coords
))
annotation_id += 1
return annotations
def __str__(self):
gt_string = "Ground truth: " + self.name + "\n"
gt_string += "=" * 80 + "\n"
for img_filename, _ in self.images.items():
gt_string += "Image: " + str(img_filename) + "\n\n"
gt_string += self.images[img_filename].__str__()
if self.annotations[img_filename]:
gt_string += "\nAnnotations for this image:\n\n"
for ann in self.annotations[img_filename]:
gt_string += ann.__str__()
else:
gt_string += "\nNo annotations for this image.\n"
gt_string += "-" * 80 + "\n"
return gt_string