import streamlit as st
from PIL import Image
from transformers import AutoImageProcessor
import torch
import joblib

with st.sidebar:
    st.subheader('Image Classifier using ResNet50')
    st.write('This is a image classification app using ResNet50. It is a state of the art model for image classification. It is a pretrained model which is trained on a large dataset of images. It can be used for classifying any image. It is a very powerful model and is very fast. It is also very accurate.')
    image = Image.open('resnet_architecture.png')
    st.image(image, caption='Bert Model')
    st.code('App Built by Ambuj Raj',language='python')


st.title('Image Classifier using ResNet50')

uploaded_file = st.file_uploader("Choose a image", type=['png', 'jpeg', 'jpg'])
if uploaded_file is not None:
    st.image(uploaded_file, width=300)
    raw_image = Image.open(uploaded_file).convert('RGB')

if st.button('Classify Image'):
    with st.spinner('Classifying Image...'):
        processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
        loaded_model = joblib.load("model.sav")
        inputs = processor(raw_image, return_tensors="pt")
        with torch.no_grad():
            logits = loaded_model(**inputs).logits
        # model predicts one of the 1000 ImageNet classes
        predicted_label = logits.argmax(-1).item()
    st.success('Image Classified!')
    st.write('Predicted Label is: ',loaded_model.config.id2label[predicted_label])