import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms, models
import numpy as np
import os
import huggingface_hub

if torch.cuda.is_available():
    device = torch.device("cuda:0")
else:
    device = torch.device("cpu")


net = models.resnet18(pretrained=False)
net.fc = nn.Linear(net.fc.in_features, 6)
net.to(device)

HF_Token = os.environ['HF_Token']

model = huggingface_hub.cached_download(huggingface_hub.hf_hub_url(
    'danyalmalik/sceneryclassifier', '1655988285.9725637_Acc0.88_modelweights.pth'), use_auth_token=HF_Token)

net.load_state_dict(torch.load(model, map_location=device))
net.eval()

mean = np.array([0.5, 0.5, 0.5])
std = np.array([0.25, 0.25, 0.25])

data_transforms = transforms.Compose([
    transforms.Resize((150, 150)),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)])

labels = ['Buildings', 'Forest', 'Glacier', 'Mountain', 'Sea', 'Street']

title = "Scenery Classifier"


def examples():
    number = 8
    for i in range(number):
        imgs = os.listdir('examples')
        egs = [os.path.join('examples/', eg) for eg in imgs]

    return egs


def predict(img):
    try:
        img = data_transforms(img)
        img = img.to(device)
        img = img.unsqueeze(0)

        with torch.no_grad():
            output = F.softmax(net(img), dim=1)

            pred = [output[0][i].item() for i in range(len(labels))]

    except:
        pred = [0 for i in range(len(labels))]

    weightage = {labels[i]: pred[i] for i in range(len(labels))}
    return weightage


gr.Interface(fn=predict, inputs=gr.Image(type='pil'),
             outputs='label', title=title,  examples=examples()).launch()