from fastai.vision.all import *
import fastai
import streamlit as st

st.header("Outfit Color Guide")

learn = load_learner('export.pkl')
labels = learn.dls.vocab

col1, col2 = st.columns([3,3])

def predict(img):
  img = PILImage.create(img)
  pred, pred_idx, probs = learn.predict(img)
  highest_prob = probs.argmax()
  file_name = f"{labels[highest_prob].split(' ')[0]}_sugg.png"
  img_sugg = PILImage.create(file_name)
  return f'Is this {labels[highest_prob]}?', img_sugg

with col1:
  st.subheader("Take a photo of your outfit or upload the photo.")
  st.write("Please stay in natural light or any light reflecting the original color of your outfit.")
  camera = st.camera_input("Take a picture")
  upload  = st.file_uploader("Upload a photo of your outfit")

with col2:
  if camera or upload:
     if camera:
        img = camera
     if upload:
        img = upload
  
     st.image(img, width = 400)
     result, sugg = predict(img)      
     st.subheader(result)
     st.subheader('Pair your outfit with these colors.')
     st.image(sugg, width = 400)