# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/data.ipynb (unless otherwise specified).

__all__ = ['RecordDataframeParser', 'RecordDataset', 'ResultsDataset', 'BboxRecordDataframeParser',
           'DataDescriptorBbox', 'StatsDescriptorBbox', 'ImageStatsDescriptorBbox', 'ClassStatsDescriptorBbox',
           'GalleryStatsDescriptorBbox', 'BboxRecordDataset', 'PrecisionRecallMetricsDescriptorObjectDetection',
           'ObjectDetectionResultsDataset', 'InstanceSegmentationRecordDataframeParser',
           'DataDescriptorInstanceSegmentation', 'StatsDescriptorInstanceSegmentation',
           'ImageStatsDescriptorInstanceSegmentation', 'ClassStatsDescriptorInstanceSegmentation',
           'GalleryStatsDescriptorInstanceSegmentation', 'InstanceSegmentationRecordDataset',
           'PrecisionRecallMetricsDescriptorInstanceSegmentation', 'InstanceSegmentationResultsDataset']

# Cell
import datetime
from typing import Union, Optional, List
import os
import json
from copy import deepcopy
from random import shuffle
from abc import ABC

import numpy as np
import pandas as pd
from PIL import Image
import panel as pn

import icevision
from icevision.core.record import BaseRecord
from icevision.core.record_defaults import ObjectDetectionRecord, InstanceSegmentationRecord
import icevision.parsers as parsers
from icevision.data.data_splitter import SingleSplitSplitter
from icevision.core.bbox import BBox
from icevision.core.mask import EncodedRLEs, MaskArray
from icevision.core.class_map import ClassMap

from pycocotools import mask as mask_utils

from fastprogress import master_bar, progress_bar

from .plotting.utils import draw_record_with_bokeh
from .metrics import APObjectDetection, APInstanceSegmentation
from .core.data import *
from .utils import erles_to_counts_to_utf8, erles_to_string, string_to_erles, correct_mask, decorrect_mask

# Cell
class RecordDataframeParser(parsers.Parser):
    """IceVision parser for pandas dataframes. This parser is mostly used by the RecordDataset to load records from a saved RecordDataset."""
    def __init__(self, record_template):
        super().__init__(record_template)

    def __iter__(self):
        for group in self.record_dataframe.groupby("filepath"):
            yield group[1]

    def __len__(self):
        return self.record_dataframe["filepath"].nunique

    def record_id(self, o):
        return o.iloc[0]["id"]

    def parse_fields(self, o, record, is_new):
        width, height = o.iloc[0]["width"], o.iloc[0]["height"]
        record.set_filepath(o.iloc[0]["filepath"])
        record.set_img_size((width, height))
        record.detection.set_class_map(self.class_map)
        record.detection.add_labels(o["label"])

# Cell
class RecordDataset(GenericDataset, ABC):
    """Base class dashboard datasets that are based on IceVision records."""
    def __init__(self, records: Union[List[BaseRecord], ObservableList, str], class_map, name=None, description=None):
        if isinstance(records, str):
            self.load_from_file(records)
        else:
            if isinstance(records, icevision.data.record_collection.RecordCollection):
                records = records._records._list
            self.records = records if isinstance(records, ObservableList) else ObservableList(records)
            self.class_map = class_map
        super().__init__(self.records, name=name, description=description)
        self.records.register_callback(self.reset_infered_data)

    def __repr__(self):
        base_string = ""
        if getattr(self, "stats_dataset", None) is not None:
            for col in self.stats_dataset.columns:
                base_string += str(col) + ": " + str(self.stats_dataset[col][0]) + " | "
            base_string = base_string[:-2]
        return base_string

    def __getitem__(self, index):
        return self.records[index]

    def __len__(self):
        return len(self.records)

    def __add__(self, other):
        """Add two RecordDatases together but sets name and description to None for the new RecordDataset."""
        if sorted(self.class_map._id2class) != sorted(other.class_map._id2class):
            raise ValueError("The two RecordDatasets don't have the same class_map")
        else:
            return RecordDataset(self.records+other.records, class_map=self.class_map)

    def split_in_train_and_val(self, train_fraction):
        records = list(self.records)
        indices = list(range(len(records)))
        shuffle(indices)
        if train_fraction > 1:
            train_fraction /= len(records)
        train = [records[index] for index in indices[:int(len(records)*train_fraction)]]
        val = [records[index] for index in indices[int(len(records)*train_fraction):]]
        return train, val

    @property
    def num_images(self):
        return len(self)

    @classmethod
    def load_from_record_dataframe(cls, record_data_df: pd.DataFrame, class_map=None, name=None, description=None):
        if class_map is None:
            class_map = ClassMap(json.loads(record_data_df["class_map"].values[0])[1:])
        records = cls.parse_df_to_records(record_data_df, class_map)
        return cls(records, class_map=class_map, name=name, description=description)

    @staticmethod
    def parse_df_to_records(record_data_df):
        raise NotImplementedError()

    def load_from_file(self, path):
        data = json.load(open(path))
        df = pd.DataFrame(data["data"])
        self.class_map = ClassMap(data["class_map"])
        self._name = data["name"]
        self._description = data["description"]
        records = self.parse_df_to_records(df, self.class_map)
        self.records = ObservableList(records)

    def save(self, save_path):
        if not os.path.isdir(save_path):
            os.makedirs(save_path, exist_ok=True)
        save_name = "dataset" if self.name == "" or self.name is None else self.name
        if not os.path.isfile(os.path.join(save_path, save_name+".json")):
            save_name = save_name+".json"
        else:
            counter = 1
            while True:
                save_name = save_name+"("+str(counter)+").json"
                if os.path.isfile(os.path.join(save_path, save_name)):
                    counter += 1
                else:
                    break

        class_map = self.class_map if self.class_map is not None else self.create_class_map_from_record_df(self.data)
        save_data = {"name": self.name, "description": self.description, "data": self.data.to_dict(), "class_map": class_map._id2class}

        json.dump(save_data, open(os.path.join(save_path, save_name), "w"), default=str)

    def get_image_by_index(self, index, width, height):
        return draw_record_with_bokeh(self[index], class_map=self.class_map, width=None, height=height, return_figure=True)

    @classmethod
    def create_new_from_mask(cls, cls_instance, mask):
        selection = cls_instance.data[mask]
        filepaths = np.unique(selection["filepath"]).tolist()
        new_records = [record for record in cls_instance.records if str(record.filepath) in filepaths]
        return cls(new_records, cls_instance.class_map)

