new version

.streamlit/config.toml

@@ -1,6 +0,0 @@
-[server]
-maxUploadSize = 10
-
-[theme]
-base="light"
-primaryColor="#0074ff"

app.py

@@ -1,80 +1,55 @@
-import streamlit as st
-import os
-from datetime import datetime
+import gradio as gr
 from PIL import Image
 from io import BytesIO
 
 from src.utils import change_background, matte
-from src.st_style import apply_prod_style
-
-# apply_prod_style(st)  # NOTE: Uncomment this for production!
 
+def process_image(uploaded_file, background_color):
+    # Add debugging statement
+    print(f"background_color received: {background_color}")
+
+    # Default value check
+    if background_color is None:
+        return "Please select a background color."
+
+    hexmap = {
+        "Transparent (PNG)": "#000000",
+        "Black": "#000000",
+        "White": "#FFFFFF",
+        "Green": "#22EE22",
+        "Red": "#EE2222",
+        "Blue": "#2222EE",
+    }
+    alpha = 0.0 if background_color == "Transparent (PNG)" else 1.0
+
+    img_input = uploaded_file  # No need to open, it's already a PIL image
+    img_matte = matte(img_input)
+    img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=hexmap[background_color])
+    
+    return img_output
 
-def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
-    if fmt not in ["jpg", "png"]:
-        raise Exception(f"Unknown image format (Available: {fmt} - case sensitive)")
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    # AI Photo Background Removal
+    
+    You want to remove your photo background, but don't have the time and effort to learn photo editing skills?
+    **This app will change or remove your photo background, in seconds.**
+    """)
+    
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload your photo here")
+        bg_color = gr.Dropdown(choices=["Transparent (PNG)", "White", "Black", "Green", "Red", "Blue"], label="Choose background color")
     
-    pil_format = "JPEG" if fmt == "jpg" else "PNG"
-    file_format = "jpg" if fmt == "jpg" else "png"
-    mime = "image/jpeg" if fmt == "jpg" else "image/png"
+    output_image = gr.Image(type="pil", label="Processed Image")
     
-    buf = BytesIO()
-    pil_image.save(buf, format=pil_format)
+    btn = gr.Button("Submit")
     
-    return st.download_button(
-        label=label,
-        data=buf.getvalue(),
-        file_name=f'{filename}.{file_format}',
-        mime=mime,
+    btn.click(fn=process_image, inputs=[image_input, bg_color], outputs=output_image)
+    
+    gr.Examples(
+        examples=[["assets/demo.jpg", "Transparent (PNG)"]],
+        inputs=[image_input, bg_color]
     )
 
-
-st.title("AI Photo Background Removal")
-st.image(Image.open("assets/demo.jpg"))
-st.write(
-    """
-    You want to remove your photo background, but don't have the time and effort to learn photo editing skills?
-    **This app will change or remove your photo background, in seconds.**
-    """
-)
-
-uploaded_file = st.file_uploader(
-    label="Upload your photo here", 
-    accept_multiple_files=False, type=["png", "jpg", "jpeg"],
-)
-
-if uploaded_file is not None:
-
-    with st.expander("Original photo", expanded=True):
-        if uploaded_file is not None:
-            st.image(uploaded_file)
-        else:
-            st.warning("You haven't uploaded any photo yet")
-
-    in_mode = st.selectbox("Choose background color", ["Transparent (PNG)", "White", "Black", "Green", "Red", "Blue"])
-    in_submit = st.button("Submit")
-
-    if uploaded_file is not None and in_submit:
-        img_input = Image.open(uploaded_file)
-        
-        with st.spinner("AI is doing magic to your photo. Please wait..."):
-            hexmap = {
-                "Transparent (PNG)": "#000000",
-                "Black": "#000000",
-                "White": "#FFFFFF",
-                "Green": "#22EE22",
-                "Red": "#EE2222",
-                "Blue": "#2222EE",
-            }
-            alpha = 0.0 if in_mode == "Transparent (PNG)" else 1.0
-            img_matte = matte(img_input)
-            img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=hexmap[in_mode])
-        
-        with st.expander("Success!", expanded=True):
-            st.image(img_output)
-            uploaded_name = os.path.splitext(uploaded_file.name)[0]
-            image_download_button(
-                pil_image=img_output,
-                filename=uploaded_name,
-                fmt="png"
-            )
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -3,5 +3,5 @@ torchvision
 numpy
 opencv-python-headless
 matplotlib
-streamlit
-altair<5
+altair<5
+gradio