import streamlit as st
import torch
from torchvision import transforms
import pickle
from PIL import Image
import io

def load_image():
    uploaded_image = st.file_uploader(label="Pick a bird image")
    if uploaded_image is not None:
        image_data = uploaded_image.getvalue()
        st.image(image_data)
        return Image.open(io.BytesIO(image_data))
    else:
        return None

def load_classes():
    with open('bird_classes_dictionary.pkl', 'rb') as f:
        class_dict = pickle.load(f)
    return class_dict

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

    model = torch.load("efficient1_re", map_location=device)
    model.eval()
    return model

def predict(model, categories, image):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    preprocess = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor()
    ])

    image_tensor = preprocess(image).to(device)
    image_batch = image_tensor.unsqueeze(0)

    with torch.inference_mode():
        output = model(image_batch)
    probabilities = torch.nn.functional.softmax(output, dim=1)

    probs, categs = torch.topk(probabilities, 1)
    st.write(list(categories.keys())[list(categories.values()).index(categs)])
    st.write(f"Prediction probability: {int(probs)}")

def main():
    st.title("Which Bird is this?")
    st.markdown("Enter the image of a bird to know its species.")
    model = load_model()
    categories = load_classes()
    image = load_image()
    result = st.button('Run on image')
    if result:
        #st.write('Calculating results...')
        predict(model, categories, image)

if __name__ == "__main__":
    main()