# Cell
class ResultsDataset(GenericDataset):
    """Dashboard dataset for the results of and object detection system."""
    metric_data_ap = None
    df_parser = None

    def __init__(self, dataframe, name=None, description=None):
        super().__init__(dataframe, name, description)
        # instanciate metric data and preload it
        self.metric_data_ap = None
        self.class_map = ClassMap([str(i) for i in self.base_data[["label", "label_num"]].drop_duplicates().sort_values("label_num")["label"].tolist()])

    def save(self, path):
        if not os.path.exists(os.path.join(*path.split("/")[:-1])):
            os.makedirs(os.path.join(*path.split("/")[:-1]))
        self.base_data.to_csv(path)

    def get_image_by_image_id(self, image_id, width=None, height=None):
        """For gallery dashboards"""
        df_pred = self.base_data[(self.base_data["filepath"] == image_id) & (self.base_data["is_prediction"] == True)]
        df_gt = self.base_data[(self.base_data["filepath"] == image_id) & (self.base_data["is_prediction"] == False)]

        parser_pred = self.df_parser(df_pred, self.class_map)
        parser_gt = self.df_parser(df_gt, self.class_map)
        res_gt = parser_gt.parse(show_pbar=False, autofix=False)[1][0]
        res_pred = parser_pred.parse(show_pbar=False, autofix=False)[1]
        if len(res_pred) == 0:
            res_pred = deepcopy(res_gt)
            res_pred.labels = []
            res_pred.bboxes = []
        else:
            res_pred = res_pred[0]
        plot_gt = draw_record_with_bokeh(res_gt, display_bbox=True, display_mask=True, return_figure=True, width=width, height=height)
        plot_pred = draw_record_with_bokeh(res_pred, display_bbox=True, display_mask=True, return_figure=True, width=width, height=height)
        return pn.Row(pn.Column(pn.Row("<b>Ground Truth</b>",  align="center"), plot_gt), pn.Column(pn.Row("<b>Prediction</b>",  align="center"), plot_pred))

    @classmethod
    def load(cls, path):
        df = pd.read_csv(path)
        return cls(df, None, None)

# Cell
class BboxRecordDataframeParser(RecordDataframeParser):
    """Extends the RecordDataframeParser for object detection"""
    def __init__(self, record_dataframe, class_map):
        super().__init__(ObjectDetectionRecord())
        self.record_dataframe = record_dataframe
        self.class_map = class_map

    def parse_fields(self, o, record, is_new):
        super().parse_fields(o, record, is_new)
        record.detection.add_bboxes([BBox(annot[1]["bbox_xmin"], annot[1]["bbox_ymin"], annot[1]["bbox_xmax"], annot[1]["bbox_ymax"]) for annot in o.iterrows()])

