import streamlit as st
# from img_classification import teachable_machine_classification
from PIL import Image, ImageOps
import streamlit_authenticator as stauth
import yaml
from yaml.loader import SafeLoader
import torch
from diffusers import StableDiffusionPipeline

torch.layer_norm(torch.tensor([1., 2., 3.], dtype=torch.float16, device='cpu'), normalized_shape=(3,))

# authentification
with open('./bla.yaml') as file:
    config = yaml.load(file, Loader=SafeLoader)
authenticator = stauth.Authenticate(
    config['credentials'],
    config['cookie']['name'],
    config['cookie']['key'],
    config['cookie']['expiry_days'],
    config['preauthorized']
)

name, authentication_status, username = authenticator.login('Login', 'main')
if authentication_status:
    authenticator.logout('Logout', 'main')
    page = st.sidebar.selectbox("探索或预测", ("苹果病分类","bla"))
    

    if page == "苹果病分类":
        st.title("使用谷歌的可教机器进行图像分类")
        st.header("苹果病")
        st.text("上传彩色苹果叶子图片")

        pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)  
        pipe = pipe.to("cpu")#("cuda")
        

        uploaded_file = st.file_uploader("选择..", type=["jpg","png","jpeg"])
        if uploaded_file is not None:
            # image = Image.open(uploaded_file).convert('RGB')
            # st.image(image, caption='上传了图片。', use_column_width=True)
            # st.write("")
            # st.write("分类...")
            # generator = torch.Generator("cuda").manual_seed(1024)
            
            # generator = torch.Generator().manual_seed(1024)
            prompt = "a photograph of an astronaut riding a horse"
            # image_gen = pipe(prompt, generator=generator).images[0]
            image = pipe(prompt).images[0]
            st.image(image)


        #     label = teachable_machine_classification(image, 'keras_model_apple.h5')
        #     if label == 0:
        #         st.write("苹果结痂")
        #     elif label == 1:
        #         st.write("黑腐病")
        #     elif label == 2:
        #         st.write("雪松苹果锈")
        #     else:
        #         st.write("健康苹果")

        # st.text("类:苹果结痂, 黑腐病, 雪松苹果锈, 健康苹果")

        # 0 apple_scrab
        # 1 black_rot
        # 2 cedar_apple_rust
        # 3 apple_healthy

    
elif authentication_status == False:
    st.error('Username/password is incorrect')
elif authentication_status == None:
    st.warning('Please enter your username and password')