# AUTOGENERATED! DO NOT EDIT! File to edit: dataset-image_object_detection.ipynb (unless otherwise specified).

__all__ = ['logger', 'ImageObjectDetectionDataset']


# Cell
import logging
from os.path import join, isfile
from functools import partial
from ..core import assign_arg_prefix, input_feedback
from ..annotation.core import AnnotationAdapter, SubsetType
from .image_classification import ImageClassificationDataset
from ..image.pillow_tools import assign_exif_orientation
from ..annotation.core import RegionShape, convert_region
from ..tensorflow.tfrecord_builder import create_tfrecord_file
from ..evaluation.core import box_area, intersection_box, union_box


# Cell

logger = logging.getLogger(__name__)


# Cell
class ImageObjectDetectionDataset(ImageClassificationDataset):
    """
    Object detection dataset.
    `name`: The name of the dataset.
    `base_path`: The data-set base-path.
    `imageset_path`: The imageset source path.
    `categories_path`: The path to the categories.txt file.
    `annotations_path`: The path to the annotations-file.
    `create_tfrecord`: Also create .tfrecord files.
    `join_overlapping_regions`: Whether overlapping regions of same category should be joined.
    `annotation_area_threshold`: Keep only annotations with minimum size (width or height) related to image size
    """

    def __init__(self, input_adapter: AnnotationAdapter, output_adapter: AnnotationAdapter, split=None, seed=None,
                 sample=None, tfrecord=False, join_overlapping_regions=False, annotation_area_threshold=None):
        super().__init__(input_adapter, output_adapter, split, seed, sample, tfrecord)
        self.join_overlapping_regions = join_overlapping_regions
        self.annotation_area_threshold = annotation_area_threshold

    @classmethod
    def argparse(cls, prefix=None):
        """
        Returns the argument parser containing argument definition for command line use.
        `prefix`: a parameter prefix to set, if needed
        return: the argument parser
        """
        parser = super(ImageObjectDetectionDataset, cls).argparse(prefix=prefix)
        parser.add_argument(assign_arg_prefix("--join_overlapping_regions", prefix),
                            dest="join_overlapping_regions",
                            help="Whether overlapping regions of same category should be joined.",
                            action="store_true")
        parser.add_argument(assign_arg_prefix("--annotation_area_threshold", prefix),
                            dest="annotation_area_threshold",
                            help="Keep only annotations with minimum size (width or height) related to image size.",
                            type=float,
                            default=None)
        return parser

    def validate(self):
        """
        Validates the annotations.
        return: The skipped annotations
        """
        # convert the annotations before doing validation
        self.convert_annotations()
        return super().validate()

    def copy(self, train_annotation_keys, val_annotation_keys, test_files=None):
        """
        Copy the images to the dataset and remove EXIF orientation information by hard-rotate the images.
        If tfrecords should be build, create tfrecords for train and val subsets and generate a labelmap.pbtxt file.
        `train_annotation_keys`: The list of training annotation keys
        `val_annotation_keys`: The list of validation annotation keys
        `test_files`: The list of test file paths
        return: A tuple containing train, val and test target file paths
        """

        train_targets, val_targets, test_targets = super().copy(train_annotation_keys, val_annotation_keys, test_files)

        # if creating a .tfrecord
        if self.tfrecord:
            # write train subset tfrecords
            if train_annotation_keys:
                annotations = dict(zip(train_annotation_keys, [self.annotations[key] for key in train_annotation_keys]))
                tfrecord_output_file = join(self.output_adapter.path, '{}.record'.format(str(SubsetType.TRAIN)))
                logger.info('Generate file {}'.format(tfrecord_output_file))
                create_tfrecord_file(tfrecord_output_file, self.categories, annotations)

            # write val subset tfrecords
            if val_annotation_keys:
                annotations = dict(zip(val_annotation_keys, [self.annotations[key] for key in val_annotation_keys]))
                tfrecord_output_file = join(self.output_adapter.path, '{}.record'.format(str(SubsetType.VAL)))
                logger.info('Generate file {}'.format(tfrecord_output_file))
                create_tfrecord_file(tfrecord_output_file, self.categories, annotations)

        return train_targets, val_targets, test_targets

    def convert_annotations(self):
        """
        Converts segmentation regions from polygon to rectangle, if exist
        """

        # only the trainval images have annotation, not the test images
        area_threshold = self.annotation_area_threshold

        steps = [
            {
                'name': 'position',
                'choices': {
                    's': 'Skip',  # just delete the annotation
                    'S': 'Skip All',
                    't': 'Trim',  # transform the annotation
                    'T': 'Trim All',
                },
                'choice': None,
                'condition': lambda p_min, p_max, size: p_min < 0 or p_max >= size,
                'message': '{} -> {} : {}Exceeds image {}. \n Box \n x: {} \n y: {} \n x_max: {} \n y_max: {}',
                'transform': lambda p, size=0: max(min(p, size - 1), 0),
            },
            {
                'name': 'size',
                'choices': {
                    's': 'Skip',  # just delete the annotation
                    'S': 'Skip All',
                    'k': 'Keep',  # transform the annotation (in this case do nothing)
                    'K': 'Keep All',
                },
                'choice': None,
                'condition': lambda p_min, p_max, _: p_max - p_min <= 1,
                'message': '{} -> {} : {}Shape {} is <= 1 pixel. \n Box \n x: {} \n y: {} \n x_max: {} \n y_max: {}',
                'transform': lambda p, size=0: p,
            },
            {
                'name': 'area',
                'choices': {
                    's': 'Skip',  # just delete the annotation
                    'S': 'Skip All',
                    'k': 'Keep',  # transform the annotation (in this case do nothing)
                    'K': 'Keep All',
                },
                'choice': None,
                'condition': lambda p_min, p_max, size: area_threshold and (p_max - p_min) / size <= area_threshold,
                'message': '{} <= {} percent. {}'.format('{} -> {} : {}Shape {} is', (area_threshold or 0) * 100,
                                                         ' \n Box \n x: {} \n y: {} \n x_max: {} \n y_max: {}'),
                'transform': lambda p, size=0: p,
            }
        ]

        logger.info('Start convert image annotations from {}'.format(self.input_adapter.path))

        for annotation in self.annotations.values():
            # skip file, if regions are empty or file do not exist
            if not (annotation.regions and isfile(annotation.file_path)):
                continue

            # convert from polygon to rect if needed
            for region in annotation.regions:
                convert_region(region, RegionShape.RECTANGLE)

            # try to join regions
            if self.join_overlapping_regions:
                self._join_regions(annotation.regions)

            image, _, __ = assign_exif_orientation(annotation.file_path)
            img_width, img_height = image.size
            delete_regions = {}
            for index, region in enumerate(annotation.regions):
                # validate the shape size
                x_min, x_max = region.points_x[:2]
                y_min, y_max = region.points_y[:2]
                for step in steps:
                    width_condition = step['condition'](x_min, x_max, img_width)
                    height_condition = step['condition'](y_min, y_max, img_height)
                    if width_condition or height_condition:
                        size_message = ['width'] if width_condition else []
                        size_message.extend(['height'] if height_condition else [])
                        message = step['message'].format(annotation.file_path, index, ' ',
                                                         ' and '.join(size_message),
                                                         x_min, y_min, x_max, y_max)

                        step['choice'] = input_feedback(message, step['choice'], step['choices'])

                        choice_op = step['choice'].lower()
                        # if skip the shapes
                        if choice_op == 's':
                            delete_regions[index] = True
                            message = step['message'].format(annotation.file_path, index,
                                                             '{} '.format(step['choices'][choice_op]),
                                                             ' and '.join(size_message),
                                                             x_min, y_min, x_max, y_max)
                            logger.info(message)

                            break
                        else:
                            region.points_x = list(map(partial(step['transform'], size=img_width), [x_min, x_max]))
                            region.points_y = list(map(partial(step['transform'], size=img_height), [y_min, y_max]))
                            message = step['message'].format(annotation.file_path, index,
                                                             '{} '.format(step['choices'][choice_op]),
                                                             ' and '.join(size_message),
                                                             x_min, y_min, x_max, y_max)
                            logger.info(message)

            # delete regions after iteration is finished
            for index in sorted(list(delete_regions.keys()), reverse=True):
                del annotation.regions[index]

        logger.info('Finished convert image annotations from {}'.format(self.input_adapter.path))

    @classmethod
    def _join_regions(cls, regions):
        """
        Join regions which overlaps.
        `regions`: the region to parse
        """
        len_before = len(regions)
        index_left = 0
        while index_left < len(regions):
            regions_joined = []
            region_left = regions[index_left]
            for index_right in range(len(regions)):
                if index_left == index_right:
                    continue
                region_right = regions[index_right]
                same_label_length = len(region_left.labels) == len(region_right.labels)
                same_label_contents = len(region_left.labels) == len(set(region_left.labels) & set(region_right.labels))
                same_labels = same_label_length and same_label_contents
                if same_labels:
                    bbox_left = (region_left.points_x, region_left.points_y)
                    bbox_right = (region_right.points_x, region_right.points_y)
                    inter_area = box_area(intersection_box(bbox_left, bbox_right))
                    if inter_area > 0:
                        points_x, points_y = union_box(bbox_left, bbox_right)
                        region_left.points_x = points_x
                        region_left.points_y = points_y
                        regions_joined.append(index_right)
            for index in regions_joined[::-1]:
                del regions[index]
            if not regions_joined:
                index_left += 1
        logger.info('Joined overlapping regions from {} -> {}.'.format(len_before, len(regions)))

    def build_info(self):
        """
        Converts annotations
        """
        super().build_info()
        logger.info('join_overlapping_regions: {}'.format(self.join_overlapping_regions))
        logger.info('annotation_area_threshold: {}'.format(self.annotation_area_threshold))