# Cell
class DataDescriptorBbox(DatasetDescriptor):
    """Dashboard dataset for object detection"""
    def calculate_description(self, obj):
        """Aggregates stats from a list of records and returns a pandas dataframe with the aggregated stats. The creation time is not necessarily the real creation time.
        This depends on the OS, for more information see: https://docs.python.org/3/library/os.html#os.stat_result."""
        data = []
        for index,record in enumerate(obj.records):
            record_commons, record_detections = record.as_dict()["common"], record.aggregate_objects()["detection"]
            for label, bbox in zip(record_detections["labels"], record_detections["bboxes"]):
                file_stats = record.filepath.stat()
                bbox_width = bbox["bbox_width"]
                bbox_height = bbox["bbox_height"]
                area = bbox_width*bbox_height
                area_normalized = area / (record.width * record.height)
                bbox_ratio = bbox_width / bbox_height
                data.append(
                    {
                        "id": record_commons["record_id"], "width": record.width, "height": record.height, "label": label, "area_square_root": bbox["bbox_sqrt_area"], "area_square_root_normalized": area_normalized**0.5,
                        "bbox_xmin": bbox["bbox_x"], "bbox_xmax": bbox["bbox_x"]+bbox["bbox_width"], "bbox_ymin": bbox["bbox_y"], "bbox_ymax": bbox["bbox_y"]+bbox["bbox_height"],
                        "bbox_xmin_normalized": bbox["bbox_x"]/record.width, "bbox_xmax_normalized": (bbox["bbox_x"]+bbox["bbox_width"])/record.width, "bbox_ymin_normalized": bbox["bbox_y"]/record.height,
                        "bbox_ymax_normalized": (bbox["bbox_y"]+bbox["bbox_height"])/record.height, "area": area, "area_normalized": area_normalized, "bbox_ratio": bbox_ratio, "bbox_ratio_normalized": bbox_ratio*(record.height/record.width),
                        "record_index": index, "bbox_width": bbox_width, "bbox_height": bbox_height, "bbox_width_normalized": bbox_width/record.width, "bbox_height_normalized": bbox_height/record.height,
                        "filepath": str(record.filepath), "creation_date": datetime.datetime.fromtimestamp(file_stats.st_ctime), "class_map": json.dumps(record.detection.class_map._id2class),
                        "modification_date": datetime.datetime.fromtimestamp(file_stats.st_mtime), "num_annotations": len(record_detections["bboxes"])
                    }
                )
        data = pd.DataFrame(data)
        data["label_num"] = data["label"]
        if obj.class_map is not None:
            data["label"] = data["label"].apply(obj.class_map.get_by_id)
        return data

# Cell
class StatsDescriptorBbox(DatasetDescriptor):
    def calculate_description(self, obj):
        stats_dict = {}
        stats_dict["no_imgs"] = [obj.data["filepath"].nunique()]
        stats_dict["no_classes"] = [obj.data["label"].nunique()]
        stats_dict["classes"] = [list(obj.data["label"].unique())]
        stats_dict["area_min"] = [obj.data["area"].min()]
        stats_dict["area_max"] = [obj.data["area"].max()]
        stats_dict["num_annotations_min"] = [obj.data["num_annotations"].min()]
        stats_dict["num_annotations_max"] = [obj.data["num_annotations"].max()]
        stats_dict["name"] = [obj._name]
        stats_dict["description"] = [obj._description]
        return pd.DataFrame(stats_dict)

# Cell
class ImageStatsDescriptorBbox(DatasetDescriptor):
    def calculate_description(self, obj):
        """Creates a dataframe containing stats about the images."""
        stats_dict = {}
        stats_dict["Num. imgs."] = obj.data["filepath"].nunique()
        stats_dict["Min Num. Objects"] = obj.data["num_annotations"].min()
        stats_dict["Max Num. Objects"] = obj.data["num_annotations"].max()
        stats_dict["Avg. Objects/Img"] = round(obj.data["num_annotations"].mean(),2)
        df = pd.DataFrame.from_dict(stats_dict, orient="index").T
        return df

# Cell
class ClassStatsDescriptorBbox(DatasetDescriptor):
    def calculate_description(self, obj):
        """Creates a dataframe containing stats about the object classes."""
        stats_dict = {}
        label_group = obj.data.groupby("label")
        for label, group in label_group:
            label_stats = {}
            label_stats["imgs"] = group["filepath"].nunique()
            label_stats["objects"] = group.shape[0]
            label_stats["avg_objects_per_img"] = label_stats["objects"]/label_stats["imgs"]
            label_stats["frac_of_labels"] = round(label_stats["objects"]/obj.data.shape[0], 2)
            stats_dict[label] = label_stats
        df = pd.DataFrame(stats_dict).T
        df = df.rename_axis('Class').reset_index()
        return df

# Cell
class GalleryStatsDescriptorBbox(DatasetDescriptor):
    def calculate_description(self, obj):
        """Creates a dataframe containing the data for a gallery."""
        df = obj.data[["id", "area", "num_annotations", "label", "bbox_ratio", "bbox_width", "bbox_height", "width", "height"]].drop_duplicates().reset_index(drop=True)
        return df

