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

__all__ = []

# Cell
from fastai.vision.all import *
from fastcore.foundation import patch
from .data import TileDataset
from scipy.stats import entropy

# Cell
@patch
def predict_tiles(self:Learner, ds_idx=1, dl=None, mc_dropout=False, n_times=1):
    "Make predictions with dropout applied."

    if dl is None: dl = self.dls[ds_idx].new(shuffled=False, drop_last=False)

    self.model.eval()
    if mc_dropout: self.apply_dropout()

    mean_list = []
    std_list = []
    for data in progress_bar(dl):
        if isinstance(data, TensorImage):
            images = data
        else:
            images, _, _ = data
        out_list = []
        for t in range(n_times):
            with torch.no_grad():
                out = self.model(images)
            out = F.softmax(out, dim=1)
            out_list.append(out)
        out_stack = torch.stack(out_list)

        out_means = torch.mean(out_stack, dim=0)
        mean_list.append(out_means)

        out_sdts = torch.std(out_stack, dim=0)
        std_list.append(out_sdts)

    softmax_pred = torch.cat(mean_list).permute(0,2,3,1)
    smx_tiles = [x for x in softmax_pred.cpu().numpy()]

    std_pred = torch.cat(std_list).permute(0,2,3,1)
    std_tiles = [x for x in std_pred.cpu().numpy()]

    smxcores = dl.reconstruct_from_tiles(smx_tiles)
    segmentations = [np.argmax(x, axis=-1) for x in smxcores]
    std_deviations = dl.reconstruct_from_tiles(std_tiles)

    return smxcores, segmentations, std_deviations

# Cell
@patch
def apply_dropout(self:Learner):
    "If a module contains 'dropout', it will be switched to .train() mode."
    for m in self.model.modules():
        if isinstance(m, nn.Dropout):
            m.train()