app.py

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")
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