# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_vision_inference.ipynb.

# %% auto 0
__all__ = ['inference', 'pred_postprocess']

# %% ../nbs/06_vision_inference.ipynb 1
import numpy as np
from pathlib import Path
from .vision_core import *
from .vision_augmentation import do_pad_or_crop, do_resize
from scipy.ndimage import label
from skimage.morphology import remove_small_objects

# %% ../nbs/06_vision_inference.ipynb 3
def inference(learn_inf, reorder, resample, fn:(Path,str), save_path:(str,Path)=None): 
    '''Predict on new data using exported model'''         

    org_img = load(fn)
    input_img, org_size = preprocess(org_img, reorder, resample)    
        
    pred, *_ = learn_inf.predict(input_img.data);
    
    pred_mask = do_pad_or_crop(pred.float(), input_img.shape[1:], padding_mode=0, mask_name=None)
    input_img.set_data(pred_mask)
    
    input_img = do_resize(input_img, org_size, image_interpolation='nearest')
    
    reoriented_array = to_original_orientation(input_img.as_sitk(), ('').join(org_img.orientation))
    
    org_img.set_data(reoriented_array)

    if save_path:
        save_fn = 'pred_' + Path(fn).parts[-1]
        save_fn = Path(save_path)/save_fn
        org_img.save(save_fn)

    return org_img

# %% ../nbs/06_vision_inference.ipynb 5
def pred_postprocess(pred_mask, remove_size=10437, percentage=0.2): 
    '''Remove small objects from predicted mask.'''
    small_objects = remove_size*percentage    
    labeled_mask, ncomponents = label(pred_mask)
    labeled_mask = remove_small_objects(labeled_mask, min_size=small_objects)
    
    return np.where(labeled_mask>0, 1., 0.)
