updated visuals

app.py

@@ -116,12 +116,16 @@ unsafe_allow_html=True)
 # Main window configuration
 st.title("Black and white colorizer")
 st.markdown("This app puts color into your black and white pictures")
-title_message = st.empty()
+st_title_message = st.empty()
+st_file_uploader = st.empty()
+st_input_img = st.empty()
+st_output_img = st.empty()
+st_download_button = st.empty()
 
-title_message.markdown("**Model loading, please wait** ⌛")
+st_title_message.markdown("**Model loading, please wait** ⌛")
 
 # # Sidebar
-color_option = st.sidebar.selectbox('Select colorizer mode',
+st_color_option = st.sidebar.selectbox('Select colorizer mode',
                                     ('Artistic', 'Stable'))
                                     
 # st.sidebar.title('Model parameters')
@@ -130,20 +134,16 @@ color_option = st.sidebar.selectbox('Select colorizer mode',
 
 # Load models
 try:
-    colorizer = load_model(model_folder, color_option)
+    colorizer = load_model(model_folder, st_color_option)
 except:
     colorizer = None
     print('Error while loading the model. Please refresh the page')
 
 if colorizer is not None:
-    title_message.markdown("**To begin, please upload an image** 👇")
+    st_title_message.markdown("**To begin, please upload an image** 👇")
 
     #Choose your own image
-    uploaded_file = st.file_uploader("Upload a black and white photo", type=['png', 'jpg', 'jpeg'])
-
-    # show = st.image(use_column_width='auto')
-    input_img_pos = st.empty()
-    output_img_pos = st.empty()
+    uploaded_file = st_file_uploader.file_uploader("Upload a black and white photo", type=['png', 'jpg', 'jpeg'])
 
     if uploaded_file is not None:
         img_name = uploaded_file.name
@@ -154,14 +154,14 @@ if colorizer is not None:
         resized_img_rgb = resize_img(img_rgb, max_img_size)
         resized_pil_img = PIL.Image.fromarray(resized_img_rgb)
 
-        title_message.markdown("**Processing your image, please wait** ⌛")
+        st_title_message.markdown("**Processing your image, please wait** ⌛")
 
         output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
         
-        title_message.markdown("**To begin, please upload an image** 👇")
+        st_title_message.markdown("**To begin, please upload an image** 👇")
 
         # Plot images
-        input_img_pos.image(resized_pil_img, 'Input image', use_column_width=True)
-        output_img_pos.image(output_pil_img, 'Output image', use_column_width=True)
+        st_input_img.image(resized_pil_img, 'Input image', use_column_width=True)
+        st_output_img.image(output_pil_img, 'Output image', use_column_width=True)
 
-        st.markdown(get_image_download_link(output_pil_img, img_name, 'Download '+img_name), unsafe_allow_html=True)
+        st_download_button.markdown(get_image_download_link(output_pil_img, img_name, 'Download Image'), unsafe_allow_html=True)