# Cell
class BboxRecordDataset(RecordDataset):
    data = DataDescriptorBbox()
    gallery_data = GalleryStatsDescriptorBbox()
    stats_dataset = StatsDescriptorBbox()
    stats_image = ImageStatsDescriptorBbox()
    stats_class = ClassStatsDescriptorBbox()
    stats = StatsDescriptorBbox()

    def __init__(self, records: Union[List[BaseRecord], ObservableList, str], class_map=None, name=None, description=None):
        super().__init__(records, class_map, name, description)
        self.record_index_image_id_map = {str(record.filepath): index for index, record in enumerate(self.records)}
        self.data = None
        self.gallery_data = None
        self.stats_dataset = None
        self.stats_image = None
        self.stats_class = None
        self.stats = None

    @staticmethod
    def parse_df_to_records(record_data_df, class_map):
        return BboxRecordDataframeParser(record_data_df, class_map).parse(SingleSplitSplitter())[0]

    def get_image_by_image_id(self, image_id, width, height):
        index = self.record_index_image_id_map[image_id]
        return draw_record_with_bokeh(self[index], display_bbox=True, display_mask=True, class_map=self.class_map, width=None, height=height, return_figure=True)

# Cell
class PrecisionRecallMetricsDescriptorObjectDetection(DatasetDescriptor):
    def __init__(self, ious=None):
        if ious is None:
            self.ious = np.arange(0.5, 1, 0.05).round(2)
        else:
            self.ious = ious

    def calculate_description(self, obj):
        return APObjectDetection(obj.base_data, self.ious).metric_data

# Cell
class ObjectDetectionResultsDataset(ResultsDataset):
    """Dashboard dataset for the results of and object detection system."""
    metric_data_ap = PrecisionRecallMetricsDescriptorObjectDetection()
    df_parser = BboxRecordDataframeParser

    @classmethod
    def init_from_preds_and_samples(cls, predictions, samples_plus_losses, padded_along_shortest=True, class_map=None, name=None, description=None):
        """The input_records are required because they are the only information source with the image stats(width, etc.) for the image on the disk."""
        data = []
        for index, (prediction, sample_plus_loss) in enumerate(zip(predictions, samples_plus_losses)):
            # TODO: At the moment only resize_and_pad or resize are handeled. Check if there are other edge cases that need to be included
            # The correction requires that the sample_plus_loss has the scaled image sizes (not the padded ones or the original ones)
            # correct the width and height to the values of the original image
            prediction = prediction.pred.as_dict()["detection"]

            if len(prediction["labels"]) > 0:
                losses = sample_plus_loss.losses
                sample_plus_loss_dict = sample_plus_loss.as_dict()
                # bbox correction
                img = Image.open(str(sample_plus_loss.common.filepath))
                width, height = img.size[0], img.size[1]
                # use bool to int for padded_along_shortest and int(sample_plus_loss["width"] < sample_plus_loss["height"]) to avoid if branches
                resize_factor = sample_plus_loss_dict["common"]["width"]/width if width > height else sample_plus_loss_dict["common"]["height"]/height
                resized_x = resize_factor * width
                resized_y = resize_factor * height
                pad_x = sample_plus_loss_dict["common"]["width"]-resized_x
                pad_y = sample_plus_loss_dict["common"]["height"]-resized_y
                correct_x = lambda x: (x-pad_x/2)/resize_factor
                correct_y = lambda y: (y-pad_y/2)/resize_factor

                for label, bbox, score in zip(prediction["labels"], prediction["bboxes"], prediction["scores"]):
                    xmin, xmax, ymin, ymax = correct_x(bbox.xmin), correct_x(bbox.xmax), correct_y(bbox.ymin), correct_y(bbox.ymax)
                    file_stats = sample_plus_loss.common.filepath.stat()
                    bbox_width = xmax - xmin
                    bbox_height = ymax - ymin
                    area = bbox_width * bbox_height
                    area_normalized = area / (width * height)
                    bbox_ratio = bbox_width / bbox_height
                    image_data = {
                            "id": sample_plus_loss_dict["common"]["record_id"], "width": width, "height": height, "label": label, "area_square_root": area**2, "area_square_root_normalized": area_normalized**2,
                            "score": score, "bbox_xmin": xmin, "bbox_xmax": xmax, "bbox_ymin": ymin, "bbox_ymax": ymax, "area": area,
                            "bbox_xmin_normalized": xmin/width, "bbox_xmax_normalized": xmax/width, "bbox_ymin_normalized": ymin/height, "bbox_ymax_normalized": ymax/height,
                            "area_normalized": area_normalized, "bbox_ratio": bbox_ratio, "bbox_ratio_normalized": bbox_ratio*(height/width), "record_index": index,
                            "bbox_width": bbox_width, "bbox_height": bbox_height, "bbox_width_normalized": bbox_width/width, "bbox_height_normalized": bbox_height/height,
                            "filepath": str(sample_plus_loss.common.filepath), "filename": str(sample_plus_loss.common.filepath).split("/")[-1], "creation_date": datetime.datetime.fromtimestamp(file_stats.st_ctime),
                            "modification_date": datetime.datetime.fromtimestamp(file_stats.st_mtime), "num_annotations": len(prediction["bboxes"]), "is_prediction": True,
                        }
                    for key, value in losses.items():
                        image_data[key] = value
                    data.append(image_data)

            if len(sample_plus_loss_dict["detection"]["labels"]) > 0:
                for label, bbox in zip(sample_plus_loss_dict["detection"]["labels"], sample_plus_loss_dict["detection"]["bboxes"]):
                    xmin, xmax, ymin, ymax = correct_x(bbox.xmin), correct_x(bbox.xmax), correct_y(bbox.ymin), correct_y(bbox.ymax)
                    file_stats = sample_plus_loss.common.filepath.stat()
                    bbox_width = xmax - xmin
                    bbox_height = ymax - ymin
                    area = bbox_width * bbox_height
                    area_normalized = area / (width * height)
                    bbox_ratio = bbox_width / bbox_height
                    image_data = {
                            "id": sample_plus_loss_dict["common"]["record_id"], "width": width, "height": height, "label": label,
                            "score": 999, "bbox_xmin": xmin, "bbox_xmax": xmax, "bbox_ymin": ymin, "bbox_ymax": ymax, "area": area, "area_square_root": area**2, "area_square_root_normalized": area_normalized**2,
                            "area_normalized": area_normalized, "bbox_ratio": bbox_ratio, "record_index": index,
                            "bbox_xmin_normalized": xmin/width, "bbox_xmax_normalized": xmax/width, "bbox_ymin_normalized": ymin/height, "bbox_ymax_normalized": ymax/height,
                            "bbox_width": bbox_width, "bbox_height": bbox_height,  "bbox_width_normalized": bbox_width/width, "bbox_height_normalized": bbox_height/height,
                            "filepath": str(sample_plus_loss.common.filepath), "filename": str(sample_plus_loss.common.filepath).split("/")[-1], "creation_date": datetime.datetime.fromtimestamp(file_stats.st_ctime),
                            "modification_date": datetime.datetime.fromtimestamp(file_stats.st_mtime), "num_annotations": len(prediction["bboxes"]), "is_prediction": False,
                        }
                    for key, value in losses.items():
                        image_data[key] = value
                    data.append(image_data)

        data = pd.DataFrame(data)

        data["label_num"] = data["label"]
        if class_map is not None:
            data["label_num"] = data["label"].apply(class_map.get_by_name)

        if not any(data["is_prediction"] == True):
            raise ValueError("No predictions found.")

        return cls(data, name, description)

