import timm
import torch
import torch.nn.functional as nnf
import gradio as gr
import numpy as np
import pandas as pd
import json
from torchvision.transforms import ToTensor, Resize
from PIL import Image


#model = torch.load("/home/user/app/model_scripted.pkl",map_location=torch.device('cpu')) 
model = torch.jit.load("/home/user/app/model_scripted.pkl")
model.eval()

with open("/home/user/app/class_mapping.json", "r") as read_file:
    classes = json.load(read_file)

def classify_image(inp):
  
  inp = Image.fromarray(inp)
  inp= Resize((300, 300))(inp)
  inp= ToTensor()(inp)
  inp = torch.unsqueeze(inp, 0)
  ##print(inp.shape)
  #inp = inp.astype(np.uint8).reshape((-1, 3, 300, 300))
  ##print(inp.shape)
  #inp = torch.from_numpy(inp).float()
  ##confidences = model(inp)

  preds = model(inp).data[0]
  #means = preds.mean(dim=0, keepdim=True)
  #stds = preds.std(dim=0, keepdim=True)
  #preds = 4 * (preds - means) / stds
  ##preds = nnf.normalize(model(inp).data[0], dim=0)
  preds = nnf.softmax(preds, dim=0)
  preds = [pred.cpu() for pred in preds]
  preds = [float(pred.detach()) for pred in preds]
  print(pd.Series(preds).describe())

  #confidences_dict = {classes[i]: float(confidences.data[0][i]) for i in range(len(confidences.data[0]))}
  confidences_dict = {classes[str(i)]: float(preds[i]) for i in range(len(preds))}

  return confidences_dict

gr.Interface(fn=classify_image, 
             inputs=gr.Image(shape=(300, 300)),
             outputs=gr.Label(num_top_classes=3)).launch(debug = True)