Spaces:
Sleeping
Sleeping
import streamlit as st | |
import torch | |
from PIL import Image | |
from read import classify | |
st.title("Pizza & Not Pizza") | |
device = torch.device("cpu") | |
checkpoint = torch.load(r"best.pth.tar", map_location=device) | |
model = checkpoint['model'] | |
classes = checkpoint['classes'] | |
tran = checkpoint['transform'] | |
# upload image | |
uploaded_file = st.file_uploader("Choose an image...", type="jpg") | |
taking_picture = st.camera_input("Take a picture...") | |
if uploaded_file is not None: | |
img = Image.open(uploaded_file) | |
st.image(img, caption='Uploaded Image.', use_column_width=True) | |
label = classify(model, img, tran, classes, device) | |
st.write(label) | |
elif taking_picture is not None: | |
img = Image.open(taking_picture) | |
st.image(img, caption='Uploaded Image.', use_column_width=True) | |
label = classify(model, img, tran, classes, device) | |
st.write(label) |