# Cell
class InstanceSegmentationRecordDataframeParser(RecordDataframeParser):
    """Extends the RecordDataframeParser for instance segmentation"""
    def __init__(self, record_dataframe, class_map):
        super().__init__(InstanceSegmentationRecord())
        self.record_dataframe = record_dataframe
        self.class_map = class_map

    def parse_fields(self, o, record, is_new):
        super().parse_fields(o, record, is_new)
        bboxes = []
        masks  = []
        for annot in o.iterrows():
            bboxes.append(BBox(annot[1]["bbox_xmin"], annot[1]["bbox_ymin"], annot[1]["bbox_xmax"], annot[1]["bbox_ymax"]))
            masks.append(EncodedRLEs([string_to_erles(annot[1]["erles_corrected"])]))
        record.detection.add_bboxes(bboxes)
        record.detection.add_masks(masks)

# Cell
class DataDescriptorInstanceSegmentation(DatasetDescriptor):
    """Dashboard dataset for object detection"""
    def calculate_description(self, obj):
        """Aggregates stats from a list of records and returns a pandas dataframe with the aggregated stats. The creation time is not necessarily the real creation time.
        This depends on the OS, for more information see: https://docs.python.org/3/library/os.html#os.stat_result."""
        data = []
        summe = 0
        for index,record in enumerate(obj.records):
            record_commons, record_detections = record.as_dict()["common"], record.as_dict()["detection"]

            if isinstance(record_detections["masks"][0], EncodedRLEs):
                masks_to_iterate_over = [entry.erles[0] for entry in record_detections["masks"]]
            else:
                masks_to_iterate_over = record_detections["masks"][0].to_erles(None,None).erles

            for label, bbox, mask in zip(record_detections["labels"], record_detections["bboxes"], masks_to_iterate_over):
                mask_array = mask_utils.decode([mask])
                file_stats = record.filepath.stat()
                area = bbox.width*bbox.height
                area_normalized = area / (record.width * record.height)
                bbox_ratio = bbox.width / bbox.height

                data.append(
                    {
                        "id": record_commons["record_id"], "width": record.width, "height": record.height, "label": label, "bbox_area_square_root": bbox.area**0.5, "bbox_area_square_root_normalized": area_normalized**0.5,
                        "bbox_xmin": bbox.xmin, "bbox_xmax": bbox.xmin+bbox.width, "bbox_ymin": bbox.ymin, "bbox_ymax": bbox.ymin+bbox.height,
                        "bbox_xmin_normalized": bbox.xmin/record.width, "bbox_xmax_normalized": (bbox.xmin+bbox.width)/record.width, "bbox_ymin_normalized": bbox.ymin/record.height,
                        "bbox_ymax_normalized": (bbox.ymin+bbox.height)/record.height, "bbox_area": area, "bbox_area_normalized": area_normalized, "bbox_ratio": bbox_ratio, "bbox_ratio_normalized": bbox_ratio*(record.height/record.width),
                        "record_index": index, "bbox_width": bbox.width, "bbox_height": bbox.height, "bbox_width_normalized": bbox.width/record.width, "bbox_height_normalized": bbox.height/record.height,
                        "filepath": str(record.filepath), "creation_date": datetime.datetime.fromtimestamp(file_stats.st_ctime),
                        "modification_date": datetime.datetime.fromtimestamp(file_stats.st_mtime), "num_annotations": len(record_detections["bboxes"]),
                        "erles_corrected": erles_to_string(mask), "erles": erles_to_string(mask), "mask_area": mask_array.sum(), "mask_area_normalized": mask_array.sum()/mask_array.size, "mask_area_normalized_by_bbox_area": mask_array.sum()/area,
                    }
                )

        data = pd.DataFrame(data)
        data["label_num"] = data["label"]
        if obj.class_map is not None:
            data["label_num"] = data["label"].apply(obj.class_map.get_by_name)
        return data

