egecandrsn's picture
Create app.py
a852064
raw
history blame contribute delete
No virus
1.26 kB
# 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}")