# app.py
from utils import load_model, predict
import torch
import json
import requests
import io
from PIL import Image
import streamlit as st

st.title('Dog Breed Classification')

# set header
st.header('Please upload a image')

# upload file
file = st.file_uploader('', type=['jpeg', 'jpg', 'png'])
url = st.text_input("Or enter image URL:")
st.button("Load Image")

# load classifier
model = load_model('model_scripted.pt')

if file is not None:
    image = Image.open(file).convert('RGB')
    st.image(image, use_column_width=True)

    # classify image
    predicted, probability = predict(model, image)

    # write classification
    st.write("## {}".format(predicted))
    st.write("### Probability: {}%".format(int(probability * 100)))

if url != '':
    try:
        response = requests.get(url)
        image = Image.open(io.BytesIO(response.content)).convert('RGB')
        st.image(image, caption="Image from URL", use_column_width=True)

        # classify image
        predicted, probability = predict(model, image)

        # write classification
        st.write("## {}".format(predicted))
        st.write("### score: {}%".format(int(probability * 100)))

    except Exception as e:
        st.error(f"Error loading image from URL: {e}")