[Streamlit] 1차 프로토타입

[baekkr95]

Background

• streamlit 화면인 app.py 만 수정했기 때문에 이슈에 app.py 코드만 올렸습니다.
• 사진 업로드하면, 중앙에 '원하는 구름 이미지를 선택해주세요' 버튼이 뜨는데, 이를 누르면 자동으로 '구름 선택 안함'이 설정되서 dehazed image가 왼쪽에 나타납니다.
• 그 후 자유롭게 어떤 구름을 선택할 지 결정하면 됩니다.

Content

import streamlit as st

import io
import os
import yaml

from PIL import Image

from my_predict import get_prediction

from confirm_button_hack import cache_on_button_press

# SETTING PAGE CONFIG TO WIDE MODE
st.set_page_config(layout="wide")


root_password = 'password'


def main():
    st.title("Dehazing Model")

    uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg","png"])

    if uploaded_file:
        image_bytes = uploaded_file.getvalue()
        image = Image.open(io.BytesIO(image_bytes))

        # st.write("Dehazing...")
        with st.container():
            col1, col2, col3 = st.columns([2,1,2])

            with col1:
                st.header("Input Image")
                st.image(image, caption='Uploaded Image')

            with col2:
                st.header("원하는 구름 이미지")
                with st.expander("원하는 구름 이미지를 선택해주세요"):
                    option = st.selectbox(
                        ' ',
                        ('선택 안 함', '작은 구름', '큰 구름', '분홍 구름'))
                    st.image("/opt/ml/final-project-level3-cv-17/PSD/cloud_image/cloud.png", caption = '작은 구름')
                    st.image("/opt/ml/final-project-level3-cv-17/PSD/cloud_image/cloud2.jpg", caption = '큰 구름')
                    st.image("/opt/ml/final-project-level3-cv-17/PSD/cloud_image/cloud3.jpg", caption = '분홍 구름')
                    dehaze_image = get_prediction(image_bytes)

            ### 원하는 구름 누를 때마다 옆에 이미지가 바뀌도록 해도 괜찮을 듯
            ### get_prediction을 나중에 백그라운드 실행으로 하면 될 듯 -> 저게 실행될 동안 다른 UI가 화면에 나타나질 않음

            if option == '선택 안 함':
                with col3:
                    st.header("Dehazed Image")
                    st.spinner("dehazing now...")
                    # dehaze_image = get_prediction(image_bytes)

                    st.image(dehaze_image, caption='그냥 dehazed')
            elif option == '작은 구름':
                with col3:
                    st.header("Dehazed Image")
                    st.spinner("dehazing now...")
                    # dehaze_image = get_prediction(image_bytes)

                    st.image(dehaze_image, caption='작은 구름 이미지')

    # 처음에 나오는 예시 사진
    # 이미지를 업로드하면 사라지게 된다.
    else:
        with st.container():
            st.header('예시 사진')
            col0, col1, col2, col3 = st.columns([1,2,2,1])
            with col1:
                st.header("Dehazing 원하는 사진")
                st.image("/opt/ml/final-project-level3-cv-17/PSD/example_image/exam.jpg")
            with col2:
                st.header("Dehazing 완료 사진")
                st.image("/opt/ml/final-project-level3-cv-17/PSD/example_image/exam2.jpg")
        

@cache_on_button_press('Authenticate')
def authenticate(password) -> bool:
    print(type(password))
    return password == root_password


# password = st.text_input('password', type="password")

# if authenticate(password):
#     st.success('You are authenticated!')
#     main()
# else:
#     st.error('The password is invalid.')

main()

Details

[baekkr95]

• 이미지 다운로드 코드 추가했습니다
• dehazed image 밑에 놓으면 될 것 같습니다

from io import BytesIO

            buf = BytesIO()
            dehaze_image.save(buf, format="JPEG")
            byte_im = buf.getvalue()

            st.download_button(
                label='Download Image',
                data=byte_im,
                file_name='dehazed_image.png',
                mime='image/jpeg'
            )