"""A module to perform data augmentation during the training and/or inference process"""

from detectron2.data import detection_utils as utils
import detectron2.data.transforms as T
import copy
import torch


def custom_mapper(dataset_dict):
    """A custom mapper to make data augmentation with the images

    Custom mapper performs resizing, flipping, color editor to artificially augment data. 

    Args:
        dataset_dict (dict): a dict mapping the images along with their labels.

    Returns:
        dict: a dict mapping the augmentated images along with their labels
    """
    dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
    image = utils.read_image(dataset_dict["file_name"], format="BGR")
    transform_list = [
        T.Resize((800,600)),
        T.RandomBrightness(0.8, 1.8),
        T.RandomContrast(0.6, 1.3),
        T.RandomSaturation(0.8, 1.4),
        T.RandomRotation(angle=[90, 90]),
        T.RandomLighting(0.7),
        T.RandomFlip(prob=0.4, horizontal=False, vertical=True),
    ]
    image, transforms = T.apply_transform_gens(transform_list, image)
    dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))

    annos = [
        utils.transform_instance_annotations(obj, transforms, image.shape[:2])
        for obj in dataset_dict.pop("annotations")
        if obj.get("iscrowd", 0) == 0
    ]
    instances = utils.annotations_to_instances(annos, image.shape[:2])
    dataset_dict["instances"] = utils.filter_empty_instances(instances)
    return dataset_dict


