# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/03_segmentation.core.ipynb.

# %% auto 0
__all__ = ['show_tfms', 'SegmentationExplorer']

# %% ../../nbs/03_segmentation.core.ipynb 3
#TODO: optimize imports
import polvo
from fastai.vision.all import *

# %% ../../nbs/03_segmentation.core.ipynb 11
# TODO: Move DataExplorer base either to common or core
class SegmentationExplorer(polvo.classification.DataExplorer):
    def __init__(self, x_tl, y_tl, id2label):
        self.x_tl = x_tl
        self.y_tl = y_tl
        self.id2label = id2label
        
        self.grouped_y = self._group_y()
        
        
    def _group_y(self):
        label_stats = []
        for item_idx, y in enumerate(progress_bar(self.y_tl)):
            labels, counts = np.unique(y, return_counts=True)
            label_stats.extend([{'item_idx': item_idx,
                                 'label_id': l,
                                 'pct': c / counts.sum(),
                                 'count': c,
                                 'label': self.id2label[l]} 
                                for l, c in zip(labels, counts)])
        
        grouped_y = pd.DataFrame(label_stats)
        grouped_y = grouped_y.groupby("label_id").agg({'label': 'first',
                                                       'item_idx': list,
                                                       'pct': list,
                                                       'count': sum}).reset_index()
        grouped_y["n_samples"] = grouped_y.item_idx.str.len()
        grouped_y = grouped_y.sort_values("count", ascending=False).reset_index(drop=True)
        return grouped_y
    
    def most_common_appearances(self):
        "Return a dafarame ordered by the number of appearances of a label in the dataset."
        return self.grouped_y.sort_values('n_samples', ascending=False).reset_index(drop=True)
    
    def most_common_pixels(self):
        "Return a dataframe ordered by the most common pixels in the dataset."
        return self.most_common()

    @classmethod
    def from_datasets(cls, dss, id2label):
        return cls(dss.tls[0], dss.tls[1], id2label)

# %% ../../nbs/03_segmentation.core.ipynb 20
show_tfms = polvo.show_tfms
