Update app.py

app.py

@@ -1,74 +1,117 @@
-import gradio as gr
+import numpy as np
+import streamlit as st
 import requests
-from PIL import Image, ImageOps
+from PIL import Image
 from io import BytesIO
-import numpy as np
+from streamlit_drawable_canvas import st_canvas
 
-# Hàm xử lý gửi ảnh và mask tới API Magic Studio
-def process_inpaint_via_api(image, mask):
-    url = "https://api.magicstudio.com/magiceraser/erase"
-    
-    headers = {
-        'accessToken': 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJjbGllbnRfaWQiOiIyRGZQSm1sY1BLalBuaGh4ajBxSVpmQktfbS12RFhTc1NMRFNOX1gtZlVVIiwiZXhwIjoxNzI3MTY0MTcyLCJhcHBfbmFtZSI6IjE3NTMwODQiLCJtZXRhbnRfdHlwZSI6ImNyZWQifQ.24oHGH9ial_gqeYrpUbYP68MwYly2ArzXbhbE1hGzQo'
-    }
-    
-    # Chuyển đổi ảnh và mask thành định dạng PNG để gửi tới API
-    image_io = BytesIO()
-    mask_io = BytesIO()
-    
-    image.save(image_io, format="PNG")
-    mask.save(mask_io, format="PNG")
-    
-    image_io.seek(0)
-    mask_io.seek(0)
-    
-    files = [
-        ('image_file', ('image.png', image_io, 'image/png')),
-        ('mask_file', ('mask.png', mask_io, 'image/png'))
-    ]
-    
-    # Gửi yêu cầu POST tới API
-    response = requests.post(url, headers=headers, files=files)
-    
-    if response.status_code == 200:
-        # Nếu API trả về thành công, xử lý và trả về ảnh kết quả
-        return Image.open(BytesIO(response.content))
-    else:
-        # Nếu có lỗi, hiển thị lỗi
-        return f"API Error: {response.status_code} - {response.text}"
-
-# Hàm xử lý kết hợp ảnh và mask
-def process(image, mask):
-    img = Image.fromarray(image).convert("RGBA")
-    
-    # Mask nhận được từ Image Sketchpad dưới dạng ảnh RGBA, cần chuyển sang ảnh nhị phân
-    mask = Image.fromarray(mask).convert("L")
-    mask = ImageOps.invert(mask)  # Invert lại mask để đảo ngược vùng vẽ
-    
-    # Gọi API với ảnh và mask
-    result = process_inpaint_via_api(img, mask)
+# API information
+url = "https://api.magicstudio.com/magiceraser/erase"
+access_token = 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJjbGllbnRfaWQiOiIyRGZQSm1sY1BLalBuaGh4ajBxSVpmQktfbS12RFhTc1NMRFNOX1gtZlVVIiwiZXhwIjoxNzI3MTY0MTcyLCJhcHBfbmFtZSI6IjE3NTMwODQiLCJtZXRhZGF0YSI6bnVsbCwiZ3JhbnRfdHlwZSI6ImNyZWQifQ.24oHGH9ial_gqeYrpUbYP68MwYly2ArzXbhbE1hGzQo'  # Your actual token
+
+def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
+    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"
     
-    # Nếu kết quả trả về là lỗi, hiển thị thông báo
-    if isinstance(result, str):
-        return None, result
+    buf = BytesIO()
+    pil_image.save(buf, format=pil_format)
     
-    # Trả về ảnh đã xử lý
-    return result, "Processing complete."
+    return st.download_button(
+        label=label,
+        data=buf.getvalue(),
+        file_name=f'{filename}.{file_format}',
+        mime=mime,
+    )
 
-# Tạo giao diện Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# AI Object Removal using Magic Studio API")
-    
-    # Tải ảnh lên từ người dùng và cho phép vẽ mask đè trực tiếp lên ảnh đã tải lên
-    image_input = gr.Image(type="numpy", tool="sketch", label="Upload an image and draw mask")
+# Set title of the app
+st.title("AI Photo Object Removal")
 
-    # Khu vực hiển thị kết quả
-    result_output = gr.Image(label="Output Image")
-    output_message = gr.Textbox(label="Message", lines=2)
-    
-    # Nút xử lý và kết nối các thành phần với hàm process
-    process_button = gr.Button("Submit")
-    process_button.click(process, inputs=[image_input, image_input], outputs=[result_output, output_message])
+# Upload the image
+uploaded_file = st.file_uploader("Upload an image to remove objects from", type=["png", "jpg", "jpeg"])
+
+if uploaded_file is not None:
+    # Load the uploaded image
+    img_input = Image.open(uploaded_file).convert("RGBA")
+
+    # Resize image if too large
+    max_size = 2000
+    img_width, img_height = img_input.size
+    if img_width > max_size or img_height > max_size:
+        if img_width > img_height:
+            new_width = max_size
+            new_height = int((max_size / img_width) * img_height)
+        else:
+            new_height = max_size
+            new_width = int((max_size / img_height) * img_width)
+        img_input = img_input.resize((new_width, new_height))
+
+    # Display the image and let the user draw the mask
+    stroke_width = st.slider("Brush size", 1, 100, 50)
+    st.write("**Draw over the parts of the image you want to remove.**")
+
+    # Create a drawing canvas for the mask
+    canvas_result = st_canvas(
+        stroke_color="rgba(255, 0, 255, 1)",  # Pink color for the brush
+        stroke_width=stroke_width,
+        background_image=img_input,
+        update_streamlit=True,
+        height=img_input.height,
+        width=img_input.width,
+        drawing_mode="freedraw",
+        key="canvas",
+    )
+
+    if canvas_result.image_data is not None:
+        # Convert the drawing into a mask
+        mask_image = Image.fromarray((canvas_result.image_data[:, :, 3] > 0).astype(int) * 255)
+        mask_image = mask_image.resize(img_input.size).convert("L")  # Convert to grayscale
+
+        # Display the mask for the user to review
+        st.write("**Generated Mask:**")
+        st.image(mask_image)
+
+        if st.button('Submit'):
+            with st.spinner("AI is processing..."):
+                try:
+                    # Convert both image and mask to binary format for API upload
+                    img_buffer = BytesIO()
+                    mask_buffer = BytesIO()
+
+                    img_input.save(img_buffer, format='PNG')
+                    mask_image.save(mask_buffer, format='PNG')
+
+                    # Prepare the files for the API
+                    files = [
+                        ('image_file', ('image.png', img_buffer.getvalue(), 'image/png')),
+                        ('mask_file', ('mask.png', mask_buffer.getvalue(), 'image/png')),
+                    ]
+                    
+                    headers = {
+                        'accessToken': access_token
+                    }
+
+                    # Send POST request to MagicEraser API
+                    response = requests.post(url, headers=headers, files=files)
+
+                    if response.status_code == 200:
+                        # Process and display the result
+                        result_image = Image.open(BytesIO(response.content)).convert("RGB")
+                        st.write("AI has finished the job!")
+                        st.image(result_image)
 
-# Khởi chạy ứng dụng Gradio
-demo.launch()
+                        # Download button for the output image
+                        uploaded_name = uploaded_file.name.split('.')[0]
+                        image_download_button(
+                            pil_image=result_image,
+                            filename=f"{uploaded_name}_output",
+                            fmt="jpg",
+                            label="Download Image"
+                        )
+                    else:
+                        st.error(f"Error: {response.status_code} - {response.text}")
+                except Exception as e:
+                    st.error(f"Error processing the image: {e}")
+if __name__ == "__main__":
+    import os
+    os.system("streamlit run app.py")