# Cell
class StatsDescriptorInstanceSegmentation(DatasetDescriptor):
    def calculate_description(self, obj):
        stats_dict = {}
        stats_dict["no_imgs"] = [obj.data["filepath"].nunique()]
        stats_dict["no_classes"] = [obj.data["label"].nunique()]
        stats_dict["classes"] = [list(obj.data["label"].unique())]
        stats_dict["mask_area_min"] = [obj.data["mask_area"].min()]
        stats_dict["mask_area_max"] = [obj.data["mask_area"].max()]
        stats_dict["num_annotations_min"] = [obj.data["num_annotations"].min()]
        stats_dict["num_annotations_max"] = [obj.data["num_annotations"].max()]
        stats_dict["name"] = [obj._name]
        stats_dict["description"] = [obj._description]
        return pd.DataFrame(stats_dict)

# Cell
class ImageStatsDescriptorInstanceSegmentation(ImageStatsDescriptorBbox):
    pass

# Cell
class ClassStatsDescriptorInstanceSegmentation(ClassStatsDescriptorBbox):
    def calculate_description(self, obj):
        """Creates a dataframe containing stats about the object classes."""
        stats_dict = {}
        label_group = obj.data.groupby("label")
        for label, group in label_group:
            label_stats = {}
            label_stats["imgs"] = group["filepath"].nunique()
            label_stats["objects"] = group.shape[0]
            label_stats["avg_objects_per_img"] = label_stats["objects"]/label_stats["imgs"]
            label_stats["frac_of_labels"] = round(label_stats["objects"]/obj.data.shape[0], 2)
            stats_dict[label] = label_stats
        df = pd.DataFrame(stats_dict).T
        df = df.rename_axis('Class').reset_index()
        return df

# Cell
class GalleryStatsDescriptorInstanceSegmentation(DatasetDescriptor):
    def calculate_description(self, obj):
        """Creates a dataframe containing the data for a gallery."""
        df = obj.data[["id", "bbox_area", "num_annotations", "label", "bbox_ratio", "bbox_width", "bbox_height", "mask_area", "mask_area_normalized", "width", "height"]].drop_duplicates().reset_index(drop=True)
        return df

