import gradio as gr
from fastai.vision.all import load_learner
from fastai import *
import torch
import os
from PIL import Image

model_path = 'multi_target_resnet18.pkl'
model = load_learner(model_path)

def result(path):    
    pred,_,probability = model.predict(path)
    arr = ['Name','Status','Disease Name']
    vals = ['', '', '']

    names = ['Maple', 'Banana', 'Cucumber', 'Mango', 'Maple', 'Pepper', 'Rose', 'Tomato']
    status = ['diseased', 'no disease found']

    for x in pred:
      if x in names:
        vals[0] = x.capitalize()
      elif x in status:
        vals[1] = x.capitalize()
      elif x == 'healthy':
        vals[2] = 'None'
      else:
        vals[2] = x.capitalize()
      
    return f'{arr[0]}:\t{vals[0]}\n{arr[1]}:\t{vals[1]}\n{arr[2]}:\t{vals[2]}\n'

path = 'test-images/'

image_path = []

for i in os.listdir(path):
  image_path.append(path+i) 

image = gr.components.Image(shape =(300,300))
label = gr.components.Label()

iface = gr.Interface(fn=result, inputs=image, outputs='text', examples = image_path)
iface.launch(inline = False)