Spaces:
Sleeping
Sleeping
File size: 1,020 Bytes
06088b6 c2fc618 06088b6 8c92bed c2fc618 06088b6 12bc61f 06088b6 12bc61f c2fc618 12bc61f c2fc618 12bc61f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import streamlit as st
import torch
from PIL import Image
import os
from read import classify
st.title("Pizza & Not Pizza")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
checkpoint = torch.load(os.path.join(os.getcwd(), "best.pth.tar"))
model = checkpoint["model"]
classes = checkpoint["classes"]
tran = checkpoint["transform"]
# upload image
while True:
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)
else:
pass
|