fix

app.py

@@ -5,75 +5,62 @@ st.set_page_config(layout="wide")
 from streamlit_drawable_canvas import st_canvas
 from PIL import Image
 from typing import Union
-import random
 import numpy as np
-import os
-import time
 
 from models import make_image_controlnet, make_inpainting
 from preprocessing import preprocess_seg_mask, get_image, get_mask
 from explanation import make_inpainting_explanation, make_regeneration_explanation, make_segmentation_explanation
 
 
-
 def on_upload() -> None:
     """Upload image to the canvas."""
     if 'input_image' in st.session_state and st.session_state['input_image'] is not None:
         image = Image.open(st.session_state['input_image']).convert('RGB')
         st.session_state['initial_image'] = image
-        if 'seg' in st.session_state:
-            del st.session_state['seg']
-        if 'unique_colors' in st.session_state:
-            del st.session_state['unique_colors']
-        if 'output_image' in st.session_state:
-            del st.session_state['output_image']
-
-def make_image_row(image_0, image_1):
-    col_0, col_1 = st.columns(2)
-    with col_0:
-        st.image(image_0, use_column_width=True)
-    with col_1:
-        st.image(image_1, use_column_width=True)
+        st.session_state.pop('seg', None)
+        st.session_state.pop('unique_colors', None)
+        st.session_state.pop('output_image', None)
+
 
 def check_reset_state() -> bool:
     """Check whether the UI elements need to be reset"""
-    if ('reset_canvas' in st.session_state and st.session_state['reset_canvas']):
+    if st.session_state.get('reset_canvas', False):
         st.session_state['reset_canvas'] = False
         return True
-    st.session_state['reset_canvas'] = False
     return False
 
 
-def move_image(source: Union[str, Image.Image],
-               dest: str,
-               rerun: bool = True,
-               remove_state: bool = True) -> None:
+def move_image(source: Union[str, Image.Image], dest: str, rerun: bool = True, remove_state: bool = True) -> None:
     """Move image from source to destination."""
-    source_image = source if isinstance(source, Image.Image) else st.session_state[source]
+    source_image = source if isinstance(source, Image.Image) else st.session_state.get(source)
 
     if remove_state:
         st.session_state['reset_canvas'] = True
-        if 'seg' in st.session_state:
-            del st.session_state['seg']
-        if 'unique_colors' in st.session_state:
-            del st.session_state['unique_colors']
-
-    st.session_state[dest] = source_image
-    if rerun:
-        st.experimental_rerun()
+        st.session_state.pop('seg', None)
+        st.session_state.pop('unique_colors', None)
 
-
-def on_change_radio() -> None:
-    """Reset the UI elements when the radio button is changed."""
-    st.session_state['reset_canvas'] = True
+    if source_image:
+        st.session_state[dest] = source_image
+        if rerun:
+            st.experimental_rerun()
 
 
 def make_canvas_dict(canvas_color, brush, paint_mode, _reset_state):
+    """Create a dictionary for the canvas settings."""
+    background_image = st.session_state.get('initial_image', None)
+    
+    if isinstance(background_image, str):  # Convert if it's a file path
+        try:
+            background_image = Image.open(background_image)
+        except Exception as e:
+            st.error(f"Failed to load background image: {e}")
+            background_image = None
+
     canvas_dict = {
         'fill_color': canvas_color,
         'stroke_color': canvas_color,
         'background_color': "#FFFFFF",
-        'background_image': st.session_state['initial_image'] if 'initial_image' in st.session_state else None,
+        'background_image': background_image if isinstance(background_image, Image.Image) else None,
         'stroke_width': brush,
         'initial_drawing': {'version': '4.4.0', 'objects': []} if _reset_state else None,
         'update_streamlit': True,
@@ -86,22 +73,25 @@ def make_canvas_dict(canvas_color, brush, paint_mode, _reset_state):
 
 
 def make_prompt_row():
+    """Create input fields for positive and negative prompts."""
     col_0_0, col_0_1 = st.columns(2)
     with col_0_0:
         st.text_input(label="Positive prompt", 
-                     value="a realistic photograph of a room, high resolution",
-                     key='positive_prompt')
+                      value="a realistic photograph of a room, high resolution",
+                      key='positive_prompt')
     with col_0_1:
         st.text_input(label="Negative prompt", 
-                     value="watermark, banner, logo, watermark, contactinfo, text, deformed, blurry, blur, lowres", 
-                     key='negative_prompt')
+                      value="watermark, banner, logo, contact info, text, deformed, blurry, lowres", 
+                      key='negative_prompt')
+
 
 def make_sidebar():
