# AUTOGENERATED! DO NOT EDIT! File to edit: ../09_fastai_model.ipynb.

# %% auto 0
__all__ = ['Fastai2_model']

# %% ../09_fastai_model.ipynb 3
from fastai.vision.all import Learner, ProgressCallback, PILImage #for loading the sample model
import torch #for loading the sample model
from PIL import Image
import numpy as np

# %% ../09_fastai_model.ipynb 4
class Fastai2_model:
    def __init__(self, github, model, force_reload=False):
        self.trainedModel = torch.hub.load(github,model, force_reload = force_reload)
        self.resize256 = lambda x: x.resize ((256,256))
        self.trainedModel.remove_cbs(ProgressCallback)

        
    def prepareSize(self, item):
        return self.resize256(item)
    
    def predict(self, image):
        image = PILImage.create(np.array(image))
        output = self.trainedModel.predict(image)
        output = PILImage.create(output [0])
        output = Image.fromarray(np.array(output)) #mode="I"
        return output #output PIL.Image.Image mode="L"
