import torch
import torch.nn as nn
import torchvision.transforms as transforms
from efficientnet_pytorch import EfficientNet
import pandas as pd
import gradio as gr


test_weight = './weight.pt'

model = EfficientNet.from_name('efficientnet-b0')
model._fc = nn.Linear(model._fc.in_features, out_features=2)
weight = torch.load(test_weight)
model.load_state_dict(weight)
model.to('cpu')


def convert_image(image):
    transformations = [transforms.ToTensor(),
                       transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                      ]
    transform = transforms.Compose(transformations)
    image = transform(image)
    return image

def test_classification(input_img):
    model.eval()
    with torch.no_grad():
        image = convert_image(input_img)
        image = image.unsqueeze(0)
        outputs = model(image)
        likelihood_ratio = nn.functional.softmax(outputs, dim=1)
    df_likelihood = pd.DataFrame(likelihood_ratio.detach().numpy()).rename(columns={1:'Likelihood'})[['Likelihood']]
    likelihood = df_likelihood.at[0,'Likelihood']
    if likelihood > 0.052569367:
        results = 'Positive'
    elif likelihood <= 0.052569367:
        results = 'Negative'
    df_results = pd.DataFrame([results],columns=['results'])
    return df_results



iface = gr.Interface(fn=test_classification, inputs=gr.inputs.Image(type="pil"), outputs="dataframe")
iface.launch()