+    """Create the sidebar with upload options and settings."""
     with st.sidebar:
         input_image = st.file_uploader("Upload Image", 
-                                     type=["png", "jpg"], 
-                                     key='input_image', 
-                                     on_change=on_upload)
+                                       type=["png", "jpg"], 
+                                       key='input_image', 
+                                       on_change=on_upload)
         generation_mode = 'Inpainting'
         paint_mode = 'freedraw'
         brush = 30 if paint_mode == "freedraw" else 5
@@ -109,23 +99,25 @@ def make_sidebar():
     return input_image, generation_mode, brush, color_chooser, paint_mode
 
 
-
 def make_output_image():
-    if 'output_image' in st.session_state:
-        output_image = st.session_state['output_image']
-        if isinstance(output_image, np.ndarray):
-            output_image = Image.fromarray(output_image)
-        if isinstance(output_image, Image.Image):
-            output_image = output_image.resize((512, 512))
-    else:
+    """Display the output image."""
+    output_image = st.session_state.get('output_image', None)
+
+    if isinstance(output_image, np.ndarray):
+        output_image = Image.fromarray(output_image)
+
+    if output_image is None:
         output_image = Image.new('RGB', (512, 512), (255, 255, 255))
 
     st.write("#### Output image")
     st.image(output_image, width=512)
+
     if st.button("Move to input image"):
         move_image('output_image', 'initial_image', remove_state=True, rerun=True)
 
+
 def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, paint_mode):
+    """Create an editable drawing canvas."""
     st.write("#### Input image")
     canvas_dict = make_canvas_dict(
         canvas_color=canvas_color,
@@ -135,35 +127,49 @@ def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, pain
     )
 
     if generation_mode == "Segmentation":
-        canvas = st_canvas(**canvas_dict)
+        st_canvas(**canvas_dict)
+
     elif generation_mode == "Inpainting":
-        image = get_image()  # Assuming this function exists in your preprocessing module
+        image = get_image()
+        if image is None:
+            st.error("Error: Could not load image for inpainting.")
+            return
+
         canvas = st_canvas(**canvas_dict)
 
         if st.button("Generate images", key='generate_button'):
-            canvas_mask = canvas.image_data
+            canvas_mask = getattr(canvas, 'image_data', None)
+
+            if canvas_mask is None:
+                st.error("Error: No mask data found on canvas.")
+                return
+
             if not isinstance(canvas_mask, np.ndarray):
                 canvas_mask = np.array(canvas_mask)
-            mask = get_mask(canvas_mask)  # Assuming this function exists in your preprocessing module
 
-            with st.spinner(text="Generating new images"):
-                print("Making image")
+            mask = get_mask(canvas_mask)
+
+            with st.spinner(text="Generating new images..."):
                 result_image = make_inpainting(
                     positive_prompt=st.session_state['positive_prompt'],
                     image=Image.fromarray(image),
                     mask_image=mask,
                     negative_prompt=st.session_state['negative_prompt'],
                 )
+
                 if isinstance(result_image, np.ndarray):
                     result_image = Image.fromarray(result_image)
+
                 st.session_state['output_image'] = result_image
 
+
 def main():
+    """Main Streamlit app function."""
     st.write("## Virtual Staging")
 
     input_image, generation_mode, brush, color_chooser, paint_mode = make_sidebar()
 
-    if not ('initial_image' in st.session_state and st.session_state['initial_image'] is not None):
+    if 'initial_image' not in st.session_state or st.session_state['initial_image'] is None:
         st.write("Please upload an image to begin")
     else:
         make_prompt_row()
@@ -182,7 +188,6 @@ def main():
         with col2:
             make_output_image()
 
+
 if __name__ == "__main__":
     main()
-    
-