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)