# Cell
class InstanceSegmentationRecordDataset(RecordDataset):
    data = DataDescriptorInstanceSegmentation()
    gallery_data = GalleryStatsDescriptorInstanceSegmentation()
    stats_dataset = StatsDescriptorInstanceSegmentation()
    stats_image = ImageStatsDescriptorInstanceSegmentation()
    stats_class = ClassStatsDescriptorInstanceSegmentation()
    stats = StatsDescriptorInstanceSegmentation()

    def __init__(self, records: Union[List[BaseRecord], ObservableList, str], class_map=None, name=None, description=None):
        super().__init__(records, class_map, name, description)
        self.record_index_image_id_map = {str(record.filepath): index for index, record in enumerate(self.records)}
        self.data = None
        self.gallery_data = None
        self.stats_dataset = None
        self.stats_image = None
        self.stats_class = None
        self.stats = None

    @staticmethod
    def parse_df_to_records(record_data_df, class_map):
        return InstanceSegmentationRecordDataframeParser(record_data_df, class_map).parse(SingleSplitSplitter())[0]

    def get_image_by_image_id(self, image_id, width, height):
        index = self.record_index_image_id_map[image_id]
        return draw_record_with_bokeh(self[index], display_bbox=True, display_mask=True, class_map=self.class_map, width=None, height=height, return_figure=True)

# Cell
class PrecisionRecallMetricsDescriptorInstanceSegmentation(DatasetDescriptor):
    def __init__(self, ious=None):
        if ious is None:
            self.ious = np.arange(0.5, 1, 0.05).round(2)
        else:
            self.ious = ious

    def calculate_description(self, obj):
        return APInstanceSegmentation(obj.base_data, self.ious).metric_data

