import torch
from torchvision import transforms
import gradio as gr
import torch
from efficientnet_pytorch import EfficientNet

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = EfficientNet.from_name('efficientnet-b0')

in_features = model._fc.in_features
model._fc = torch.nn.Linear(in_features, 2) 
model.load_state_dict(torch.load('model_transfer.pt', map_location=torch.device('cpu')))

model.to(device)

model.eval()


labels = ["Organic Waste","Recyclable Waste"]

def predict(inp):
    inp = transforms.ToTensor()(inp).unsqueeze(0)
    
    inp = inp.to(device)
    
    with torch.no_grad():
        prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
        
        confidences = {labels[i]: float(prediction[i]) for i in range(len(prediction))}
    
    return confidences


gr.Interface(
    fn=predict, 
    inputs=gr.components.Image(type="pil"),
    outputs=gr.components.Label(num_top_classes=2),
    examples=["tissue.jpg", "carrots.jpg"],
    theme="default",
    css=".footer{display:none !important}"
).launch()