import torch
from torch import nn
from torch.nn import functional as F
import torchvision
from torchvision import transforms

from PIL import Image

from captum.attr import IntegratedGradients
from captum.attr import GradientShap
from captum.attr import Occlusion
from captum.attr import NoiseTunnel
from captum.attr import visualization as viz
import jovian 

from matplotlib.colors import LinearSegmentedColormap
import json
import numpy as np
import os

model = torchvision.models.resnet34(pretrained=False);
model.fc = torch.nn.Linear(in_features=512, out_features=8)
checkpoint = torch.load('IC_ResNet34_95.pth',map_location = torch.device('cpu'))
model.load_state_dict(checkpoint['state_dict'])

labels = ['Rs 10','Rs 20','Rs 50','Rs 100','Rs 200','Rs 500','Rs 2000','No Indian Currency']

import matplotlib.pyplot as plt
def evaluate(file_path):
    labels = ['Rs 10','Rs 20','Rs 50','Rs 100','Rs 200','Rs 500','Rs 2000','No Indian Currency']
    img = Image.open(file_path)
    img = transforms.Resize(size=(224, 224))(img)
    img = transforms.ToTensor()(img)
    img = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(img)
    img = img.unsqueeze(0)
    
    with torch.no_grad():
        model.eval()
        prediction = model(img)
        indx = torch.argmax(prediction)
        #print(f'Predicted scenery : {labels[indx]}')
        
       # return labels[indx]
    im = plt.imread(file_path)
    plt.imshow(im)
    plt.axis('off')
    plt.title(f'Predicted Indian Currency : {labels[indx]}')

evaluate('50.jpg')