# Cell
class InstanceSegmentationResultsDataset(ResultsDataset):
    """Dashboard dataset for the results of and object detection system."""
    metric_data_ap = PrecisionRecallMetricsDescriptorInstanceSegmentation()
    df_parser = InstanceSegmentationRecordDataframeParser

    @staticmethod
    def get_masks_to_iterate_over(prediction):
        if len(prediction["masks"]) > 0:
            if isinstance(prediction["masks"][0], EncodedRLEs):
                masks_to_iterate_over = [entry.erles[0] for entry in prediction["masks"]]
            elif isinstance(prediction["masks"][0], MaskArray):
                masks_to_iterate_over = [entry.to_erles(None, None).erles[0] for entry in prediction["masks"]]
            else:
                masks_to_iterate_over = prediction["masks"][0].to_erles(None,None).erles
        else:
            masks_to_iterate_over = prediction["mask_array"].to_erles(None,None).erles
        return masks_to_iterate_over

    @classmethod
    def init_from_preds_and_samples(cls, predictions, samples_plus_losses, padded_along_shortest=True, class_map=None, name=None, description=None):
        """The input_records are required because they are the only information source with the image stats(width, etc.) for the image on the disk."""
        data = []
        for index, (prediction, sample_plus_loss) in enumerate(zip(predictions, samples_plus_losses)):
            # TODO: At the moment only resize_and_pad or resize are handeled. Check if there are other edge cases that need to be included
            # The correction requires that the sample_plus_loss has the scaled image sizes (not the padded ones or the original ones)
            # correct the width and height to the values of the original image
            prediction = prediction.pred.as_dict()["detection"]
            losses = sample_plus_loss.losses
            sample_plus_loss_dict = sample_plus_loss.as_dict()
            # bbox correction
            img = Image.open(sample_plus_loss.common.filepath)
            width, height = img.size[0], img.size[1]
            # use bool to int for padded_along_shortest and int(sample_plus_loss["width"] < sample_plus_loss["height"]) to avoid if branches
            resize_factor = sample_plus_loss_dict["common"]["width"]/width if width > height else sample_plus_loss_dict["common"]["height"]/height
            resized_x = resize_factor * width
            resized_y = resize_factor * height
            pad_x = sample_plus_loss_dict["common"]["width"]-resized_x
            pad_y = sample_plus_loss_dict["common"]["height"]-resized_y
            correct_x = lambda x: (x-pad_x/2)/resize_factor
            correct_y = lambda y: (y-pad_y/2)/resize_factor

            if len(prediction["labels"]) > 0:
                masks_to_iterate_over = cls.get_masks_to_iterate_over(prediction)

                for label, bbox, score, mask in zip(prediction["labels"], prediction["bboxes"], prediction["scores"], masks_to_iterate_over):
                    mask_array = mask_utils.decode([mask])
                    xmin, xmax, ymin, ymax = correct_x(bbox.xmin), correct_x(bbox.xmax), correct_y(bbox.ymin), correct_y(bbox.ymax)
                    file_stats = sample_plus_loss.common.filepath.stat()
                    bbox_width = xmax - xmin
                    bbox_height = ymax - ymin
                    area = bbox_width * bbox_height
                    area_normalized = area / (width * height)
                    bbox_ratio = bbox_width / bbox_height
                    # correct mask
                    corrected_mask = correct_mask(mask_array, pad_x, pad_y, width, height)
                    corrected_mask = corrected_mask.to_erles(None, None).erles[0]

                    image_data = {
                        "id": sample_plus_loss_dict["common"]["record_id"], "width": width, "height": height, "label": label, "bbox_area_square_root": area**2, "bbox_area_square_root_normalized": area_normalized**2,
                        "score": score, "bbox_xmin": xmin, "bbox_xmax": xmax, "bbox_ymin": ymin, "bbox_ymax": ymax, "bbox_area": area,
                        "bbox_xmin_normalized": xmin/width, "bbox_xmax_normalized": xmax/width, "bbox_ymin_normalized": ymin/height, "bbox_ymax_normalized": ymax/height,
                        "bbox_area_normalized": area_normalized, "bbox_ratio": bbox_ratio, "bbox_ratio_normalized": bbox_ratio*(height/width), "record_index": index,
                        "bbox_width": bbox_width, "bbox_height": bbox_height, "bbox_width_normalized": bbox_width/width, "bbox_height_normalized": bbox_height/height,
                        "erles_corrected": erles_to_string(corrected_mask), "erles": erles_to_string(mask), "mask_area": mask_array.sum(), "mask_area_normalized": mask_array.sum()/mask_array.size, "mask_area_normalized_by_bbox_area": mask_array.sum()/area,
                        "filepath": str(sample_plus_loss.common.filepath), "filename": str(sample_plus_loss.common.filepath).split("/")[-1], "creation_date": datetime.datetime.fromtimestamp(file_stats.st_ctime),
                        "modification_date": datetime.datetime.fromtimestamp(file_stats.st_mtime), "num_annotations": len(prediction["bboxes"]), "is_prediction": True,
                    }
                    for key, value in losses.items():
                        image_data[key] = value
                    data.append(image_data)

            if len(sample_plus_loss_dict["detection"]["labels"]) > 0:
                masks_to_iterate_over = cls.get_masks_to_iterate_over(sample_plus_loss_dict["detection"])
                for label, bbox, mask in zip(sample_plus_loss_dict["detection"]["labels"], sample_plus_loss_dict["detection"]["bboxes"], masks_to_iterate_over):
                    mask_array = mask_utils.decode([mask])
                    xmin, xmax, ymin, ymax = correct_x(bbox.xmin), correct_x(bbox.xmax), correct_y(bbox.ymin), correct_y(bbox.ymax)
                    file_stats = sample_plus_loss.common.filepath.stat()
                    bbox_width = xmax - xmin
                    bbox_height = ymax - ymin
                    area = bbox_width * bbox_height
                    area_normalized = area / (width * height)
                    bbox_ratio = bbox_width / bbox_height
                    # correct mask
                    decorrected_mask = decorrect_mask(mask_array, int(pad_x/2), int(pad_y/2), width, height)
                    decorrected_mask = decorrected_mask.to_erles(None, None).erles[0]

                    # Below the erles_corrected is set to the mask and not the mask_corrected because atm icevision does not transform the masks in the record. Only in the output of the data loader is transformed
                    image_data = {
                        "id": sample_plus_loss_dict["common"]["record_id"], "width": width, "height": height, "label": label,
                        "score": 999, "bbox_xmin": xmin, "bbox_xmax": xmax, "bbox_ymin": ymin, "bbox_ymax": ymax, "bbox_area": area, "bbox_area_square_root": area**2, "bbox_area_square_root_normalized": area_normalized**2,
                        "bbox_area_normalized": area_normalized, "bbox_ratio": bbox_ratio, "record_index": index,
                        "bbox_xmin_normalized": xmin/width, "bbox_xmax_normalized": xmax/width, "bbox_ymin_normalized": ymin/height, "bbox_ymax_normalized": ymax/height,
                        "bbox_width": bbox_width, "bbox_height": bbox_height,  "bbox_width_normalized": bbox_width/width, "bbox_height_normalized": bbox_height/height,
                        "erles_corrected": erles_to_string(mask), "erles": erles_to_string(decorrected_mask), "mask_area": mask_array.sum(), "mask_area_normalized": mask_array.sum()/mask_array.size, "mask_area_normalized_by_bbox_area": mask_array.sum()/area,
                        "filepath": str(sample_plus_loss.common.filepath), "filename": str(sample_plus_loss.common.filepath).split("/")[-1], "creation_date": datetime.datetime.fromtimestamp(file_stats.st_ctime),
                        "modification_date": datetime.datetime.fromtimestamp(file_stats.st_mtime), "num_annotations": len(prediction["bboxes"]), "is_prediction": False,
                    }
                    for key, value in losses.items():
                        image_data[key] = value
                    data.append(image_data)

        data = pd.DataFrame(data)
        data["label_num"] = data["label"]
        if class_map is not None:
            data["label_num"] = data["label"].apply(class_map.get_by_name)

        if not any(data["is_prediction"] == True):
            raise ValueError("No predictions found.")

        return cls(data, name, description)